[inference]update pir argmax and bilinear trt op convert (#70325)

bukejiyu · web-flow · commit f0687bd746e2 · 2024-12-19T22:25:33.000-08:00
* update pir argmax and bilinear op convert

* update
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1402,13 +1402,6 @@ class BilinearInterpV2Pattern
         return false;
       }
     }
-    pir::Value size_tensor = op.operand_source(2);
-    if (size_tensor != nullptr) {
-      VLOG(3) << "The Paddle-TRT doesn't support the SizeTensor for "
-                 "BilinearInterpV2";
-      return false;
-    }
-
     auto data_format =
         op->attribute<pir::StrAttribute>("data_format").AsString();
     if (data_format != "NCHW" && data_format != "NHWC") {
@@ -1421,7 +1414,26 @@ class BilinearInterpV2Pattern
       VLOG(3) << "The interp_method of BilinearInterpV2 is not bilinear";
       return false;
     }
-
+#if IS_TRT_VERSION_GE(8200)
+    auto size_tensor = op.operand_source(2);
+    if (size_tensor.impl()) {
+      auto size_tensor_type = size_tensor.type();
+      if (size_tensor_type.isa<pir::VectorType>()) {
+        auto vector_type = size_tensor.type().dyn_cast<pir::VectorType>();
+        if (vector_type.size() == 2) {
+          op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+          return true;
+        }
+      }
+    }
+#else
+    auto size_tensor = op.operand_source(2);
+    if (size_tensor.impl() != nullptr) {
+      VLOG(3) << "The Paddle-TRT doesn't support the SizeTensor for "
+                 "BilinearInterpV2";
+      return false;
+    }
+#endif
     pir::Value scale_tensor = op.operand_source(3);
 
     bool has_scale_input = false;
diff --git a/python/paddle/tensorrt/impls/common.py b/python/paddle/tensorrt/impls/common.py
@@ -16,6 +16,7 @@
 import numpy as np
 import tensorrt as trt
 
+from paddle import pir
 from paddle.tensorrt.converter_utils import get_shape_tensor_element
 from paddle.tensorrt.register import converter_registry
 from paddle.tensorrt.util import get_trt_version_list
@@ -103,10 +104,39 @@ def bilinear_interp_converter(network, paddle_op, inputs):
 
     outsize_tensor = None
     if trt_version_float >= 8.2:
-        if len(inputs) > 1 and inputs[1] is not None:
-            output_tensor_operand = paddle_op.operands()[1].source()
-            outsize_tensor = inputs[1]
-
+        if not pir.is_fake_value(paddle_op.operands()[1].source()):
+            size_tensor_operand = paddle_op.operands()[1].source()
+            if len(inputs) > 1 and inputs[1] is not None:
+                output_tensor_operand = paddle_op.operands()[1].source()
+                outsize_tensor = inputs[1]
+        elif not pir.is_fake_value(paddle_op.operands()[2].source()):
+            size_tensor_operand = paddle_op.operands()[2].source()
+            size_tensor = inputs[2]
+            if size_tensor_operand.is_combine():
+                size_tensors = []
+                if not isinstance(size_tensor, list):
+                    size_tensors = [size_tensor]
+                else:
+                    size_tensors = size_tensor
+                if len(size_tensors) >= 2:
+                    # Extract the first two elements representing height and width
+                    outsize_h = size_tensors[0]
+                    outsize_w = size_tensors[1]
+                    outsize_tensor = network.add_concatenation(
+                        [outsize_h, outsize_w]
+                    ).get_output(0)
+            else:
+                size_tensor_shape = size_tensor_operand.source().shape
+                if size_tensor_shape.size >= 2:
+                    outsize_h = network.add_slice(
+                        size_tensor, start=[0], shape=[1], stride=[1]
+                    ).get_output(0)
+                    outsize_w = network.add_slice(
+                        size_tensor, start=[1], shape=[1], stride=[1]
+                    ).get_output(0)
+                    outsize_tensor = network.add_concatenation(
+                        [outsize_h, outsize_w]
+                    ).get_output(0)
     use_scales = True
     if outsize_tensor is not None:
         use_scales = False
diff --git a/python/paddle/tensorrt/impls/search.py b/python/paddle/tensorrt/impls/search.py
@@ -19,6 +19,7 @@
     get_shape_tensor_element,
     squeeze_trt,
     trt_cast,
+    trt_gather,
     trt_reshape,
     trt_shape,
     trt_unsqueeze,
@@ -59,14 +60,23 @@ def argmax_converter(network, paddle_op, inputs):
     if keepdims:
         return topk_layer.get_output(1)
     else:
