[TorchToLinalg] Add lowering of torch.aten.pixel_unshuffle op #4278

alaa-ali · 2025-07-18T04:10:48Z

This PR will fix the following issue:
Add lowering of torch.aten.pixel_unshuffle op to linalg dialect

This code snippet can reproduce the issue:

func.func @pixel_unshuffle(%arg0: !torch.vtensor<[1,8,4,4],f32>) -> !torch.vtensor<[1,32,2,2],f32> attributes {torch.assume_strict_symbolic_shapes} {
  %int2 = torch.constant.int 2
  %0 = torch.aten.pixel_unshuffle %arg0, %int2 : !torch.vtensor<[1,8,4,4],f32>, !torch.int -> !torch.vtensor<[1,32,2,2],f32>
  return %0 : !torch.vtensor<[1,32,2,2],f32>
}

The decomposition is based on this specification:
https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.pixel_unshuffle.html
and PyTorch implementation could be found in main/aten/src/ATen/native/PixelShuffle.cpp:
https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/PixelShuffle.cpp

With code changes, torch.aten.pixel_unshuffle will be lowered to the following:

module {
  func.func @main(%arg0: !torch.vtensor<[1,8,4,4],f32>) -> !torch.vtensor<[1,32,2,2],f32> attributes {torch.assume_strict_symbolic_shapes} {
    %int2 = torch.constant.int 2
    %int0 = torch.constant.int 0
    %int1 = torch.constant.int 1
    %int3 = torch.constant.int 3
    %int4 = torch.constant.int 4
    %int5 = torch.constant.int 5
    %0 = torch.prim.ListConstruct %int0, %int1, %int3, %int5, %int2, %int4 : (!torch.int, !torch.int, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
    %1 = torch.prims.split_dim %arg0, %int2, %int2 : !torch.vtensor<[1,8,4,4],f32>, !torch.int, !torch.int -> !torch.vtensor<[1,8,2,2,4],f32>
    %2 = torch.prims.split_dim %1, %int4, %int2 : !torch.vtensor<[1,8,2,2,4],f32>, !torch.int, !torch.int -> !torch.vtensor<[1,8,2,2,2,2],f32>
    %3 = torch.aten.permute %2, %0 : !torch.vtensor<[1,8,2,2,2,2],f32>, !torch.list<int> -> !torch.vtensor<[1,8,2,2,2,2],f32>
    %4 = torch.prims.collapse %3, %int2, %int3 : !torch.vtensor<[1,8,2,2,2,2],f32>, !torch.int, !torch.int -> !torch.vtensor<[1,8,4,2,2],f32>
    %5 = torch.prims.collapse %4, %int1, %int2 : !torch.vtensor<[1,8,4,2,2],f32>, !torch.int, !torch.int -> !torch.vtensor<[1,32,2,2],f32>
    return %5 : !torch.vtensor<[1,32,2,2],f32>
  }
}

alaa-ali · 2025-07-18T19:50:45Z

@silvasean @rsuderman @zjgarvey @penguin-wwy

ivangarcia44

In general it looks good to me. My main feedback is:

Can this be generalized for ND if the pixel shuffle operator allows ND?
What is the expected batch dimension behavior? Can this be tested?

Thanks!

projects/pt1/python/torch_mlir_e2e_test/test_suite/basic.py

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

ivangarcia44 · 2025-07-24T19:05:04Z

Are there tests for batch batch dimension omission and more than one batch dimensions? Thanks

alaa-ali · 2025-08-01T20:56:21Z

Are there tests for batch batch dimension omission and more than one batch dimensions? Thanks

This has been captured in e2e tests. Thanks for your feedback.

ivangarcia44 · 2025-08-01T20:58:33Z

LGTM

This PR has tests for rank 3, 4, and 5 which cover the case of no batch dimension, 1 batch dimension and 2 batch dimensions. All my concerns are addressed in the PR. Looks good to me.

alaa-ali · 2025-08-04T14:38:24Z

Hi everyone, a kind reminder to provide feedback, please. This PR adds support of torch.aten.pixel_unshuffle op.
Thank you

@rsuderman @zjgarvey @penguin-wwy @newling @sahas3 @ramiro050 @qedawkins @vivekkhandelwal1

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

sahas3 · 2025-08-06T12:37:54Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+//     (*leading_dims, C, H*r, W*r),
+//
+// where leading_dims is of size N, then
+//    X = pixel_unshuffle(input, downscale_factor)


Using r here instead of downscale_factor will be better for consistency. You can also move line 3731 mentioning r is the downscale_factor above to help readability.

sahas3 · 2025-08-06T12:38:52Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+//    X = X.collapse(...)       # shape (*leading_dims, C, r*r, H, W)
+//    X = X.collapse(...)       # shape (*leading_dims, C*r*r, H, W)


Do we need two collapses -- isn't collapsing directly to C*r*r sufficient?

sahas3 · 2025-08-06T13:39:47Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+    auto getTypeFromShape = [inOptionalDType](auto &&vals) {
+      // Get a vector of integers from a vector of Values.
+      auto getIntShape = [](auto &&vals) {
+        SmallVector<int64_t> shape;
+        shape.reserve(vals.size());
+        for (auto v : vals) {
+          int64_t cst_val;
+          if (matchPattern(v, m_TorchConstantInt(&cst_val))) {
+            shape.push_back(cst_val);
+          } else {
+            shape.push_back(kUnknownSize);
+          }
+        }
+        return shape;
+      };
+
+      const auto intShape = getIntShape(vals);
+      return ValueTensorType::get(vals[0].getContext(),
+                                  llvm::ArrayRef(intShape), inOptionalDType);
+    };


These methods were added as utilities in #4259. Once that is merged, can you update your code to reuse the utilities?

sahas3 · 2025-08-06T13:43:16Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+    // Get the size of the dimension 'i'. Note the use of 'createOrFold' instead
+    // of 'create': if the dimension size is known, then the AtenSizeIntOp is
+    // folded to a ConstantOp.
+    auto getDimSize = [&](uint64_t i) -> Value {
+      Value dim =
+          rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(i));
+      return rewriter.createOrFold<AtenSizeIntOp>(loc, inValue, dim);
+    };


