[Paddle TensorRT] Add pd_op.index_put、pd_op.pow converter (#71157)

PolaKuma · web-flow · commit b12c77e3fb3f · 2025-02-19T10:45:38.000+08:00
* add pow and index_put

* fix codestyle

* Update test_converter_math.py

* update

* fix codestyle
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -266,6 +266,28 @@ using FloorDivideOpPattern =
 using RemainderOpPattern =
     ElementwiseCommonOpPattern<paddle::dialect::RemainderOp>;
 
+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;
+  }
+};
+
 template <typename OpType>
 class ActOpPattern : public pir::OpRewritePattern<OpType> {
  public:
@@ -1977,6 +1999,41 @@ class FullWithTensorPattern
   }
 };
 
+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;
+    }
+    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."
+              << value_num;
+      return false;
+    }
+    pir::Value indices = op.operand_source(1);
+    pir::VectorType vec_type = indices.type().dyn_cast<pir::VectorType>();
+    size_t output_num = vec_type.size();
+    for (size_t j = 0; j < output_num; j++) {
+      auto dtype =
+          vec_type[j].dyn_cast<paddle::dialect::DenseTensorType>().dtype();
+      if (!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 TakeAlongAxisOpPattern
     : public pir::OpRewritePattern<paddle::dialect::TakeAlongAxisOp> {
  public:
@@ -2748,6 +2805,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<ElementwisePowOpPattern>(context));
     ps.Add(std::make_unique<MeanOpPattern>(context));
     ps.Add(std::make_unique<RemainderOpPattern>(context));
+    ps.Add(std::make_unique<PowOpPattern>(context));
     ps.Add(std::make_unique<MulticlassNms3OpPattern>(context));
     ps.Add(std::make_unique<ArgmaxOpPattern>(context));
     ps.Add(std::make_unique<ArgminOpPattern>(context));
@@ -2780,6 +2838,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<CeluOpPattern>(context));
     ps.Add(std::make_unique<OneHotOpPattern>(context));
     ps.Add(std::make_unique<TemporalShiftOpPattern>(context));
+    ps.Add(std::make_unique<IndexPutOpPattern>(context));
     ps.Add(std::make_unique<InstanceNormOpPattern>(context));
     ps.Add(std::make_unique<EinsumOpPattern>(context));
     ps.Add(std::make_unique<PNormOpPattern>(context));
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/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,
@@ -968,3 +969,76 @@ def numel_converter(network, paddle_op, inputs):
         shape_tensor, trt.ReduceOperation.PROD, axes=1, keep_dims=False
     )
     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, indices_list, value_tensor = inputs
+    indices_tensor = indices_list[0]
+    input_shape_tensor = trt_shape(network, input_tensor)
+    input_dims = input_tensor.shape
+    indices_dims = indices_tensor.shape
+    rank = len(input_dims)
+
+    # indices
+    indices_shape_vec = [
+        add_1D_constant_layer(
+            network, indices_dims[i] if i < len(indices_dims) else 1
+        )
+        for i in range(rank)
+    ]
+    start_tensor_vec = [add_1D_constant_layer(network, 0)] * rank
+    stride_tensor_vec = [add_1D_constant_layer(network, 1)] * rank
+    indices_tensor_temp = trt_reshape(
+        network,
+        indices_tensor,
+        trt_concat(network, indices_shape_vec),
+        is_shape_tensor=True,
+    )
+    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)
diff --git a/python/paddle/tensorrt/impls/math.py b/python/paddle/tensorrt/impls/math.py
@@ -18,6 +18,7 @@
 from paddle.tensorrt.converter_utils import (
     add_1D_constant_layer,
     add_cast_reduce_layer,
+    add_constant_layer,
     add_elementwise_layer,
     add_reduce_layer,
     broadcast,
@@ -218,6 +219,21 @@ def _get_constant_or_expand_tensor(
     return layer.get_output(0)
 
 
+@converter_registry.register("pd_op.pow", trt_version="trt_version_ge=8.0")
+def pow_converter(network, paddle_op, inputs):
+    from paddle.tensorrt.util import support_fp32_mix_precision
+
+    x = inputs[0]
+    factor = paddle_op.attrs()["y"]
+    dims_x = x.shape
+    trt_dims_y = trt.Dims([1] * len(dims_x))
+    w_data = [factor]
+    y = add_constant_layer(network, w_data, trt_dims_y, np.float32)
+    layer = network.add_elementwise(x, y, trt.ElementWiseOperation.POW)
+    support_fp32_mix_precision(paddle_op.name(), layer)
+    return layer.get_output(0)
+
+
 @converter_registry.register("pd_op.remainder", trt_version="8.x")
 @converter_registry.register("pd_op.remainder_", trt_version="8.x")
 def remainder_converter(network, paddle_op, inputs):
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -776,5 +776,45 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestIndexPutTRTCasePattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.index_put
+        self.api_args = {
+            "x": np.zeros([2, 3]).astype("float32"),
+            "indices": [np.array([1, 0]).astype("bool")],
+            "value": np.array([1]).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x", "indices", "value"]}
+        self.min_shape = {"x": [2, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [4, 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, 1]).astype("bool")],
+            "value": np.array([2]).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()
+
+    def test_fp16_result(self):
+        self.check_trt_result(precision_mode="fp16")
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/tensorrt/test_converter_math.py b/test/tensorrt/test_converter_math.py
@@ -119,6 +119,44 @@ def test_trt_result_fp32(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": 2.5,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+
+class TestPowCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.pow
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "y": 2,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+
 class TestRemainderFloatTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.remainder
