add pd_op.min converter for PIR TRT and base ReduceCommonOpPattern (#68783)

zhink · web-flow · commit 17d68c9934f1 · 2024-10-21T10:55:45.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
@@ -1220,68 +1220,6 @@ class ArgmaxOpPattern
     return true;
   }
 };
-class MaxOpPattern : public pir::OpRewritePattern<paddle::dialect::MaxOp> {
- public:
-  using pir::OpRewritePattern<paddle::dialect::MaxOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::MaxOp op,
-                       pir::PatternRewriter &rewriter) const override {
-    if (op->HasAttribute(kCanRunTrtAttr) &&
-        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
-      return false;
-    }
-    if (!op->HasAttribute("keepdim")) {
-      VLOG(3) << "the max does not have attr keep_dim ";
-      return false;
-    }
-    pir::Value x = op.operand_source(0);
-    auto x_dtype = pir::GetDataTypeFromValue(x);
-    if (!(x_dtype.isa<pir::Float32Type>() || x_dtype.isa<pir::Float64Type>() ||
-          x_dtype.isa<pir::Int32Type>() || x_dtype.isa<pir::Int64Type>())) {
-      VLOG(3) << "max input data type must be int32 or int64 or "
-                 "float32 or "
-                 "float64";
-      return false;
-    }
-    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
-    return true;
-  }
-};
-
-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 BilinearInterpV2Pattern
     : public pir::OpRewritePattern<paddle::dialect::BilinearInterpOp> {
@@ -1459,6 +1397,59 @@ class NearestInterV2Pattern
   }
 };
 
+// Add ReduceCommonOpPattern base class to simplify code
+template <typename OpType>
+class ReduceCommonOpPattern : 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;
+    }
+    if (!op->HasAttribute("keepdim")) {
+      VLOG(3) << "the max does not have attr keep_dim ";
+      return false;
+    }
+
+    if constexpr (std::is_same_v<OpType, paddle::dialect::MeanOp>) {
+      if (!op->HasAttribute("axis")) {
+        VLOG(3) << "The axis attribute does not exist";
+        return false;
+      }
+    }
+
+    pir::Value x = op.operand_source(0);
+    auto x_dtype = pir::GetDataTypeFromValue(x);
+    if (!(x_dtype.isa<pir::Float32Type>() || x_dtype.isa<pir::Float64Type>() ||
+          x_dtype.isa<pir::Int32Type>() || x_dtype.isa<pir::Int64Type>())) {
+      if constexpr (std::is_same_v<OpType, paddle::dialect::MinOp>) {
+        VLOG(3) << "min input data type must be int32 or int64 or "
+                   "float32 or "
+                   "float64";
+      } else if constexpr (std::is_same_v<OpType, paddle::dialect::MaxOp>) {
+        VLOG(3) << "max input data type must be int32 or int64 or "
+                   "float32 or "
+                   "float64";
+      } else if constexpr (std::is_same_v<OpType, paddle::dialect::MeanOp>) {
+        VLOG(3) << "mean input data type must be int32 or int64 or "
+                   "float32 or "
+                   "float64";
+      }
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
+// use type aliases to simplify usage
+using MinOpPattern = ReduceCommonOpPattern<paddle::dialect::MinOp>;
+using MaxOpPattern = ReduceCommonOpPattern<paddle::dialect::MaxOp>;
+using MeanOpPattern = ReduceCommonOpPattern<paddle::dialect::MeanOp>;
+
 class TanhOpPattern : public pir::OpRewritePattern<paddle::dialect::TanhOp> {
  public:
   using pir::OpRewritePattern<paddle::dialect::TanhOp>::OpRewritePattern;
@@ -1557,6 +1548,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<MulticlassNms3OpPattern>(context));
     ps.Add(std::make_unique<ArgmaxOpPattern>(context));
     ps.Add(std::make_unique<MaxOpPattern>(context));
+    ps.Add(std::make_unique<MinOpPattern>(context));
     ps.Add(std::make_unique<BilinearInterpV2Pattern>(context));
     ps.Add(std::make_unique<NearestInterV2Pattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
diff --git a/python/paddle/tensorrt/impls/math.py b/python/paddle/tensorrt/impls/math.py
@@ -118,3 +118,28 @@ def multiply_converter(network, paddle_op, inputs):
     return add_elementwise_layer(
         network, paddle_op, inputs, trt.ElementWiseOperation.PROD
     )
+
+
+@converter_registry.register("pd_op.min", trt_version="8.x")
+def min_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    axis = paddle_op.operands()[1].source().get_defining_op().attrs()["value"]
+    input_shape = paddle_op.operands()[0].source().shape
+    keepdim = paddle_op.attrs()["keepdim"]
+    if network.has_implicit_batch_dimension:
+        assert (
+            axis != 0
+        ), "can't reduce on axis == 0 when network has implicit batch dimension"
+    output_shape = []
+    if len(axis) == 0:
+        axis = list(range(len(input_shape)))
+    for i in range(len(axis)):
+        if axis[i] < 0:
+            axis[i] = len(input_shape) + axis[i]
+    layer = network.add_reduce(
+        input_tensor,
+        trt.ReduceOperation.MIN,
+        axes=get_axes_for_reduce_op(axis),
+        keep_dims=keepdim,
+    )
+    return layer.get_output(0)
diff --git a/test/tensorrt/test_converter_math.py b/test/tensorrt/test_converter_math.py
@@ -24,14 +24,12 @@ class TestMaxTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.max
         self.api_args = {
-            "x": np.array([[0.2, 0.3, 0.5, 0.9], [0.1, 0.2, 0.6, 0.7]]).astype(
-                "float32"
-            ),
+            "x": np.random.randn(2, 4).astype(np.float32),
             "axis": [0, 1],
         }
         self.program_config = {"feed_list": ["x"]}
-        self.min_shape = {"x": [1, 4], "y": [1, 4]}
-        self.max_shape = {"x": [5, 4], "y": [5, 4]}
+        self.min_shape = {"x": [1, 4]}
+        self.max_shape = {"x": [5, 4]}
 
     def test_trt_result(self):
         self.check_trt_result()
@@ -97,5 +95,20 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestMinTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.min
+        self.api_args = {
+            "x": np.random.randn(2, 4).astype(np.float32),
+            "axis": [0, 1],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 4]}
+        self.max_shape = {"x": [5, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()