This is shared with #4259 as well. It'll be good to move this into an utility method and probably move all these utilities to lib/Conversion/Utils/Utils.cpp as the shared location to be used elsewhere in the code base too.

sahas3 · 2025-08-06T13:47:07Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+    SmallVector<Value> partiallyExpandedShape = leadingDims;
+    partiallyExpandedShape.append({inC, outH, factor, inW});


Can you move this before it's use also rename to heightSplitShape for readability?

alaa-ali added 4 commits July 17, 2025 23:18

add lowering torch.aten.pixel_unshuffle op to linalg

9068bbc

Merge branch 'llvm:main' into pixel.unshuffle

be02c1f

fix decompose-complex-ops.mlir filecheck

9b44fc7

remove pixel_unshuffle e2e tests from fx_importer_stablehlo failed tests

f409948

alaa-ali added 2 commits July 22, 2025 00:07

Merge branch 'llvm:main' into pixel.unshuffle

111dfda

add pixel_unshuffle_fulldynamic test

82215f3

ivangarcia44 reviewed Jul 23, 2025

View reviewed changes

projects/pt1/python/torch_mlir_e2e_test/test_suite/basic.py Show resolved Hide resolved

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py Show resolved Hide resolved

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Show resolved Hide resolved

sahas3 reviewed Aug 6, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[TorchToLinalg] Add lowering of torch.aten.pixel_unshuffle op #4278

[TorchToLinalg] Add lowering of torch.aten.pixel_unshuffle op #4278

Uh oh!

alaa-ali commented Jul 18, 2025 •

edited

Loading

Uh oh!

alaa-ali commented Jul 18, 2025

Uh oh!

ivangarcia44 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ivangarcia44 commented Jul 24, 2025

Uh oh!

alaa-ali commented Aug 1, 2025

Uh oh!

ivangarcia44 commented Aug 1, 2025

Uh oh!

alaa-ali commented Aug 4, 2025

Uh oh!

Uh oh!

sahas3 Aug 6, 2025

Uh oh!

sahas3 Aug 6, 2025

Uh oh!

sahas3 Aug 6, 2025

Uh oh!

sahas3 Aug 6, 2025

Uh oh!

sahas3 Aug 6, 2025

Uh oh!

Uh oh!

		// X = X.collapse(...) # shape (leading_dims, C, rr, H, W)
		// X = X.collapse(...) # shape (leading_dims, Cr*r, H, W)

		SmallVector<Value> partiallyExpandedShape = leadingDims;
		partiallyExpandedShape.append({inC, outH, factor, inW});

[TorchToLinalg] Add lowering of torch.aten.pixel_unshuffle op #4278

Are you sure you want to change the base?

[TorchToLinalg] Add lowering of torch.aten.pixel_unshuffle op #4278

Uh oh!

Conversation

alaa-ali commented Jul 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

alaa-ali commented Jul 18, 2025

Uh oh!

ivangarcia44 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ivangarcia44 commented Jul 24, 2025

Uh oh!

alaa-ali commented Aug 1, 2025

Uh oh!

ivangarcia44 commented Aug 1, 2025

Uh oh!

alaa-ali commented Aug 4, 2025

Uh oh!

Uh oh!

sahas3 Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

sahas3 Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

sahas3 Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

sahas3 Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

sahas3 Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

alaa-ali commented Jul 18, 2025 •

edited

Loading