[Inference]Add expand and expand_as converter (#68546)

YuanRisheng · web-flow · commit f25bfbb4c633 · 2024-10-09T17:34:55.000+08:00
* add expand and expand_as

* fix slice
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -71,6 +71,7 @@ DEFINE_GENERAL_PATTERN(FusedConv2dAddAct, paddle::dialect::FusedConv2dAddActOp)
 DEFINE_GENERAL_PATTERN(DepthwiseConv2d, paddle::dialect::DepthwiseConv2dOp)
 DEFINE_GENERAL_PATTERN(Shape, paddle::dialect::ShapeOp)
 DEFINE_GENERAL_PATTERN(Expand, paddle::dialect::ExpandOp)
+DEFINE_GENERAL_PATTERN(ExpandAs, paddle::dialect::ExpandAsOp)
 DEFINE_GENERAL_PATTERN(Sigmoid, paddle::dialect::SigmoidOp)
 DEFINE_GENERAL_PATTERN(Sqrt, paddle::dialect::SqrtOp)
 DEFINE_GENERAL_PATTERN(Hardsigmoid, paddle::dialect::HardsigmoidOp)
@@ -1308,6 +1309,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Gelu)
     ADD_PATTERN(Shape)
     ADD_PATTERN(Expand)
+    ADD_PATTERN(ExpandAs)
     ADD_PATTERN(Sigmoid)
     ADD_PATTERN(Sqrt)
     ADD_PATTERN(Hardsigmoid)
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -180,17 +180,19 @@ def convert_subgraph_to_trt(self, program, group_op):
                     continue
                 define_op_name = source.get_defining_op().name()
                 if define_op_name == "builtin.combine":
+                    operand_list = []
                     for combined_operand in source.get_defining_op().operands():
                         combined_source = combined_operand.source()
                         combined_source_id = combined_source.id
                         if combined_source_id in value_to_trt_tensor:
-                            operands.append(
+                            operand_list.append(
                                 value_to_trt_tensor[combined_source_id]
                             )
                         else:
                             raise RuntimeError(
                                 f'{combined_source_id} not found in value_to_trt_tensor'
                             )
+                    operands.append(operand_list)
                 else:
                     source_id = source.id
                     if source_id in value_to_trt_tensor:
diff --git a/python/paddle/tensorrt/converter_utils.py b/python/paddle/tensorrt/converter_utils.py
@@ -31,6 +31,9 @@
     __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
 )
 
+version = trt.__version__
+version_list = list(map(int, version.split('.')))
+
 
 def has_dynamic_shape(shape):
     return any(s == -1 for s in shape)
@@ -159,11 +162,46 @@ 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):
-    constant_data = np.array([data], dtype=dtype)
+    if not isinstance(data, list):
+        data = [data]
+    constant_data = np.array(data, dtype=dtype)
     constant_layer = network.add_constant(constant_data.shape, constant_data)
     return constant_layer.get_output(0)
 
 
+# Concat not make rank changed
+def trt_concat(network, inputs, axis=0):
+    concat_layer = network.add_concatenation(inputs=inputs)
+    if axis != 0:
+        concat_layer.axis = axis
+    return concat_layer.get_output(0)
+
+
+def trt_cast(network, input, dtype):
+    identity_layer = network.add_identity(input)
+    identity_layer.set_output_type(0, dtype)
+    identity_layer.get_output(0).dtype = dtype
+    return identity_layer.get_output(0)
+
+
+def trt_shape(network, input):
+    shape_layer = network.add_shape(input)
+    if version_list[0] >= 10:  # trt_version >=10
+        return trt_cast(network, shape_layer.get_output(0), trt.int32)
+    return shape_layer.get_output(0)
+
+
+def trt_reshape(network, input, new_shape, name="", is_shape_tensor=False):
+    reshape_layer = network.add_shuffle(input)
+    if is_shape_tensor:
+        reshape_layer.set_input(1, new_shape)
+    else:
+        reshape_layer.reshape_dims = new_shape
+    if name != "":
+        reshape_layer.name = name
+    return reshape_layer.get_output(0)
+
+
 # Get element tensor of 1D shape tensor
 def get_shape_tensor_element(network, x, index):
     assert index >= 0, (
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -22,8 +22,11 @@
     get_positive_dim,
     get_shape_tensor_element,
     has_dynamic_shape,
+    trt_concat,
     trt_max,
     trt_min,
+    trt_reshape,
+    trt_shape,
     trt_sub,
     trt_sum,
 )
@@ -156,8 +159,8 @@ def flatten_converter(network, paddle_op, inputs):
 # In the converter, pd_op.concat has three inputs, because builtin.combine has two inputs.
 @converter_registry.register("pd_op.concat", trt_version="8.x")
 def concat_converter(network, paddle_op, inputs):
-    input_tensors = inputs[:-1]
-    axis_tensor = inputs[-1]
+    input_tensors = inputs[0]
+    axis_tensor = inputs[1]
     concat_layer = network.add_concatenation(inputs=input_tensors)
 
     axis = paddle_op.operands()[1].source().get_defining_op().attrs()["value"]
