make padding layer converter more efficient (#1470)

Author: nvpohanh

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -406,46 +406,34 @@ def acc_ops_pad_with_slice_layer(
     )
 
     input_shape = input_val.shape
-    pre_start = tuple(i - 1 for i in input_shape)
     prefix_len = len(input_shape) - len(pad) // 2
-    pre_shape = tuple(
-        input_shape[i] + (pad[-(i - prefix_len) * 2 - 2] if i >= prefix_len else 0)
+    start = tuple(
+        -pad[-(i - prefix_len) * 2 - 2] if i >= prefix_len else 0
         for i in range(0, len(input_shape))
     )
-    pre_stride = [-1] * len(input_shape)
+
+    shape = tuple(
+        input_shape[i]
+        + (
+            pad[-(i - prefix_len) * 2 - 1] + pad[-(i - prefix_len) * 2 - 2]
+            if i >= prefix_len
+            else 0
+        )
+        for i in range(0, len(input_shape))
+    )
+    stride = tuple([1] * len(shape))
 
     layer = network.add_slice(
         input_val,
-        pre_start,
-        pre_shape,
-        pre_stride,
+        start,
+        shape,
+        stride,
     )
-    layer.set_input(4, value_const)
-    layer.mode = trt.SliceMode.FILL
-    set_layer_name(layer, target, f"pre_{name}")
-    half_pad_output = layer.get_output(0)
 
-    shape = half_pad_output.shape
-    mid_start = tuple(i - 1 for i in shape)
-    mid_stride = [-1] * len(shape)
-    layer = network.add_slice(half_pad_output, mid_start, shape, mid_stride)
     layer.set_input(4, value_const)
     layer.mode = trt.SliceMode.FILL
-    set_layer_name(layer, target, f"transpose_{name}")
-    transpose_output = layer.get_output(0)
-
-    shape = transpose_output.shape
-    post_start = tuple([0] * len(shape))
-    post_shape = tuple(
-        shape[i] + (pad[-(i - prefix_len) * 2 - 1] if i >= prefix_len else 0)
-        for i in range(0, len(shape))
-    )
-    post_stride = tuple([1] * len(shape))
+    set_layer_name(layer, target, name)
 
-    layer = network.add_slice(transpose_output, post_start, post_shape, post_stride)
-    layer.set_input(4, value_const)
-    layer.mode = trt.SliceMode.FILL
-    set_layer_name(layer, target, f"post_{name}")
     return layer.get_output(0)
 
 

py/torch_tensorrt/fx/test/passes/test_fuse_permute_matmul_trt.py

@@ -37,8 +37,20 @@ class TestFusePermuteMatmul(AccTestCase):
                 lambda x: x.permute(0, 1, 3, 2),
                 torch.matmul,
             ),
-            param("transpose_lhs_bmm_broadcast", (3, 2), (3, 3, 4), tranpose_last_two_dims, op=torch.matmul),
-            param("transpose_rhs_bmm_broadcast", (3, 3, 4), (3, 4), rhs_op=tranpose_last_two_dims, op=torch.matmul),
+            param(
+                "transpose_lhs_bmm_broadcast",
+                (3, 2),
+                (3, 3, 4),
+                tranpose_last_two_dims,
+                op=torch.matmul,
+            ),
+            param(
+                "transpose_rhs_bmm_broadcast",
+                (3, 3, 4),
+                (3, 4),
+                rhs_op=tranpose_last_two_dims,
+                op=torch.matmul,
+            ),
         ]
     )
     def test_fuse_permute_matmul(