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

This PR will fix the following issue:
[Add lowering of torch.aten.pixel_unshuffle op to linalg
dialect](#4260)

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>
  }
}
```

Author: alaa-ali

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -8668,6 +8668,30 @@ def Torch_AtenPixelShuffleOp : Torch_Op<"aten.pixel_shuffle", [
   }];
 }
 
+def Torch_AtenPixelUnshuffleOp : Torch_Op<"aten.pixel_unshuffle", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::pixel_unshuffle : (Tensor, int) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    Torch_IntType:$downscale_factor
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenPixelUnshuffleOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenPixelUnshuffleOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+}
+
 def Torch_AtenChannelShuffleOp : Torch_Op<"aten.channel_shuffle", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -7613,6 +7613,56 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %15 = torch.aten.append.t %6, %14 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
 "    return %6 : !torch.list<int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.pixel_unshuffle\"(%arg0: !torch.list<int>, %arg1: !torch.int) -> !torch.list<int> {\n"
+"    %int1 = torch.constant.int 1\n"
+"    %int-3 = torch.constant.int -3\n"
+"    %str = torch.constant.str \"AssertionError: width  must be divisible by downscale_factor in pixel_unshuffle\"\n"
+"    %int-1 = torch.constant.int -1\n"
+"    %str_0 = torch.constant.str \"AssertionError: height  must be divisible by downscale_factor in pixel_unshuffle\"\n"
+"    %int-2 = torch.constant.int -2\n"
+"    %none = torch.constant.none\n"
+"    %str_1 = torch.constant.str \"AssertionError: input must be at least rank-3 in pixel_unshuffle\"\n"
+"    %int3 = torch.constant.int 3\n"
+"    %int0 = torch.constant.int 0\n"
+"    %0 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
+"    %1 = torch.aten.ge.int %0, %int3 : !torch.int, !torch.int -> !torch.bool\n"
+"    torch.prim.If %1 -> () {\n"
+"      torch.prim.If.yield\n"
+"    } else {\n"
+"      torch.prim.RaiseException %str_1, %none : !torch.str, !torch.none\n"
+"      torch.prim.If.yield\n"
+"    }\n"
+"    %2 = torch.aten.mul.int %arg1, %arg1 : !torch.int, !torch.int -> !torch.int\n"
+"    %3 = torch.aten.__getitem__.t %arg0, %int-2 : !torch.list<int>, !torch.int -> !torch.int\n"
+"    %4 = torch.aten.remainder.int %3, %arg1 : !torch.int, !torch.int -> !torch.int\n"
+"    %5 = torch.aten.eq.int %4, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"    torch.prim.If %5 -> () {\n"
+"      torch.prim.If.yield\n"
+"    } else {\n"
+"      torch.prim.RaiseException %str_0, %none : !torch.str, !torch.none\n"
+"      torch.prim.If.yield\n"
+"    }\n"
+"    %6 = torch.aten.__getitem__.t %arg0, %int-1 : !torch.list<int>, !torch.int -> !torch.int\n"
+"    %7 = torch.aten.remainder.int %6, %arg1 : !torch.int, !torch.int -> !torch.int\n"
+"    %8 = torch.aten.eq.int %7, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"    torch.prim.If %8 -> () {\n"
+"      torch.prim.If.yield\n"
+"    } else {\n"
+"      torch.prim.RaiseException %str, %none : !torch.str, !torch.none\n"
+"      torch.prim.If.yield\n"
+"    }\n"
+"    %9 = torch.aten.slice.t %arg0, %int0, %int-3, %int1 : !torch.list<int>, !torch.int, !torch.int, !torch.int -> !torch.list<int>\n"
+"    %10 = torch.aten.__getitem__.t %arg0, %int-3 : !torch.list<int>, !torch.int -> !torch.int\n"
+"    %11 = torch.aten.mul.int %10, %2 : !torch.int, !torch.int -> !torch.int\n"
+"    %12 = torch.aten.append.t %9, %11 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
+"    %13 = torch.aten.__getitem__.t %arg0, %int-2 : !torch.list<int>, !torch.int -> !torch.int\n"
+"    %14 = torch.aten.floordiv.int %13, %arg1 : !torch.int, !torch.int -> !torch.int\n"
+"    %15 = torch.aten.append.t %9, %14 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
+"    %16 = torch.aten.__getitem__.t %arg0, %int-1 : !torch.list<int>, !torch.int -> !torch.int\n"
+"    %17 = torch.aten.floordiv.int %16, %arg1 : !torch.int, !torch.int -> !torch.int\n"
+"    %18 = torch.aten.append.t %9, %17 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
+"    return %9 : !torch.list<int>\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten.channel_shuffle\"(%arg0: !torch.list<int>, %arg1: !torch.int) -> !torch.list<int> {\n"
 "    %none = torch.constant.none\n"
 "    %str = torch.constant.str \"AssertionError: input must be at least rank-3 in channel_shuffle\"\n"
@@ -12411,6 +12461,10 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
 "    return %0#1 : !torch.int\n"
 "  }\n"
+"  func.func @\"__torch_mlir_dtype_fn.aten.pixel_unshuffle\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.int) -> !torch.int {\n"
+"    %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
+"    return %0#1 : !torch.int\n"
+"  }\n"
 "  func.func @\"__torch_mlir_dtype_fn.aten.channel_shuffle\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.int) -> !torch.int {\n"
 "    %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
 "    return %0#1 : !torch.int\n"

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -3536,30 +3536,6 @@ class DecomposeAten_LinalgDetOp : public OpRewritePattern<Aten_LinalgDetOp> {
 };
 } // namespace
 
-namespace { // Start of rearrangement ops utility functions
-// Extracts shape as vector of int64_t from vector of Value
-SmallVector<int64_t> getIntShapeFromValues(ArrayRef<Value> vals) {
-  SmallVector<int64_t> shape;
-  shape.reserve(vals.size());
-  for (Value 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;
-}
-
-// Converts a vector of Value (shape dimensions) into a ValueTensorType
-ValueTensorType getTypeFromShape(ArrayRef<Value> vals, Type inOptionalDType) {
-  SmallVector<int64_t> intShape = getIntShapeFromValues(vals);
-  return ValueTensorType::get(vals[0].getContext(), llvm::ArrayRef(intShape),
-                              inOptionalDType);
-}
-} // namespace
-
 // Decompose aten.pixel_shuffle into: prims.split_dim, aten.permute, and
 // prims.collapse operations.
 //
@@ -3609,18 +3585,9 @@ class DecomposeAtenPixelShuffleOp
 
     auto nLeadingDims = inRank - 3;
 
-    // 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);
-    };
-
-    auto inC = getDimSize(inRank - 3);
-    auto inH = getDimSize(inRank - 2);
-    auto inW = getDimSize(inRank - 1);
+    auto inC = getTensorDimSize(rewriter, inValue, inRank - 3);
+    auto inH = getTensorDimSize(rewriter, inValue, inRank - 2);
+    auto inW = getTensorDimSize(rewriter, inValue, inRank - 1);
 
     auto factor = op.getUpscaleFactor();
 
@@ -3678,23 +3645,26 @@ class DecomposeAtenPixelShuffleOp
     auto partiallyExpanded =
         rewriter
             .create<PrimsSplitDimOp>(
-                loc, getTypeFromShape(partiallyExpandedShape, inOptionalDType),
+                loc,
+                getTensorTypeFromShapeValues(partiallyExpandedShape,
+                                             inOptionalDType),
                 inValue, dimensionConstants[nLeadingDims], outC)
             .getResult();
 
     // Split new dimension factorSquared -> (factor, factor)
     auto fullyExpanded = rewriter.create<PrimsSplitDimOp>(
-        loc, getTypeFromShape(prePermuteShape, inOptionalDType),
+        loc, getTensorTypeFromShapeValues(prePermuteShape, inOptionalDType),
         partiallyExpanded, dimensionConstants[nLeadingDims + 1], factor);
 
     // Perform the permutation
     auto permuted = rewriter.create<AtenPermuteOp>(
-        loc, getTypeFromShape(postPermuteShape, inOptionalDType), fullyExpanded,
-        permuteDimsOrder);
+        loc, getTensorTypeFromShapeValues(postPermuteShape, inOptionalDType),
+        fullyExpanded, permuteDimsOrder);
 
     // Collapse final 2 dimension
     auto partiallyCollapsed = rewriter.create<PrimsCollapseOp>(
-        loc, getTypeFromShape(partiallyCollapsedShape, inOptionalDType),
+        loc,
+        getTensorTypeFromShapeValues(partiallyCollapsedShape, inOptionalDType),
         permuted, dimensionConstants[nLeadingDims + 3],
         dimensionConstants[nLeadingDims + 4]);
 
@@ -3709,6 +3679,142 @@ class DecomposeAtenPixelShuffleOp
 };
 } // namespace
 
+// Decompose aten.pixel_unshuffle into: prims.split_dim, aten.permute, and
+// prims.collapse operations.
+//
+// We want to do the exact opposite of aten.pixel_shuffle
+//
+// 'r' is referred to as the 'downscale factor' or just 'factor' below.
+//
+// If input is a tensor of shape
+//     (*leading_dims, C, H*r, W*r),
+//
+// where leading_dims is of size N, then
+//    X = pixel_unshuffle(input, downscale_factor)
+//
+// gets replaced with
+//    X = input.split_dim(...)  # shape (*leading_dims, C, H, r, W*r)
+//    X = X.split_dim(...)      # shape (*leading_dims, C, H, r, W, r)
+//    X = X.permute(0, ..., N, N+2, N+4, N+1, N+3)
+//                              # shape (*leading_dims, C, r, r, H, W)
+//    X = X.collapse(...)       # shape (*leading_dims, C*r*r, H, W)
+//
+namespace {
+class DecomposeAtenPixelUnshuffleOp
+    : public OpRewritePattern<AtenPixelUnshuffleOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenPixelUnshuffleOp op,
+                                PatternRewriter &rewriter) const override {
+
+    Location loc = op.getLoc();
+    Value inValue = op.getSelf();
+    auto inType = cast<BaseTensorType>(inValue.getType());
+    auto maybeSizes = inType.getOptionalSizes();
+    if (!maybeSizes) {
+      return rewriter.notifyMatchFailure(
+          op, "Expected input tensor to have known rank.");
+    }
+    auto inShape = maybeSizes.value();
+    auto inRank = inShape.size();
+
+    // The input tensor must have at least 3 dimensions: (1) the channel
+    // dimension which gets bigger by 'factor*factor', (2) the H channel which
+    // gets smaller by 'factor' and (3) the W channel which get smaller by
+    // 'factor'. The total number of dimensions is 3 + N, where N is the number
+    // of leading dimensions, and N >= 0 so the input must have rank at least 3.
+    if (inRank < 3)
+      return rewriter.notifyMatchFailure(
+          op, "Expected input tensor to have rank greater than 2.");
+
+    const auto inOptionalDType = inType.getOptionalDtype();
+
+    auto nLeadingDims = inRank - 3;
+
+    auto inC = getTensorDimSize(rewriter, inValue, inRank - 3);
+    auto inH = getTensorDimSize(rewriter, inValue, inRank - 2);
+    auto inW = getTensorDimSize(rewriter, inValue, inRank - 1);
+
+    auto factor = op.getDownscaleFactor();
+
+    Value factorSquared =
+        rewriter.createOrFold<AtenMulIntOp>(loc, factor, factor);
+
+    Value outC = rewriter.createOrFold<AtenMulIntOp>(loc, inC, factorSquared);
+
+    Value outH = rewriter.createOrFold<AtenFloordivIntOp>(loc, inH, factor);
+    Value outW = rewriter.createOrFold<AtenFloordivIntOp>(loc, inW, factor);
+
+    SmallVector<Value> dimensionConstants;
+    dimensionConstants.reserve(inRank + 2);
+    for (unsigned i = 0; i < inRank + 2; ++i) {
+      dimensionConstants.push_back(
+          rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(i)));
+    }
+
+    SmallVector<Value> leadingDims;
+    leadingDims.reserve(nLeadingDims);
+    for (unsigned i = 0; i < nLeadingDims; ++i) {
+      Value leadingDimSize = rewriter.createOrFold<AtenSizeIntOp>(
+          loc, inValue, dimensionConstants[i]);
+      leadingDims.push_back(leadingDimSize);
+    }
+
+    SmallVector<Value> prePermuteShape = leadingDims;
+    prePermuteShape.append({inC, outH, factor, outW, factor});
+
+    SmallVector<Value> postPermuteShape = leadingDims;
+    postPermuteShape.append({inC, factor, factor, outH, outW});
+
+    SmallVector<Value> partiallyCollapsedShape = leadingDims;
+    partiallyCollapsedShape.append({inC, factorSquared, outH, outW});
+
+    SmallVector<Value> outShape = leadingDims;
+    outShape.append({outC, outH, outW});
+
+    SmallVector<Value> permutation{dimensionConstants.begin(),
+                                   dimensionConstants.begin() + nLeadingDims};
+    SmallVector<uint64_t> permutationTail{0, 2, 4, 1, 3};
+    for (uint64_t d : permutationTail) {
+      permutation.push_back(dimensionConstants[nLeadingDims + d]);
+    }
+
+    Value permuteDimsOrder = rewriter.create<PrimListConstructOp>(
+        loc, Torch::ListType::get(Torch::IntType::get(op->getContext())),
+        permutation);
+
+    SmallVector<Value> heightSplitShape = leadingDims;
+    heightSplitShape.append({inC, outH, factor, inW});
+
+    // Split input channel inH -> (outH, factor)
+    auto partiallyExpanded =
+        rewriter
+            .create<PrimsSplitDimOp>(
+                loc,
+                getTensorTypeFromShapeValues(heightSplitShape, inOptionalDType),
+                inValue, dimensionConstants[nLeadingDims + 1], outH)
+            .getResult();
+
+    // Split new dimension inW -> (outW, factor)
+    auto fullyExpanded = rewriter.create<PrimsSplitDimOp>(
+        loc, getTensorTypeFromShapeValues(prePermuteShape, inOptionalDType),
+        partiallyExpanded, dimensionConstants[nLeadingDims + 3], outW);
+
+    // Perform the permutation
+    auto permuted = rewriter.create<AtenPermuteOp>(
+        loc, getTensorTypeFromShapeValues(postPermuteShape, inOptionalDType),
+        fullyExpanded, permuteDimsOrder);
+
+    // Collapse final 2 dimensions back to original rank
+    rewriter.replaceOpWithNewOp<PrimsCollapseOp>(
+        op, op.getType(), permuted, dimensionConstants[nLeadingDims],
+        dimensionConstants[nLeadingDims + 2]);
+
+    return success();
+  }
+};
+} // namespace
+
 // Decompose aten.channel_shuffle into: prims.split_dim, aten.permute, and
 // prims.collapse operations.
 //
@@ -3763,23 +3869,14 @@ class DecomposeAtenChannelShuffleOp
 
     auto numOfSpatialDims = inRank - 2;
 
-    // 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 = [&rewriter, &inValue, loc](uint64_t i) -> Value {
-      Value dim =
-          rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(i));
-      return rewriter.createOrFold<AtenSizeIntOp>(loc, inValue, dim);
-    };
-
     // The channel dimension is always the second dimension. PyTorch errors out
     // if the batch dimension (first dimension) is not present. See comment at
     // the top of this class for details.
-    auto inC = getDimSize(1);
+    auto inC = getTensorDimSize(rewriter, inValue, 1);
     SmallVector<Value> inSpatialDims;
     inSpatialDims.reserve(numOfSpatialDims);
     for (unsigned i = 2; i < (2 + numOfSpatialDims); ++i) {
-      inSpatialDims.push_back(getDimSize(i));
+      inSpatialDims.push_back(getTensorDimSize(rewriter, inValue, i));
     }
 
     auto groups = op.getGroups();
@@ -3832,14 +3929,14 @@ class DecomposeAtenChannelShuffleOp
     auto expandedTensor =
         rewriter
             .create<PrimsSplitDimOp>(
-                loc, getTypeFromShape(splitShape, inOptionalDType), inValue,
-                dimC, tempC)
+                loc, getTensorTypeFromShapeValues(splitShape, inOptionalDType),
+                inValue, dimC, tempC)
             .getResult();
 
     // Perform the permutation
     auto permuted = rewriter.create<AtenPermuteOp>(
-        loc, getTypeFromShape(permuteShape, inOptionalDType), expandedTensor,
-        permuteDimsOrder);
+        loc, getTensorTypeFromShapeValues(permuteShape, inOptionalDType),
+        expandedTensor, permuteDimsOrder);
 
     // Collapse (C, groups) back into a single channel dimension
     rewriter.replaceOpWithNewOp<PrimsCollapseOp>(op, op.getType(), permuted,
@@ -12909,6 +13006,7 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposeAtenRenormOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenLinalgCrossOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenPixelShuffleOp>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenPixelUnshuffleOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenChannelShuffleOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenTOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAten_LogSoftmaxBackwardDataOp>(

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -421,6 +421,7 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context,
   target.addIllegalOp<Aten_LinalgDetOp>();
   target.addIllegalOp<AtenLinalgSlogdetOp>();
   target.addIllegalOp<AtenPixelShuffleOp>();
+  target.addIllegalOp<AtenPixelUnshuffleOp>();
   target.addIllegalOp<AtenChannelShuffleOp>();
   target.addIllegalOp<AtenTOp>();
   target.addIllegalOp<Aten_LogSoftmaxBackwardDataOp>();

lib/Dialect/Torch/Utils/Utils.cpp

@@ -327,7 +327,8 @@ bool Torch::isViewLikeOp(Operation *op) {
       AtenNarrowTensorOp, AtenToDeviceOp, PrimsSqueezeOp, AtenMovedimIntOp,
       PrimsViewOfOp, AtenRealOp, AtenImagOp, PrimsSplitDimOp,
       AtenViewAsComplexOp, AtenViewAsRealOp, AtenPixelShuffleOp,
-      AtenChannelShuffleOp, AtenDiagonalOp, AtenUnfoldOp>(op);
+      AtenPixelUnshuffleOp, AtenChannelShuffleOp, AtenDiagonalOp, AtenUnfoldOp>(
+      op);
 }
 
 Value Torch::getConstantWithGivenDtypeAndValue(PatternRewriter &rewriter,