diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc index ae2d09a827c7f6..b1dc188dad3b5b 100644 --- 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) @@ -2194,6 +2195,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 index 20e8cfe6fb9611..73681f02589b5f 100644 --- a/python/paddle/tensorrt/impls/activation.py +++ b/python/paddle/tensorrt/impls/activation.py @@ -135,6 +135,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 index c3f077364c14b2..6b37c776dc0ed1 100644 --- 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