add pd_op.mean and pd_op.sqrt converter for PIR TRT (#67974)

zhink · web-flow · commit ed2e83dc2fbe · 2024-09-12T17:24:01.000+08:00
* add pd_op.mean pd_op.sqrt converter for PIR TRT

* fix code style and delete some code

* 怎么会被合入 ~ ~

* chufa

* 触发

* 触发

* 触发

* codestyle

* codestyle

* codestyle
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -72,6 +72,7 @@ DEFINE_GENERAL_PATTERN(DepthwiseConv2d, paddle::dialect::DepthwiseConv2dOp)
 DEFINE_GENERAL_PATTERN(Shape, paddle::dialect::ShapeOp)
 DEFINE_GENERAL_PATTERN(Expand, paddle::dialect::ExpandOp)
 DEFINE_GENERAL_PATTERN(Sigmoid, paddle::dialect::SigmoidOp)
+DEFINE_GENERAL_PATTERN(Sqrt, paddle::dialect::SqrtOp)
 DEFINE_GENERAL_PATTERN(Hardsigmoid, paddle::dialect::HardsigmoidOp)
 DEFINE_GENERAL_PATTERN(Hardswish, paddle::dialect::HardswishOp)
 
@@ -1214,6 +1215,42 @@ class MaxOpPattern : public pir::OpRewritePattern<paddle::dialect::MaxOp> {
   }
 };
 
+class MeanOpPattern : public pir::OpRewritePattern<paddle::dialect::MeanOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::MeanOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::MeanOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+
+    if (!op->HasAttribute("axis")) {
+      VLOG(3) << "The axis attribute does not exist";
+      return false;
+    }
+
+    if (!op->HasAttribute("keepdim")) {
+      VLOG(3) << "The keepdim attribute does not exist";
+      return false;
+    }
+
+    pir::Value input = op.operand_source(0);
+    auto input_type = pir::GetDataTypeFromValue(input);
+    if (!input_type.isa<pir::Int32Type>() &&
+        !input_type.isa<pir::Int64Type>() &&
+        !input_type.isa<pir::Float32Type>() &&
+        !input_type.isa<pir::Float64Type>()) {
+      VLOG(3) << "The input type of pd_op.mean is not int32 or int64 or "
+                 "float32 or float64";
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
@@ -1245,6 +1282,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Shape)
     ADD_PATTERN(Expand)
     ADD_PATTERN(Sigmoid)
+    ADD_PATTERN(Sqrt)
     ADD_PATTERN(Hardsigmoid)
     ADD_PATTERN(Hardswish)
 #if IS_TRT_VERSION_GE(8600)
@@ -1280,6 +1318,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<MinimumOpPattern>(context));
     ps.Add(std::make_unique<MaximumOpPattern>(context));
     ps.Add(std::make_unique<FloorDivideOpPattern>(context));
+    ps.Add(std::make_unique<MeanOpPattern>(context));
     ps.Add(std::make_unique<RemainderOpPattern>(context));
     ps.Add(std::make_unique<MulticlassNms3OpPattern>(context));
     ps.Add(std::make_unique<ArgmaxOpPattern>(context));
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -15,6 +15,7 @@
 import ctypes
 import hashlib
 import logging
+import os
 
 import numpy as np
 import tensorrt as trt
@@ -39,6 +40,7 @@
 from .impls.others import *  # noqa: F403
 from .impls.pooling import *  # noqa: F403
 from .impls.search import *  # noqa: F403
+from .impls.stat import *  # noqa: F403
 from .register import converter_registry
 from .util import map_dtype
 
diff --git a/python/paddle/tensorrt/impls/activation.py b/python/paddle/tensorrt/impls/activation.py
@@ -11,28 +11,14 @@
 # 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 logging
-import os
-import sys
 
 import numpy as np
-
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
-
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
-from paddle.tensorrt.register import converter_registry
-
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
 from paddle.tensorrt.converter_utils import (
     get_trt_plugin,
 )
+from paddle.tensorrt.register import converter_registry
 
 
 @converter_registry.register("pd_op.relu", trt_version="8.x")
diff --git a/python/paddle/tensorrt/impls/conv.py b/python/paddle/tensorrt/impls/conv.py
@@ -11,24 +11,12 @@
 # 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 logging
-import os
-import sys
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
 
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
 from paddle.tensorrt.register import converter_registry
 
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
-
 
 @converter_registry.register("pd_op.conv2d", trt_version="8.x")
 def conv2d_converter(network, paddle_op, inputs):
diff --git a/python/paddle/tensorrt/impls/creation.py b/python/paddle/tensorrt/impls/creation.py
@@ -11,26 +11,12 @@
 # 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 logging
-import os
-import sys
 
 import numpy as np
-
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
-
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
 from paddle.tensorrt.register import converter_registry
 
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
-
 
 @converter_registry.register("pd_op.full_int_array", trt_version="8.x")
 def full_int_array_converter(network, paddle_op, inputs):
