【SCU】【Paddle TensorRT No.11、12、65】Add pd_op.bitwise_and、pd_op.bitwise_or、pd_op.bitwise_not converter (#69599)

PolaKuma · web-flow · commit db2acbd43a17 · 2024-12-31T14:07:15.000+08:00
* add_bitwise_logic

* fix codestyle

* fix codestyle

* fix codestyle

* fix again

* fix

* update_not_markerpass

* fix codestyle

* Update trt_op_marker_pass.cc

* fix

* update

* fix codestyle

* fix codestyle

* fix codestyle

* Update trt_op_marker_pass.cc

* Update trt_op_marker_pass.cc

* fix codestyle

* add_fp16

* fix codestyle

* fix codestyle

* Update test_converter_logic.py

* update

* fix codestyle

* Update test_converter_logic.py

* Update logic.py

* Update logic.py
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1744,6 +1744,68 @@ class NotEqualOpPattern
   }
 };
 
+template <typename OpType>
+class BitwiseCommonOpPattern : 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 input_operand = op.operand_source(0);
+    auto input_type = pir::GetDataTypeFromValue(input_operand);
+    if (!input_type.isa<pir::BoolType>()) {
+      if constexpr (std::is_same_v<OpType, paddle::dialect::BitwiseAndOp>) {
+        VLOG(3) << "the bitwise_and only supports input of BOOL.";
+      } else {
+        VLOG(3) << "the bitwise_or only supports input of BOOL.";
+      }
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+using BitwiseAndOpPattern =
+    BitwiseCommonOpPattern<paddle::dialect::BitwiseAndOp>;
+using BitwiseOrOpPattern = BitwiseCommonOpPattern<paddle::dialect::BitwiseOrOp>;
+
+class BitwiseNotOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::BitwiseNotOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::BitwiseNotOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(paddle::dialect::BitwiseNotOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+
+    pir::Value input_operand = op.operand_source(0);
+    auto input_type = pir::GetDataTypeFromValue(input_operand);
+    if (input_type.isa<pir::Int8Type>() || input_type.isa<pir::UInt8Type>()) {
+      VLOG(3) << "INT8 / UINT8 type convert to TRT is not supported.";
+      return false;
+    }
+#if IS_TRT_VERSION_LT(8600)
+    pir::Value x = op.operand_source(0);
+    auto x_type = x.type().dyn_cast<paddle::dialect::DenseTensorType>();
+    auto x_shape = x_type.dims();
+    int dims = x_shape.size();
+    if (dims < 1) {
+      VLOG(3) << "BOOL type does not support 0-dim input when TensorRT < 8.6.";
+      return false;
+    }
+#endif
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class FullLikeOpPattern
     : public pir::OpRewritePattern<paddle::dialect::FullLikeOp> {
  public:
@@ -2168,6 +2230,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<BitwiseAndOpPattern>(context));
+    ps.Add(std::make_unique<BitwiseOrOpPattern>(context));
+    ps.Add(std::make_unique<BitwiseNotOpPattern>(context));
     ps.Add(std::make_unique<LogicalNot_OpPattern>(context));
     ps.Add(std::make_unique<LogicalOrOpPattern>(context));
     ps.Add(std::make_unique<LogicalOr_OpPattern>(context));
diff --git a/python/paddle/tensorrt/impls/logic.py b/python/paddle/tensorrt/impls/logic.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import tensorrt as trt
 
 from paddle.tensorrt.converter_utils import (
@@ -24,6 +25,8 @@
     "pd_op.greater_than": trt.ElementWiseOperation.GREATER,
     "pd_op.less_than": trt.ElementWiseOperation.LESS,
     "pd_op.equal": trt.ElementWiseOperation.EQUAL,
+    "pd_op.bitwise_and": trt.ElementWiseOperation.AND,
+    "pd_op.bitwise_or": trt.ElementWiseOperation.OR,
     "pd_op.logical_xor": trt.ElementWiseOperation.XOR,
     "pd_op.logical_or": trt.ElementWiseOperation.OR,
     "pd_op.logical_or_": trt.ElementWiseOperation.OR,
@@ -34,6 +37,8 @@
 @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.bitwise_and", trt_version="8.x")
+@converter_registry.register("pd_op.bitwise_or", trt_version="8.x")
 @converter_registry.register("pd_op.logical_xor", 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")
@@ -55,6 +60,30 @@ def not_equal_converter(network, paddle_op, inputs):
     return layer_output
 
 
+@converter_registry.register("pd_op.bitwise_not", trt_version="8.x")
+def bitwise_not_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    if input_tensor.dtype == trt.bool:
+        bitwise_not_layer = network.add_unary(
+            input_tensor, trt.UnaryOperation.NOT
+        )
+        layer_output = bitwise_not_layer.get_output(0)
+    else:
+        neg_one_tensor_dims = trt.Dims([1] * len(input_tensor.shape))
+        neg_one_value = np.array([-1], dtype=np.int32)
+        neg_one_weights = trt.Weights(neg_one_value)
+        neg_one_tensor = network.add_constant(
+            neg_one_tensor_dims, neg_one_weights
+        ).get_output(0)
+        mul_neg_one = network.add_elementwise(
+            input_tensor, neg_one_tensor, trt.ElementWiseOperation.PROD
+        ).get_output(0)
+        layer_output = network.add_elementwise(
+            mul_neg_one, neg_one_tensor, trt.ElementWiseOperation.SUM
+        ).get_output(0)
+    return layer_output
+
+
 @converter_registry.register("pd_op.logical_not", trt_version="8.x")
 @converter_registry.register("pd_op.logical_not_", trt_version="8.x")
 def logic_not_converter(network, paddle_op, inputs):
diff --git a/test/tensorrt/test_converter_logic.py b/test/tensorrt/test_converter_logic.py
@@ -140,6 +140,130 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestAndRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_and
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("bool"),
+            "y": np.random.randn(3).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [3]}
+        self.max_shape = {"x": [5, 3], "y": [3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestAndRTPatternDifferentShapes(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_and
+        self.api_args = {
+            "x": np.random.randn(4, 5).astype("bool"),
+            "y": np.random.randn(1, 5).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 5], "y": [1, 5]}
+        self.max_shape = {"x": [10, 5], "y": [1, 5]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestAndRTPatternDifferentShapes1(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_and
+        self.api_args = {
+            "x": np.random.randint(0, 2, (2, 3)).astype("bool"),
+            "y": np.random.randint(0, 2, (2, 3)).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [1, 3]}
+        self.max_shape = {"x": [5, 3], "y": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestOrRTPatternBroadcast(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_or
+        self.api_args = {
+            "x": np.random.randn(2, 1).astype("bool"),
+            "y": np.random.randn(2, 3).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [2, 1], "y": [2, 3]}
+        self.max_shape = {"x": [2, 1], "y": [2, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestOrRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_or
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("bool"),
+            "y": np.random.randn(3).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [3]}
+        self.max_shape = {"x": [5, 3], "y": [3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestNotRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_not
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestNotRTPatternEdgeCase(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_not
+        self.api_args = {
+            "x": np.zeros((2, 3)).astype("bool"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
 class TestLogicalOrTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.logical_or
@@ -165,6 +289,67 @@ def test_trt_diff_shape_result(self):
         self.check_trt_result()
 
 
+class TestAndRTPatternErrorType(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_and
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("int32"),
+            "y": np.random.randn(3).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [3]}
+        self.max_shape = {"x": [5, 3], "y": [3]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
+class TestOrRTPatternErrorType(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_or
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("int32"),
+            "y": np.random.randn(3).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3], "y": [3]}
+        self.max_shape = {"x": [5, 3], "y": [3]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
+class TestNotRTINT8(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_not
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("int8"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
+class TestNotRTINT64(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.bitwise_not
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
 class TestLogicalOrMarker(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.logical_or
