Repository: pytorch/TensorRT (language: python)

Issue:
Add support for padding aten operators
```[tasklist]
### Tasks
- [x] torch.ops.aten.pad.default
- [x] torch.ops.aten.constant_pad_nd.default
- [x] torch.ops.aten.reflection_pad1d.default
- [x] torch.ops.aten.reflection_pad2d.default
- [x] torch.ops.aten.reflection_pad3d.default
- [x] torch.ops.aten.replication_pad1d.default
- [x] torch.ops.aten.replication_pad2d.default
- [x] torch.ops.aten.replication_pad3d.default
- [x] torch.ops.aten._pad_circular.default
```


Discussion / hints:

Fix patch:
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index e578ec2fc7..fb031c4f19 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -2054,3 +2054,125 @@ def aten_ops_addmm(
         beta=kwargs.get("beta", 1),
         alpha=kwargs.get("alpha", 1),
     )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.constant_pad_nd.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_constant_pad(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.pad.constant_padNd(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+        args_bounds_check(args, 2, 0),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.reflection_pad1d.default)
+@dynamo_tensorrt_converter(torch.ops.aten.reflection_pad2d.default)
+@dynamo_tensorrt_converter(torch.ops.aten.reflection_pad3d.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_reflection_pad(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.pad.reflection_padNd(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.replication_pad1d.default)
+@dynamo_tensorrt_converter(torch.ops.aten.replication_pad2d.default)
+@dynamo_tensorrt_converter(torch.ops.aten.replication_pad3d.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_replication_pad(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.pad.replication_padNd(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten._pad_circular.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_circular_pad(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.pad.circular_padNd(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.pad.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_pad(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.pad.pad(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        pad=args[1],
+        mode=args_bounds_check(args, 2, "constant"),
+        value=args_bounds_check(args, 3, None),
+    )
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/__init__.py b/py/torch_tensorrt/dynamo/conversion/impl/__init__.py
index b448b40bc3..62bf556beb 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/__init__.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/__init__.py
@@ -16,6 +16,7 @@
     linear,
     matmul,
     normalization,
+    pad,
     permutation,
     pool,
     reduce,
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/cat.py b/py/torch_tensorrt/dynamo/conversion/impl/cat.py
index 24149d01b0..d6ffc77377 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/cat.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/cat.py
@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Sequence, Union
+from typing import Optional, Sequence, Union
 
 import numpy as np
 import torch
@@ -6,12 +6,11 @@
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 from torch_tensorrt.dynamo.conversion.converter_utils import (
-    SourceIR,
     get_positive_dim,
     get_trt_tensor,
 )
 from torch_tensorrt.fx.converters.converter_utils import set_layer_name
-from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+from torch_tensorrt.fx.types import TRTTensor
 
 
 def cat(
@@ -23,12 +22,12 @@ def cat(
     dim: int,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     trt_inputs = []
-    for each_input in input:
+    for i, each_input in enumerate(input):
         if not isinstance(each_input, TRTTensor):
-            each_input = get_trt_tensor(ctx, each_input, name + "_tensor_{i}")
+            each_input = get_trt_tensor(ctx, each_input, f"{name}_tensor_{i}")
         trt_inputs.append(each_input)
     concat_layer = ctx.net.add_concatenation(trt_inputs)
     dim = get_positive_dim(dim, len(input[0].shape))
     concat_layer.axis = dim
-    set_layer_name(concat_layer, target, name + "_gather", source_ir)
+    set_layer_name(concat_layer, target, f"{name}_gather", source_ir)
     return concat_layer.get_output(0)
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/pad.py b/py/torch_tensorrt/dynamo/conversion/impl/pad.py
new file mode 100644
index 0000000000..3764667ffb
--- /dev/null
+++ b/py/torch_tensorrt/dynamo/conversion/impl/pad.py
@@ -0,0 +1,205 @@
+from typing import Optional, Sequence, Union
+
+import tensorrt as trt
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor
+from torch_tensorrt.fx.converters.converter_utils import (
+    has_dynamic_shape,
+    set_layer_name,
+)
+from torch_tensorrt.fx.types import TRTTensor
+
+"""
+Note: IPaddingLayer is deprecated in TensorRT 8.2 and will be removed in TensorRT 10.0.
+Use ISliceLayer to pad the tensor, which supports new non-constant, reflects padding
+mode and clamp, and supports padding output with dynamic shape.
+"""
+
+
+def constant_padNd(
+    ctx: ConversionContext,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    pad: Sequence[int],
+    value: Union[int, float] = 0,
+) -> TRTTensor:
+    if has_dynamic_shape(input.shape):
+        assert input.shape[1] != -1, "Channel dim can't be dynamic for padding."
+
+    rank = len(input.shape)
+
+    if len(pad) // 2 > rank:
+        raise RuntimeError(
+            f"Trying to pad last {len(pad) // 2} dimension but the input only has {rank} dimension."
+        )
+
+    start_list = [0] * rank
+    new_shape = list(input.shape)
+
+    for i in range(0, len(pad) // 2):
+        start_list[-i - 1] = -pad[i * 2]
+        new_shape[-i - 1] += pad[i * 2] + pad[i * 2 + 1]
+
+    stride_list = [1] * rank
+    layer = ctx.net.add_slice(
+        input,
+        start=tuple(start_list),
+        shape=tuple(new_shape),
+        stride=tuple(stride_list),
+    )
+    value_const = get_trt_tensor(ctx, value, f"{name}_value", input.dtype)
+    layer.set_input(4, value_const)
+    layer.mode = trt.SliceMode.FILL
+
+    set_layer_name(layer, target, name, source_ir)
+    return layer.get_output(0)
+
+
+def reflection_padNd(
+    ctx: ConversionContext,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    padding: Sequence[int],
+) -> TRTTensor:
+    if has_dynamic_shape(input.shape):
+        assert input.shape[1] != -1, "Channel dim can't be dynamic for padding."
+
+    rank = len(input.shape)
+
+    if len(padding) // 2 > rank:
+        raise RuntimeError(
+            f"Trying to pad last {len(padding) // 2} dimension but the input only has {rank} dimension."
+        )
+
+    start_list = [0] * rank
+    new_shape = list(input.shape)
+
+    for i in range(0, len(padding) // 2):
+        start_list[-i - 1] = -padding[i * 2]
+        new_shape[-i - 1] += padding[i * 2] + padding[i * 2 + 1]
+
+    stride_list = [1] * rank
+    layer = ctx.net.add_slice(
+        input,
+        start=tuple(start_list),
+        shape=tuple(new_shape),
+        stride=tuple(stride_list),
+    )
+    layer.mode = trt.SliceMode.REFLECT
+
+    set_layer_name(layer, target, name, source_ir)
+    return layer.get_output(0)
+
+
+def replication_padNd(
+    ctx: ConversionContext,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    padding: Sequence[int],
+) -> TRTTensor:
+    if has_dynamic_shape(input.shape):
+        assert input.shape[1] != -1, "Channel dim can't be dynamic for padding."
+
+    rank = len(input.shape)
+
+    if len(padding) // 2 > rank:
+        raise RuntimeError(
+            f"Trying to pad last {len(padding) // 2} dimension but the input only has {rank} dimension."
+        )
+
+    start_list = [0] * rank
+    new_shape = list(input.shape)
+
+    for i in range(0, len(padding) // 2):
+        start_list[-i - 1] = -padding[i * 2]
+        new_shape[-i - 1] += padding[i * 2] + padding[i * 2 + 1]
+
+    stride_list = [1] * rank
+    layer = ctx.net.add_slice(
+        input,
+        start=tuple(start_list),
+        shape=tuple(new_shape),
+        stride=tuple(stride_list),
+    )
+    layer.mode = trt.SliceMode.CLAMP
+
+    set_layer_name(layer, target, name, source_ir)
+    return layer.get_output(0)
+
+
+def circular_padNd(
+    ctx: ConversionContext,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    pad: Sequence[int],
+) -> TRTTensor:
+    if has_dynamic_shape(input.shape):
+        assert input.shape[1] != -1, "Channel dim can't be dynamic for padding."
+
+    rank = len(input.shape)
+
+    if len(pad) // 2 > rank:
+        raise RuntimeError(
+            f"Trying to pad last {len(pad) // 2} dimension but the input only has {rank} dimension."
+        )
+
+    start_list = [0] * rank
+    new_shape = list(input.shape)
+
+    for i in range(0, len(pad) // 2):
+        start_list[-i - 1] = -pad[i * 2]
+        new_shape[-i - 1] += pad[i * 2] + pad[i * 2 + 1]
+
+    stride_list = [1] * rank
+    layer = ctx.net.add_slice(
+        input,
+        start=tuple(start_list),
+        shape=tuple(new_shape),
+        stride=tuple(stride_list),
+    )
+    layer.mode = trt.SliceMode.WRAP
+
+    set_layer_name(layer, target, name, source_ir)
+    return layer.get_output(0)
+
+
+def pad(
+    ctx: ConversionContext,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    pad: Sequence[int],
+    mode: str = "constant",
+    value: Optional[float] = None,
+) -> TRTTensor:
+    if mode == "constant":
+        return constant_padNd(
+            ctx,
+            target,
+            source_ir,
+            f"{name}_{mode}",
+            input,
+            pad,
+            value if value is not None else 0,
+        )
+    elif mode == "reflect":
+        return reflection_padNd(ctx, target, source_ir, f"{name}_{mode}", input, pad)
+    elif mode == "replicate":
+        return replication_padNd(ctx, target, source_ir, f"{name}_{mode}", input, pad)
+    elif mode == "circular":
+        return circular_padNd(ctx, target, source_ir, f"{name}_{mode}", input, pad)
+    else:
+        raise RuntimeError(
+            f'We currently only support for `mode` in ["constant", "reflect", "replicate", "circular"], but got {mode}'
+        )


Test patch:
diff --git a/tests/py/dynamo/conversion/test_pad_aten.py b/tests/py/dynamo/conversion/test_pad_aten.py
new file mode 100644
index 0000000000..2803736ad0
--- /dev/null
+++ b/tests/py/dynamo/conversion/test_pad_aten.py
@@ -0,0 +1,241 @@
+import torch
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestConstantPadConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ((1, 2), (1, 1), 0),
+            ((2, 1), (2, 1), 1),
+            ((3, 4, 2), (1, 2), 2),
+            ((3, 4, 2), (1, 2, 3, 1, 2, 3), 0),
+            ((3, 3, 4, 2), (1, 2, 3, 4), 0),
+            ((3, 3, 4, 2), (1, 2, 3, 4), 2),
+            ((3, 3, 4, 2, 1), (1, 2, 3, 4, 5, 1, 2, 3, 4, 5), 0),
+            ((3, 3, 4, 2, 1, 2), (1, 2, 3, 4, 1, 2, 3, 4), 4),
+        ]
+    )
+    def test_constant_pad(self, shape, pad, value):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.constant_pad_nd.default(input, pad, value)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+
+class TestReflectionPadConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 2D or 3D
+            ((3, 3), (1, 1)),
+            ((3, 3), (2, 2)),
+            ((2, 2, 2), (1, 1)),
+            ((2, 2, 4), (2, 3)),
+        ]
+    )
+    def test_reflection_pad1d(self, shape, padding):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.reflection_pad1d.default(input, padding)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 3D or 4D
+            ((2, 2, 2), (1, 1, 1, 1)),
+            ((1, 2, 4), (2, 2, 1, 1)),
+            ((2, 2, 3, 3), (1, 1, 2, 2)),
+            ((2, 3, 4, 5), (4, 3, 0, 1)),
+        ]
+    )
+    def test_reflection_pad2d(self, shape, padding):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.reflection_pad2d.default(input, padding)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 4D or 5D
+            ((2, 2, 2, 2), (1, 1, 1, 1, 1, 1)),
+            ((1, 2, 3, 4), (3, 2, 2, 1, 1, 1)),
+            ((2, 2, 3, 3, 4), (3, 3, 2, 1, 1, 2)),
+            ((2, 3, 4, 5, 6), (4, 3, 2, 1, 1, 0)),
+        ]
+    )
+    def test_reflection_pad3d(self, shape, padding):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.reflection_pad3d.default(input, padding)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+
+class TestReplicationPadConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 2D or 3D
+            ((3, 3), (1, 1)),
+            ((3, 3), (2, 2)),
+            ((2, 2, 2), (1, 1)),
+            ((2, 2, 4), (2, 3)),
+        ]
+    )
+    def test_replication_pad1d(self, shape, padding):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.replication_pad1d.default(input, padding)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 3D or 4D
+            ((2, 2, 2), (1, 1, 1, 1)),
+            ((1, 2, 4), (2, 2, 1, 1)),
+            ((2, 2, 3, 3), (1, 1, 2, 2)),
+            ((2, 3, 4, 5), (4, 3, 0, 1)),
+        ]
+    )
+    def test_replication_pad2d(self, shape, padding):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.replication_pad2d.default(input, padding)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 4D or 5D
+            ((2, 2, 2, 2), (1, 1, 1, 1, 1, 1)),
+            ((1, 2, 3, 4), (3, 2, 2, 1, 1, 1)),
+            ((2, 2, 3, 3, 4), (3, 3, 2, 1, 1, 2)),
+            ((2, 3, 4, 5, 6), (4, 3, 2, 1, 1, 0)),
+        ]
+    )
+    def test_replication_pad3d(self, shape, padding):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.replication_pad3d.default(input, padding)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+
+class TestCircularPadConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 2D or 3D
+            ((3, 3), (1, 1)),
+            ((3, 3), (2, 2)),
+            ((2, 2, 2), (1, 1)),
+            ((2, 2, 4), (2, 3)),
+        ]
+    )
+    def test_circular_pad1d(self, shape, pad):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten._pad_circular.default(input, pad)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 3D or 4D
+            ((2, 2, 2), (1, 1, 1, 1)),
+            ((1, 2, 4), (2, 2, 1, 1)),
+            ((2, 2, 3, 3), (1, 1, 2, 2)),
+            ((2, 3, 4, 5), (4, 3, 0, 1)),
+        ]
+    )
+    def test_circular_pad2d(self, shape, pad):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten._pad_circular.default(input, pad)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+    @parameterized.expand(
+        [
+            # Per pytorch doc, the input should be 4D or 5D
+            ((2, 2, 2, 2), (1, 1, 1, 1, 1, 1)),
+            ((1, 2, 3, 4), (3, 2, 2, 1, 1, 1)),
+            ((2, 2, 3, 3, 4), (3, 3, 2, 1, 1, 2)),
+            ((2, 3, 4, 5, 6), (4, 3, 2, 1, 1, 0)),
+        ]
+    )
+    def test_circular_pad3d(self, shape, pad):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten._pad_circular.default(input, pad)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+
+class TestPadConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ((3, 3), (2, 2), "constant"),
+            ((2, 2, 4), (2, 3, 1, 0), "reflect"),
+            ((1, 2, 3, 4), (3, 2, 2, 1, 1, 1), "replicate"),
+            ((2, 3, 4, 5), (3, 2, 1, 0), "circular"),
+        ]
+    )
+    def test_pad(self, shape, pad, mode, value=None):
+        class TestModule(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.pad.default(input, pad, mode, value)
+
+        input = [torch.randn(shape)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()