[Paddle TensorRT] Support isnan, group_norm and take_along_axis (#70817)

ckl117 · web-flow · commit 646f6c64deef · 2025-01-16T16:00:57.000+08:00
* add isnan, group_norm and take_along_axis

* check isnan

* remove unused files
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -109,6 +109,7 @@ DEFINE_GENERAL_PATTERN(Ceil, paddle::dialect::CeilOp)
 DEFINE_GENERAL_PATTERN(Rsqrt, paddle::dialect::RsqrtOp)
 DEFINE_GENERAL_PATTERN(Reciprocal, paddle::dialect::ReciprocalOp)
 DEFINE_GENERAL_PATTERN(Erf, paddle::dialect::ErfOp)
+DEFINE_GENERAL_PATTERN(Isnan, paddle::dialect::IsnanOp)
 DEFINE_GENERAL_PATTERN(Sign, paddle::dialect::SignOp)
 DEFINE_GENERAL_PATTERN(Round, paddle::dialect::RoundOp)
 DEFINE_GENERAL_PATTERN(Numel, paddle::dialect::NumelOp)
@@ -1910,6 +1911,34 @@ class FullWithTensorPattern
   }
 };
 
+class TakeAlongAxisOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::TakeAlongAxisOp> {
+ public:
+  using pir::OpRewritePattern<
+      paddle::dialect::TakeAlongAxisOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(paddle::dialect::TakeAlongAxisOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    pir::Value x_value = op.operand_source(0);
+    auto x_shape = pir::GetShapeFromValue(x_value);
+    pir::Value index_value = op.operand_source(1);
+    auto index_shape = pir::GetShapeFromValue(index_value);
+    if (x_shape.size() != index_shape.size()) {
+      VLOG(3) << "TakeAlongAxis op Index input dims size ["
+              << index_shape.size() << "] is not equal to input dims size ["
+              << x_shape.size() << "].";
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class StridedSliceOpPattern
     : public pir::OpRewritePattern<paddle::dialect::StridedSliceOp> {
  public:
@@ -2394,6 +2423,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Rsqrt)
     ADD_PATTERN(Reciprocal)
     ADD_PATTERN(Erf)
+    ADD_PATTERN(Isnan)
     ADD_PATTERN(Sign)
     ADD_PATTERN(Round)
     ADD_PATTERN(Numel)
@@ -2458,6 +2488,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<WherePattern>(context));
     ps.Add(std::make_unique<FullLikeOpPattern>(context));
     ps.Add(std::make_unique<FullWithTensorPattern>(context));
+    ps.Add(std::make_unique<TakeAlongAxisOpPattern>(context));
     ps.Add(std::make_unique<StridedSliceOpPattern>(context));
     ps.Add(std::make_unique<TopkOpPattern>(context));
     ps.Add(std::make_unique<CumsumOpPattern>(context));
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -791,6 +791,28 @@ def tile_converter(network, paddle_op, inputs):
     return slice_layer.get_output(0)
 
 
+@converter_registry.register(
+    "pd_op.take_along_axis", trt_version="trt_version_ge=8.2"
+)
+def take_along_axis_converter(network, paddle_op, inputs):
+    axis = paddle_op.attrs().get("axis", 0)
+    input_tensor = inputs[0]
+    index_tensor = inputs[1]
+
+    input_dims = input_tensor.shape
+    if axis < 0:
+        axis += len(input_dims)
+
+    gather_layer = network.add_gather_v2(
+        input_tensor, index_tensor, trt.GatherMode.ELEMENT
+    )
+    gather_layer.axis = axis
+
+    output_tensor = gather_layer.get_output(0)
+
+    return output_tensor
+
+
 @converter_registry.register("pd_op.strided_slice", trt_version="8.x")
 def strided_slice_converter(network, paddle_op, inputs):
     input_tensor = inputs[0]
diff --git a/python/paddle/tensorrt/impls/math.py b/python/paddle/tensorrt/impls/math.py
@@ -29,6 +29,7 @@
     get_shape_tensor_element,
     trt_cast,
     trt_concat,
+    trt_equal,
     trt_expand,
     trt_max,
     trt_reshape,
@@ -400,6 +401,14 @@ def elementwise_pow_converter(network, paddle_op, inputs):
     )
 
 
+@converter_registry.register("pd_op.isnan", trt_version="8.x")
+def isnan_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    equal_tensor = trt_equal(network, input_tensor, input_tensor)
+    layer = network.add_unary(equal_tensor, trt.UnaryOperation.NOT)
+    return layer.get_output(0)
+
+
 @converter_registry.register("pd_op.minimum", trt_version="8.x")
 def minimum_converter(network, paddle_op, inputs):
     min_layer = add_elementwise_layer(
diff --git a/python/paddle/tensorrt/impls/norm.py b/python/paddle/tensorrt/impls/norm.py
@@ -16,11 +16,16 @@
 import tensorrt as trt
 
 from paddle.tensorrt.converter_utils import (
+    add_1D_constant_layer,
     append_ones,
     get_axes_for_reduce_op,
     get_dynamic_dims,
     get_trt_plugin,
     has_dynamic_shape,
+    trt_expand,
+    trt_prod,
+    trt_reshape,
+    trt_sum,
 )
 from paddle.tensorrt.register import converter_registry
 
@@ -152,3 +157,39 @@ def instance_norm_converter(network, paddle_op, inputs):
     )
     instance_norm_layer = network.add_plugin_v2(instance_norm_inputs, plugin)
     return instance_norm_layer.get_output(0)
+
+
+@converter_registry.register(
+    "pd_op.group_norm", trt_version="trt_version_ge=8.6"
+)
+def group_norm_converter(network, paddle_op, inputs):
+    x, scale, bias = inputs
+    groups = paddle_op.attrs().get("groups", 1)
+    eps = paddle_op.attrs().get("epsilon", 1e-05)
+
+    axes_mask = 0
+    x_shape = paddle_op.operands()[0].source().shape
+    rank_x = len(x_shape)
+
+    fake_shape = [1, groups] + [1] * (rank_x - 2)
+    broadcast_shape = [1, x_shape[1]] + [1] * (rank_x - 2)
+    for d in range(2, rank_x):
+        axes_mask |= 1 << d
+
+    weight_one = add_1D_constant_layer(network, 1.0, np.float32)
+    bias_zero = add_1D_constant_layer(network, 0.0, np.float32)
+    fake_shape = add_1D_constant_layer(network, fake_shape, np.int32)
+    weight_one = trt_expand(network, weight_one, 1, fake_shape, rank_x)
+    bias_zero = trt_expand(network, bias_zero, 1, fake_shape, rank_x)
+    layer = network.add_normalization(x, weight_one, bias_zero, axes_mask)
+    layer.num_groups = groups
+    layer.epsilon = eps
+    output = layer.get_output(0)
+    if scale is not None:
+        scale = trt_reshape(network, scale, broadcast_shape)
+        output = trt_prod(network, output, scale)
+    if bias is not None:
+        bias = trt_reshape(network, bias, broadcast_shape)
+        output = trt_sum(network, output, bias)
+
+    return output
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -460,6 +460,57 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestTakeAlongAxisCase0TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.take_along_axis
+        self.api_args = {
+            "X": np.random.random([3, 4, 10]).astype("float32"),
+            "Index": np.random.randint(0, 2, size=(3, 4, 10)).astype("int64"),
+            "axis": 1,
+        }
+        self.program_config = {"feed_list": ["X", "Index"]}
+        self.min_shape = {"X": [1, 4, 10], "Index": [1, 4, 10]}
+        self.opt_shape = {"X": [3, 4, 10], "Index": [3, 4, 10]}
+        self.max_shape = {"X": [5, 4, 10], "Index": [5, 4, 10]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestTakeAlongAxisCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.take_along_axis
+        self.api_args = {
+            "X": np.random.random([3, 4, 10]).astype("float32"),
+            "Index": np.random.randint(0, 2, size=(3, 4, 10)).astype("int64"),
+            "axis": -1,
+        }
+        self.program_config = {"feed_list": ["X", "Index"]}
+        self.min_shape = {"X": [1, 4, 10], "Index": [1, 4, 10]}
+        self.opt_shape = {"X": [3, 4, 10], "Index": [3, 4, 10]}
+        self.max_shape = {"X": [5, 4, 10], "Index": [5, 4, 10]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestTakeAlongAxisFP16TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.take_along_axis
+        self.api_args = {
+            "X": np.random.random([3, 4, 10]).astype("float32"),
+            "Index": np.random.randint(0, 2, size=(3, 4, 10)).astype("int64"),
+            "axis": 1,
+        }
+        self.program_config = {"feed_list": ["X", "Index"]}
+        self.min_shape = {"X": [1, 4, 10], "Index": [1, 4, 10]}
+        self.opt_shape = {"X": [3, 4, 10], "Index": [3, 4, 10]}
+        self.max_shape = {"X": [5, 4, 10], "Index": [5, 4, 10]}
+
+    def test_trt_result(self):
+        self.check_trt_result(precision_mode="fp16")
+
+
 class TestStrideSliceCase1TRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.strided_slice
diff --git a/test/tensorrt/test_converter_math.py b/test/tensorrt/test_converter_math.py
@@ -507,6 +507,51 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestIsnanFP32TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.isnan
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        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 TestIsnanFP16TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.isnan
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        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(precision_mode="fp16")
+
+
+class TestIsnanIntTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.isnan
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        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 TestMaximumTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.maximum
diff --git a/test/tensorrt/test_converter_norm.py b/test/tensorrt/test_converter_norm.py
@@ -95,5 +95,45 @@ def test_trt_result(self):
         self.check_marker(expected_result=False)
 
 
+class TestGroupNormNCHWFP32TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.group_norm
+        self.api_args = {
+            "x": np.random.random([4, 32, 64, 64]).astype(np.float32),
+            "num_groups": 2,
+            "epsilon": 1e-05,
+            "weight": np.random.randn(32).astype(np.float32),
+            "bias": np.random.randn(32).astype(np.float32),
+            "data_format": "NCHW",
+        }
+        self.program_config = {"feed_list": ["x", "weight", "bias"]}
+        self.min_shape = {"x": [1, 32, 64, 64]}
+        self.opt_shape = {"x": [4, 32, 64, 64]}
+        self.max_shape = {"x": [6, 32, 64, 64]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestGroupNormNCHWFP16TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.group_norm
+        self.api_args = {
+            "x": np.random.random([4, 32, 64, 64]).astype(np.float32),
+            "num_groups": 2,
+            "epsilon": 1e-05,
+            "weight": np.random.randn(32).astype(np.float32),
+            "bias": np.random.randn(32).astype(np.float32),
+            "data_format": "NCHW",
+        }
+        self.program_config = {"feed_list": ["x", "weight", "bias"]}
+        self.min_shape = {"x": [1, 32, 64, 64]}
+        self.opt_shape = {"x": [4, 32, 64, 64]}
+        self.max_shape = {"x": [6, 32, 64, 64]}
+
+    def test_trt_result(self):
+        self.check_trt_result(precision_mode="fp16")
+
+
 if __name__ == '__main__':
     unittest.main()
