PaddlePaddle
diff --git a/‎paddle/fluid/inference/tensorrt/convert/op_converter.h
Copy file name to clipboard
+77-1Lines changed: 77 additions & 1 deletion b/‎paddle/fluid/inference/tensorrt/convert/op_converter.h
Copy file name to clipboard
+77-1Lines changed: 77 additions & 1 deletion
diff --git a/‎paddle/fluid/inference/tensorrt/convert/range_op.cc
Copy file name to clipboard
+7-2Lines changed: 7 additions & 2 deletions b/‎paddle/fluid/inference/tensorrt/convert/range_op.cc
Copy file name to clipboard
+7-2Lines changed: 7 additions & 2 deletions
diff --git a/‎paddle/fluid/inference/tensorrt/convert/set_value_op.cc
Copy file name to clipboard
+11-38Lines changed: 11 additions & 38 deletions b/‎paddle/fluid/inference/tensorrt/convert/set_value_op.cc
Copy file name to clipboard
+11-38Lines changed: 11 additions & 38 deletions
diff --git a/‎paddle/fluid/inference/tensorrt/convert/top_k_op.cc
Copy file name to clipboard
+53-59Lines changed: 53 additions & 59 deletions b/‎paddle/fluid/inference/tensorrt/convert/top_k_op.cc
Copy file name to clipboard
+53-59Lines changed: 53 additions & 59 deletions
@@ -373,6 +373,13 @@ class OpConverter {
     engine->ClearWeights();
   }
 
