【SCU】【Paddle TensorRT No.4】Add pd_op.stanh converter (#69539)

PolaKuma · web-flow · commit 461fbd138e4b · 2024-12-11T16:58:55.000+08:00
* add_tanh

* fix codestyle

* fix codestyle

* fix

* fix codestyle

* fix codestyle

* 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
@@ -84,6 +84,7 @@ DEFINE_GENERAL_PATTERN(Swish, paddle::dialect::SwishOp)
 DEFINE_GENERAL_PATTERN(Log, paddle::dialect::LogOp)
 DEFINE_GENERAL_PATTERN(Floor, paddle::dialect::FloorOp)
 DEFINE_GENERAL_PATTERN(Roll, paddle::dialect::RollOp)
+DEFINE_GENERAL_PATTERN(Stanh, paddle::dialect::StanhOp)
 DEFINE_GENERAL_PATTERN(Softplus, paddle::dialect::SoftplusOp)
 DEFINE_GENERAL_PATTERN(ThresholdedRelu, paddle::dialect::ThresholdedReluOp)
 DEFINE_GENERAL_PATTERN(Flip, paddle::dialect::FlipOp)
@@ -2166,6 +2167,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Log)
     ADD_PATTERN(Floor)
     ADD_PATTERN(Roll)
+    ADD_PATTERN(Stanh)
     ADD_PATTERN(Softplus)
     ADD_PATTERN(ThresholdedRelu)
     ADD_PATTERN(Flip)
diff --git a/python/paddle/tensorrt/impls/activation.py b/python/paddle/tensorrt/impls/activation.py
@@ -139,6 +139,17 @@ def swish_silu_converter(network, paddle_op, inputs):
     return trt_prod(network, inputs[0], layer_output)
 
 
+@converter_registry.register("pd_op.stanh", trt_version="8.x")
+def stanh_converter(network, paddle_op, inputs):
+    x = inputs[0]
+    scale_a = paddle_op.attrs()["scale_a"]
+    scale_b = paddle_op.attrs()["scale_b"]
+    stanh_layer = network.add_activation(x, trt.ActivationType.SCALED_TANH)
+    stanh_layer.alpha = scale_b
+    stanh_layer.beta = scale_a
+    return stanh_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.mish", trt_version="8.x")
 def mish_converter(network, paddle_op, inputs):
     x = inputs[0]
diff --git a/test/tensorrt/test_converter_activation.py b/test/tensorrt/test_converter_activation.py
@@ -128,6 +128,22 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestStanhFloatTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.stanh
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "scale_a": 0.67,
+            "scale_b": 1.7159,
+        }
+        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 TestCeluTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.nn.functional.celu
