[HEU][Paddle TensorRT No.62、64] Add (pd_op.index_put、pd_op.pow) converter

paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc

@@ -2264,6 +2264,28 @@ class InstanceNormOpPattern
   }
 };
 
+class PowOpPattern : public pir::OpRewritePattern<paddle::dialect::PowOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::PowOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::PowOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    pir::Value x = op.operand_source(0);
+    auto x_dtype = pir::GetDataTypeFromValue(x);
+    if (x_dtype.isa<pir::Int32Type>()) {
+      VLOG(3) << "These operations (pow) do not support int32 "
+                 "datatype.";
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class EinsumOpPattern
     : public pir::OpRewritePattern<paddle::dialect::EinsumOp> {
  public:
@@ -2298,6 +2320,39 @@ class EinsumOpPattern
   }
 };
 
+class IndexPutOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::IndexPutOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::IndexPutOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::IndexPutOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+#if IS_TRT_VERSION_LT(8510)
+    VLOG(3) << "index_put is not supported when TensorRT < 8.5.1";
+    return false;
+#endif
+    pir::Value value = op.operand_source(2);
+    auto value_shape = pir::GetShapeFromValue(value);
+    int value_num = std::accumulate(
+        value_shape.begin(), value_shape.end(), 1, std::multiplies<int>());
+    if (value_num != 1) {
+      VLOG(3) << " index_put op only support value_num = 1 in tensorrt.";
+      return false;
+    }
+    pir::Value indices = op.operand_source(1);
+    auto indices_dtype = pir::GetDataTypeFromValue(indices);
+    if (!indices_dtype.isa<pir::BoolType>()) {
+      VLOG(3) << " index_put op only support bool indices in tensorrt.";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class PNormOpPattern : public pir::OpRewritePattern<paddle::dialect::PNormOp> {
  public:
   using pir::OpRewritePattern<paddle::dialect::PNormOp>::OpRewritePattern;
@@ -2502,6 +2557,8 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<TanhOpPattern>(context));
     ps.Add(std::make_unique<CeluOpPattern>(context));
     ps.Add(std::make_unique<OneHotOpPattern>(context));
+    ps.Add(std::make_unique<PowOpPattern>(context));
+    ps.Add(std::make_unique<IndexPutOpPattern>(context));
     ps.Add(std::make_unique<TemporalShiftOpPattern>(context));
     ps.Add(std::make_unique<InstanceNormOpPattern>(context));
     ps.Add(std::make_unique<EinsumOpPattern>(context));

python/paddle/tensorrt/impls/manipulation.py

@@ -26,6 +26,7 @@
     get_shape_tensor_element,
     has_dynamic_shape,
     resize_to_1d,
+    trt_cast,
     trt_concat,
     trt_expand,
     trt_floor_div,
@@ -945,6 +946,91 @@ def roll_converter(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
+@converter_registry.register("pd_op.index_put", trt_version="8.x")
+def index_put_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+
+    indices_op = paddle_op.operands()[1].source().get_defining_op()
+    value_op = paddle_op.operands()[2].source().get_defining_op()
+
+    if indices_op.name() == "pd_op.full_int_array":
+        indices = indices_op.attrs()["value"]
+        indices_tensor = add_1D_constant_layer(network, indices)
+    else:
+        indices_tensor = inputs[1]
+
+    if value_op.name() == "pd_op.full_int_array":
+        value = value_op.attrs()["value"]
+        value_tensor = add_1D_constant_layer(network, value)
+    else:
+        value_tensor = inputs[2]
+
+    input_shape_tensor = trt_shape(network, input_tensor)
+    input_shape = input_tensor.shape
+    rank = len(input_shape)
+    indices_shape = indices_tensor.shape
+    indices_dims = len(indices_shape)
+
+    # indices
+    indices_shape_vec = []
+    start_tensor_vec = []
+    stride_tensor_vec = []
+    for i in range(rank):
+        indices_one = indices_shape[i] if i < len(indices_dims) else 1
+        indices_shape_vec.append(add_1D_constant_layer(indices_one))
+        start_tensor_vec.append(add_1D_constant_layer(0))
+        stride_tensor_vec.append(add_1D_constant_layer(1))
+    indices_tensor_temp = trt_reshape(
+        network, indices_tensor, trt_concat(network, indices_shape_vec)
+    )
+    start_tensor = trt_concat(network, start_tensor_vec)
+    stride_tensor = trt_concat(network, stride_tensor_vec)
+
+    # slice
+    stride = [1] * rank
+    indices_slice_layer = network.add_slice(
+        trt_cast(network, indices_tensor_temp, trt.float32),
+        stride,
+        stride,
+        stride,
+    )
+    indices_slice_layer.set_input(1, start_tensor)
+    indices_slice_layer.set_input(2, input_shape_tensor)
+    indices_slice_layer.set_input(3, stride_tensor)
+    indices_slice_layer.mode = trt.SampleMode.CLAMP
+    bool_indices_tensor = trt_cast(
+        network, indices_slice_layer.get_output(0), trt.bool
+    )
+
+    # nonzero
+    nonzero_layer = network.add_non_zero(bool_indices_tensor)
+    indices_tensor = nonzero_layer.get_output(0)
+    permutation = trt.Permutation([1, 0])
+    trans_layer = network.add_shuffle(indices_tensor)
+    trans_layer.first_transpose = permutation
+    indices_tensor = trans_layer.get_output(0)
+    indices_new_shape_tensor = trt_shape(network, indices_tensor)
+    indices_count_tensor = get_shape_tensor_element(
+        network, indices_new_shape_tensor, 0
+    )
+
+    # value
+    value_stride = [1]
+    value_slice_layer = network.add_slice(
+        value_tensor, value_stride, value_stride, value_stride
+    )
+    value_slice_layer.set_input(1, add_1D_constant_layer(network, 0))
+    value_slice_layer.set_input(2, indices_count_tensor)
+    value_slice_layer.set_input(3, add_1D_constant_layer(network, 1))
+    value_slice_layer.mode = trt.SampleMode.CLAMP
+    value_tensor = value_slice_layer.get_output(0)
+
+    layer = network.add_scatter(
+        input_tensor, indices_tensor, value_tensor, trt.ScatterMode.ND
+    )
+    return layer.get_output(0)
+
+
 @converter_registry.register("pd_op.numel", trt_version="8.x")
 def numel_converter(network, paddle_op, inputs):
     input_tensor = inputs[0]

python/paddle/tensorrt/impls/math.py

@@ -394,6 +394,13 @@ def sqrt_converter(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
+@converter_registry.register("pd_op.pow", trt_version="8.x")
+def pow_op_converter(network, paddle_op, inputs):
+    return add_elementwise_layer(
+        network, paddle_op, inputs, trt.ElementWiseOperation.POW
+    )
+
+
 @converter_registry.register("pd_op.elementwise_pow", trt_version="8.x")
 def elementwise_pow_converter(network, paddle_op, inputs):
     return add_elementwise_layer(

test/tensorrt/tensorrt_test_base.py

@@ -93,6 +93,28 @@ def create_fake_program(self):
                         )
                         new_list_args.append(input_data)
                     api_args[feed_name] = new_list_args
+                elif isinstance(self.api_args[feed_name], tuple):
+                    new_tuple_args = []
+                    for index, elem in enumerate(self.api_args[feed_name]):
+                        if (
+                            feed_name in self.min_shape.keys()
+                            and feed_name in self.max_shape.keys()
+                        ):
+                            input_shape_without_dynamic_dim = elem.shape[1:]
+                            input_shape = [-1]
+                            input_shape.extend(input_shape_without_dynamic_dim)
+                        else:
+                            input_shape = elem.shape
+
+                        input_dtype = elem.dtype
+
+                        input_data = paddle.static.data(
+                            name=f"{feed_name}{index}",
+                            shape=input_shape,
+                            dtype=input_dtype,
+                        )
+                        new_tuple_args.append(input_data)
+                    api_args[feed_name] = tuple(new_tuple_args)
                 else:
                     empty_min_max_shape = (
                         self.min_shape is None
@@ -150,6 +172,9 @@ def run_program(self, main_program, fetch_list):
                     feed_name
                 ].items():
                     feed_data[sub_arg_name] = sub_arg_value
+            elif isinstance(self.api_args[feed_name], tuple):
+                for index, elem in enumerate(self.api_args[feed_name]):
+                    feed_data[f"{feed_name}{index}"] = elem
             else:
                 feed_data[feed_name] = self.api_args[feed_name]
         ret = exe.run(main_program, feed=feed_data, fetch_list=fetch_list)
@@ -229,6 +254,23 @@ def check_trt_result(self, rtol=1e-4, atol=1e-4, precision_mode="fp32"):
                             ).astype(
                                 self.api_args[feed_name][sub_feed_name].dtype
                             )
+                elif isinstance(self.api_args[feed_name], tuple):
+                    # shape_tensor is tuple
+                    for index, elem in enumerate(self.api_args[feed_name]):
+                        sub_feed_name = f"{feed_name}{index}"
+                        if (
+                            feed_name not in self.min_shape.keys()
+                            and feed_name not in self.max_shape.keys()
+                        ):
+                            min_shape_data[sub_feed_name] = elem
+                            max_shape_data[sub_feed_name] = elem
+                        else:
+                            min_shape_data[sub_feed_name] = np.random.randn(
+                                *self.min_shape[feed_name][index]
+                            ).astype(elem.dtype)
+                            max_shape_data[sub_feed_name] = np.random.randn(
+                                *self.max_shape[feed_name][index]
+                            ).astype(elem.dtype)
                 else:
                     # shape_tensor is list
                     if (

test/tensorrt/test_converter_manipulation.py

@@ -692,6 +692,60 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestIndexPutCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.index_put
+        self.api_args = {
+            "x": np.random.random([3, 3]).astype("float32"),
+            "indices": (np.array([0, 1, 2]), np.array([1, 2, 1])),
+            "value": np.ones(3).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x", "indices", "value"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+    def test_fp16_result(self):
+        self.check_trt_result(precision_mode="fp16")
+
+
+class TestIndexPutCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.index_put
+        self.api_args = {
+            "x": np.random.random([3, 3]).astype("int64"),
+            "indices": (np.array([0, 1, 2]), np.array([1, 2, 1])),
+            "value": np.ones(3).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["x", "indices", "value"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestIndexPutCase3TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.index_put
+        self.api_args = {
+            "x": np.random.random([3, 3]).astype("float32"),
+            "indices": (np.array([0, 1, 2]), np.array([1, 2, 1])),
+            "value": np.ones(3).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x", "indices", "value"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
 class TestNumelTRTCase1Pattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.numel

test/tensorrt/test_converter_math.py

@@ -507,6 +507,56 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestPowCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.pow
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "y": float(np.random.randn()),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp32(self):
+        self.enable_fp16 = False
+        self.check_trt_result()
+
+    def test_trt_result_fp16(self):
+        self.enable_fp16 = True
+        self.check_trt_result()
+
+
+class TestPowCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.pow
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("int64"),
+            "y": int(np.random.randn()),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestPowCase3TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.pow
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "y": float(np.random.randn()),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
 class TestIsnanFP32TRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.isnan