[Paddle Tensorrt] add pd_op.topk,pd_op.roll converter (#69117)

lizexu123 · web-flow · commit cd4063b22b06 · 2024-11-05T14:38:23.000+08:00
* pd_op.strided_slice

* merge develop

* pd_op.topk

* fix

* merge develop

* pd_op.roll

* add pir::value,topk

* fix

* fix

* 增加单测时长
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -80,6 +80,7 @@ DEFINE_GENERAL_PATTERN(Hardswish, paddle::dialect::HardswishOp)
 DEFINE_GENERAL_PATTERN(Assign, paddle::dialect::AssignOp)
 DEFINE_GENERAL_PATTERN(AssignValue_, paddle::dialect::AssignValue_Op)
 DEFINE_GENERAL_PATTERN(AssignOut, paddle::dialect::AssignOut_Op)
+DEFINE_GENERAL_PATTERN(Roll, paddle::dialect::RollOp)
 
 #undef DEFINE_GENERAL_PATTERN
 
@@ -1434,7 +1435,7 @@ class StackOpPattern : public pir::OpRewritePattern<paddle::dialect::StackOp> {
     }
 
     pir::Value x = op.operand_source(0);
-    int rank;
+    int rank = 1;
     auto x_type = x.type();
     if (x_type.isa<pir::VectorType>()) {
       rank = x_type.dyn_cast<pir::VectorType>().size();
@@ -1443,7 +1444,7 @@ class StackOpPattern : public pir::OpRewritePattern<paddle::dialect::StackOp> {
       rank = x_shape.size();
     }
 
-    int axis;
+    int axis = 1;
     if (op->HasAttribute("axis")) {
       axis = op->attribute<pir::Int32Attribute>("axis").data();
     } else {
@@ -1561,6 +1562,32 @@ class StridedSliceOpPattern
   }
 };
 
+class TopkOpPattern : public pir::OpRewritePattern<paddle::dialect::TopkOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::TopkOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::TopkOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    if (!op->HasAttribute("axis")) {
+      VLOG(3) << "pd_op.topk must has axis attribute";
+      return false;
+    }
+    if (op->HasAttribute("sorted")) {
+      bool sorted = op->attribute<pir::BoolAttribute>("sorted").data();
+      if (!sorted) {
+        VLOG(3)
+            << "pd_op.topk does not support results not sorted in tensorrt.";
+        return false;
+      }
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
@@ -1599,6 +1626,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(AssignOut)
     ADD_PATTERN(Assign)
     ADD_PATTERN(AssignValue_)
+    ADD_PATTERN(Roll)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
 #endif
@@ -1649,6 +1677,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<TanhOpPattern>(context));
     ps.Add(std::make_unique<FullWithTensorPattern>(context));
     ps.Add(std::make_unique<StridedSliceOpPattern>(context));
+    ps.Add(std::make_unique<TopkOpPattern>(context));
     return ps;
   }
 };
diff --git a/python/paddle/tensorrt/converter_utils.py b/python/paddle/tensorrt/converter_utils.py
@@ -162,11 +162,12 @@ def add_elementwise_layer(network, paddle_op, inputs, op_type):
 
 
 # Create and add 1D constant layer
-def add_1D_constant_layer(network, data, dtype=np.int32):
+def add_1D_constant_layer(network, data, dtype=np.int32, is_scalar=False):
     if not isinstance(data, list):
         data = [data]
     constant_data = np.array(data, dtype=dtype)
-    constant_layer = network.add_constant(constant_data.shape, constant_data)
+    shape = () if is_scalar else (len(data),)
+    constant_layer = network.add_constant(shape, constant_data)
     return constant_layer.get_output(0)
 
 
@@ -204,13 +205,12 @@ def trt_reshape(network, input, new_shape, name="", is_shape_tensor=False):
 
 
 # Get element tensor of 1D shape tensor
-def get_shape_tensor_element(network, x, index):
+def get_shape_tensor_element(network, x, index, is_scalar=False):
     assert index >= 0, (
         "The index should be greater or equal than 0, but got %d" % index
     )
-    gather_layer = network.add_gather(
-        input=x, indices=add_1D_constant_layer(network, index), axis=0
-    )
+    index_tensor = add_1D_constant_layer(network, index, is_scalar=is_scalar)
+    gather_layer = network.add_gather(input=x, indices=index_tensor, axis=0)
     return gather_layer.get_output(0)
 
 
@@ -514,3 +514,55 @@ def fix_negative_indices(network, input_shape, indices):
     sub = trt_mul(network, sign, input_shape)
     fixed_indices = trt_sub(network, indices, sub)
     return fixed_indices
