[mlir][linalg] Add support for lowering pack with outer_dims_perm.

[hanchung]

Simplify the logic in computePackingMetadata

[hanchung]

rebase on https://reviews.llvm.org/D148867

[qcolombet]

LGTM, please add a test for the lowering to "simple" pad.

mlir/test/Dialect/Linalg/transform-lower-pack.mlir

249

Could you add a test with the lowering to a "simple" pad?
I.e., outer dimensions should be 1s and inner_dim_pos should not perform any transpose.
In other words, shuffling dims of size one should still lower as a pad with no transpose.

Use func.func @pack_as_pad as template and just add an outer_dims field.

[hanchung]

add a test and rebase to main

[chelini]

thanks!

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

497 ↗

(On Diff #516195)

nit: obained -> obtained

[hanchung]

fix typo

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

 /// ```
 FailureOr<CollapseShapeRankReducingSliceSimplificationInfo>
 getSimplifyCollapseShapeWithRankReducingSliceInfo(
     RankedTensorType sourceType,
     ArrayRef<ReassociationIndices> reassociationIndices);
 
 struct PackingMetadata {
   SmallVector<int64_t> insertPositions;
+  SmallVector<int64_t> outerPositions;
   SmallVector<ReassociationIndices> reassociations;
 };
 
 /// Given a vector of `positions` indices representing desired packing insertion
 /// points into a target vector (i.e. pack/unpack.inner_dim_pos), compute the
 /// final positions in the target shape as well as the reshape reassociations.
 // Note: This should not be called with a large positions array (or the
 // implementation needs to be updated to use an N.log N sort instead of
 // repeated N^2 counts).
 PackingMetadata computePackingMetadata(int64_t packedRank,
                                        ArrayRef<int64_t> innerDimPos);
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_RESHAPEOPSUTILS_H

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

struct LowerPackResult {
   tensor::ExpandShapeOp expandShapeOp;
   linalg::TransposeOp transposeOp;
 };
 
 /// Rewrite pack as pad + reshape + transpose.
 static FailureOr<LowerPackResult> lowerPack(RewriterBase &rewriter,
                                             tensor::PackOp packOp) {
   // 1. Filter out NYI cases.
-  if (!packOp.getOuterDimsPerm().empty())
-    return rewriter.notifyMatchFailure(packOp, "outer dims perm NYI");
-
   auto packedTensorType =
       packOp->getResultTypes().front().cast<RankedTensorType>();
   if (!packedTensorType.hasStaticShape()) {
     return rewriter.notifyMatchFailure(
         packOp,
         "non-static shape NYI, needs a more powerful tensor.expand_shape op");
   }
 
   Location loc = packOp->getLoc();
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(packOp);
 
-  // 2. Compute the permutation vector to move the last `numPackedDims` into the
-  // `innerPosDims` of a shape of rank `packedRank`.
+  // 2. Compute the permutation vector to shuffle packed shape into the shape
+  // before any outer or inner permutations have been applied. The permutation
+  // can be obained from two permutations:
+  //   a) Compute the permutation vector to move the last `numPackedDims` into
+  //      the `innerPosDims` of a shape of rank `packedRank`.
+  //   b) Compute the permutation vector to move outer dims if the pack op
+  //      has outer_dims_perm.
+  // Apply (b) permutation on (a) permutation to get the final permutation.
   int64_t numPackedDims = packOp.getInnerDimsPos().size();
   int64_t packedRank = packedTensorType.getRank();
   auto lastDims = llvm::to_vector(
       llvm::seq<int64_t>(packedRank - numPackedDims, packedRank));
   PackingMetadata packingMetadata = computePackingMetadata(
       packedTensorType.getRank(), packOp.getInnerDimsPos());
-  SmallVector<int64_t> lastDimsToInsertPositionsPerm = computePermutationVector(
+  SmallVector<int64_t> innerPositionsPerm = computePermutationVector(
       packedRank, lastDims, packingMetadata.insertPositions);
 
+  SmallVector<int64_t> outerPos = packingMetadata.outerPositions;
+  ArrayRef<int64_t> outerPerm = packOp.getOuterDimsPerm();
+  if (!outerPerm.empty())
+    applyPermutationToVector(outerPos, outerPerm);
+  SmallVector<int64_t> outerPositionPerm = computePermutationVector(
+      packedRank, packingMetadata.outerPositions, outerPos);
+
+  SmallVector<int64_t> packedToStripMinedShapePerm = innerPositionsPerm;
+  applyPermutationToVector(packedToStripMinedShapePerm, outerPositionPerm);
+
   // 3. Compute the stripMinedShape: this is the packed shape before any outer
   // or inner permutations have been applied.
   SmallVector<int64_t> stripMinedShape(packedTensorType.getShape());
-  applyPermutationToVector(stripMinedShape, lastDimsToInsertPositionsPerm);
+  applyPermutationToVector(stripMinedShape, packedToStripMinedShapePerm);
 
   // 4. Pad the source of packOp to a shape we can expand into stripMinedShape.
   RankedTensorType collapsed = tensor::CollapseShapeOp::inferCollapsedType(
       RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape),
       packingMetadata.reassociations);
   Value paddingValue = packOp.getPaddingValue();
   if (!paddingValue) {
     paddingValue = rewriter.create<arith::ConstantOp>(
         loc, rewriter.getZeroAttr(getElementTypeOrSelf(collapsed)));
   }
   auto padOp =
       tensor::createPadHighOp(collapsed, packOp.getSource(), paddingValue,
                               /*nofold=*/false, loc, rewriter);
 
   LLVM_DEBUG(
       DBGSNL(); DBGSNL(); llvm::interleaveComma(packingMetadata.insertPositions,
                                                 DBGS() << "insertPositions: ");
+      DBGSNL(); llvm::interleaveComma(packingMetadata.outerPositions,
+                                      DBGS() << "outerPositions: ");
       DBGSNL(); llvm::interleaveComma(packedTensorType.getShape(),
                                       DBGS() << "packedShape: ");
       DBGSNL();
-      llvm::interleaveComma(lastDimsToInsertPositionsPerm,
-                            DBGS() << "lastDimsToInsertPositionsPerm: ");
+      llvm::interleaveComma(outerPositionPerm, DBGS() << "outerPositionPerm: ");
+      DBGSNL(); llvm::interleaveComma(innerPositionsPerm,
+                                      DBGS() << "innerPositionsPerm: ");
+      DBGSNL();
+      llvm::interleaveComma(packedToStripMinedShapePerm,
+                            DBGS() << "packedToStripMinedShapePerm: ");
       DBGSNL(); llvm::interleaveComma(
           packingMetadata.reassociations, DBGS() << "reassociations: ",
           [&](ReassociationIndices ri) {
             llvm::interleaveComma(ri, llvm::dbgs() << "|");
           });
       DBGSNL();
       llvm::interleaveComma(stripMinedShape, DBGS() << "stripMinedShape: ");
       DBGSNL(); DBGS() << "collapsed type: " << collapsed; DBGSNL(););
static FailureOr<LowerPackResult> lowerPack(RewriterBase &rewriter,
   }
   // 5. Expand from the padded result to the stripMinedShape.
   auto reshapeOp = rewriter.create<tensor::ExpandShapeOp>(
       loc,
       RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape),
       padOp.getResult(), packingMetadata.reassociations);
 
   // 6. Transpose stripMinedShape to packedShape.
-  SmallVector<int64_t> insertPositionsToLastDimsPerm = computePermutationVector(
-      packedRank, packingMetadata.insertPositions, lastDims);
+  SmallVector<int64_t> transpPerm =
+      invertPermutationVector(packedToStripMinedShapePerm);
   auto transposeOp = rewriter.create<linalg::TransposeOp>(
-      loc, reshapeOp.getResult(), packOp.getDest(),
-      insertPositionsToLastDimsPerm);
+      loc, reshapeOp.getResult(), packOp.getDest(), transpPerm);
 
   LLVM_DEBUG(DBGSNL(); DBGSNL(); DBGSNL();
              DBGS() << "reshape op: " << reshapeOp; DBGSNL();
-             llvm::interleaveComma(insertPositionsToLastDimsPerm,
-                                   DBGS() << "insertPositionsToLastDimsPerm: ");
+             llvm::interleaveComma(transpPerm, DBGS() << "transpPerm: ");
              DBGSNL(); DBGS() << "transpose op: " << transposeOp; DBGSNL(););
 
   // 7. Replace packOp by transposeOp.
   rewriter.replaceOp(packOp, transposeOp->getResults());
 
   return LowerPackResult{padOp, reshapeOp, transposeOp};
 }
 

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

