Fixup: Create and use vector::createUnrollIterator() until

Instead of progressively unrolling a leading dimension at a time, this
now uses `vector::createUnrollIterator()` which returns an iterator for
all leading dimensions of a vector type (until a target rank).

Author: MacDue

mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h

@@ -266,10 +266,11 @@ void populateVectorMaskedLoadStoreEmulationPatterns(RewritePatternSet &patterns,
 
 /// Populate the pattern set with the following patterns:
 ///
-/// [InterleaveOpLowering]
-/// Progressive lowering of InterleaveOp to ExtractOp + InsertOp + lower-D
-/// InterleaveOp until dim 1.
+/// [UnrollInterleaveOp]
+/// A one-shot unrolling of InterleaveOp to (one or more) ExtractOp +
+/// InterleaveOp (of `targetRank`) + InsertOp.
 void populateVectorInterleaveLoweringPatterns(RewritePatternSet &patterns,
+                                              int64_t targetRank = 1,
                                               PatternBenefit benefit = 1);
 
 } // namespace vector

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

@@ -9,6 +9,7 @@
 #ifndef MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
 #define MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
 
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Support/LLVM.h"
@@ -75,6 +76,28 @@ FailureOr<std::pair<int, int>> isTranspose2DSlice(vector::TransposeOp op);
 ///         vector<2x1x2x2xi32> from memref<5x4x3x2xi32>)
 bool isContiguousSlice(MemRefType memrefType, VectorType vectorType);
 
+/// Returns an iterator for all positions in the leading dimensions of `vType`
+/// up to the `targetRank`. If any leading dimension before the `targetRank` is
+/// scalable (so cannot be unrolled), it will return an iterator for positions
+/// up to the first scalable dimension.
+///
+/// If no leading dimensions can be unrolled an empty optional will be returned.
+///
+/// Examples:
+///
+///   For vType = vector<2x3x4> and targetRank = 1
+///
+///   The resulting iterator will yield:
+///     [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]
+///
+///   For vType = vector<3x[4]x5> and targetRank = 0
+///
+///   The scalable dimension blocks unrolling so the iterator yields only:
+///     [0], [1], [2]
+///
+std::optional<StaticTileOffsetRange>
+createUnrollIterator(VectorType vType, int64_t targetRank = 1);
+
 } // namespace vector
 
 /// Constructs a permutation map of invariant memref indices to vector

mlir/lib/Dialect/Vector/Transforms/LowerVectorInterleave.cpp

@@ -13,6 +13,7 @@
 
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/PatternMatch.h"
 
