[Paddle Tensorrt] add pd_op.tanh marker and converter (#68614)

Hanyonggong · web-flow · commit 079677e1e72e · 2024-10-18T10:37:00.000+08:00
* add tanh_marker

* add math_converter

* mod marker

* add version_judge

* add func_activation_converter

* mod file

* add if

* mod file

* mod converter

* mod map
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1459,6 +1459,31 @@ class NearestInterV2Pattern
   }
 };
 
+class TanhOpPattern : public pir::OpRewritePattern<paddle::dialect::TanhOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::TanhOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::TanhOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+#if IS_TRT_VERSION_LT(8600)
+    pir::Value x = op.operand_source(0);
+    auto x_type = x.type().dyn_cast<paddle::dialect::DenseTensorType>();
+    auto x_shape = x_type.dims();
+    int dims = x_shape.size();
+    if (dims < 1) {
+      VLOG(3) << "Tanh op does not support 0 dim input when TensorRT < 8.6.";
+      return false;
+    }
+#endif
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
@@ -1534,6 +1559,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<MaxOpPattern>(context));
     ps.Add(std::make_unique<BilinearInterpV2Pattern>(context));
     ps.Add(std::make_unique<NearestInterV2Pattern>(context));
+    ps.Add(std::make_unique<TanhOpPattern>(context));
     return ps;
   }
 };
diff --git a/python/paddle/tensorrt/impls/activation.py b/python/paddle/tensorrt/impls/activation.py
@@ -20,11 +20,21 @@
 )
 from paddle.tensorrt.register import converter_registry
 
+activation_type_map = {
+    "pd_op.tanh": trt.ActivationType.TANH,
+    "pd_op.relu": trt.ActivationType.RELU,
+    "pd_op.sigmoid": trt.ActivationType.SIGMOID,
+}
+
 
 @converter_registry.register("pd_op.relu", trt_version="8.x")
-def relu_converter(network, paddle_op, inputs):
-    relu_layer = network.add_activation(inputs[0], trt.ActivationType.RELU)
-    return relu_layer.get_output(0)
+@converter_registry.register("pd_op.tanh", trt_version="8.x")
+@converter_registry.register("pd_op.sigmoid", trt_version="8.x")
+def activation_converter(network, paddle_op, inputs):
+    layer = network.add_activation(
+        inputs[0], activation_type_map[paddle_op.name()]
+    )
+    return layer.get_output(0)
 
 
 @converter_registry.register("pd_op.softmax", trt_version="8.x")
diff --git a/python/paddle/tensorrt/impls/ops.py b/python/paddle/tensorrt/impls/ops.py
@@ -23,11 +23,3 @@ def sqrt_converter(network, paddle_op, inputs):
 
     sqrt_layer = network.add_unary(input_tensor, trt.UnaryOperation.SQRT)
     return sqrt_layer.get_output(0)
-
-
-@converter_registry.register("pd_op.sigmoid", trt_version="8.x")
-def sigmoid_converter(network, paddle_op, inputs):
-    sigmoid_layer = network.add_activation(
-        inputs[0], trt.ActivationType.SIGMOID
-    )
-    return sigmoid_layer.get_output(0)
diff --git a/test/tensorrt/test_converter_activation.py b/test/tensorrt/test_converter_activation.py
@@ -48,5 +48,43 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+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.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1]}
+        self.max_shape = {"x": [5]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestTANHTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.tanh
+        self.api_args = {"x": np.random.randn(3).astype(np.float32)}
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1]}
+        self.max_shape = {"x": [5]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+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),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3], "y": [1, 3]}
+        self.max_shape = {"x": [5, 3], "y": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/tensorrt/test_converter_ops.py b/test/tensorrt/test_converter_ops.py
@@ -34,19 +34,5 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
-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),
-        }
-        self.program_config = {"feed_list": ["x"]}
-        self.min_shape = {"x": [1, 3], "y": [1, 3]}
-        self.max_shape = {"x": [5, 3], "y": [5, 3]}
-
-    def test_trt_result(self):
-        self.check_trt_result()
-
-
 if __name__ == '__main__':
     unittest.main()