int64_t numInsertedBefore = llvm::count_if(
         innerDimPos, [&pos](int64_t pos2) { return pos > pos2; });
     res.insertPositions.push_back(pos + numInsertedBefore + offset);
   }
 
   DenseSet<int64_t> posSet(res.insertPositions.begin(),
                            res.insertPositions.end());
   res.reassociations.reserve(packedRank);
   for (int64_t i = 1; i <= packedRank; ++i) {
+    res.outerPositions.push_back(i - 1);
     if (!posSet.contains(i)) {
       res.reassociations.push_back(ReassociationIndices{i - 1});
       continue;
     }
     res.reassociations.push_back(ReassociationIndices{i - 1, i});
     ++i;
   }
   return res;
 }

mlir/test/Dialect/Linalg/transform-lower-pack.mlir

^bb1(%module_op: !pdl.operation):
   %unpack = transform.structured.match ops{["tensor.unpack"]} in %module_op
     : (!pdl.operation) -> !transform.op<"tensor.unpack">
   transform.structured.lower_unpack %unpack : (!transform.op<"tensor.unpack">)
     -> (!transform.op<"tensor.empty">,
         !transform.op<"linalg.transpose">,
         !transform.op<"tensor.collapse_shape">,
         !transform.op<"tensor.extract_slice">)
 }
