PaddlePaddle
diff --git a/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+75-167Lines changed: 75 additions & 167 deletions b/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+75-167Lines changed: 75 additions & 167 deletions
diff --git a/‎python/paddle/tensorrt/converter_utils.py
Copy file name to clipboard
+55Lines changed: 55 additions & 0 deletions b/‎python/paddle/tensorrt/converter_utils.py
Copy file name to clipboard
+55Lines changed: 55 additions & 0 deletions
@@ -910,7 +910,6 @@ class SplitWithNumOpPattern
     return true;
   }
 };
-
 class GreaterEqualOpPattern
     : public pir::OpRewritePattern<paddle::dialect::GreaterEqualOp> {
  public:
@@ -996,46 +995,43 @@ class LessThanOpPattern
   }
 };
 
-class MultiplyOpPattern
-    : public pir::OpRewritePattern<paddle::dialect::MultiplyOp> {
+// Add ElementwiseCommonOpPattern base class to simplify code
+template <typename OpType>
+class ElementwiseCommonOpPattern : public pir::OpRewritePattern<OpType> {
  public:
-  using pir::OpRewritePattern<paddle::dialect::MultiplyOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::MultiplyOp 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()) {
-      return false;
-    }
-    pir::Value x = op.operand_source(0);
-    pir::Value y = op.operand_source(1);
-    auto x_dtype = pir::GetDataTypeFromValue(x);
-    auto y_dtype = pir::GetDataTypeFromValue(y);
-    if (x_dtype.isa<pir::BoolType>() || y_dtype.isa<pir::BoolType>()) {
-      VLOG(3) << "elementwise_mul do not support boolean datatype.";
+        op->template attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
       return false;
     }
 
-    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
-    return true;
-  }
-};
-
-class SubtractOpPattern
-    : public pir::OpRewritePattern<paddle::dialect::SubtractOp> {
- public:
-  using pir::OpRewritePattern<paddle::dialect::SubtractOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::SubtractOp op,
-                       pir::PatternRewriter &rewriter) const override {
-    if (op->HasAttribute(kCanRunTrtAttr) &&
-        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
-      return false;
-    }
     pir::Value x = op.operand_source(0);
     pir::Value y = op.operand_source(1);
     auto x_dtype = pir::GetDataTypeFromValue(x);
     auto y_dtype = pir::GetDataTypeFromValue(y);
     if (x_dtype.isa<pir::BoolType>() || y_dtype.isa<pir::BoolType>()) {
-      VLOG(3) << "elementwise_sub do not support boolean datatype.";
+      if constexpr (std::is_same_v<OpType, paddle::dialect::MultiplyOp>) {
+        VLOG(3) << "elementwise_mul do not support boolean datatype.";
+      } else if constexpr (std::is_same_v<OpType,  // NOLINT
+                                          paddle::dialect::SubtractOp>) {
+        VLOG(3) << "elementwise_sub do not support boolean datatype.";
+      } else if constexpr (std::is_same_v<OpType, paddle::dialect::DivideOp>) {
+        VLOG(3) << "elementwise_div do not support boolean datatype.";
+      } else if constexpr (std::is_same_v<OpType, paddle::dialect::MinimumOp>) {
+        VLOG(3) << "elementwise_min do not support boolean datatype.";
+      } else if constexpr (std::is_same_v<OpType, paddle::dialect::MaximumOp>) {
+        VLOG(3) << "elementwise_max do not support boolean datatype.";
+      } else if constexpr (std::is_same_v<OpType,  // NOLINT
+                                          paddle::dialect::FloorDivideOp>) {
+        VLOG(3) << "elementwise_floordiv do not support boolean datatype.";
+      } else if constexpr (std::is_same_v<OpType,  // NOLINT
+                                          paddle::dialect::RemainderOp>) {
+        VLOG(3) << "elementwise_mod do not support boolean datatype.";
+      } else {
+        VLOG(3) << "elementwise other do not support boolean datatype.";
+      }
       return false;
     }
 
@@ -1044,29 +1040,17 @@ class SubtractOpPattern
   }
 };
 
