[HEU][Paddle TensorRT No.69-72,74-85] Add UnaryOp converter (#70535)

fangfangssj · web-flow · commit 29c3d915e086 · 2025-01-08T11:05:00.000+08:00
* add converter

* fix

* add marker

* fix

* 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
@@ -94,7 +94,24 @@ DEFINE_GENERAL_PATTERN(Flip, paddle::dialect::FlipOp)
 DEFINE_GENERAL_PATTERN(Mish, paddle::dialect::MishOp)
 DEFINE_GENERAL_PATTERN(AssignValue, paddle::dialect::AssignValueOp)
 DEFINE_GENERAL_PATTERN(AssignValue_, paddle::dialect::AssignValue_Op)
+DEFINE_GENERAL_PATTERN(Exp, paddle::dialect::ExpOp)
+DEFINE_GENERAL_PATTERN(Abs, paddle::dialect::AbsOp)
+DEFINE_GENERAL_PATTERN(Abs_, paddle::dialect::Abs_Op)
+DEFINE_GENERAL_PATTERN(Sin, paddle::dialect::SinOp)
+DEFINE_GENERAL_PATTERN(Cos, paddle::dialect::CosOp)
+DEFINE_GENERAL_PATTERN(Sinh, paddle::dialect::SinhOp)
+DEFINE_GENERAL_PATTERN(Cosh, paddle::dialect::CoshOp)
+DEFINE_GENERAL_PATTERN(Asinh, paddle::dialect::AsinhOp)
+DEFINE_GENERAL_PATTERN(Acosh, paddle::dialect::AcoshOp)
+DEFINE_GENERAL_PATTERN(Atanh, paddle::dialect::AtanhOp)
+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(Sign, paddle::dialect::SignOp)
+DEFINE_GENERAL_PATTERN(Round, paddle::dialect::RoundOp)
 DEFINE_GENERAL_PATTERN(Numel, paddle::dialect::NumelOp)
+
 #undef DEFINE_GENERAL_PATTERN
 
 // Add ReduceCommonOpPattern base class to simplify code
@@ -267,8 +284,30 @@ class ActOpPattern : public pir::OpRewritePattern<OpType> {
 using TanhOpPattern = ActOpPattern<paddle::dialect::TanhOp>;
 using CeluOpPattern = ActOpPattern<paddle::dialect::CeluOp>;
 using TanhShrinkOpPattern = ActOpPattern<paddle::dialect::TanhShrinkOp>;
-using LogicalNotOpPattern = ActOpPattern<paddle::dialect::LogicalNotOp>;
-using LogicalNot_OpPattern = ActOpPattern<paddle::dialect::LogicalNot_Op>;
+
+template <typename OpType>
+class Logical_NotOpPattern : public pir::OpRewritePattern<OpType> {
+ public:
+  using pir::OpRewritePattern<OpType>::OpRewritePattern;
+  bool MatchAndRewrite(OpType op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->template attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    pir::Value x = op.operand_source(0);
+    auto x_dtype = pir::GetDataTypeFromValue(x);
+    if (!x_dtype.isa<pir::BoolType>()) {
+      VLOG(3) << " logical_not op only support bool input in tensorrt.";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+using LogicalNotOpPattern = Logical_NotOpPattern<paddle::dialect::LogicalNotOp>;
+using LogicalNot_OpPattern =
+    Logical_NotOpPattern<paddle::dialect::LogicalNot_Op>;
 
 class Pool2dOpPattern
     : public pir::OpRewritePattern<paddle::dialect::Pool2dOp> {
@@ -538,24 +577,6 @@ class ArangeOpPattern
   }
 };
 
-class SignOpPattern : public pir::OpRewritePattern<paddle::dialect::SignOp> {
- public:
-  using pir::OpRewritePattern<paddle::dialect::SignOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::SignOp op,
-                       pir::PatternRewriter &rewriter) const override {
-    if (op->HasAttribute(kCanRunTrtAttr) &&
-        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
-      return false;
-    }
-#if IS_TRT_VERSION_LT(8200)
-    VLOG(3) << "sign op is only supported by tensorrt8.2 above ";
-    return false;
-#endif
-    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
-    return true;
-  }
-};
-
 class GroupNormOpPattern
     : public pir::OpRewritePattern<paddle::dialect::GroupNormOp> {
  public:
@@ -2273,6 +2294,23 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Mish)
     ADD_PATTERN(AssignValue)
     ADD_PATTERN(AssignValue_)
+    ADD_PATTERN(Exp)
+    ADD_PATTERN(Abs)
+    ADD_PATTERN(Abs_)
+    ADD_PATTERN(Cos)
+    ADD_PATTERN(Sin)
+    ADD_PATTERN(Cos)
+    ADD_PATTERN(Sinh)
+    ADD_PATTERN(Cosh)
+    ADD_PATTERN(Asinh)
+    ADD_PATTERN(Acosh)
+    ADD_PATTERN(Atanh)
+    ADD_PATTERN(Ceil)
+    ADD_PATTERN(Rsqrt)
+    ADD_PATTERN(Reciprocal)
+    ADD_PATTERN(Erf)
+    ADD_PATTERN(Sign)
+    ADD_PATTERN(Round)
     ADD_PATTERN(Numel)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
@@ -2283,7 +2321,6 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<DepthwiseConv2dTransposeOpPattern>(context));
     ps.Add(std::make_unique<DeformableConvOpPattern>(context));
     ps.Add(std::make_unique<ArangeOpPattern>(context));
-    ps.Add(std::make_unique<SignOpPattern>(context));
     ps.Add(std::make_unique<LogicalNotOpPattern>(context));
     ps.Add(std::make_unique<BitwiseAndOpPattern>(context));
     ps.Add(std::make_unique<BitwiseOrOpPattern>(context));
diff --git a/python/paddle/tensorrt/converter_utils.py b/python/paddle/tensorrt/converter_utils.py
@@ -686,6 +686,29 @@ def squeeze_trt(network, input_tensor, axes):
 def unary_op_converter(network, paddle_op, inputs):
     from paddle.tensorrt import PrecisionMode
 
+    ops_type_map = {
+        "pd_op.sqrt": [trt.UnaryOperation.SQRT],
+        "pd_op.sqrt_": [trt.UnaryOperation.SQRT],
+        "pd_op.floor": [trt.UnaryOperation.FLOOR],
+        "pd_op.exp": [trt.UnaryOperation.EXP],
+        "pd_op.abs": [trt.UnaryOperation.ABS],
+        "pd_op.abs_": [trt.UnaryOperation.ABS],
+        "pd_op.sin": [trt.UnaryOperation.SIN],
+        "pd_op.cos": [trt.UnaryOperation.COS],
+        "pd_op.sinh": [trt.UnaryOperation.SINH],
+        "pd_op.cosh": [trt.UnaryOperation.COSH],
+        "pd_op.asinh": [trt.UnaryOperation.ASINH],
+        "pd_op.acosh": [trt.UnaryOperation.ACOSH],
+        "pd_op.atanh": [trt.UnaryOperation.ATANH],
+        "pd_op.ceil": [trt.UnaryOperation.CEIL],
+        "pd_op.reciprocal": [trt.UnaryOperation.RECIP],
+        "pd_op.erf": [trt.UnaryOperation.ERF],
+        "pd_op.sign": [trt.UnaryOperation.SIGN],
+        "pd_op.round": [trt.UnaryOperation.ROUND],
+        "pd_op.logical_not": [trt.UnaryOperation.NOT],
+        "pd_op.rsqrt": [trt.UnaryOperation.SQRT, trt.UnaryOperation.RECIP],
+    }
+
     input_tensor = inputs[0]
     layer = None
     org_type = input_tensor.dtype
@@ -707,9 +730,10 @@ def unary_op_converter(network, paddle_op, inputs):
             identity_layer.set_output_type(0, trt.float16)
         input_tensor = identity_layer.get_output(0)
 
-    if paddle_op.name() in ["pd_op.logical_not", "pd_op.logical_not_"]:
-        layer = network.add_unary(input_tensor, trt.UnaryOperation.NOT)
-        input_tensor = layer.get_output(0)
+    if paddle_op.name() in ops_type_map:
+        for trt_op in ops_type_map[paddle_op.name()]:
+            layer = network.add_unary(input_tensor, trt_op)
+            input_tensor = layer.get_output(0)
     else:
         raise NotImplementedError(
             f"Unsupported unary operation: {paddle_op.name()}"
diff --git a/python/paddle/tensorrt/impls/ops.py b/python/paddle/tensorrt/impls/ops.py
@@ -11,21 +11,32 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import tensorrt as trt
 
+from paddle.tensorrt.converter_utils import unary_op_converter
 from paddle.tensorrt.register import converter_registry
 
-ops_type_map = {
-    "pd_op.sqrt": trt.UnaryOperation.SQRT,
-    "pd_op.sqrt_": trt.UnaryOperation.SQRT,
-    "pd_op.floor": trt.UnaryOperation.FLOOR,
-}
-
 
 @converter_registry.register("pd_op.sqrt", trt_version="trt_version_ge=8.0")
 @converter_registry.register("pd_op.sqrt_", trt_version="trt_version_ge=8.0")
-@converter_registry.register("pd_op.floor", trt_version="8.x")
-def sqrt_converter(network, paddle_op, inputs):
-    input_tensor = inputs[0]
-    layer = network.add_unary(input_tensor, ops_type_map[paddle_op.name()])
-    return layer.get_output(0)
+@converter_registry.register("pd_op.floor", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.exp", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.abs", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.abs_", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.sin", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.cos", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.sinh", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.cosh", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.asinh", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.acosh", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.atanh", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.ceil", trt_version="trt_version_ge=8.0")
+@converter_registry.register(
+    "pd_op.reciprocal", trt_version="trt_version_ge=8.0"
+)
+@converter_registry.register("pd_op.erf", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.rsqrt", trt_version="trt_version_ge=8.0")
+@converter_registry.register("pd_op.sign", trt_version="trt_version_ge=8.2")
+@converter_registry.register("pd_op.round", trt_version="trt_version_ge=8.2")
+def UnaryOpConverter(network, paddle_op, inputs):
+    layer_output = unary_op_converter(network, paddle_op, inputs)
+    return layer_output
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -14,7 +14,7 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_conv PROPERTIES TIMEOUT "300")
   set_tests_properties(test_export PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_norm PROPERTIES TIMEOUT "300")
-  set_tests_properties(test_converter_ops PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_ops PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_stat PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_math PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_activation PROPERTIES TIMEOUT "300")
diff --git a/test/tensorrt/test_converter_ops.py b/test/tensorrt/test_converter_ops.py