+
+
+def unsqueeze_trt(network, input_tensor, axes):
+    input_shape = network.add_shape(input_tensor).get_output(0)
+
+    axis_set = set(axes)
+
+    subscripts = list(range(len(input_tensor.shape)))
+
+    for axis in sorted(axis_set):
+        subscripts.insert(axis, len(input_tensor.shape))
+
+    one_tensor = network.add_constant(
+        (1,), np.array([1], dtype=np.int32)
+    ).get_output(0)
+    extended_shape = network.add_concatenation(
+        [input_shape, one_tensor]
+    ).get_output(0)
+
+    gather_layer = network.add_gather(
+        extended_shape,
+        network.add_constant(
+            (len(subscripts),), np.array(subscripts, dtype=np.int32)
+        ).get_output(0),
+        axis=0,
+    )
+    new_shape_tensor = gather_layer.get_output(0)
+
+    reshaped_tensor = network.add_shuffle(input_tensor)
+    reshaped_tensor.set_input(1, new_shape_tensor)
+
+    return reshaped_tensor.get_output(0)
+
+
+def squeeze_trt(network, input_tensor, axes):
+    input_shape = network.add_shape(input_tensor).get_output(0)
+    input_shape = input_tensor.shape
+    all_dims = list(range(len(input_shape)))
+    remaining_dims = [dim for dim in all_dims if dim not in axes]
+
+    input_shape_tensor = network.add_shape(input_tensor).get_output(0)
+
+    remaining_dims_tensor = network.add_constant(
+        (len(remaining_dims),), np.array(remaining_dims, dtype=np.int32)
+    ).get_output(0)
+
+    new_shape_tensor = network.add_gather(
+        input_shape_tensor, remaining_dims_tensor, axis=0
+    ).get_output(0)
+    reshape_layer = network.add_shuffle(input_tensor)
+    reshape_layer.set_input(1, new_shape_tensor)
+    return reshape_layer.get_output(0)
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -829,3 +829,71 @@ def strided_slice_converter(network, paddle_op, inputs):
     layer.set_input(3, step_tensor)
 
     return layer.get_output(0)
+
+
+@converter_registry.register("pd_op.roll", trt_version="8.x")
+def roll_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    axis = paddle_op.attrs()["axis"]
+
+    shifts_op = paddle_op.operands()[1].source().get_defining_op()
+    if shifts_op.name() == "pd_op.full_int_array":
+        shifts = shifts_op.attrs()["value"]
+    else:
+        shifts = inputs[1]
+
+    axis_size = len(axis)
+    input_shape_tensor = trt_shape(network, input_tensor)
+
+    for i in range(axis_size):
+        axi = axis[i]
+        if isinstance(shifts, trt.ITensor):
+            shift = get_shape_tensor_element(network, shifts, i)
+            input_shift = shift
+        else:
+            shift = shifts[i]
+            input_shift = add_1D_constant_layer(network, shift)
+        input_axis = get_shape_tensor_element(network, input_shape_tensor, axi)
+
+        # 1.sub_value mod input_axis
+        input1 = trt_sub(network, input_axis, input_shift)
+        tmp_div_res = trt_floor_div(network, input1, input_axis)
+        tmp_prod_res = trt_mul(network, tmp_div_res, input_axis)
+        start = trt_sub(network, input1, tmp_prod_res)
+        # 2.avoid start less than 0,start mod input_axis
+        start = trt_sum(network, start, input_axis)
+        tmp_div_res1 = trt_floor_div(network, start, input_axis)
+        tmp_prod_res1 = trt_mul(network, tmp_div_res1, input_axis)
+        start = trt_sub(network, start, tmp_prod_res1)
+        zero_tensor = add_1D_constant_layer(network, 0)
+        step = add_1D_constant_layer(network, 1)
+        # 3.make index_tensor0
+        sub_qutient = trt_sub(network, input_axis, start)
+        quotient_tensor = trt_floor_div(network, sub_qutient, step)
+        start1 = get_shape_tensor_element(network, start, 0, is_scalar=True)
+        fill_layer0 = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
+        fill_layer0.set_input(0, quotient_tensor)
+        fill_layer0.set_input(1, start1)
+        fill_layer0.set_input(2, step)
+        index_tensor0 = fill_layer0.get_output(0)
+        # 4.make index_tensor1
+        sub_qutient_tensor = trt_sub(network, start, zero_tensor)
+        quotient_tensor = trt_floor_div(network, sub_qutient_tensor, step)
+        start2 = add_1D_constant_layer(network, 0, is_scalar=True)
+        fill_layer1 = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
+        fill_layer1.set_input(0, quotient_tensor)
+        fill_layer1.set_input(1, start2)
+        fill_layer1.set_input(2, step)
+        index_tensor1 = fill_layer1.get_output(0)
+        itensors = [index_tensor0, index_tensor1]
+        concat_input_tensor = trt_concat(network, itensors)
+        if i == 0:
+            layer = network.add_gather(
+                input=input_tensor, indices=concat_input_tensor, axis=axi
+            )
+        else:
+            layer = network.add_gather(
+                input=layer.get_output(0), indices=concat_input_tensor, axis=axi
+            )
+
+    return layer.get_output(0)
diff --git a/python/paddle/tensorrt/impls/search.py b/python/paddle/tensorrt/impls/search.py
@@ -15,6 +15,11 @@
 
 import tensorrt as trt
 