diff --git a/python/paddle/tensorrt/impls/linalg.py b/python/paddle/tensorrt/impls/linalg.py
@@ -11,26 +11,14 @@
 # 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 logging
-import os
-import sys
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
 
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
-from paddle.tensorrt.register import converter_registry
-
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
 from paddle.tensorrt.converter_utils import (
     broadcast,
 )
+from paddle.tensorrt.register import converter_registry
 
 
 @converter_registry.register("pd_op.matmul", trt_version="8.x")
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -11,28 +11,16 @@
 # 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 logging
-import os
-import sys
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
 
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
-from paddle.tensorrt.register import converter_registry
-
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
 from paddle.tensorrt.converter_utils import (
     get_axes_for_reduce_op,
     get_positive_dim,
     has_dynamic_shape,
 )
+from paddle.tensorrt.register import converter_registry
 
 
 @converter_registry.register("pd_op.reshape", trt_version="8.x")
diff --git a/python/paddle/tensorrt/impls/math.py b/python/paddle/tensorrt/impls/math.py
@@ -11,30 +11,16 @@
 # 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 logging
-import os
-import sys
 
 import numpy as np
-
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
-
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
-from paddle.tensorrt.register import converter_registry
-
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
 from paddle.tensorrt.converter_utils import (
     add_elementwise_layer,
     broadcast,
     get_axes_for_reduce_op,
 )
+from paddle.tensorrt.register import converter_registry
 
 
 @converter_registry.register("pd_op.add", trt_version="8.x")
diff --git a/python/paddle/tensorrt/impls/norm.py b/python/paddle/tensorrt/impls/norm.py
@@ -11,31 +11,17 @@
 # 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 logging
-import os
-import sys
 
 import numpy as np
-
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
-
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
-from paddle.tensorrt.register import converter_registry
-
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
 from paddle.tensorrt.converter_utils import (
     append_ones,
     get_axes_for_reduce_op,
     get_dynamic_dims,
     has_dynamic_shape,
 )
+from paddle.tensorrt.register import converter_registry
 
 
 @converter_registry.register(
diff --git a/python/paddle/tensorrt/impls/ops.py b/python/paddle/tensorrt/impls/ops.py
@@ -16,6 +16,15 @@
 from paddle.tensorrt.register import converter_registry
 
 
+@converter_registry.register("pd_op.sqrt", trt_version="8.x")
+@converter_registry.register("pd_op.sqrt_", trt_version="8.x")
+def sqrt_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+
+    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(
diff --git a/python/paddle/tensorrt/impls/pooling.py b/python/paddle/tensorrt/impls/pooling.py
@@ -11,24 +11,12 @@
 # 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 logging
-import os
-import sys
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
 
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
 from paddle.tensorrt.register import converter_registry
 
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
-
 
 @converter_registry.register("pd_op.pool2d", trt_version="8.x")
 def pool2d_converter(network, paddle_op, inputs):
diff --git a/python/paddle/tensorrt/impls/search.py b/python/paddle/tensorrt/impls/search.py
@@ -11,24 +11,12 @@
 # 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 logging
-import os
-import sys
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
-if parent_dir not in sys.path:
-    sys.path.append(parent_dir)
 
 import tensorrt as trt
 
-from paddle.base.log_helper import get_logger
 from paddle.tensorrt.register import converter_registry
 
-_logger = get_logger(
-    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
-)
-
 
 @converter_registry.register("pd_op.nonzero", trt_version="8.x")
 def non_zero_converter(network, paddle_op, inputs):
diff --git a/python/paddle/tensorrt/impls/stat.py b/python/paddle/tensorrt/impls/stat.py
@@ -0,0 +1,33 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 get_axes_for_reduce_op
+from paddle.tensorrt.register import converter_registry
+
+
+@converter_registry.register("pd_op.mean", trt_version="8.x")
+def mean_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    keep_dim = paddle_op.attrs().get("keepdim")
+    dim = paddle_op.attrs().get("axis")
+
+    mean_layer = network.add_reduce(
+        input_tensor,
+        trt.ReduceOperation.AVG,
+        axes=get_axes_for_reduce_op(dim, network.has_implicit_batch_dimension),
+        keep_dims=keep_dim,
+    )
+    return mean_layer.get_output(0)
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -13,9 +13,10 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_model_resnet50 PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_conv2d PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_add PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_ops PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_stat PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_batch_norm PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_math PROPERTIES TIMEOUT "100")
-  set_tests_properties(test_converter_ops PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_activation PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_others PROPERTIES TIMEOUT "100")
 endif()
diff --git a/test/tensorrt/test_converter_ops.py b/test/tensorrt/test_converter_ops.py
diff --git a/test/tensorrt/test_converter_stat.py b/test/tensorrt/test_converter_stat.py