+
+// -----
+
+// CHECK-LABEL: func.func @pack_with_outer_dims_perm(
+func.func @pack_with_outer_dims_perm(%src: tensor<100x200x128x256xi32>,
+                                     %dest: tensor<200x4x16x100x16x32xi32>)
+    -> tensor<200x4x16x100x16x32xi32> {
+  //      CHECK: %[[C0:.*]] = arith.constant 0 : index
+  //      CHECK: tensor.pad {{.*}} low[%[[C0]], %[[C0]], %[[C0]], %[[C0]]]
+  //      CHECK:   : tensor<100x200x128x256xi32> to tensor<100x200x128x256xi32>
+  //      CHECK: tensor.expand_shape %{{.*}} [{{.*}}[0], [1], [2, 3], [4, 5]]
+  // CHECK-SAME:   : tensor<100x200x128x256xi32> into tensor<100x200x4x32x16x16xi32>
+  //      CHECK: linalg.transpose
+  // CHECK-SAME:   ins(%{{.*}} : tensor<100x200x4x32x16x16xi32>)
+  // CHECK-SAME:   outs(%{{.*}} : tensor<200x4x16x100x16x32xi32>)
+  // CHECK-SAME:   permutation = [1, 2, 4, 0, 5, 3]
+  %0 = tensor.pack %src
+    outer_dims_perm = [1, 2, 3, 0]
+    inner_dims_pos = [3, 2]
+    inner_tiles = [16, 32]
+    into %dest : tensor<100x200x128x256xi32> -> tensor<200x4x16x100x16x32xi32>
+  return %0 : tensor<200x4x16x100x16x32xi32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !pdl.operation):
+  %pack = transform.structured.match ops{["tensor.pack"]} in %module_op
+    : (!pdl.operation) -> !transform.op<"tensor.pack">
+  transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">)
+    -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_with_pad_and_outer_dims_perm(
+func.func @pack_with_pad_and_outer_dims_perm(%src: tensor<100x200x127x255xi32>,
+                                             %dest: tensor<200x4x16x100x16x32xi32>)
+    -> tensor<200x4x16x100x16x32xi32> {
+  //      CHECK: %[[C0:.*]] = arith.constant 0 : index
+  //      CHECK: tensor.pad {{.*}} low[%[[C0]], %[[C0]], %[[C0]], %[[C0]]]
+  //      CHECK:   : tensor<100x200x127x255xi32> to tensor<100x200x128x256xi32>
+  //      CHECK: tensor.expand_shape %{{.*}} [{{.*}}[0], [1], [2, 3], [4, 5]]
+  // CHECK-SAME:   : tensor<100x200x128x256xi32> into tensor<100x200x4x32x16x16xi32>
+  //      CHECK: linalg.transpose
+  // CHECK-SAME:   ins(%{{.*}} : tensor<100x200x4x32x16x16xi32>)
+  // CHECK-SAME:   outs(%{{.*}} : tensor<200x4x16x100x16x32xi32>)
+  // CHECK-SAME:   permutation = [1, 2, 4, 0, 5, 3]
+  %cst_0 = arith.constant 0 : i32
+  %0 = tensor.pack %src
+    padding_value(%cst_0 : i32)
+    outer_dims_perm = [1, 2, 3, 0]
+    inner_dims_pos = [3, 2]
+    inner_tiles = [16, 32]
+    into %dest : tensor<100x200x127x255xi32> -> tensor<200x4x16x100x16x32xi32>
+  return %0 : tensor<200x4x16x100x16x32xi32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !pdl.operation):
+  %pack = transform.structured.match ops{["tensor.pack"]} in %module_op
+    : (!pdl.operation) -> !transform.op<"tensor.pack">
+  transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">)
+    -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
+}