feat: support circular padding dynamo converters for 1D, 2D, and 3D

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -2090,3 +2090,26 @@ def aten_ops_replication_pad(
         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],
+    )

py/torch_tensorrt/dynamo/conversion/impl/pad.py

@@ -210,3 +210,62 @@ def replication_padNd(
         raise RuntimeError(
             f"We currently only support for padding 1D, 2D, and 3D, but got {padding_dims}D"
         )
+
+
+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."
+
+    padding_dims = len(pad) // 2
+
+    if padding_dims == 1 or padding_dims == 2 or padding_dims == 3:
+        for i in range(padding_dims):
+            dim = -1 - i
+            pre_pad, post_pad = pad[2 * i], pad[2 * i + 1]
+            pre_pad_tensor = impl.slice.slice_op(
+                ctx,
+                target,
+                source_ir,
+                f"{name}_slice_pre{i}",
+                input,
+                dim=dim,
+                start=input.shape[dim] - pre_pad,
+                stop=input.shape[dim],
+                step=1,
+            )
+
+            post_pad_tensor = impl.slice.slice_op(
+                ctx,
+                target,
+                source_ir,
+                f"{name}_slice_post{i}",
+                input,
+                dim=dim,
+                start=0,
+                stop=post_pad,
+                step=1,
+            )
+
+            output = impl.cat.cat(
+                ctx,
+                target,
+                source_ir,
+                f"{name}_concat_dim{dim}",
+                input=(pre_pad_tensor, input, post_pad_tensor),
+                dim=dim,
+            )
+            input = output
+
+        return output
+
+    else:
+        raise RuntimeError(
+            f"We currently only support for padding 1D, 2D, and 3D, but got {padding_dims}D"
+        )

tests/py/dynamo/conversion/test_pad_aten.py

@@ -154,5 +154,67 @@ def forward(self, 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,
+        )
+
+
 if __name__ == "__main__":
     run_tests()