【TRT_Converter】add pd_op.gather converter (#70875)

chang-wenbin · web-flow · commit 4aaa08da4ae8 · 2025-02-06T21:01:43.000+08:00
* add pd_op.gather converter

* gather_op

* gather_op
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -650,6 +650,7 @@ class TransposeOpPattern
     return true;
   }
 };
+
 class GatherOpPattern
     : public pir::OpRewritePattern<paddle::dialect::GatherOp> {
  public:
@@ -660,9 +661,10 @@ class GatherOpPattern
         op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
       return false;
     }
-    pir::Value axis = op.operand_source(2);
-    if (!axis) {
-      VLOG(3) << "axis is empty. Skipping rewrite.";
+    if (!op.axis().defining_op()->isa<paddle::dialect::FullOp>()) {
+      VLOG(3) << "When axis is not a constant "
+                 "Skip to convert into TRT.";
+
       return false;
     }
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -76,6 +76,21 @@ def reshape_converter(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
+@converter_registry.register("pd_op.gather", trt_version="8.x")
+def gather_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    index_tensor = inputs[1]
+    axis_value = get_input_constant_value(paddle_op, inputs, 2)[0]
+    axis_value = int(axis_value)
+
+    reshape_layer = network.add_shuffle(index_tensor)
+    reshape_layer.reshape_dims = (-1,)
+    gather_layer = network.add_gather(
+        input_tensor, reshape_layer.get_output(0), axis_value
+    )
+    return gather_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.gather_nd", trt_version="8.x")
 def gather_nd_converter(network, paddle_op, inputs):
     input_tensor, indices_tensor = inputs
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -232,6 +232,57 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestGatherCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.gather
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("float32"),
+            "index": np.array([0, 2]).astype("int64"),
+            "axis": 1,
+        }
+        self.program_config = {"feed_list": ["x", "index"]}
+        self.min_shape = {"x": [1, 4, 10], "index": [1]}
+        self.opt_shape = {"x": [1, 4, 10], "index": [1]}
+        self.max_shape = {"x": [5, 4, 10], "index": [5]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestGatherCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.gather
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("int64"),
+            "index": np.array([0, 2]).astype("int64"),
+            "axis": 1,
+        }
+        self.program_config = {"feed_list": ["x", "index"]}
+        self.min_shape = {"x": [1, 4, 10], "index": [1]}
+        self.opt_shape = {"x": [1, 4, 10], "index": [1]}
+        self.max_shape = {"x": [5, 4, 10], "index": [5]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestGatherCase3TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.gather
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("int64"),
+            "index": np.array([0, 2]).astype("int64"),
+            "axis": np.array([2]).astype("int64"),
+        }
+        self.program_config = {"feed_list": ["x", "index", "axis"]}
+        self.min_shape = {"x": [1, 4, 10], "index": [1]}
+        self.opt_shape = {"x": [1, 4, 10], "index": [1]}
+        self.max_shape = {"x": [5, 4, 10], "index": [5]}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
 class TestSplitWithNumTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.split
