Repository: pytorch/TensorRT

Issue:
↔ [Converter] Add support for `torch.ops.aten.pixel_unshuffle.default` in Torch-TensorRT
##

- **Function Schema**: https://github.com/pytorch/pytorch/blob/6f74b7607207054c1f5331183ab725640c99fe8d/aten/src/ATen/native/native_functions.yaml#L4483

- **Original PyTorch API**: https://pytorch.org/docs/stable/generated/torch.nn.PixelUnshuffle.html

## Additional context

```py
import torch
import torch.nn as nn
import torch_tensorrt


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.m = nn.PixelUnshuffle(3)

    def forward(self, x):
        return self.m(x)


model = MyModule().eval().cuda().half()

inputs = [
    torch.randn((1, 1, 12, 12), dtype=torch.half, device="cuda"),
]

optimized_model = torch_tensorrt.compile(
    model,
    ir="dynamo",
    inputs=inputs,
    enabled_precisions={torch.half},
    debug=True,
    min_block_size=1,
)
```

```
DEBUG:torch_tensorrt.dynamo._compiler:Input graph: graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %pixel_unshuffle : [num_users=1] = call_function[target=torch.ops.aten.pixel_unshuffle.default](args = (%arg0_1, 3), kwargs = {})
    return (pixel_unshuffle,)
DEBUG:torch_tensorrt.dynamo.lowering.passes.constant_folding:Graph after constant folding:
graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %pixel_unshuffle : [num_users=1] = call_function[target=torch.ops.aten.pixel_unshuffle.default](args = (%arg0_1, 3), kwargs = {})
    return (pixel_unshuffle,)
DEBUG:torch_tensorrt.dynamo._compiler:Lowered Input graph: graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %pixel_unshuffle : [num_users=1] = call_function[target=torch.ops.aten.pixel_unshuffle.default](args = (%arg0_1, 3), kwargs = {})
    return (pixel_unshuffle,)
INFO:torch_tensorrt.dynamo._compiler:Compilation Settings: CompilationSettings(precision=torch.float16, debug=True, workspace_size=0, min_block_size=1, torch_executed_ops=set(), pass_through_build_failures=False, max_aux_streams=None, version_compatible=False, optimization_level=None, use_python_runtime=False, truncate_long_and_double=False, use_fast_partitioner=True, enable_experimental_decompositions=False, device=Device(type=DeviceType.GPU, gpu_id=0), require_full_compilation=False, disable_tf32=False, sparse_weights=False, refit=False, engine_capability=<EngineCapability.DEFAULT: 0>, num_avg_timing_iters=1, dla_sram_size=1048576, dla_local_dram_size=1073741824, dla_global_dram_size=536870912, dryrun=False, hardware_compatible=False, output_format='exported_program')

DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
Supported Nodes:

DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
Unsupported or Excluded Nodes:
- torch.ops.aten.pixel_unshuffle.default + Operator Count: 1

WARNING:torch_tensorrt.dynamo._compiler:0 supported operations detected in subgraph containing 1 computational nodes. Skipping this subgraph, since min_block_size was detected to be 1
```


Discussion:

Resolving patch:
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -2319,6 +2319,29 @@ def aten_ops_pixel_shuffle(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.pixel_unshuffle.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_pixel_unshuffle(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.shuffle.pixel_unshuffle(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
 @enforce_tensor_types({0: (TRTTensor,)})
 @dynamo_tensorrt_converter(torch.ops.aten.argmax.default)
 def aten_ops_argmax(
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/shuffle.py b/py/torch_tensorrt/dynamo/conversion/impl/shuffle.py
--- a/py/torch_tensorrt/dynamo/conversion/impl/shuffle.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/shuffle.py
@@ -60,3 +60,47 @@ def pixel_shuffle(
         permuted_tensor,
         shape[:-3] + (out_channels, out_height, out_width),
     )
+
+
+def pixel_unshuffle(
+    ctx: ConversionContext,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    downscale_factor: int,
+) -> TRTTensor:
+    shape = input.shape
+    in_channels, in_height, in_width = shape[-3:]
+    out_channels = in_channels * (downscale_factor**2)
+    out_height = in_height // downscale_factor
+    out_width = in_width // downscale_factor
+    new_shape = shape[:-3] + (
+        in_channels,
+        out_height,
+        downscale_factor,
+        out_width,
+        downscale_factor,
+    )
+    reshaped_tensor = reshape(
+        ctx, target, source_ir, f"{name}_reshape1", input, new_shape
+    )
+    rank = len(new_shape)
+    permute_shape = tuple(range(rank - 5)) + (
+        rank - 5,  # in_channels
+        rank - 3,  # downscale_factor
+        rank - 1,  # downscale_factor
+        rank - 4,  # out_height
+        rank - 2,  # out_width
+    )
+    permuted_tensor = impl.permutation.permute(
+        ctx, target, source_ir, f"{name}_permute", reshaped_tensor, permute_shape
+    )
+    return reshape(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_reshape2",
+        permuted_tensor,
+        shape[:-3] + (out_channels, out_height, out_width),
+    )


Test patch:
diff --git a/tests/py/dynamo/conversion/test_pixel_unshuffle_aten.py b/tests/py/dynamo/conversion/test_pixel_unshuffle_aten.py
new file mode 100644
--- /dev/null
+++ b/tests/py/dynamo/conversion/test_pixel_unshuffle_aten.py
@@ -0,0 +1,29 @@
+import torch
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestPixelUnshuffleConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ((1, 1, 1), 1),
+            ((1, 1, 12, 12), 3),
+            ((2, 3, 4, 25, 30), 5),
+        ]
+    )
+    def test_pixel_unshuffle(self, shape, downscale_factor):
+        class PixelUnshuffle(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.pixel_unshuffle.default(x, downscale_factor)
+
+        inputs = [torch.randn(shape)]
+        self.run_test(
+            PixelUnshuffle(),
+            inputs,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()