+from paddle.tensorrt.converter_utils import (
+    squeeze_trt,
+    trt_cast,
+    unsqueeze_trt,
+)
 from paddle.tensorrt.register import converter_registry
 
 
@@ -59,3 +64,40 @@ def argmax_converter(network, paddle_op, inputs):
             output_dims.append(input_dims[i])
         squeeze_layer.reshape_dims = tuple(output_dims)
         return squeeze_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]
+
+    input_shape = paddle_op.operands()[0].source().shape
+
+    k = paddle_op.attrs().get("k", 1)
+    axis = paddle_op.attrs().get("axis", -1)
+    largest = paddle_op.attrs().get("largest", True)
+    flag = trt.TopKOperation.MAX if largest else trt.TopKOperation.MIN
+
+    input_rank = len(input_shape)
+
+    expand_to_2d = input_rank == 1
+    if expand_to_2d:
+        input_tensor = unsqueeze_trt(network, input_tensor, [1])
+
+    input_type = input_tensor.dtype
+    if input_type == trt.DataType.INT32:
+        input_tensor = trt_cast(network, input_tensor, trt.DataType.FLOAT)
+
+    if axis < 0:
+        axis += input_rank
+    layer = network.add_topk(input_tensor, flag, k, 1 << axis)
+    values = layer.get_output(0)
+    indices = layer.get_output(1)
+
+    if expand_to_2d:
+        values = squeeze_trt(network, values, [1])
+        indices = squeeze_trt(network, indices, [1])
+
+    if input_type == trt.DataType.INT32:
+        values = trt_cast(network, values, trt.DataType.INT32)
+
+    return values, indices
diff --git a/python/paddle/tensorrt/util.py b/python/paddle/tensorrt/util.py
@@ -23,6 +23,7 @@
 
 
 def map_dtype(pd_dtype):
+    version_list = get_trt_version_list()
     if pd_dtype == "FLOAT32":
         return trt.float32
     elif pd_dtype == "FLOAT16":
@@ -33,6 +34,9 @@ def map_dtype(pd_dtype):
         return trt.int8
     elif pd_dtype == "BOOL":
         return trt.bool
+    # trt version<10.0 not support int64,so convert int64 to int32
+    elif pd_dtype == "INT64":
+        return trt.int64 if version_list[0] >= 10 else trt.int32
     # Add other dtype mappings as needed
     else:
         raise TypeError(f"Unsupported dtype: {pd_dtype}")
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -13,15 +13,15 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_model_resnet50 PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_conv PROPERTIES TIMEOUT "300")
   set_tests_properties(test_export PROPERTIES TIMEOUT "500")
-  set_tests_properties(test_converter_norm PROPERTIES TIMEOUT "100")
-  set_tests_properties(test_converter_ops PROPERTIES TIMEOUT "100")
-  set_tests_properties(test_converter_stat PROPERTIES TIMEOUT "100")
-  set_tests_properties(test_converter_math PROPERTIES TIMEOUT "100")
-  set_tests_properties(test_converter_activation PROPERTIES TIMEOUT "100")
-  set_tests_properties(test_converter_others PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_norm PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_ops PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_stat PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_math PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_activation PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_others PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_manipulation PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "300")
-  set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_common PROPERTIES TIMEOUT "300")
-  set_tests_properties(test_converter_logic PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_logic PROPERTIES TIMEOUT "300")
 endif()
diff --git a/test/tensorrt/tensorrt_test_base.py b/test/tensorrt/tensorrt_test_base.py
@@ -53,7 +53,6 @@ def create_fake_program(self):
                     for sub_arg_name, sub_arg_value in self.api_args[
                         feed_name
                     ].items():
-
                         if (
                             feed_name in self.min_shape.keys()
                             and feed_name in self.max_shape.keys()
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -471,5 +471,53 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestRollCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.roll
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("float32"),
+            "shift": 1,
+            "axis": 0,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 4, 10]}
+        self.max_shape = {"x": [5, 4, 10]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestRollCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.roll
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("int64"),
+            "shift": 1,
+            "axis": 1,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 4, 10]}
+        self.max_shape = {"x": [5, 4, 10]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestRollCase3TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.roll
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("float32"),
+            "shift": np.array([1]).astype("int64"),
+            "axis": 1,
+        }
+        self.program_config = {"feed_list": ["x", "shift"]}
+        self.min_shape = {"x": [1, 4, 10]}
+        self.max_shape = {"x": [5, 4, 10]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/tensorrt/test_converter_search.py b/test/tensorrt/test_converter_search.py
