Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

… dev/unify_gpu_context3

paddle/fluid/framework/ir/remove_padding_recover_padding_pass.cc

@@ -359,6 +359,7 @@ void RemovePaddingRecoverPaddingPass::ApplyImpl(ir::Graph* graph) const {
 
     std::vector<int64_t> skip_layernorm_x_shape =
         skip_layernorm_x->Var()->GetShape();
+    check_flag = true;
     if (skip_layernorm_x_shape.size() != multihead_matmul_input_shape.size()) {
       check_flag = false;
       VLOG(3) << "Transformer model remove_padding shape check failed, return "
@@ -395,6 +396,7 @@ void RemovePaddingRecoverPaddingPass::ApplyImpl(ir::Graph* graph) const {
     GET_IR_NODE_FROM_SUBGRAPH(fc_op, fc_op, fc);
 
     std::vector<int64_t> fc_input_shape = fc_input->Var()->GetShape();
+    check_flag = true;
     if ((fc_input_shape.size() != multihead_matmul_input_shape.size()) ||
         (fc_input_shape.size() != 3)) {
       check_flag = false;
@@ -446,11 +448,13 @@ void RemovePaddingRecoverPaddingPass::ApplyImpl(ir::Graph* graph) const {
 
     std::vector<int64_t> activation_input_shape =
         activation_input->Var()->GetShape();
+    check_flag = true;
     if ((activation_input_shape.size() !=
          multihead_matmul_input_shape.size()) ||
         (activation_input_shape.size() != 3)) {
       check_flag = false;
-      VLOG(3) << "Transformer model remove_padding shape check failed, return "
+      VLOG(3) << "Activation: Transformer model remove_padding "
+                 "shape(activation_input_shape.size()) check failed, return "
                  "remove_padding pass.";
       return;
     }
@@ -465,7 +469,8 @@ void RemovePaddingRecoverPaddingPass::ApplyImpl(ir::Graph* graph) const {
       check_flag = false;
     }
     if (!check_flag) {
-      VLOG(3) << "Transformer model remove_padding shape check failed, return "
+      VLOG(3) << "Activation: Transformer model remove_padding "
+                 "shape(activation_input_shape[i]) check failed, return "
                  "remove_padding pass.";
       return;
     }
@@ -530,6 +535,7 @@ void RemovePaddingRecoverPaddingPass::ApplyImpl(ir::Graph* graph) const {
 
     std::vector<int64_t> skip_layernorm_x_shape =
         preln_skip_layernorm_x->Var()->GetShape();
+    check_flag = true;
     if (skip_layernorm_x_shape.size() != multihead_matmul_input_shape.size()) {
       check_flag = false;
       VLOG(3) << "Transformer model remove_padding shape check failed, return "

paddle/fluid/inference/api/details/zero_copy_tensor.cc

@@ -88,6 +88,11 @@ void Tensor::ReshapeStrings(const size_t &shape) {
 
 template <typename T>
 T *Tensor::mutable_data(PlaceType place) {
+#ifdef PADDLE_WITH_ONNXRUNTIME
+  if (is_ort_tensor_) {
+    return ORTGetMutableData<T>();
+  }
+#endif
   EAGER_GET_TENSOR(paddle::framework::LoDTensor);
   PADDLE_ENFORCE_GT(
       tensor->numel(),
@@ -720,6 +725,17 @@ void Tensor::SetOrtBinding(const std::shared_ptr<Ort::IoBinding> binding) {
   binding_ = binding;
 }
 
+template <typename T>
+T *Tensor::ORTGetMutableData() {
+  auto binding = binding_.lock();
+  PADDLE_ENFORCE_NOT_NULL(binding,
+                          paddle::platform::errors::PreconditionNotMet(
+                              "output tensor [%s] no binding ptr", name_));
+  std::vector<Ort::Value> outputs = binding->GetOutputValues();
+  Ort::Value &value = outputs[idx_];
+  return value.GetTensorMutableData<T>();
+}
+
 template <typename T>
 void Tensor::ORTCopyToCpu(T *data) const {
   auto binding = binding_.lock();

paddle/fluid/inference/api/paddle_tensor.h

@@ -198,6 +198,9 @@ class PD_INFER_DECL Tensor {
 
   void SetOrtBinding(const std::shared_ptr<Ort::IoBinding> binding);
 
+  template <typename T>
+  T* ORTGetMutableData();
+
   template <typename T>
   void ORTCopyFromCpu(const T* data);
 

paddle/fluid/inference/tensorrt/convert/emb_eltwise_layernorm.cc

@@ -60,6 +60,50 @@ class EmbEltwiseLayerNormOpConverter : public OpConverter {
           std::vector<std::string>{word_id_name, pos_id_name, sent_id_name};
       emb_names =
           std::vector<std::string>{word_emb_name, pos_emb_name, sent_emb_name};
+
+      auto mask_id_tensor = engine_->GetITensor("mask_id");
+      auto mask_dims = mask_id_tensor->getDimensions();
+      auto slice_start_dims = mask_dims;
+      auto slice_stride_dims = mask_dims;
+
+      for (int i = 0; i < mask_dims.nbDims; i++) {
+        slice_start_dims.d[i] = 0;
+        slice_stride_dims.d[i] = 1;
+      }
+
+      auto* shape_tensor = Shape(mask_id_tensor);
+      std::vector<nvinfer1::ITensor*> size_vec_tensor;
+      for (int i = 0; i < mask_dims.nbDims; i++) {
+        size_vec_tensor.push_back(Add1DConstantLayer(1));
+      }
+      size_vec_tensor[1] = GetEleTensorOfShape(shape_tensor, 1);
+      auto size_tensor = Concat(size_vec_tensor);
+
+      auto slice_layer =
+          TRT_ENGINE_ADD_LAYER(engine_,
+                               Slice,
+                               *mask_id_tensor,
+                               slice_start_dims,
+                               slice_start_dims,
+                               slice_stride_dims);  // unuseful slice_start_dims
+      slice_layer->setInput(2, *size_tensor);
+      slice_layer->setName(
+          ("Embeltwise_slice_layer (Output: slice_max_seqlen " +
+           op_desc.Output("Out")[0] + ")")
+              .c_str());
+      engine_->SetTensorDynamicRange(slice_layer->getOutput(0), 1.0f);
+
+      auto* reshape_layer =
+          TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *slice_layer->getOutput(0));
+      nvinfer1::Dims shape_dim;
+      shape_dim.nbDims = 1;
+      shape_dim.d[0] = -1;
+      reshape_layer->setReshapeDimensions(shape_dim);
+      reshape_layer->setName(("Embeltwise_reshape_layer (Output: max_seqlen " +
+                              op_desc.Output("Out")[0] + ")")
+                                 .c_str());
+      engine_->SetTensorDynamicRange(reshape_layer->getOutput(0), 1.0f);
+      engine_->SetITensor("max_seqlen_tensor", reshape_layer->getOutput(0));
     } else {
       id_names = op_desc.Input("Ids");
       emb_names = op_desc.Input("Embs");
@@ -192,20 +236,8 @@ class EmbEltwiseLayerNormOpConverter : public OpConverter {
       plugin_inputs.emplace_back(
           engine_->GetITensor(pos_id_name));  // cu_seqlens,
                                               // eval_placeholder_2
-      auto max_seqlen_tensor = engine_->GetITensor(mask_id_name);
-      auto* shuffle_layer =
-          TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *max_seqlen_tensor);
-      nvinfer1::Dims shape_dim;
-      shape_dim.nbDims = 1;
-      shape_dim.d[0] = -1;
-      shuffle_layer->setReshapeDimensions(shape_dim);
-      shuffle_layer->setName(
-          ("Embeltwise_Shuffle_reshape (Output: max_seqlen " +
-           op_desc.Output("Out")[0] + ")")
-              .c_str());
-      engine_->SetTensorDynamicRange(shuffle_layer->getOutput(0), 1.0f);
-      plugin_inputs.emplace_back(
-          shuffle_layer->getOutput(0));  // max_seqlen, eval_placeholder_3
+      plugin_inputs.emplace_back(engine_->GetITensor(
+          "max_seqlen_tensor"));  // max_seqlen, eval_placeholder_3
 
       auto creator = GetPluginRegistry()->getPluginCreator(
           "CustomEmbLayerNormPluginDynamic", "2");

paddle/fluid/inference/tensorrt/convert/fc_op.cc

@@ -1,11 +1,8 @@
 /* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
-
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
-
 http://www.apache.org/licenses/LICENSE-2.0
-
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

paddle/fluid/inference/tensorrt/convert/fused_token_prune_op.cc

@@ -23,7 +23,6 @@ class FusedTokenPruneOpConverter : public OpConverter {
                   bool test_mode) override {
     framework::OpDesc op_desc(op, nullptr);
     nvinfer1::ILayer* layer = nullptr;
-
     auto* Attn = engine_->GetITensor(op_desc.Input("Attn").front());
     auto* X = engine_->GetITensor(op_desc.Input("X").front());
     auto* Mask = engine_->GetITensor(op_desc.Input("Mask").front());
@@ -36,37 +35,61 @@ class FusedTokenPruneOpConverter : public OpConverter {
         op_desc.HasAttr("keep_order")
             ? PADDLE_GET_CONST(bool, op_desc.GetAttr("keep_order"))
             : false;
-
-    std::vector<nvinfer1::ITensor*> itensors = {Attn, X, Mask, NewMask};
-
     auto output_name = op_desc.Output("SlimmedX")[0];
     auto out_inds_name = op_desc.Output("CLSInds")[0];
     if (engine_->with_dynamic_shape()) {
-#if IS_TRT_VERSION_GE(6000)
       bool with_fp16 =
           engine_->WithFp16() && !engine_->disable_trt_plugin_fp16();
 
       if (engine_->precision() == AnalysisConfig::Precision::kInt8) {
         with_fp16 = true;
       }
+      bool flag_varseqlen = engine_->use_varseqlen();
       plugin::FusedTokenPrunePluginDynamic* plugin =
           new plugin::FusedTokenPrunePluginDynamic(
-              with_fp16, keep_first_token, keep_order);
-      layer = engine_->AddDynamicPlugin(itensors.data(), 4, plugin);
-#else
-      PADDLE_THROW(platform::errors::Fatal(
-          "You are running the TRT Dynamic Shape mode, need to confirm that "
-          "your TRT version is no less than 6.0"));
-#endif
+              with_fp16, keep_first_token, keep_order, flag_varseqlen);
+      if (flag_varseqlen) {
+        auto* word_id = engine_->GetITensor("word_id");
+        auto* pos_id = engine_->GetITensor("pos_id");
+        auto* mask_id = engine_->GetITensor("mask_id");
+        std::vector<nvinfer1::ITensor*> itensors = {
+            Attn, X, Mask, NewMask, word_id, pos_id, mask_id};
+        layer = engine_->AddDynamicPlugin(itensors.data(), 7, plugin);
+
+        layer->getOutput(0)->setName(output_name.c_str());
+        engine_->SetITensor(output_name, layer->getOutput(0));
+
+        layer->getOutput(1)->setName(out_inds_name.c_str());
+        engine_->SetITensor(out_inds_name, layer->getOutput(1));
+
+        engine_->DeleteITensor("word_id", word_id);
+        layer->getOutput(2)->setName("word_id_after_token_prune");
+        engine_->SetITensor("word_id", layer->getOutput(2));
+
+        engine_->DeleteITensor("pos_id", pos_id);
+        layer->getOutput(3)->setName("pos_id_after_token_prune");
+        engine_->SetITensor("pos_id", layer->getOutput(3));
+
+        engine_->DeleteITensor("mask_id", mask_id);
+        layer->getOutput(4)->setName("mask_id_after_token_prune");
+        engine_->SetITensor("mask_id", layer->getOutput(4));
+      } else {
+        std::vector<nvinfer1::ITensor*> itensors = {Attn, X, Mask, NewMask};
+        layer = engine_->AddDynamicPlugin(itensors.data(), 4, plugin);
+        layer->getOutput(0)->setName(output_name.c_str());
+        engine_->SetITensor(output_name, layer->getOutput(0));
+        layer->getOutput(1)->setName(out_inds_name.c_str());
+        engine_->SetITensor(out_inds_name, layer->getOutput(1));
+      }
+      layer->setName(
+          ("fused_token_prune(Output: " + output_name + ")").c_str());
     } else {
       PADDLE_THROW(platform::errors::Fatal(
           "You are running the Ernie(Bert) model in static shape mode, which "
           "is not supported for the time being.\n"
           "You can use the config.SetTRTDynamicShapeInfo(...) interface to set "
           "the shape information to run the dynamic shape mode."));
     }
-    RreplenishLayerAndOutput(
-        layer, "fused_token_prune", {output_name, out_inds_name}, test_mode);
   }
 };
 

paddle/fluid/inference/tensorrt/convert/multihead_matmul_op.cc

@@ -94,6 +94,8 @@ class MultiheadMatMulOpConverter : public OpConverter {
         nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT,
                                static_cast<void*>(bias_data),
                                static_cast<int32_t>(bias_t->numel())};
+        auto max_seqlen_tensor = engine_->GetITensor("max_seqlen_tensor");
+        auto pos_id_tensor = engine_->GetITensor("pos_id");
         if (engine_->with_interleaved()) {
           VLOG(4) << "fused multihead_matmul op: use_varseqlen and "
                      "with_interleaved";
@@ -154,31 +156,9 @@ class MultiheadMatMulOpConverter : public OpConverter {
 
           std::vector<nvinfer1::ITensor*> plugin_inputs;
           plugin_inputs.emplace_back(fc_layer->getOutput(0));
-          if (engine_->Has("ernie_pos_name")) {
-            plugin_inputs.emplace_back(engine_->GetITensor(
-                engine_->Get<std::string>("ernie_pos_name")));
-          } else {
-            plugin_inputs.emplace_back(engine_->GetITensor(
-                engine_->network()
-                    ->getInput(2)
-                    ->getName()));  // cu_seqlens, eval_placeholder_2
-          }
-          auto max_seqlen_tensor =
-              engine_->GetITensor(engine_->network()->getInput(3)->getName());
-          engine_->SetTensorDynamicRange(max_seqlen_tensor, 1.0f);
-          auto* shuffle_layer = TRT_ENGINE_ADD_LAYER(
-              engine_,
-              Shuffle,
-              *const_cast<nvinfer1::ITensor*>(max_seqlen_tensor));
-          nvinfer1::Dims shape_dim;
-          shape_dim.nbDims = 1;
-          shape_dim.d[0] = -1;
-          shuffle_layer->setReshapeDimensions(shape_dim);
-          engine_->SetTensorDynamicRange(shuffle_layer->getOutput(0), 1.0f);
+          plugin_inputs.emplace_back(pos_id_tensor);
           plugin_inputs.emplace_back(
-              shuffle_layer->getOutput(0));  // max_seqlen, eval_placeholder_3
-          shuffle_layer->setName(
-              ("Multihead: Shuffle: (Output: " + output_name + ")").c_str());
+              max_seqlen_tensor);  // max_seqlen, eval_placeholder_3
           auto plugin_layer = engine_->network()->addPluginV2(
               plugin_inputs.data(), plugin_inputs.size(), *plugin);
           layer = plugin_layer;
@@ -299,20 +279,9 @@ class MultiheadMatMulOpConverter : public OpConverter {
           std::vector<nvinfer1::ITensor*> plugin_inputs;
           plugin_inputs.emplace_back(fc_layer->getOutput(0));
           plugin_inputs.emplace_back(engine_->GetITensor("qkv_plugin_mask"));
-          plugin_inputs.emplace_back(engine_->GetITensor("pos_id"));
-
-          auto max_seqlen_tensor = engine_->GetITensor("mask_id");
-          auto* shuffle_layer = TRT_ENGINE_ADD_LAYER(
-              engine_,
-              Shuffle,
-              *const_cast<nvinfer1::ITensor*>(max_seqlen_tensor));
-          nvinfer1::Dims shape_dim;
-          shape_dim.nbDims = 1;
-          shape_dim.d[0] = -1;
-          shuffle_layer->setReshapeDimensions(shape_dim);
-          engine_->SetTensorDynamicRange(shuffle_layer->getOutput(0), 1.0f);
+          plugin_inputs.emplace_back(pos_id_tensor);
           plugin_inputs.emplace_back(
-              shuffle_layer->getOutput(0));  // max_seqlen, eval_placeholder_3
+              max_seqlen_tensor);  // max_seqlen, eval_placeholder_3
 
           auto plugin_layer = engine_->network()->addPluginV2(
               plugin_inputs.data(), plugin_inputs.size(), *plugin);

paddle/fluid/inference/tensorrt/convert/preln_emb_eltwise_layernorm.cc

@@ -157,20 +157,47 @@ class PrelnEmbEltwiseLayerNormOpConverter : public OpConverter {
     plugin_inputs.emplace_back(
         engine_->GetITensor(pos_id_name));  // cu_seqlens,
                                             // eval_placeholder_2
-    auto max_seqlen_tensor = engine_->GetITensor(mask_id_name);
-    auto* shuffle_layer =
-        TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *max_seqlen_tensor);
+    auto mask_id_tensor = engine_->GetITensor("mask_id");
+    auto mask_dims = mask_id_tensor->getDimensions();
+    auto slice_start_dims = mask_dims;
+    auto slice_size_dims = mask_dims;
+    auto slice_stride_dims = mask_dims;
+
+    for (int i = 0; i < mask_dims.nbDims; i++) {
+      slice_start_dims.d[i] = 0;
+      slice_size_dims.d[i] = 1;
+      slice_stride_dims.d[i] = 1;
+    }
+    slice_size_dims.d[1] = mask_dims.d[1];
+    auto* slice_size_tensor = Add1DConstantLayer(slice_size_dims);
+    auto slice_layer =
+        TRT_ENGINE_ADD_LAYER(engine_,
+                             Slice,
+                             *mask_id_tensor,
+                             slice_start_dims,
+                             slice_start_dims,
+                             slice_stride_dims);  // unuseful slice_start_dims
+    slice_layer->setInput(2, *slice_size_tensor);
+    slice_layer->setName(
+        ("PrelnEmbeltwise_slice_layer (Output: slice_max_seqlen " +
+         op_desc.Output("Out")[0] + ")")
+            .c_str());
+    engine_->SetTensorDynamicRange(slice_layer->getOutput(0), 1.0f);
+
+    auto* reshape_layer =
+        TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *slice_layer->getOutput(0));
     nvinfer1::Dims shape_dim;
     shape_dim.nbDims = 1;
     shape_dim.d[0] = -1;
-    shuffle_layer->setReshapeDimensions(shape_dim);
-    shuffle_layer->setName(
-        ("PrelnEmbeltwise_Shuffle_reshape (Output: max_seqlen " +
-         op_desc.Output("Out_0")[0] + ")")
+    reshape_layer->setReshapeDimensions(shape_dim);
+    reshape_layer->setName(
+        ("PrelnEmbeltwise_reshape_layer (Output: max_seqlen " +
+         op_desc.Output("Out")[0] + ")")
             .c_str());
-    engine_->SetTensorDynamicRange(shuffle_layer->getOutput(0), 1.0f);
+    engine_->SetTensorDynamicRange(reshape_layer->getOutput(0), 1.0f);
+    engine_->SetITensor("max_seqlen_tensor", reshape_layer->getOutput(0));
     plugin_inputs.emplace_back(
-        shuffle_layer->getOutput(0));  // max_seqlen, eval_placeholder_3
+        reshape_layer->getOutput(0));  // max_seqlen, eval_placeholder_3
 
     auto creator = GetPluginRegistry()->getPluginCreator(
         "CustomEmbLayerNormPluginDynamic", "3");