@@ -23,66 +24,62 @@ using namespace mlir::vector;
 
 namespace {
 
-/// Progressive lowering of InterleaveOp.
-///
-/// Each leading dimension is unrolled until the result of the interleave is
-/// rank 1 (or the dimension is scalable, so can't be unrolled).
+/// A one-shot unrolling of vector.interleave to the `targetRank`.
 ///
 /// Example:
 ///
+/// ```mlir
+/// vector.interleave %a, %b : vector<1x2x3x4xi64>
 /// ```
-/// %0 = vector.interleave %lhs, %rhs : vector<2x...8xty>
-/// ```
-/// Becomes:
-/// ```
-/// %lhs_0 = vector.extract %lhs[0]
-/// %rhs_0 = vector.extract %rhs[0]
-/// %lhs_1 = vector.extract %lhs[1]
-/// %rhs_1 = vector.extract %rhs[1]
-/// %zip_0 = vector.interleave %lhs_0, %rhs_0
-/// %zip_1 = vector.interleave %lhs_1, %rhs_1
-/// %res_0 = vector.insert %zip_0, %undef[0]
-///     %0 = vector.insert %zip_1, %res_0[1]
+/// Would be unrolled to:
+/// ```mlir
+/// %result = arith.constant dense<0> : vector<1x2x3x8xi64>
+/// %0 = vector.extract %a[0, 0, 0]                 ─┐
+///        : vector<4xi64> from vector<1x2x3x4xi64>  |
+/// %1 = vector.extract %b[0, 0, 0]                  |
+///        : vector<4xi64> from vector<1x2x3x4xi64>  | - Repeated 6x for
+/// %2 = vector.interleave %0, %1 : vector<4xi64>    |   all leading positions
+/// %3 = vector.insert %2, %result [0, 0, 0]         |
+///        : vector<8xi64> into vector<1x2x3x8xi64>  ┘
 /// ```
 ///
-/// If %zip_0 and %zip_1 still have a rank > 1 they will be unrolled again
-/// following the same pattern.
-class InterleaveOpLowering : public OpRewritePattern<vector::InterleaveOp> {
+/// Note: If any leading dimension before the `targetRank` is scalable the
+/// unrolling will stop before the scalable dimension.
+class UnrollInterleaveOp : public OpRewritePattern<vector::InterleaveOp> {
 public:
-  using OpRewritePattern::OpRewritePattern;
+  UnrollInterleaveOp(int64_t targetRank, MLIRContext *context,
+                     PatternBenefit benefit = 1)
+      : OpRewritePattern(context, benefit), targetRank(targetRank){};
 
   LogicalResult matchAndRewrite(vector::InterleaveOp op,
                                 PatternRewriter &rewriter) const override {
     VectorType resultType = op.getResultVectorType();
-    // 1-D vector.interleave ops can be directly lowered to LLVM (later).
-    if (resultType.getRank() == 1)
+    auto unrollIterator = vector::createUnrollIterator(resultType, targetRank);
+    if (!unrollIterator)
       return failure();
 
-    // Below we unroll the leading (or front) dimension. If that dimension is
-    // scalable we can't unroll it.
-    if (resultType.getScalableDims().front())
-      return failure();
-
-    // n-D case: Unroll the leading dimension.
     auto loc = op.getLoc();
     Value result = rewriter.create<arith::ConstantOp>(
         loc, resultType, rewriter.getZeroAttr(resultType));
-    for (int idx = 0, end = resultType.getDimSize(0); idx < end; ++idx) {
-      Value extractLhs = rewriter.create<ExtractOp>(loc, op.getLhs(), idx);
-      Value extractRhs = rewriter.create<ExtractOp>(loc, op.getRhs(), idx);
+    for (auto position : *unrollIterator) {
+      Value extractLhs = rewriter.create<ExtractOp>(loc, op.getLhs(), position);
+      Value extractRhs = rewriter.create<ExtractOp>(loc, op.getRhs(), position);
       Value interleave =
           rewriter.create<InterleaveOp>(loc, extractLhs, extractRhs);
-      result = rewriter.create<InsertOp>(loc, interleave, result, idx);
+      result = rewriter.create<InsertOp>(loc, interleave, result, position);
     }
 
     rewriter.replaceOp(op, result);
     return success();
   }
+
+private:
+  int64_t targetRank = 1;
 };
 
 } // namespace
 
 void mlir::vector::populateVectorInterleaveLoweringPatterns(
-    RewritePatternSet &patterns, PatternBenefit benefit) {
-  patterns.add<InterleaveOpLowering>(patterns.getContext(), benefit);
+    RewritePatternSet &patterns, int64_t targetRank, PatternBenefit benefit) {
+  patterns.add<UnrollInterleaveOp>(targetRank, patterns.getContext(), benefit);
 }

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

@@ -303,3 +303,25 @@ bool vector::isContiguousSlice(MemRefType memrefType, VectorType vectorType) {
 
   return llvm::all_of(leadingDims, [](auto x) { return x == 1; });
 }
+
+std::optional<StaticTileOffsetRange>
+vector::createUnrollIterator(VectorType vType, int64_t targetRank) {
+  if (vType.getRank() <= targetRank)
+    return {};
+  // Attempt to unroll until targetRank or the first scalable dimension (which
+  // cannot be unrolled).
+  auto shapeToUnroll = vType.getShape().drop_back(targetRank);
+  auto scalableDimsToUnroll = vType.getScalableDims().drop_back(targetRank);
+  auto it =
+      std::find(scalableDimsToUnroll.begin(), scalableDimsToUnroll.end(), true);
+  auto firstScalableDim = it - scalableDimsToUnroll.begin();
+  if (firstScalableDim == 0)
+    return {};
+  // All scalable dimensions should be removed now.
+  scalableDimsToUnroll = scalableDimsToUnroll.slice(0, firstScalableDim);
+  assert(!llvm::is_contained(scalableDimsToUnroll, true) &&
+         "unexpected leading scalable dimension");
+  // Create an unroll iterator for leading dimensions.
+  shapeToUnroll = shapeToUnroll.slice(0, firstScalableDim);
+  return StaticTileOffsetRange(shapeToUnroll, /*unrollStep=*/1);
+}

mlir/test/Dialect/Vector/vector-interleave-lowering-transforms.mlir

@@ -36,6 +36,28 @@ func.func @vector_interleave_2d_scalable(%a: vector<2x[8]xi16>, %b: vector<2x[8]
   return %0 : vector<2x[16]xi16>
 }
 
+// CHECK-LABEL: @vector_interleave_4d
+//  CHECK-SAME:     %[[LHS:.*]]: vector<1x2x3x4xi64>, %[[RHS:.*]]: vector<1x2x3x4xi64>)
+func.func @vector_interleave_4d(%a: vector<1x2x3x4xi64>, %b: vector<1x2x3x4xi64>) -> vector<1x2x3x8xi64>
+{
+  // CHECK: %[[LHS_0:.*]] = vector.extract %[[LHS]][0, 0, 0] : vector<4xi64> from vector<1x2x3x4xi64>
+  // CHECK: %[[RHS_0:.*]] = vector.extract %[[RHS]][0, 0, 0] : vector<4xi64> from vector<1x2x3x4xi64>
+  // CHECK: %[[ZIP_0:.*]] = vector.interleave %[[LHS_0]], %[[RHS_0]] : vector<4xi64>
+  // CHECK: %[[RES_0:.*]] = vector.insert %[[ZIP_0]], %{{.*}} [0, 0, 0] : vector<8xi64> into vector<1x2x3x8xi64>
+  // CHECK-COUNT-5: vector.interleave %{{.*}}, %{{.*}} : vector<4xi64>
+  %0 = vector.interleave %a, %b : vector<1x2x3x4xi64>
+  return %0 : vector<1x2x3x8xi64>
+}
+
+// CHECK-LABEL: @vector_interleave_nd_with_scalable_dim
+func.func @vector_interleave_nd_with_scalable_dim(%a: vector<1x3x[2]x2x3x4xf16>, %b: vector<1x3x[2]x2x3x4xf16>) -> vector<1x3x[2]x2x3x8xf16>
+{
+  // The scalable dim blocks unrolling so only the first two dims are unrolled.
+  // CHECK-COUNT-3: vector.interleave %{{.*}}, %{{.*}} : vector<[2]x2x3x4xf16>
+  %0 = vector.interleave %a, %b : vector<1x3x[2]x2x3x4xf16>
+  return %0 : vector<1x3x[2]x2x3x8xf16>
+}
+
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
     %f = transform.structured.match ops{["func.func"]} in %module_op