-class DivideOpPattern
-    : public pir::OpRewritePattern<paddle::dialect::DivideOp> {
- public:
-  using pir::OpRewritePattern<paddle::dialect::DivideOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::DivideOp op,
-                       pir::PatternRewriter &rewriter) const override {
-    if (op->HasAttribute(kCanRunTrtAttr) &&
-        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
-      return false;
-    }
-    pir::Value x = op.operand_source(0);
-    pir::Value y = op.operand_source(1);
-    auto x_dtype = pir::GetDataTypeFromValue(x);
-    auto y_dtype = pir::GetDataTypeFromValue(y);
-    if (x_dtype.isa<pir::BoolType>() || y_dtype.isa<pir::BoolType>()) {
-      VLOG(3) << "elementwise_div do not support boolean datatype.";
-      return false;
-    }
-
-    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
-    return true;
-  }
-};
+using MultiplyOpPattern =
+    ElementwiseCommonOpPattern<paddle::dialect::MultiplyOp>;
+using SubtractOpPattern =
+    ElementwiseCommonOpPattern<paddle::dialect::SubtractOp>;
+using DivideOpPattern = ElementwiseCommonOpPattern<paddle::dialect::DivideOp>;
+using MinimumOpPattern = ElementwiseCommonOpPattern<paddle::dialect::MinimumOp>;
+using MaximumOpPattern = ElementwiseCommonOpPattern<paddle::dialect::MaximumOp>;
+using FloorDivideOpPattern =
+    ElementwiseCommonOpPattern<paddle::dialect::FloorDivideOp>;
+using RemainderOpPattern =
+    ElementwiseCommonOpPattern<paddle::dialect::RemainderOp>;
 
 class ElementwisePowOpPattern
     : public pir::OpRewritePattern<paddle::dialect::ElementwisePowOp> {
@@ -1093,99 +1077,6 @@ class ElementwisePowOpPattern
     return true;
   }
 };
-class MinimumOpPattern
-    : public pir::OpRewritePattern<paddle::dialect::MinimumOp> {
- public:
-  using pir::OpRewritePattern<paddle::dialect::MinimumOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::MinimumOp op,
-                       pir::PatternRewriter &rewriter) const override {
-    if (op->HasAttribute(kCanRunTrtAttr) &&
-        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
-      return false;
-    }
-    pir::Value x = op.operand_source(0);
-    pir::Value y = op.operand_source(1);
-    auto x_dtype = pir::GetDataTypeFromValue(x);
-    auto y_dtype = pir::GetDataTypeFromValue(y);
-    if (x_dtype.isa<pir::BoolType>() || y_dtype.isa<pir::BoolType>()) {
-      VLOG(3) << "elementwise_min do not support boolean datatype.";
-      return false;
-    }
-
-    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
-    return true;
-  }
-};
-class MaximumOpPattern
-    : public pir::OpRewritePattern<paddle::dialect::MaximumOp> {
- public:
-  using pir::OpRewritePattern<paddle::dialect::MaximumOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::MaximumOp op,
-                       pir::PatternRewriter &rewriter) const override {
-    if (op->HasAttribute(kCanRunTrtAttr) &&
-        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
-      return false;
-    }
-    pir::Value x = op.operand_source(0);
-    pir::Value y = op.operand_source(1);
-    auto x_dtype = pir::GetDataTypeFromValue(x);
-    auto y_dtype = pir::GetDataTypeFromValue(y);
-    if (x_dtype.isa<pir::BoolType>() || y_dtype.isa<pir::BoolType>()) {
-      VLOG(3) << "elementwise_max do not support boolean datatype.";
-      return false;
-    }
-
-    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
-    return true;
-  }
-};
-
-class FloorDivideOpPattern
-    : public pir::OpRewritePattern<paddle::dialect::FloorDivideOp> {
- public:
-  using pir::OpRewritePattern<paddle::dialect::FloorDivideOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::FloorDivideOp op,
-                       pir::PatternRewriter &rewriter) const override {
-    if (op->HasAttribute(kCanRunTrtAttr) &&
-        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
-      return false;
-    }
-    pir::Value x = op.operand_source(0);
-    pir::Value y = op.operand_source(1);
-    auto x_dtype = pir::GetDataTypeFromValue(x);
-    auto y_dtype = pir::GetDataTypeFromValue(y);
-    if (x_dtype.isa<pir::BoolType>() || y_dtype.isa<pir::BoolType>()) {
-      VLOG(3) << "elementwise_floordiv do not support boolean datatype.";
-      return false;
-    }
-    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
-    return true;
-  }
-};
-
-class RemainderOpPattern
-    : public pir::OpRewritePattern<paddle::dialect::RemainderOp> {
- public:
-  using pir::OpRewritePattern<paddle::dialect::RemainderOp>::OpRewritePattern;
-  bool MatchAndRewrite(paddle::dialect::RemainderOp op,
-                       pir::PatternRewriter &rewriter) const override {
-    if (op->HasAttribute(kCanRunTrtAttr) &&
-        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
-      return false;
-    }
-    pir::Value x = op.operand_source(0);
-    pir::Value y = op.operand_source(1);
-    auto x_dtype = pir::GetDataTypeFromValue(x);
-    auto y_dtype = pir::GetDataTypeFromValue(y);
-    if (x_dtype.isa<pir::BoolType>() || y_dtype.isa<pir::BoolType>()) {
-      VLOG(3) << "elementwise_mod do not support boolean datatype.";
-      return false;
-    }
-
-    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
-    return true;
-  }
-};
 
 class MulticlassNms3OpPattern
     : public pir::OpRewritePattern<paddle::dialect::MulticlassNms3Op> {
@@ -1506,38 +1397,49 @@ class ReduceCommonOpPattern : public pir::OpRewritePattern<OpType> {
         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;
-      }
+    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";
+    if constexpr (std::is_same_v<OpType, paddle::dialect::AnyOp>) {
+      if (!x_dtype.isa<pir::BoolType>()) {
+        VLOG(3) << "any op input data type must be bool";
+        return false;
+      }
+    } else if constexpr (std::is_same_v<OpType, paddle::dialect::AllOp>) {
+      if (!x_dtype.isa<pir::BoolType>()) {
+        VLOG(3) << "all op input data type must be bool";
+        return false;
+      }
+    } else {
+      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;
       }
-      return false;
     }