-        squeeze_layer = network.add_shuffle(topk_layer.get_output(1))
-        output_dims = []
-        for i in range(len(input_dims)):
-            if i == axis:
-                continue
-            output_dims.append(input_dims[i])
-        squeeze_layer.reshape_dims = tuple(output_dims)
-        return squeeze_layer.get_output(0)
+        topk_out = topk_layer.get_output(1)
+        topk_out_shape_size = len(topk_out.shape)
+        # Mark which dimensions to squeeze
+        should_squeeze = [False] * topk_out_shape_size
+        should_squeeze[axis] = True
+
+        # Get dimensions to keep
+        gather_indices = [
+            i for i, squeeze in enumerate(should_squeeze) if not squeeze
+        ]
+
+        # Add Shuffle layer
+        layer = network.add_shuffle(topk_out)
+        shape_tensor = trt_shape(network, topk_out)
+        real_shape_tensor = trt_gather(network, shape_tensor, gather_indices)
+        layer.set_input(1, real_shape_tensor)
+        return layer.get_output(0)
 
 
 @converter_registry.register("pd_op.argmin", trt_version="8.x")
diff --git a/test/tensorrt/tensorrt_test_base.py b/test/tensorrt/tensorrt_test_base.py
@@ -53,6 +53,8 @@ def create_fake_program(self):
         with paddle.static.program_guard(main_program, startup_program):
             api_args = copy.deepcopy(self.api_args)
             for feed_name in self.program_config["feed_list"]:
+                if self.api_args[feed_name] is None:
+                    continue
                 if isinstance(self.api_args[feed_name], dict):
                     new_list_args = []
                     for sub_arg_name, sub_arg_value in self.api_args[
@@ -131,6 +133,8 @@ def run_program(self, main_program, fetch_list):
         exe = paddle.static.Executor(place)
         feed_data = dict()  # noqa: C408
         for feed_name in self.program_config["feed_list"]:
+            if self.api_args[feed_name] is None:
+                continue
             if isinstance(self.api_args[feed_name], dict):
                 for sub_arg_name, sub_arg_value in self.api_args[
                     feed_name
@@ -154,7 +158,7 @@ def prepare_feed(self):
                     new_list_args[sub_arg_name] = self.api_args[arg_name][i]
                 self.api_args[arg_name] = new_list_args
 
-    def check_trt_result(self, rtol=1e-5, atol=1e-5, precision_mode=None):
+    def check_trt_result(self, rtol=1e-5, atol=1e-5, precision_mode="fp32"):
         paddle.framework.set_flags({"FLAGS_trt_min_group_size": 1})
         with paddle.pir_utils.IrGuard():
             self.prepare_feed()
@@ -179,6 +183,8 @@ def check_trt_result(self, rtol=1e-5, atol=1e-5, precision_mode=None):
             min_shape_data = dict()  # noqa: C408
             max_shape_data = dict()  # noqa: C408
             for feed_name in self.program_config["feed_list"]:
+                if self.api_args[feed_name] is None:
+                    continue
                 if isinstance(self.api_args[feed_name], dict):
                     # shape_tensor
                     if (
@@ -229,7 +235,6 @@ def check_trt_result(self, rtol=1e-5, atol=1e-5, precision_mode=None):
                             max_shape_data[feed_name] = np.random.randn(
                                 *self.max_shape[feed_name]
                             ).astype(self.api_args[feed_name].dtype)
-
             scope = paddle.static.global_scope()
             main_program = warmup_shape_infer(
                 main_program,
@@ -286,7 +291,6 @@ def check_trt_result(self, rtol=1e-5, atol=1e-5, precision_mode=None):
                 raise ValueError(
                     "The last op of convert pir Program in test must be split op that is the next op of pd_op.engine."
                 )
-
             output_trt = self.run_program(program_with_trt, trt_fetch_list)
 
         # Check that the results are close to each other within a tolerance of 1e-3
diff --git a/test/tensorrt/test_converter_common.py b/test/tensorrt/test_converter_common.py
@@ -188,6 +188,38 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestBilinearSizeTensorTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = bilinear_python_api
+        self.api_args = {
+            "x": np.random.random([2, 3, 6, 10]).astype("float32"),
+            "OutSize": None,
+            "SizeTensor": [
+                np.array([2], dtype="int64"),
+                np.array([2], dtype="int64"),
+            ],
+            "Scale": None,
+            "attrs": {
+                "data_layout": "NCHW",
+                "scale": [],
+                "out_h": -1,
+                "out_w": -1,
+                "out_d": -1,
+                "interp_method": "bilinear",
+                "align_corners": False,
+                "align_mode": 0,
+            },
+        }
+        self.program_config = {
+            "feed_list": ["x", "OutSize", "SizeTensor", "Scale"]
+        }
+        self.min_shape = {"x": [2, 3, 6, 10]}
+        self.max_shape = {"x": [12, 3, 6, 10]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestNearestNHWCTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = nearest_python_api