@@ -214,6 +217,76 @@ def squeeze_converter(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
+def get_expand_output(network, input, rank, shape_tensor, shape_rank):
+    if rank < shape_rank:
+        one_rank_tensor = add_1D_constant_layer(
+            network, [1] * (shape_rank - rank)
+        )
+        in_shape_tensor = trt_shape(network, input)
+        itensors = [one_rank_tensor, in_shape_tensor]
+        input_shape_tensor = trt_concat(network, itensors)
+    else:
+        input_shape_tensor = trt_shape(network, input)
+
+    new_input_tensor = trt_reshape(network, input, input_shape_tensor, "", True)
+
+    start = [0] * shape_rank
+    starts_tensor = add_1D_constant_layer(network, start)
+    one_tensor = add_1D_constant_layer(network, 1)
+    sizes_tensor = trt_max(network, input_shape_tensor, shape_tensor)
+    input_sub_tensor = trt_sub(network, input_shape_tensor, one_tensor)
+    strides_tensor = trt_min(network, one_tensor, input_sub_tensor)
+
+    slice_layer = network.add_slice(
+        new_input_tensor, start, [0] * len(start), [0] * len(start)
+    )
+    slice_layer.set_input(1, starts_tensor)
+    slice_layer.set_input(2, sizes_tensor)
+    slice_layer.set_input(3, strides_tensor)
+
+    return slice_layer.get_output(0)
+
+
+@converter_registry.register("pd_op.expand", trt_version="8.x")
+def expand_converter(network, paddle_op, inputs):
+    input = inputs[0]
+    input_dims = input.shape
+    rank = len(input_dims)
+    paddle_shape_tensor = paddle_op.operands()[1].source()
+
+    shape_tensor_source_op = paddle_shape_tensor.get_defining_op()
+    if shape_tensor_source_op.name() == "pd_op.full_int_array":
+        shape = shape_tensor_source_op.attrs()["value"]
+        shape_tensor = add_1D_constant_layer(network, shape)
+        shape_rank = len(shape)
+    elif paddle_shape_tensor.type().as_vec_type():
+        shape_tensors = inputs[1]
+        shape_rank = len(shape_tensors)
+        shape_tensor = trt_concat(network, shape_tensors)
+    else:
+        shape_tensor = inputs[1]
+        shape_rank = shape_tensor.shape[0]
+    return get_expand_output(network, input, rank, shape_tensor, shape_rank)
+
+
+@converter_registry.register("pd_op.expand_as", trt_version="8.x")
+def expand_as_converter(network, paddle_op, inputs):
+    input = inputs[0]
+    input_dims = input.shape
+    rank = len(input_dims)
+    y = paddle_op.operands()[1].source()
+
+    if y.initialized():
+        y_t = inputs[1]
+        shape_tensor = trt_shape(network, y_t)
+        shape_rank = len(y_t.shape)
+    else:
+        shape = paddle_op.attrs().get("target_shape")
+        shape_tensor = add_1D_constant_layer(network, shape)
+        shape_rank = len(shape)
+    return get_expand_output(network, input, rank, shape_tensor, shape_rank)
+
+
 @converter_registry.register("pd_op.cast", trt_version="8.x")
 @converter_registry.register("pd_op.cast_", trt_version="8.x")
 def cast_converter(network, paddle_op, inputs):
@@ -270,7 +343,7 @@ def slice_converter(network, paddle_op, inputs):
             len(starts),
             len(axes),
         )
-        for idx in axes:
+        for idx in range(len(axes)):
             if starts[idx] < 0:
                 starts_tensor[axes[idx]] = trt_max(
                     network,
@@ -293,7 +366,7 @@ def slice_converter(network, paddle_op, inputs):
                 )
     else:
         starts = inputs[1]
-        for idx in axes:
+        for idx in range(len(axes)):
             starts_tensor[axes[idx]] = get_shape_tensor_element(
                 network, starts, idx
             )
@@ -306,7 +379,7 @@ def slice_converter(network, paddle_op, inputs):
             len(ends),
             len(axes),
         )
-        for idx in axes:
+        for idx in range(len(axes)):
             if ends[idx] < 0:
                 ends_tensor[axes[idx]] = trt_max(
                     network,
@@ -329,7 +402,7 @@ def slice_converter(network, paddle_op, inputs):
                 )
     else:
         ends = inputs[2]
-        for idx in axes:
+        for idx in range(len(axes)):
             ends_tensor[axes[idx]] = get_shape_tensor_element(
                 network, ends, idx
             )
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -92,6 +92,36 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestExpandTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.expand
+        self.api_args = {
+            "x": np.random.randn(1, 3).astype("float32"),
+            "shape": [6, 3],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [6, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestExpandWithShapeTensorTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.expand
+        self.api_args = {
+            "x": np.random.randn(1, 3).astype("float32"),
+            "shape": np.array([6, 3]).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "shape"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [6, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 def slice_api(x, axes, starts, ends, infer_flags, decrease_axis):
     return _C_ops.slice(x, axes, starts, ends, infer_flags, decrease_axis)
 
@@ -115,6 +145,21 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestExpandWithDiffRankTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.expand
+        self.api_args = {
+            "x": np.array([1, 2, 3]).astype("float32"),
+            "shape": [2, 3],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {}
+        self.max_shape = {}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestSliceTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.slice
@@ -132,6 +177,23 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestExpandAsTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.expand_as
+        self.api_args = {
+            "x": np.array([[1, 2, 3]]).astype("float32"),
+            "y": np.array([[1, 2, 3], [4, 5, 6], [1, 2, 3], [4, 5, 6]]).astype(
+                "int32"
+            ),
+        }
+        self.program_config = {"feed_list": ["x", "y"]}
+        self.min_shape = {"x": [1, 3]}
+        self.max_shape = {"x": [4, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestSliceWithInputStartTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.slice
