feat: support pad dynamo converter

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -2151,3 +2151,28 @@ def aten_ops_circular_pad(
         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),
+    )

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

@@ -20,7 +20,7 @@ def constant_padNd(
     name: str,
     input: TRTTensor,
     pad: Sequence[int],
-    value: int = 0,
+    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."
@@ -269,3 +269,35 @@ def circular_padNd(
         raise RuntimeError(
             f"We currently only support for padding 1D, 2D, and 3D, but got {padding_dims}D"
         )
+
+
+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}'
+        )

tests/py/dynamo/conversion/test_pad_aten.py

@@ -216,5 +216,26 @@ def forward(self, 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()