[SCU][Paddle TensorRT No.26、29] Add (pd_op.logical_or、pd_op.logical_or、pd_op.logical_and) converter (#69706)

mori0umi · web-flow · commit faa0339fcdd9 · 2024-12-13T16:02:47.000+08:00
* add logical_or &amp; logical_or_

* fix bug

* fix marker

* fix

* add logical_and
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1204,6 +1204,35 @@ class LessThanOpPattern
     return true;
   }
 };
+
+template <typename OpType>
+class LogicalCommonOpPattern : 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);
+    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) << op->name() << " op only supports bool datatype";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+using LogicalOrOpPattern = LogicalCommonOpPattern<paddle::dialect::LogicalOrOp>;
+using LogicalOr_OpPattern =
+    LogicalCommonOpPattern<paddle::dialect::LogicalOr_Op>;
+using LogicalAndOpPattern =
+    LogicalCommonOpPattern<paddle::dialect::LogicalAndOp>;
+
 class MulticlassNms3OpPattern
     : public pir::OpRewritePattern<paddle::dialect::MulticlassNms3Op> {
  public:
@@ -2185,6 +2214,9 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     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<LogicalOrOpPattern>(context));
+    ps.Add(std::make_unique<LogicalOr_OpPattern>(context));
+    ps.Add(std::make_unique<LogicalAndOpPattern>(context));
     ps.Add(std::make_unique<GroupNormOpPattern>(context));
     ps.Add(std::make_unique<TransposeOpPattern>(context));
     ps.Add(std::make_unique<GatherOpPattern>(context));
diff --git a/python/paddle/tensorrt/impls/logic.py b/python/paddle/tensorrt/impls/logic.py
@@ -23,12 +23,18 @@
     "pd_op.greater_than": trt.ElementWiseOperation.GREATER,
     "pd_op.less_than": trt.ElementWiseOperation.LESS,
     "pd_op.equal": trt.ElementWiseOperation.EQUAL,
+    "pd_op.logical_or": trt.ElementWiseOperation.OR,
+    "pd_op.logical_or_": trt.ElementWiseOperation.OR,
+    "pd_op.logical_and": trt.ElementWiseOperation.AND,
 }
 
 
 @converter_registry.register("pd_op.greater_than", trt_version="8.x")
 @converter_registry.register("pd_op.less_than", trt_version="8.x")
 @converter_registry.register("pd_op.equal", trt_version="8.x")
+@converter_registry.register("pd_op.logical_or", trt_version="8.x")
+@converter_registry.register("pd_op.logical_or_", trt_version="8.x")
+@converter_registry.register("pd_op.logical_and", trt_version="8.x")
 def logic_converter(network, paddle_op, inputs):
     layer_output = add_elementwise_layer(
         network, paddle_op, inputs, logic_type_map[paddle_op.name()]
diff --git a/test/tensorrt/test_converter_logic.py b/test/tensorrt/test_converter_logic.py
@@ -140,5 +140,122 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestLogicalOrTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.logical_or
+
+    def test_trt_result(self):
+        self.api_args = {
+            "x": np.random.choice([True, False], size=(3,)).astype("bool"),
+            "y": np.random.choice([True, False], size=(3,)).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1], "y": [1]}
+        self.max_shape = {"x": [5], "y": [5]}
+        self.check_trt_result()
+
+    def test_trt_diff_shape_result(self):
+        self.api_args = {
+            "x": np.random.choice([True, False], size=(2, 3)).astype("bool"),
+            "y": np.random.choice([True, False], size=(3)).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [3]}
+        self.max_shape = {"x": [4, 3], "y": [3]}
+        self.check_trt_result()
+
+
+class TestLogicalOrMarker(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.logical_or
+        self.api_args = {
+            "x": np.random.randn(3).astype("int64"),
+            "y": np.random.randn(3).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.target_marker_op = "pd_op.logical_or"
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
+class TestLogicalAndTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.logical_and
+
+    def test_trt_result(self):
+        self.api_args = {
+            "x": np.random.choice([True, False], size=(3,)).astype("bool"),
+            "y": np.random.choice([True, False], size=(3,)).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1], "y": [1]}
+        self.max_shape = {"x": [5], "y": [5]}
+        self.check_trt_result()
+
+    def test_trt_diff_shape_result(self):
+        self.api_args = {
+            "x": np.random.choice([True, False], size=(2, 3)).astype("bool"),
+            "y": np.random.choice([True, False], size=(3)).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [3]}
+        self.max_shape = {"x": [4, 3], "y": [3]}
+        self.check_trt_result()
+
+
+class TestLogicalAndMarker(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.logical_and
+        self.api_args = {
+            "x": np.random.randn(3).astype("int64"),
+            "y": np.random.randn(3).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.target_marker_op = "pd_op.logical_and"
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
+class TestLogicalOr_TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.logical_or_
+
+    def test_trt_result(self):
+        self.api_args = {
+            "x": np.random.choice([True, False], size=(3,)).astype("bool"),
+            "y": np.random.choice([True, False], size=(3,)).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1], "y": [1]}
+        self.max_shape = {"x": [5], "y": [5]}
+        self.check_trt_result()
+
+    def test_trt_diff_shape_result(self):
+        self.api_args = {
+            "x": np.random.choice([True, False], size=(2, 3)).astype("bool"),
+            "y": np.random.choice([True, False], size=(3)).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [3]}
+        self.max_shape = {"x": [4, 3], "y": [3]}
+        self.check_trt_result()
+
+
+class TestLogicalOr_Marker(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.logical_or_
+        self.api_args = {
+            "x": np.random.randn(3).astype("int64"),
+            "y": np.random.randn(3).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.target_marker_op = "pd_op.logical_or_"
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
 if __name__ == '__main__':
     unittest.main()
