pd_op.softplus (#69315)

lizexu123 · web-flow · commit 83738fb0b4fd · 2024-11-14T18:43:30.000+08:00
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1490,13 +1490,14 @@ class StackOpPattern : public pir::OpRewritePattern<paddle::dialect::StackOp> {
   }
 };
 
-class TanhOpPattern : public pir::OpRewritePattern<paddle::dialect::TanhOp> {
+template <typename OpType>
+class ActOpPattern : public pir::OpRewritePattern<OpType> {
  public:
-  using pir::OpRewritePattern<paddle::dialect::TanhOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::TanhOp op,
+  using pir::OpRewritePattern<OpType>::OpRewritePattern;
+  bool MatchAndRewrite(OpType op,
                        pir::PatternRewriter &rewriter) const override {
     if (op->HasAttribute(kCanRunTrtAttr) &&
-        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+        op->template attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
       return false;
     }
 #if IS_TRT_VERSION_LT(8600)
@@ -1514,6 +1515,8 @@ class TanhOpPattern : public pir::OpRewritePattern<paddle::dialect::TanhOp> {
     return true;
   }
 };
+using TanhOpPattern = ActOpPattern<paddle::dialect::TanhOp>;
+using SoftplusOpPatten = ActOpPattern<paddle::dialect::SoftplusOp>;
 
 class WherePattern : public pir::OpRewritePattern<paddle::dialect::WhereOp> {
  public:
@@ -1821,6 +1824,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<FullWithTensorPattern>(context));
     ps.Add(std::make_unique<StridedSliceOpPattern>(context));
     ps.Add(std::make_unique<TopkOpPattern>(context));
+    ps.Add(std::make_unique<SoftplusOpPatten>(context));
     ps.Add(std::make_unique<EqualOpPattern>(context));
     ps.Add(std::make_unique<NotEqualOpPattern>(context));
     return ps;
diff --git a/python/paddle/tensorrt/impls/activation.py b/python/paddle/tensorrt/impls/activation.py
@@ -104,6 +104,23 @@ def hardswish_converter(network, paddle_op, inputs):
     return hardswish_layer.get_output(0)
 
 
+@converter_registry.register("pd_op.softplus", trt_version="8.x")
+def softplus_converter(network, paddle_op, inputs):
+    x = inputs[0]
+    beta = paddle_op.attrs()["beta"]
+    threshold = paddle_op.attrs()["threshold"]
+    layer_clip = network.add_activation(x, trt.ActivationType.CLIP)
+    layer_clip.alpha = -3.40282e038
+    layer_clip.beta = threshold / beta
+
+    softplus_layer = network.add_activation(
+        layer_clip.get_output(0), trt.ActivationType.SOFTPLUS
+    )
+    softplus_layer.alpha = 1.0 / beta
+    softplus_layer.beta = beta
+    return softplus_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.swish", trt_version="8.x")
 @converter_registry.register("pd_op.silu", trt_version="8.x")
 def swish_silu_converter(network, paddle_op, inputs):
diff --git a/test/tensorrt/test_converter_activation.py b/test/tensorrt/test_converter_activation.py
@@ -48,10 +48,10 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
-class TestRELUTRTPattern(TensorRTBaseTest):
+class TestReluTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.nn.functional.relu
-        self.api_args = {"x": np.random.randn(3).astype(np.float32)}
+        self.api_args = {"x": np.random.randn(3).astype("float32")}
         self.program_config = {"feed_list": ["x"]}
         self.min_shape = {"x": [1]}
         self.max_shape = {"x": [5]}
@@ -60,10 +60,10 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
-class TestTANHTRTPattern(TensorRTBaseTest):
+class TestTanhTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.tanh
-        self.api_args = {"x": np.random.randn(3).astype(np.float32)}
+        self.api_args = {"x": np.random.randn(3).astype("float32")}
         self.program_config = {"feed_list": ["x"]}
         self.min_shape = {"x": [1]}
         self.max_shape = {"x": [5]}
@@ -76,11 +76,25 @@ class TestSigmoidTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.nn.functional.sigmoid
         self.api_args = {
-            "x": np.random.randn(2, 3).astype(np.float32),
+            "x": np.random.randn(2, 3).astype("float32"),
         }
         self.program_config = {"feed_list": ["x"]}
-        self.min_shape = {"x": [1, 3], "y": [1, 3]}
-        self.max_shape = {"x": [5, 3], "y": [5, 3]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestSoftplusTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.Softplus()
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+        }
+        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()