-
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
     return true;
   }
@@ -1547,6 +1449,9 @@ class ReduceCommonOpPattern : public pir::OpRewritePattern<OpType> {
 using MinOpPattern = ReduceCommonOpPattern<paddle::dialect::MinOp>;
 using MaxOpPattern = ReduceCommonOpPattern<paddle::dialect::MaxOp>;
 using MeanOpPattern = ReduceCommonOpPattern<paddle::dialect::MeanOp>;
+using AnyOpPattern = ReduceCommonOpPattern<paddle::dialect::AnyOp>;
+using AllOpPattern = ReduceCommonOpPattern<paddle::dialect::AllOp>;
+using SumOpPattern = ReduceCommonOpPattern<paddle::dialect::SumOp>;
 
 class TanhOpPattern : public pir::OpRewritePattern<paddle::dialect::TanhOp> {
  public:
@@ -1652,6 +1557,9 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     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<AllOpPattern>(context));
+    ps.Add(std::make_unique<AnyOpPattern>(context));
+    ps.Add(std::make_unique<SumOpPattern>(context));
     ps.Add(std::make_unique<BilinearInterpV2Pattern>(context));
     ps.Add(std::make_unique<NearestInterV2Pattern>(context));
     ps.Add(std::make_unique<StackOpPattern>(context));
 
@@ -447,3 +447,58 @@ def convert_conv2d(network, paddle_op, inputs):
     layer.dilation_nd = nv_dilations
 
     return layer.get_output(0)
+
+
+def add_reduce_layer(network, paddle_op, inputs, op_type):
+    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,
+        op_type,
+        axes=get_axes_for_reduce_op(axis),
+        keep_dims=keepdim,
+    )
+    layer.get_output(0).dtype = layer.get_input(0).dtype
+    return layer.get_output(0)
+
+
+def add_cast_reduce_layer(network, paddle_op, inputs, op_type):
+    input_tensor = inputs[0]
+    cast_layer = network.add_identity(input_tensor)
+    cast_layer.set_output_type(0, trt.int32)
+    cast_layer.get_output(0).dtype = trt.int32
+
+    axis = paddle_op.attrs().get("axis")
+    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(
+        cast_layer.get_output(0),
+        op_type,
+        axes=get_axes_for_reduce_op(axis),
+        keep_dims=keepdim,
+    )
+    layer.set_output_type(0, trt.bool)
+    layer.get_output(0).dtype = cast_layer.get_output(0).dtype
+    return layer.get_output(0)