+  nvinfer1::ITensor* Cast(nvinfer1::ITensor* input, nvinfer1::DataType dtype) {
+    auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Identity, *input);
+    layer->setOutputType(0, dtype);
+    layer->getOutput(0)->setType(dtype);
+    return layer->getOutput(0);
+  }
+
   // rank(result) = rank(input)
   nvinfer1::ITensor* Gather(nvinfer1::ITensor* input,
                             const std::vector<int32_t> indices,
@@ -384,6 +391,59 @@ class OpConverter {
     return result;
   }
 
+  nvinfer1::ITensor* Unsqueeze(nvinfer1::ITensor* input,
+                               const std::vector<int32_t> axis) {
+    const auto dims = input->getDimensions();
+    const std::unordered_set<int32_t> axis_data(axis.begin(), axis.end());
+    std::vector<int32_t> subscripts(dims.nbDims);
+    std::iota(subscripts.begin(), subscripts.end(), 0);
+    for (const auto& axis_value : axis_data) {
+      subscripts.insert(subscripts.begin() + axis_value, dims.nbDims);
+    }
+    nvinfer1::ITensor* input_shape{nullptr};
+    if (engine_->with_dynamic_shape()) {
+      input_shape = Shape(input);
+    } else {
+      input_shape = Add1DConstantLayer(dims);
+    }
+    auto* new_dim =
+        TRT_ENGINE_ADD_LAYER(engine_,
+                             Gather,
+                             *Concat(std::vector<nvinfer1::ITensor*>{
+                                 input_shape, Add1DConstantLayer(1)}),
+                             *Add1DConstantLayer(subscripts),
+                             0)
+            ->getOutput(0);
+    auto result = Reshape(input, new_dim);
+    return result;
+  }
+
+  nvinfer1::ITensor* Squeeze(nvinfer1::ITensor* input,
+                             const std::vector<int32_t> axis) {
+    const auto dims = input->getDimensions();
+    std::vector<int32_t> subscripts(dims.nbDims);
+    std::iota(subscripts.begin(), subscripts.end(), 0);
+    auto p =
+        std::remove_if(subscripts.begin(), subscripts.end(), [axis](int x) {
+          return std::find(axis.begin(), axis.end(), x) != axis.end();
+        });
+    subscripts.resize(p - subscripts.begin());
+
+    nvinfer1::ITensor* input_shape{nullptr};
+    if (engine_->with_dynamic_shape()) {
+      input_shape = Shape(input);
+    } else {
+      input_shape = Add1DConstantLayer(dims);
+    }
+
+    auto* new_dim =
+        TRT_ENGINE_ADD_LAYER(
+            engine_, Gather, *input_shape, *Add1DConstantLayer(subscripts), 0)
+            ->getOutput(0);
+    auto result = Reshape(input, new_dim);
+    return result;
+  }
+
   // paddle allows negative index
   // for axis length = 5, paddle allows [-5, 4]
   nvinfer1::ITensor* FixNegIndices(nvinfer1::ITensor* input_shape,
@@ -406,7 +466,23 @@ class OpConverter {
                              nvinfer1::ITensor* newShape,
                              const std::string& name = "") {
     auto* shuffle = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
-    shuffle->setInput(1, *newShape);
+    if (engine_->with_dynamic_shape()) {
+      shuffle->setInput(1, *newShape);
+    } else {
+      auto shape = newShape->getDimensions();
+      shuffle->setReshapeDimensions(shape);
+    }
+    if (name != "") {
+      shuffle->setName(name.c_str());
+    }
+    return shuffle->getOutput(0);
+  }
+
+  nvinfer1::ITensor* Reshape(nvinfer1::ITensor* input,
+                             nvinfer1::Dims shape,
+                             const std::string& name = "") {
+    auto* shuffle = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
+    shuffle->setReshapeDimensions(shape);
     if (name != "") {
       shuffle->setName(name.c_str());
     }
 
@@ -46,8 +46,13 @@ class RangeOpConverter : public OpConverter {
       quotient_tensor = fquotient_tensor;
     }
     auto number_tensor = Max(Sub(zero_tensor, quotient_tensor), zero_tensor);
-    auto* start1 = engine_->GetITensor(op_desc.Input("Start")[0], true);
-
+    auto* start1 = engine_->GetITensor(op_desc.Input("Start")[0]);
+#if IS_TRT_VERSION_LT(8000)
+    nvinfer1::Dims start_dims{0, {1}, { nvinfer1::DimensionType::kSPATIAL }};
+#else
+    nvinfer1::Dims start_dims{0, {1}};
+#endif
+    start1 = Reshape(start1, start_dims);
     layer = TRT_ENGINE_ADD_LAYER(
         engine_, Fill, nvinfer1::Dims{}, nvinfer1::FillOperation::kLINSPACE);
     layer->setInput(0, *number_tensor);
 
@@ -24,16 +24,6 @@ limitations under the License. */
     }                                                                       \
   } while (0)
 
-namespace paddle {
-namespace framework {
-class Scope;
-
-namespace proto {
-class OpDesc;
-}  // namespace proto
-}  // namespace framework
-}  // namespace paddle
-
 namespace paddle {
 namespace inference {
 namespace tensorrt {
@@ -55,6 +45,14 @@ class SetValueConverter : public OpConverter {
 
     auto* inputs = engine_->GetITensor(op_desc.Input("Input")[0]);
     auto* updates = engine_->GetITensor(op_desc.Input("ValueTensor")[0]);
+    const auto decrease_axes = PADDLE_GET_CONST(
+        std::vector<int64_t>, op_desc.GetAttr("decrease_axes"));
+    std::vector<int32_t> decr_axes{decrease_axes.begin(), decrease_axes.end()};
+    auto value_rank = updates->getDimensions().nbDims;
+    auto input_rank = inputs->getDimensions().nbDims;
+    if (decrease_axes.size() > 0 && value_rank != input_rank) {
+      updates = Unsqueeze(updates, decr_axes);
+    }
 
     int64_t axes = 0;
     int64_t starts = 0;
@@ -115,39 +113,14 @@ class SetValueConverter : public OpConverter {
         indices.insert(indices.end(), axes_index.begin(), axes_index.end());
       }
 
-      nvinfer1::Dims indice_dims = update_dims;
-
-      // create a tensor to store data
-      std::vector<int> indice_dim_vec;
-      for (int i = 0; i < update_dims.nbDims; i++) {
-        indice_dim_vec.emplace_back(update_dims.d[i]);
-      }
-      auto indice_tensor_dims = phi::make_ddim(indice_dim_vec);
-      std::unique_ptr<phi::DenseTensor> indice_tensor(
-          std::make_unique<phi::DenseTensor>());
-      indice_tensor->Resize(indice_tensor_dims);
-
-      auto* dev_ctx = static_cast<phi::CPUContext*>(
-          platform::DeviceContextPool::Instance().Get(platform::CPUPlace()));
-      auto* weight_data = dev_ctx->template HostAlloc<int>(indice_tensor.get());
-
-      memcpy(weight_data, indices.data(), sizeof(int) * indice_tensor->numel());
-
-      TensorRTEngine::Weight weight{
-          nvinfer1::DataType::kINT32,
-          static_cast<void*>(weight_data),
-          static_cast<size_t>(indice_tensor->numel())};
       auto output_name = op_desc.Output("Out")[0];
-      engine_->SetWeights("set_value_index_" + output_name,
-                          std::move(indice_tensor));
-
-      auto const_layer =
-          TRT_ENGINE_ADD_LAYER(engine_, Constant, indice_dims, weight.get());
+      const auto const_layer = AddConstantLayer(
+          indices.data(), update_dims, "set_value_index_" + output_name);
 
       auto* layer = TRT_ENGINE_ADD_LAYER(engine_,
                                          Scatter,
                                          *inputs,
-                                         *const_layer->getOutput(0),
+                                         *const_layer,
                                          *updates,
                                          nvinfer1::ScatterMode::kELEMENT);
 
 
@@ -33,82 +33,76 @@ class TopKOpConverter : public OpConverter {
   void operator()(const framework::proto::OpDesc& op,
                   const framework::Scope& scope,
                   bool test_mode) override {
-    VLOG(3) << "convert a top_k op to tensorrt TopK layer";
-    // Here the two nullptr looks strange, that's because the
-    // framework::OpDesc's constructor is strange.
+    VLOG(3) << "convert a top_k op to tensorrt layer";
     framework::OpDesc op_desc(op, nullptr);
 
     auto* input_tensor = engine_->GetITensor(op_desc.Input("X")[0]);
 
-    const int k = op_desc.HasAttr("k")
-                      ? PADDLE_GET_CONST(int, op_desc.GetAttr("k"))
-                      : 1.0f;
-
-    nvinfer1::Dims input_dims = input_tensor->getDimensions();
-    int axis = input_dims.nbDims;
-    nvinfer1::ITopKLayer* layer =
-        TRT_ENGINE_ADD_LAYER(engine_,
-                             TopK,
-                             *input_tensor,
-                             nvinfer1::TopKOperation::kMAX,
-                             k,
-                             1 << (axis - 1));
-
-    std::vector<std::string> output_names;
-    output_names.push_back(op_desc.Output("Out").front());
-    output_names.push_back(op_desc.Output("Indices").front());
-
-    RreplenishLayerAndOutput(layer, "top_k", output_names, test_mode);
-  }
-};
-class TopKv2OpConverter : public OpConverter {
- public:
-  TopKv2OpConverter() {}
-  void operator()(const framework::proto::OpDesc& op,
-                  const framework::Scope& scope,
-                  bool test_mode) override {
-    // Here the two nullptr looks strange, that's because the
-    // framework::OpDesc's constructor is strange.
-    framework::OpDesc op_desc(op, nullptr);
-
-    auto* input_tensor = engine_->GetITensor(op_desc.Input("X")[0]);
-
-    const int k = op_desc.HasAttr("k")
-                      ? PADDLE_GET_CONST(int, op_desc.GetAttr("k"))
-                      : 1.0f;
-    const int axis = op_desc.HasAttr("axis")
-                         ? PADDLE_GET_CONST(int, op_desc.GetAttr("axis"))
-                         : 1.0f;
+    const int k =
+        op_desc.HasAttr("k") ? PADDLE_GET_CONST(int, op_desc.GetAttr("k")) : 1;
+    int axis = op_desc.HasAttr("axis")
+                   ? PADDLE_GET_CONST(int, op_desc.GetAttr("axis"))
+                   : -1;
     const bool largest =
         op_desc.HasAttr("largest")
             ? PADDLE_GET_CONST(bool, op_desc.GetAttr("largest"))
             : true;
     auto flag =
         largest ? nvinfer1::TopKOperation::kMAX : nvinfer1::TopKOperation::kMIN;
+
+    auto input_dims = input_tensor->getDimensions();
+    auto input_rank = input_dims.nbDims;
+    // 1d needs expand to 2d
+    bool expand_to_2d = (input_rank == 1);
+    if (engine_->with_dynamic_shape() && expand_to_2d) {
+      input_tensor = Unsqueeze(input_tensor, std::vector<int32_t>{1});
+    }
+
+    // INT32 only, other data type should to casted to INT32.
+    nvinfer1::DataType type = input_tensor->getType();
+    bool cast = (type == nvinfer1::DataType::kINT32);
+    if (cast) {
+      input_tensor = Cast(input_tensor, nvinfer1::DataType::kFLOAT);
+    }
+
     nvinfer1::ITopKLayer* layer = nullptr;
-    if (axis == -1) {
-      nvinfer1::Dims input_dims = input_tensor->getDimensions();
-      layer = TRT_ENGINE_ADD_LAYER(
-          engine_, TopK, *input_tensor, flag, k, 1 << (input_dims.nbDims - 1));
-    } else {
-      if (engine_->with_dynamic_shape()) {
-        layer = TRT_ENGINE_ADD_LAYER(
-            engine_, TopK, *input_tensor, flag, k, 1 << axis);
-      } else {
-        layer = TRT_ENGINE_ADD_LAYER(
-            engine_, TopK, *input_tensor, flag, k, 1 << (axis - 1));
-      }
+    if (axis > 0 && !engine_->with_dynamic_shape()) {
+      axis -= 1;
     }
-    std::vector<std::string> output_names;
-    output_names.push_back(op_desc.Output("Out").front());
-    output_names.push_back(op_desc.Output("Indices").front());
+    if (axis < 0) axis += input_rank;
+
+    layer =
+        TRT_ENGINE_ADD_LAYER(engine_, TopK, *input_tensor, flag, k, 1 << axis);
+
+    nvinfer1::ITensor* values = layer->getOutput(0);
+    nvinfer1::ITensor* indices = layer->getOutput(1);
+
+    // un-expand to 1d
+    if (engine_->with_dynamic_shape() && expand_to_2d) {
+      values = Squeeze(values, std::vector<int32_t>{1});
+      indices = Squeeze(indices, std::vector<int32_t>{1});
+    }
+
+    // cast back
+    if (cast) {
+      values = Cast(values, nvinfer1::DataType::kINT32);
+    }
+
+    auto out_name = op_desc.Output("Out").front();
+    auto indices_name = op_desc.Output("Indices").front();
+    values->setName(out_name.c_str());
+    engine_->SetITensor(out_name.c_str(), values);
+
+    indices->setName(indices_name.c_str());
+    engine_->SetITensor(indices_name.c_str(), indices);
 
-    RreplenishLayerAndOutput(layer, "top_k_v2", output_names, test_mode);
+    layer->setName(
+        ("top_k (Output: " + out_name + "," + indices_name + ")").c_str());
   }
 };
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle
 
 REGISTER_TRT_OP_CONVERTER(top_k, TopKOpConverter);
-REGISTER_TRT_OP_CONVERTER(top_k_v2, TopKv2OpConverter);
+REGISTER_TRT_OP_CONVERTER(top_k_v2, TopKOpConverter);