【TRT_Converter】add where TRT_Converter (#68876)

chang-wenbin · web-flow · commit 1b9e2462a081 · 2024-11-08T18:19:38.000+08:00
* TRT_Converter

* where

* add test_converter_search

* add int32 test &amp; LT

* merge

* update where

* update where converter

* updata where del non_zero

* updata where del non_zero

* update where int64
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1475,6 +1475,31 @@ class TanhOpPattern : public pir::OpRewritePattern<paddle::dialect::TanhOp> {
   }
 };
 
+class WherePattern : public pir::OpRewritePattern<paddle::dialect::WhereOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::WhereOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::WhereOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    pir::Value x = op.operand_source(1);
+    pir::Value y = op.operand_source(2);
+    if (x == nullptr || y == nullptr) {
+      VLOG(3) << "pd_op.where x or y tensor value is null";
+      return false;
+    }
+#if IS_TRT_VERSION_LT(8400)
+    VLOG(3) << "where is not supported when TensorRT < 8.4";
+    return false;
+#endif
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class FullWithTensorPattern
     : public pir::OpRewritePattern<paddle::dialect::FullWithTensorOp> {
  public:
@@ -1666,6 +1691,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<NearestInterV2Pattern>(context));
     ps.Add(std::make_unique<StackOpPattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
+    ps.Add(std::make_unique<WherePattern>(context));
     ps.Add(std::make_unique<FullWithTensorPattern>(context));
     ps.Add(std::make_unique<StridedSliceOpPattern>(context));
     ps.Add(std::make_unique<TopkOpPattern>(context));
diff --git a/python/paddle/tensorrt/impls/search.py b/python/paddle/tensorrt/impls/search.py
@@ -66,6 +66,17 @@ def argmax_converter(network, paddle_op, inputs):
         return squeeze_layer.get_output(0)
 
 
+@converter_registry.register("pd_op.where", trt_version="8.x")
+def where_converter(network, paddle_op, inputs):
+    condition = inputs[0]
+    x = inputs[1]
+    y = inputs[2]
+
+    select_layer = network.add_select(condition, x, y)
+
+    return select_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.topk", trt_version="8.x")
 def topk_converter(network, paddle_op, inputs):
     input_tensor = inputs[0]
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -23,5 +23,7 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_common PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_search PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_logic PROPERTIES TIMEOUT "300")
+
 endif()
diff --git a/test/tensorrt/test_converter_search.py b/test/tensorrt/test_converter_search.py
@@ -35,6 +35,38 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestWhereTRTPatternCase1(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.where
+        self.api_args = {
+            "condition": np.random.choice([True, False], size=(2, 3)),
+            "x": np.random.randn(2, 3).astype("float32"),
+            "y": np.random.randn(2, 3).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["condition", "x", "y"]}
+        self.min_shape = {"condition": [1, 3], "x": [1, 3], "y": [1, 3]}
+        self.max_shape = {"condition": [5, 3], "x": [5, 3], "y": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestWhereTRTPatternCase2(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.where
+        self.api_args = {
+            "condition": np.random.choice([True, False], size=(2, 3)),
+            "x": np.random.randn(2, 3).astype("int64"),
+            "y": np.random.randn(2, 3).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["condition", "x", "y"]}
+        self.min_shape = {"condition": [1, 3], "x": [1, 3], "y": [1, 3]}
+        self.max_shape = {"condition": [5, 3], "x": [5, 3], "y": [5, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestTopkCase1TRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.topk
