[Paddle Tensorrt] pd_op.full_with_tensor converter (#68941)

lizexu123 · web-flow · commit 36e3d1835459 · 2024-10-28T14:03:52.000+08:00
* merge develop

* merge develop

* merge develop

* pd_op.full_with_tensor

* fix

* fix
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1482,6 +1482,56 @@ class TanhOpPattern : public pir::OpRewritePattern<paddle::dialect::TanhOp> {
   }
 };
 
+class FullWithTensorPattern
+    : public pir::OpRewritePattern<paddle::dialect::FullWithTensorOp> {
+ public:
+  using pir::OpRewritePattern<
+      paddle::dialect::FullWithTensorOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::FullWithTensorOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    pir::Value value = op.operand_source(0);
+    if (value == nullptr) {
+      VLOG(3) << "pd_op.full_with_tensor value is null";
+      return false;
+    }
+#if IS_TRT_VERSION_LT(8500)
+    if (pir::GetDefiningOpForInput(op, 1)
+            ->isa<paddle::dialect::FullIntArrayOp>()) {
+      paddle::dialect::FullIntArrayOp full_int_array =
+          pir::GetDefiningOpForInput(op, 1)
+              ->dyn_cast<paddle::dialect::FullIntArrayOp>();
+      auto shape_attr = full_int_array->attribute<pir::ArrayAttribute>("value");
+      if (shape_attr.size() == 1) {
+        VLOG(3) << "pd_op.full_with_tensor shape is not support when TensorRT "
+                   "< 8.5.0";
+        return false;
+      }
+    } else {
+      pir::Value shape = op.operand_source(1);
+      if (shape != nullptr) {
+        VLOG(3) << "pd_op.full_with_tensor shape is not support when TensorRT "
+                   "< 8.5.0";
+        return false;
+      }
+    }
+#endif
+    auto dtype =
+        op->attribute<paddle::dialect::DataTypeAttribute>("dtype").data();
+    if (dtype != phi::DataType::INT32 && dtype != phi::DataType::INT64 &&
+        dtype != phi::DataType::FLOAT32) {
+      VLOG(3) << "pd_op.full_with_tensor only support int32, int64, float32";
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
@@ -1568,6 +1618,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<NearestInterV2Pattern>(context));
     ps.Add(std::make_unique<StackOpPattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
+    ps.Add(std::make_unique<FullWithTensorPattern>(context));
     return ps;
   }
 };
diff --git a/python/paddle/tensorrt/impls/creation.py b/python/paddle/tensorrt/impls/creation.py
@@ -96,3 +96,91 @@ def arange_converter(network, paddle_op, inputs):
     fill_layer.set_input(2, step)
 
     return fill_layer.get_output(0)
+
+
+@converter_registry.register("pd_op.full_with_tensor", trt_version="8.x")
+def full_with_tensor_converter(network, paddle_op, inputs):
+    value_input = inputs[0]
+
+    shape_tensor = None
+    dtype = paddle_op.attrs()["dtype"]
+
+    operands = paddle_op.operands()
+    num_operands = len(operands)
+
+    if num_operands >= 2:
+        shape_tensor = inputs[1]
+        if isinstance(shape_tensor, list):
+            shape_tensor_list = shape_tensor
+        else:
+            shape_tensor_list = [shape_tensor]
+
+    shape_op = paddle_op.operands()[1].source().get_defining_op()
+    if shape_op.name() == "pd_op.full_int_array":
+        shape_tensor = shape_op.attrs()["value"]
+        is_static_shape = True
+    else:
+        shape_tensor = inputs[1]
+        is_static_shape = False
+
+    shape_nbDims = 0
+    tensor_rank = 0
+    if isinstance(shape_tensor, trt.ITensor):
+        shape_x = shape_tensor.shape
+        shape_nbDims = len(shape_x)
+        shapes_tensor = shape_tensor
+    elif isinstance(shape_tensor, (list, tuple)):
+        shape_nbDims = len(shape_tensor)
+        shapes_tensor = shape_tensor
+    else:
+        raise TypeError(f"Unsupported shape_tensor type: {type(shape_tensor)}")
+
+    if shape_tensor is not None and len(shape_tensor_list) == 1:
+        is_dynamic_shape = True
+    elif len(shape_tensor_list) >= 1:
+        is_dynamic_shape = True
+    else:
+        is_dynamic_shape = False
+
+    if is_dynamic_shape:
+        if len(shape_tensor_list) == 1:
+            shape_tensor = shape_tensor_list[0]
+            if not isinstance(shape_tensor, trt.ITensor):
+                raise TypeError("shape_tensor must be an ITensor")
+            if len(shape_tensor.shape) != 1:
+                raise ValueError("The rank of shape_tensor must be 1")
+            tensor_rank = shape_tensor.shape[0]
+            shapes_tensor = shape_tensor
+        else:
+            shape_tensors = []
+            for tensor in shape_tensor_list:
+                if len(tensor.shape) == 0:
+                    tensor = trt_reshape(network, tensor, (1,))
+                shape_tensors.append(tensor)
+
+            concat_layer = network.add_concatenation(shape_tensors)
+            shapes_tensor = concat_layer.get_output(0)
+            tensor_rank = len(shape_tensors)
+
+        fill_layer = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
+        fill_layer.set_input(0, shapes_tensor)
+
+    if dtype == paddle.int32 or dtype == paddle.int64:
+        beta_vec = [0] * tensor_rank
+        value_input = trt_reduce_to_scalar(network, value_input)
+        fill_layer.set_input(1, value_input)
+        fill_layer.set_input(
+            2, add_1D_constant_layer(network, beta_vec, np.int32)
+        )
+    elif dtype == paddle.float32:
+        beta_vec = [0.0] * tensor_rank
+        value_input = trt_reduce_to_scalar(network, value_input)
+        fill_layer.set_input(1, value_input)
+        fill_layer.set_input(
+            2, add_1D_constant_layer(network, beta_vec, np.float32)
+        )
+    else:
+        raise ValueError(f"Unsupported dtype for full_with_tensor: {dtype}")
+
+    output_tensor = fill_layer.get_output(0)
+    return output_tensor
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -20,7 +20,7 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_activation PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_others PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_manipulation PROPERTIES TIMEOUT "300")
-  set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_common PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_logic PROPERTIES TIMEOUT "100")
diff --git a/test/tensorrt/test_converter_creation.py b/test/tensorrt/test_converter_creation.py
@@ -91,5 +91,37 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestFullWithTensorTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.tensor.fill_constant
+        self.api_args = {
+            "shape": np.array([1]).astype("int32"),
+            "dtype": "float32",
+            "value": np.array([0.0]).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["value", "shape"]}
+        self.min_shape = {}
+        self.max_shape = {}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestFullWithTensorCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.tensor.fill_constant
+        self.api_args = {
+            "shape": [1, 1],
+            "dtype": np.float32,
+            "value": np.array([1.0]).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["value"]}
+        self.min_shape = {}
+        self.max_shape = {}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == "__main__":
     unittest.main()
