[mlir][vector] Cleanup VectorUnroll and create a generic tile iteration utility

[christopherbate]

Summary

This change refactors some of the utilities used to unroll larger vector
computations into smaller vector computations. In fact, the indexing
computations used here are rather generic and are useful in other dialects or
downstream projects. Therefore, a utility for iterating over all possible tile
offsets for a particular pair of static (shape, tiled shape) is introduced in
IndexingUtils and replaces the existing computations in the vector unrolling
transformations. This builds off of the refactoring of IndexingUtils introduced
in 203fad476b7e.

[nicolasvasilache]

sorry for missing this

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

277

Some basic comments to describe what each entry is and how it is stored (i.e. I believe tileShape is in the original order while sliceStrides is in loopOrder) ?

mlir/lib/Dialect/Utils/IndexingUtils.cpp

280

Does the padding actually happen ?

286

I find this way of iterating followed by a permutation somewhat confusing.

Would it be possible to rewrite as something like:

SmallVector<int64_t> basis;
basis.reserve(...);
for ... { basis.push_back(...); }
sliceStrides = computeStrides(basis);

?

It is arguably a little less efficient to create a temporary basis vector but I would take the piece of mind of not having to triple check the logic :)

[christopherbate]

Address comments, add unittest

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

 
 #ifndef MLIR_DIALECT_UTILS_INDEXINGUTILS_H
 #define MLIR_DIALECT_UTILS_INDEXINGUTILS_H
 
 #include "mlir/IR/Builders.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator.h"
 #include <optional>
+#include <utility>
 
 namespace mlir {
 class ArrayAttr;
 
 //===----------------------------------------------------------------------===//
 // Utils that operate on static integer values.
 //===----------------------------------------------------------------------===//
 
SmallVector<AffineExpr> delinearize(AffineExpr linearIndex,
                                     ArrayRef<AffineExpr> strides);
 SmallVector<AffineExpr> delinearize(AffineExpr linearIndex,
                                     ArrayRef<int64_t> strides);
 
 //===----------------------------------------------------------------------===//
 // Permutation utils.
 //===----------------------------------------------------------------------===//
 
+template <typename T>
+SmallVector<T> applyPermutation(ArrayRef<T> input,
+                                ArrayRef<int64_t> permutation) {
+  assert(input.size() == permutation.size() &&
+         "expected input rank to equal permutation rank");
+  auto permutationRange = llvm::map_range(
+      llvm::seq<unsigned>(0, input.size()),
+      [&](int64_t idx) -> T { return input[permutation[idx]]; });
+  return llvm::to_vector(permutationRange);
+}
+
+template <typename T>
+SmallVector<T> applyPermutation(const SmallVectorImpl<T> &input,
+                                ArrayRef<int64_t> permutation) {
+  return applyPermutation(ArrayRef(input), permutation);
+}
+
 /// Apply the permutation defined by `permutation` to `inVec`.
 /// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
 /// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation
 /// vector `permutation = [2, 0, 1]`, this function leaves `inVec = ['c', 'a',
 /// 'b']`.
 template <typename T, unsigned N>
 void applyPermutationToVector(SmallVector<T, N> &inVec,
                               ArrayRef<int64_t> permutation) {
-  SmallVector<T, N> auxVec(inVec.size());
-  for (const auto &en : enumerate(permutation))
-    auxVec[en.index()] = inVec[en.value()];
-  inVec = auxVec;
+  inVec = applyPermutation(inVec, permutation);
 }
 
 /// Helper method to apply to inverse a permutation.
 SmallVector<int64_t> invertPermutationVector(ArrayRef<int64_t> permutation);
 
 /// Method to check if an interchange vector is a permutation.
 bool isPermutationVector(ArrayRef<int64_t> interchange);
 

 /// Returns AffineExpr and list of values to apply to it, e.g.:
 ///
 /// auto &&[expr, values] = computeLinearIndex(...);
 /// offset = affine::makeComposedFoldedAffineApply(builder, loc, expr, values);
 std::pair<AffineExpr, SmallVector<OpFoldResult>>
 computeLinearIndex(OpFoldResult sourceOffset, ArrayRef<OpFoldResult> strides,
                    ArrayRef<OpFoldResult> indices);
 
+//===----------------------------------------------------------------------===//
+// Utilities for decomposing larger shapes
+//===----------------------------------------------------------------------===//
+
+namespace detail {
+/// Encapsulates the set of parameters that are used to make tile offset
+/// calculations in the TileOffsetRangeIterator.
+class TileOffsetRangeImpl {
+public:
+  TileOffsetRangeImpl(ArrayRef<int64_t> shape, ArrayRef<int64_t> tileShape,
+                      ArrayRef<int64_t> loopOrder);
+
+  int64_t getMaxLinearIndex() const { return maxLinearIndex; }
+
+  SmallVector<int64_t> getStaticTileOffsets(int64_t linearIndex) const;
+
+  SmallVector<AffineExpr> getDynamicTileOffsets(AffineExpr linearIndex) const;
+
+  template <typename T>
+  SmallVector<T> getTileOffsets(T linearIndex) const {
+    if constexpr (std::is_same_v<T, int64_t>)
+      return getStaticTileOffsets(linearIndex);
+    else
+      return getDynamicTileOffsets(linearIndex);
+  }
+
+private:
+  /// The sub-shape that divides the larger outer shape (which is provided to
+  /// the constructor).
+  SmallVector<int64_t> tileShape;
+  /// The inverse permutation to the `loopOrder` permutation provided in the
+  /// constructor.
+  SmallVector<int64_t> inverseLoopOrder;
+  /// The strides for the basis 'div(shape, tileShape)' permuted by `loopOrder`.
+  SmallVector<int64_t> sliceStrides;
+  /// The maximum linear index in the iteration space given by basis 'div(shape,
+  /// tileShape)'.
+  int64_t maxLinearIndex;
+};
+
+/// The STL-style iterator implementation for StaticTileOffsetRange.
+template <typename ElementType>
+class TileOffsetRangeIterator
+    : public llvm::iterator_facade_base<TileOffsetRangeIterator<ElementType>,
+                                        std::forward_iterator_tag,
+                                        SmallVector<ElementType>> {
+public:
+  TileOffsetRangeIterator(const TileOffsetRangeImpl &params, ElementType index)
+      : params(params), index(index) {}
+
+  void operator++() { incrementIndex(1); }
+  TileOffsetRangeIterator operator++(int) {
+    const auto copy = *this;
+    ++*this;
+    return copy;
+  }
+
+  bool operator==(const TileOffsetRangeIterator &other) const {
+    return index == other.index;
+  }
+  bool operator!=(const TileOffsetRangeIterator &other) const {
+    return index != other.index;
+  }
+
+  SmallVector<ElementType> operator*() const {
+    return params.getTileOffsets(index);
+  }
+  void operator+=(int64_t offset) { incrementIndex(offset); }
+
+private:
+  void incrementIndex(int64_t offset) { index = index + offset; }
+  const TileOffsetRangeImpl params;
+  int64_t index;
+};
+} // namespace detail
+
+/// A range-style iterator that allows for iterating over the offsets of all
+/// potential tiles of size `tileShape` within the larger shape `shape`, using
+/// an ordering specified by `loopOrder`. The `loopOrder` specifies the order of
+/// unrolling by numbering the dimensions in order from "outer most for loop"
+/// (slowest changing) to "inner most for loop" (fastest changing).
+///
+/// For example, for `shape = {10, 20, 30}`, `tileShape = {5, 10, 15}`, and
+/// `loopOrder={2, 0, 1}`, the iterating over this range will yield offsets:
+///
+/// ```
+/// {0, 0,  0}, {0, 10,  0}, {5, 0,  0}, {5, 10,  0}, {0, 0, 15},
+/// {0, 10, 15}, {5, 0, 15}, {0, 10, 15}, {5, 10, 15}
+/// ```
+///
+/// This is useful in contexts where a vector computation over a larger shape
+/// needs to be unrolled to a set of operations on subsets of the original
+/// operands, such as during the "vector unrolling" transformations.
+///
+/// The size of `tileShape` must be less-than-or-equal-to the size of `shape`.a
+/// If the rank of `tileShape` is smaller than `shape`, then `tileShape`
+/// elements correspond to the trailing dimensions of `shape`, and the leading
+/// dimensions are considered untiled and `tileShape` is effectively prepended
+/// with the leading dims of `shape`.
+class StaticTileOffsetRange {
+public:
+  using IteratorTy = detail::TileOffsetRangeIterator<int64_t>;
+  using ParamsTy = detail::TileOffsetRangeImpl;
+
+  StaticTileOffsetRange(ArrayRef<int64_t> shape, ArrayRef<int64_t> tileShape,
+                        ArrayRef<int64_t> loopOrder)
+      : params(shape, tileShape, loopOrder), beginValue(params, 0),
+        pastEndValue(params, params.getMaxLinearIndex()) {
+    assert(shape.size() >= tileShape.size());
+    assert(loopOrder.size() == shape.size());
+  }
+
+  /// Create the range with identity loop order.
+  StaticTileOffsetRange(ArrayRef<int64_t> shape, ArrayRef<int64_t> tileShape)
+      : params(shape, tileShape,
+               llvm::to_vector(llvm::seq<int64_t>(0, shape.size()))),
+        beginValue(params, 0),
+        pastEndValue(params, params.getMaxLinearIndex()) {
+    assert(shape.size() >= tileShape.size());
+  }
+
+  IteratorTy begin() const { return beginValue; }
+  IteratorTy end() const { return pastEndValue; }
+
+  /// Returns the total number of tiles that fit in the larger shape.
+  size_t size() const { return params.getMaxLinearIndex(); }
+
+private:
+  const ParamsTy params;
+  IteratorTy beginValue;
+  IteratorTy pastEndValue;
+};
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_INDEXINGUTILS_H

mlir/include/mlir/IR/AffineExpr.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_AFFINEEXPR_H
 #define MLIR_IR_AFFINEEXPR_H
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 #include <functional>
 #include <type_traits>
 
 namespace mlir {
 
 class MLIRContext;
 class AffineMap;

 inline AffineExpr operator-(int64_t val, AffineExpr expr) {
   return expr * (-1) + val;
 }
 
 /// These free functions allow clients of the API to not use classes in detail.
 AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context);
 AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context);
 AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context);
+SmallVector<AffineExpr> getAffineConstantExprs(ArrayRef<int64_t> constants,
+                                               MLIRContext *context);
 AffineExpr getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs,
                                  AffineExpr rhs);
 
 /// Constructs an affine expression from a flat ArrayRef. If there are local
 /// identifiers (neither dimensional nor symbolic) that appear in the sum of
 /// products expression, 'localExprs' is expected to have the AffineExpr
 /// for it, and is substituted into. The ArrayRef 'eq' is expected to be in the
 /// format [dims, symbols, locals, constant term].

mlir/lib/Dialect/Utils/IndexingUtils.cpp

 AffineExpr mlir::linearize(MLIRContext *ctx, ArrayRef<AffineExpr> offsets,
                            ArrayRef<AffineExpr> basis) {
   AffineExpr zero = getAffineConstantExpr(0, ctx);
   return linearizeImpl(offsets, basis, zero);
 }
 
 AffineExpr mlir::linearize(MLIRContext *ctx, ArrayRef<AffineExpr> offsets,
                            ArrayRef<int64_t> basis) {
-  SmallVector<AffineExpr> basisExprs = llvm::to_vector(llvm::map_range(
-      basis, [ctx](int64_t v) { return getAffineConstantExpr(v, ctx); }));
-  return linearize(ctx, offsets, basisExprs);
+
+  return linearize(ctx, offsets, getAffineConstantExprs(basis, ctx));
 }
 
 SmallVector<AffineExpr> mlir::delinearize(AffineExpr linearIndex,
                                           ArrayRef<AffineExpr> strides) {
   return delinearizeImpl(
       linearIndex, strides,
       [](AffineExpr e1, AffineExpr e2) { return e1.floorDiv(e2); });
 }
 
 SmallVector<AffineExpr> mlir::delinearize(AffineExpr linearIndex,
                                           ArrayRef<int64_t> strides) {
   MLIRContext *ctx = linearIndex.getContext();
-  SmallVector<AffineExpr> basisExprs = llvm::to_vector(llvm::map_range(
-      strides, [ctx](int64_t v) { return getAffineConstantExpr(v, ctx); }));
-  return delinearize(linearIndex, ArrayRef<AffineExpr>{basisExprs});
+  return delinearize(linearIndex, getAffineConstantExprs(strides, ctx));
 }
 
 //===----------------------------------------------------------------------===//
 // Permutation utils.
 //===----------------------------------------------------------------------===//
 
 SmallVector<int64_t>
 mlir::invertPermutationVector(ArrayRef<int64_t> permutation) {

 std::pair<AffineExpr, SmallVector<OpFoldResult>>
 mlir::computeLinearIndex(OpFoldResult sourceOffset,
                          ArrayRef<OpFoldResult> strides,
                          ArrayRef<OpFoldResult> indices) {
   assert(strides.size() == indices.size());
   auto sourceRank = static_cast<unsigned>(strides.size());
 
   // Hold the affine symbols and values for the computation of the offset.
   SmallVector<OpFoldResult> values(2 * sourceRank + 1);
   SmallVector<AffineExpr> symbols(2 * sourceRank + 1);
 
   bindSymbolsList(getContext(sourceOffset), MutableArrayRef{symbols});
   AffineExpr expr = symbols.front();
   values[0] = sourceOffset;
 
   for (unsigned i = 0; i < sourceRank; ++i) {
     // Compute the stride.
     OpFoldResult origStride = strides[i];
 
     // Build up the computation of the offset.
     unsigned baseIdxForDim = 1 + 2 * i;
     unsigned subOffsetForDim = baseIdxForDim;
     unsigned origStrideForDim = baseIdxForDim + 1;
     expr = expr + symbols[subOffsetForDim] * symbols[origStrideForDim];
     values[subOffsetForDim] = indices[i];
     values[origStrideForDim] = origStride;
   }
 
   return {expr, values};
 }
+
+//===----------------------------------------------------------------------===//
+// TileOffsetRange
+//===----------------------------------------------------------------------===//
+
+/// Apply left-padding by 1 to the tile shape if required.
+static SmallVector<int64_t> padTileShapeToSize(ArrayRef<int64_t> tileShape,
+                                               unsigned paddedSize) {
+  assert(tileShape.size() <= paddedSize &&
+         "expected tileShape to <= paddedSize");
+  if (tileShape.size() == paddedSize)
+    return to_vector(tileShape);
+  SmallVector<int64_t> result(paddedSize - tileShape.size(), 1);
+  llvm::append_range(result, tileShape);
+  return result;
+}
+
+mlir::detail::TileOffsetRangeImpl::TileOffsetRangeImpl(
+    ArrayRef<int64_t> shape, ArrayRef<int64_t> tileShape,
+    ArrayRef<int64_t> loopOrder)
+    : tileShape(padTileShapeToSize(tileShape, shape.size())),
+      inverseLoopOrder(invertPermutationVector(loopOrder)),
+      sliceStrides(shape.size()) {
+  // Divide the shape by the tile shape.
+  std::optional<SmallVector<int64_t>> shapeRatio =
+      mlir::computeShapeRatio(shape, tileShape);
+  assert(shapeRatio && shapeRatio->size() == shape.size() &&
+         "target shape does not evenly divide the original shape");
+  assert(isPermutationVector(loopOrder) && loopOrder.size() == shape.size() &&
+         "expected loop order to be a permutation of rank equal to outer "
+         "shape");
+
+  maxLinearIndex = mlir::computeMaxLinearIndex(*shapeRatio);
+  mlir::applyPermutationToVector(*shapeRatio, loopOrder);
+  sliceStrides = mlir::computeStrides(*shapeRatio);
+}
+
+SmallVector<int64_t> mlir::detail::TileOffsetRangeImpl::getStaticTileOffsets(
+    int64_t linearIndex) const {
+  SmallVector<int64_t> tileCoords = applyPermutation(
+      delinearize(linearIndex, sliceStrides), inverseLoopOrder);
+  return computeElementwiseMul(tileCoords, tileShape);
+}
+
+SmallVector<AffineExpr>
+mlir::detail::TileOffsetRangeImpl::getDynamicTileOffsets(
+    AffineExpr linearIndex) const {
+  MLIRContext *ctx = linearIndex.getContext();
+  SmallVector<AffineExpr> tileCoords = applyPermutation(
+      delinearize(linearIndex, sliceStrides), inverseLoopOrder);
+  return mlir::computeElementwiseMul(tileCoords,
+                                     getAffineConstantExprs(tileShape, ctx));
+}

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

 
 #define DEBUG_TYPE "vector-unroll"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
 using namespace mlir;
 using namespace mlir::vector;
 
-/// During unrolling from `originalShape` to `targetShape` return the offset for
-/// the slice `index`.
-static SmallVector<int64_t> getVectorOffset(ArrayRef<int64_t> ratioStrides,
-                                            int64_t index,
-                                            ArrayRef<int64_t> targetShape) {
-  return computeElementwiseMul(delinearize(index, ratioStrides), targetShape);
-}
-
-/// A functor that accomplishes the same thing as `getVectorOffset` but
-/// allows for reordering the traversal of the dimensions. The order of
-/// traversal is given in "for loop order" (outer to inner).
-namespace {
-class DecomposeShapeIterator {
-private:
-  SmallVector<int64_t> vectorShape;
-  SmallVector<int64_t> loopOrder;
-  SmallVector<int64_t> sliceStrides;
-  int64_t maxIndexVal{1};
-
-public:
-  DecomposeShapeIterator(ArrayRef<int64_t> originalShape,
-                         ArrayRef<int64_t> targetShape,
-                         ArrayRef<int64_t> loopOrder)
-      : vectorShape(targetShape.begin(), targetShape.end()),
-        loopOrder(loopOrder.begin(), loopOrder.end()),
-        sliceStrides(originalShape.size()) {
-    assert(originalShape.size() >= targetShape.size());
-    assert(loopOrder.size() == originalShape.size());
-
-    // Compute the count for each dimension.
-    auto maybeShapeRatio = computeShapeRatio(originalShape, targetShape);
-    assert(maybeShapeRatio && "Shape does not evenly divide");
-    // Pad `sliceDimCounts` with leading 1s so that all sizes match.
-    SmallVector<int64_t> sliceDimCounts = *maybeShapeRatio;
-    maxIndexVal = computeMaxLinearIndex(sliceDimCounts);
-
-    // Reversing "loop order" gives dimensions from fastest varying to slowest
-    // varying (smallest stride to largest stride).
-    int64_t accum = 1;
-    for (auto idx : llvm::reverse(loopOrder)) {
-      sliceStrides[idx] = accum;
-      accum *= sliceDimCounts[idx];
-    }
-  }
-
-  // Turn the linear index into a d-tuple based on units of vectors of size
-  // `vectorShape`. The linear index is assumed to represent traversal of the
-  // dimensions based on `order`.
-  SmallVector<int64_t> delinearize(int64_t index) const {
-    // Traverse in for loop order (largest stride to smallest stride).
-    SmallVector<int64_t> vectorOffsets(sliceStrides.size());
-    for (auto idx : loopOrder) {
-      vectorOffsets[idx] = index / sliceStrides[idx];
-      index %= sliceStrides[idx];
-    }
-    return vectorOffsets;
-  }
-
-  int64_t maxIndex() const { return maxIndexVal; }
-
-  /// Return the offset within d-tuple based on the ordering given by
-  /// `loopOrder`.
-  SmallVector<int64_t> getVectorOffset(int64_t index) const {
-    SmallVector<int64_t> vectorOffsets = delinearize(index);
-    SmallVector<int64_t> elementOffsets =
-        computeElementwiseMul(vectorShape, vectorOffsets);
-    return elementOffsets;
-  }
-};
-} // namespace
-
 /// Compute the indices of the slice `index` for a tranfer op.
 static SmallVector<Value> sliceTransferIndices(ArrayRef<int64_t> elementOffsets,
                                                ArrayRef<Value> indices,
                                                AffineMap permutationMap,
                                                Location loc,
                                                OpBuilder &builder) {
   MLIRContext *ctx = builder.getContext();
   auto isBroadcast = [](AffineExpr expr) {
LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
 
     // Prepare the result vector;
     Value result = rewriter.create<arith::ConstantOp>(
         loc, sourceVectorType, rewriter.getZeroAttr(sourceVectorType));
     auto targetType =
         VectorType::get(*targetShape, sourceVectorType.getElementType());
     SmallVector<Value> originalIndices(readOp.getIndices().begin(),
                                        readOp.getIndices().end());
-
     SmallVector<int64_t> loopOrder =
         getUnrollOrder(originalSize.size(), readOp, options);
-    DecomposeShapeIterator indexToOffsets(originalSize, *targetShape,
-                                          loopOrder);
-    for (int64_t i = 0; i < indexToOffsets.maxIndex(); i++) {
-      SmallVector<int64_t> elementOffsets = indexToOffsets.getVectorOffset(i);
+    for (SmallVector<int64_t> elementOffsets :
+         StaticTileOffsetRange(originalSize, *targetShape, loopOrder)) {
       SmallVector<Value> indices =
           sliceTransferIndices(elementOffsets, originalIndices,
                                readOp.getPermutationMap(), loc, rewriter);
       auto slicedRead = rewriter.create<vector::TransferReadOp>(
           loc, targetType, readOp.getSource(), indices,
           readOp.getPermutationMapAttr(), readOp.getPadding(), readOp.getMask(),
           readOp.getInBoundsAttr());
 
LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
     if (!targetShape)
       return failure();
     auto sourceVectorType = writeOp.getVectorType();
     SmallVector<int64_t> strides(targetShape->size(), 1);
     Location loc = writeOp.getLoc();
     ArrayRef<int64_t> originalSize = sourceVectorType.getShape();
     SmallVector<Value> originalIndices(writeOp.getIndices().begin(),
                                        writeOp.getIndices().end());
-
     SmallVector<int64_t> loopOrder =
         getUnrollOrder(originalSize.size(), writeOp, options);
-    DecomposeShapeIterator indexToOffsets(originalSize, *targetShape,
-                                          loopOrder);
     Value resultTensor;
-    for (int64_t i = 0; i < indexToOffsets.maxIndex(); i++) {
-      SmallVector<int64_t> elementOffsets = indexToOffsets.getVectorOffset(i);
+    for (SmallVector<int64_t> elementOffsets :
+         StaticTileOffsetRange(originalSize, *targetShape, loopOrder)) {
       Value slicedVector = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, writeOp.getVector(), elementOffsets, *targetShape, strides);
       SmallVector<Value> indices =
           sliceTransferIndices(elementOffsets, originalIndices,
                                writeOp.getPermutationMap(), loc, rewriter);
       Operation *slicedWrite = rewriter.create<vector::TransferWriteOp>(
           loc, slicedVector, resultTensor ? resultTensor : writeOp.getSource(),
           indices, writeOp.getPermutationMapAttr(), writeOp.getInBoundsAttr());
LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
     AffineMap dstAffineMap = contractOp.getIndexingMapsArray()[accIndex];
     llvm::MapVector<
         SmallVector<int64_t>, Value,
         llvm::DenseMap<SmallVector<int64_t>, unsigned, OffsetMapInfo>>
         accCache;
 
     SmallVector<int64_t> loopOrder = getUnrollOrder(
         contractOp.getIteratorTypes().size(), contractOp, options);
-    DecomposeShapeIterator indexToOffsets(originalSize, *targetShape,
-                                          loopOrder);
-    const int64_t sliceCount = indexToOffsets.maxIndex();
-    for (int64_t i = 0; i < sliceCount; i++) {
-      SmallVector<int64_t> offsets = indexToOffsets.getVectorOffset(i);
+
+    for (SmallVector<int64_t> offsets :
+         StaticTileOffsetRange(originalSize, *targetShape, loopOrder)) {
       SmallVector<Value> slicesOperands(contractOp.getNumOperands());
 
       // Helper to compute the new shape of each operand and extract the slice.
       auto extractOperand = [&](unsigned index, Value operand,
                                 AffineMap permutationMap,
                                 ArrayRef<int64_t> operandOffets) {
         SmallVector<int64_t> operandShape = applyPermutationMap(
             permutationMap, ArrayRef<int64_t>(*targetShape));
struct UnrollMultiReductionPattern
 
   LogicalResult matchAndRewrite(vector::MultiDimReductionOp reductionOp,
                                 PatternRewriter &rewriter) const override {
     std::optional<SmallVector<int64_t>> targetShape =
         getTargetShape(options, reductionOp);
     if (!targetShape)
       return failure();
     SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
-    SmallVector<int64_t> ratio = *computeShapeRatio(originalSize, *targetShape);
     llvm::MapVector<
         SmallVector<int64_t>, Value,
         llvm::DenseMap<SmallVector<int64_t>, unsigned, OffsetMapInfo>>
         accCache;
-    // Compute shape ratio of 'shape' and 'sizes'.
-    int64_t sliceCount = computeMaxLinearIndex(ratio);
     Location loc = reductionOp.getLoc();
 
     // Stride of the ratios, this gives us the offsets of sliceCount in a basis
     // of multiples of the targetShape.
-    auto ratioStrides = computeStrides(ratio);
-    for (int64_t i = 0; i < sliceCount; i++) {
-      SmallVector<int64_t> offsets =
-          getVectorOffset(ratioStrides, i, *targetShape);
-
+    for (SmallVector<int64_t> offsets :
+         StaticTileOffsetRange(originalSize, *targetShape)) {
       SmallVector<Value> operands;
       SmallVector<int64_t> operandStrides(offsets.size(), 1);
       Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, reductionOp.getSource(), offsets, *targetShape, operandStrides);
       operands.push_back(slicedOperand);
       SmallVector<int64_t> dstShape;
       SmallVector<int64_t> destOffset;
       for (size_t i : llvm::seq(size_t(0), targetShape->size())) {
LogicalResult matchAndRewrite(Operation *op,
     if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
       return failure();
     auto targetShape = getTargetShape(options, op);
     if (!targetShape)
       return failure();
     auto dstVecType = cast<VectorType>(op->getResult(0).getType());
     SmallVector<int64_t> originalSize =
         *cast<VectorUnrollOpInterface>(op).getShapeForUnroll();
-    SmallVector<int64_t> ratio = *computeShapeRatio(originalSize, *targetShape);
-    int64_t sliceCount = computeMaxLinearIndex(ratio);
     Location loc = op->getLoc();
     // Prepare the result vector.
     Value result = rewriter.create<arith::ConstantOp>(
         loc, dstVecType, rewriter.getZeroAttr(dstVecType));
     SmallVector<int64_t> strides(targetShape->size(), 1);
     VectorType newVecType =
         VectorType::get(*targetShape, dstVecType.getElementType());
 
-    // Stride of the ratios, this gives us the offsets of sliceCount in a basis
-    // of multiples of the targetShape.
-    auto ratioStrides = computeStrides(ratio);
-    for (int64_t i = 0; i < sliceCount; i++) {
-      SmallVector<int64_t> offsets =
-          getVectorOffset(ratioStrides, i, *targetShape);
+    // Create the unrolled computation.
+    for (SmallVector<int64_t> offsets :
+         StaticTileOffsetRange(originalSize, *targetShape)) {
       SmallVector<Value> extractOperands;
       for (OpOperand &operand : op->getOpOperands()) {
         auto vecType = dyn_cast<VectorType>(operand.get().getType());
         if (!vecType) {
           extractOperands.push_back(operand.get());
           continue;
         }
         extractOperands.push_back(
struct UnrollReductionPattern : public OpRewritePattern<vector::ReductionOp> {
 
   LogicalResult matchAndRewrite(vector::ReductionOp reductionOp,
                                 PatternRewriter &rewriter) const override {
     std::optional<SmallVector<int64_t>> targetShape =
         getTargetShape(options, reductionOp);
     if (!targetShape)
       return failure();
     SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
-    auto ratio = *computeShapeRatio(originalSize, *targetShape);
-    int64_t sliceCount = ratio[0];
 
     // Create unrolled vector reduction.
     Location loc = reductionOp.getLoc();
     Value accumulator = nullptr;
-
-    // Stride of the ratios, this gives us the offsets of sliceCount in a basis
-    // of multiples of the targetShape.
-    auto ratioStrides = computeStrides(ratio);
-    for (int64_t i = 0; i < sliceCount; ++i) {
-      SmallVector<int64_t> offsets =
-          getVectorOffset(ratioStrides, i, *targetShape);
+    for (SmallVector<int64_t> offsets :
+         StaticTileOffsetRange(originalSize, *targetShape)) {
       SmallVector<int64_t> strides(offsets.size(), 1);
       Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, reductionOp.getVector(), offsets, *targetShape, strides);
       Operation *newOp = cloneOpWithOperandsAndTypes(
           rewriter, loc, reductionOp, slicedOperand, reductionOp.getType());
       Value result = newOp->getResult(0);
 
       if (!accumulator) {
if (transposeOp.getResultVectorType().getRank() == 0)
       return failure();
     auto targetShape = getTargetShape(options, transposeOp);
     if (!targetShape)
       return failure();
     auto originalVectorType = transposeOp.getResultVectorType();
     SmallVector<int64_t> strides(targetShape->size(), 1);
     Location loc = transposeOp.getLoc();
     ArrayRef<int64_t> originalSize = originalVectorType.getShape();
-    SmallVector<int64_t> ratio = *computeShapeRatio(originalSize, *targetShape);
-    int64_t sliceCount = computeMaxLinearIndex(ratio);
+
     // Prepare the result vector;
     Value result = rewriter.create<arith::ConstantOp>(
         loc, originalVectorType, rewriter.getZeroAttr(originalVectorType));
     SmallVector<int64_t> permutation;
     transposeOp.getTransp(permutation);
 
-    // Stride of the ratios, this gives us the offsets of sliceCount in a basis
-    // of multiples of the targetShape.
-    auto ratioStrides = computeStrides(ratio);
-    for (int64_t i = 0; i < sliceCount; i++) {
-      SmallVector<int64_t> elementOffsets =
-          getVectorOffset(ratioStrides, i, *targetShape);
+    // Unroll the computation.
+    for (SmallVector<int64_t> elementOffsets :
+         StaticTileOffsetRange(originalSize, *targetShape)) {
       SmallVector<int64_t> permutedOffsets(elementOffsets.size());
       SmallVector<int64_t> permutedShape(elementOffsets.size());
       // Compute the source offsets and shape.
       for (auto indices : llvm::enumerate(permutation)) {
         permutedOffsets[indices.value()] = elementOffsets[indices.index()];
         permutedShape[indices.value()] = (*targetShape)[indices.index()];
       }
       Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
LogicalResult matchAndRewrite(vector::GatherOp gatherOp,
     // Prepare the result vector;
     Value result = rewriter.create<arith::ConstantOp>(
         loc, sourceVectorType, rewriter.getZeroAttr(sourceVectorType));
     auto targetType =
         VectorType::get(*targetShape, sourceVectorType.getElementType());
 
     SmallVector<int64_t> loopOrder =
         getUnrollOrder(originalSize.size(), gatherOp, options);
-    DecomposeShapeIterator indexToOffsets(originalSize, *targetShape,
-                                          loopOrder);
-    for (int64_t i = 0, e = indexToOffsets.maxIndex(); i < e; ++i) {
+    for (SmallVector<int64_t> elementOffsets :
+         StaticTileOffsetRange(originalSize, *targetShape, loopOrder)) {
       // To get the unrolled gather, extract the same slice based on the
       // decomposed shape from each of the index, mask, and pass-through
       // vectors.
-      SmallVector<int64_t> elementOffsets = indexToOffsets.getVectorOffset(i);
       Value indexSubVec = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, gatherOp.getIndexVec(), elementOffsets, *targetShape, strides);
       Value maskSubVec = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, gatherOp.getMask(), elementOffsets, *targetShape, strides);
       Value passThruSubVec = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, gatherOp.getPassThru(), elementOffsets, *targetShape, strides);
       auto slicedGather = rewriter.create<vector::GatherOp>(
           loc, targetType, gatherOp.getBase(), gatherOp.getIndices(),

mlir/lib/IR/AffineExpr.cpp

AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *context) {
   auto assignCtx = [context](AffineConstantExprStorage *storage) {
     storage->context = context;
   };
 
   StorageUniquer &uniquer = context->getAffineUniquer();
   return uniquer.get<AffineConstantExprStorage>(assignCtx, constant);
 }
 
+SmallVector<AffineExpr>
+mlir::getAffineConstantExprs(ArrayRef<int64_t> constants,
+                             MLIRContext *context) {
+  return llvm::to_vector(llvm::map_range(constants, [&](int64_t constant) {
+    return getAffineConstantExpr(constant, context);
+  }));
+}
+
 /// Simplify add expression. Return nullptr if it can't be simplified.
 static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
   // Fold if both LHS, RHS are a constant.
   if (lhsConst && rhsConst)
     return getAffineConstantExpr(lhsConst.getValue() + rhsConst.getValue(),
                                  lhs.getContext());

mlir/unittests/Dialect/Utils/CMakeLists.txt

 add_mlir_unittest(MLIRDialectUtilsTests
   StructuredOpsUtilsTest.cpp
+  IndexingUtilsTest.cpp
 )
 target_link_libraries(MLIRDialectUtilsTests
   PRIVATE
   MLIRDialectUtils)

mlir/unittests/Dialect/Utils/IndexingUtilsTest.cpp

+//===- IndexingUtilsTest.cpp - IndexingUtils unit tests -------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "llvm/ADT/STLExtras.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+TEST(StaticTileOffsetRange, checkIteratorCanonicalOrder) {
+  // Tile <4x8> by <2x4> with canonical row-major order.
+  std::vector<SmallVector<int64_t>> expected = {{0, 0}, {0, 4}, {2, 0}, {2, 4}};
+  for (auto [idx, tileOffset] :
+       llvm::enumerate(StaticTileOffsetRange({4, 8}, {2, 4}, {0, 1})))
+    EXPECT_EQ(tileOffset, expected[idx]);
+
+  // Check the constructor for default order and test use with zip iterator.
+  for (auto [tileOffset, tileOffsetDefault] :
+       llvm::zip(StaticTileOffsetRange({4, 8}, {2, 4}, {0, 1}),
+                 StaticTileOffsetRange({4, 8}, {2, 4})))
+    EXPECT_EQ(tileOffset, tileOffsetDefault);
+}
+
+TEST(StaticTileOffsetRange, checkIteratorRowMajorOrder) {
+  // Tile <4x8> by <2x4> with canonical row-major order.
+  std::vector<SmallVector<int64_t>> expected = {{0, 0}, {2, 0}, {0, 4}, {2, 4}};
+  for (auto [idx, tileOffset] :
+       llvm::enumerate(StaticTileOffsetRange({4, 8}, {2, 4}, {1, 0})))
+    EXPECT_EQ(tileOffset, expected[idx]);
+}
+
+TEST(StaticTileOffsetRange, checkLeadingOneFill) {
+  // Tile <4x8> by <4>. A smaller tile shape gets right-aligned to the shape.
+  for (auto [idx, tileOffset] :
+       llvm::enumerate(StaticTileOffsetRange({4, 8}, {4}))) {
+    SmallVector<int64_t> expected = {static_cast<int64_t>(idx) / 2,
+                                     static_cast<int64_t>(idx) % 2 * 4};
+    EXPECT_EQ(tileOffset, expected);
+  }
+  for (auto [idx, tileOffset] :
+       llvm::enumerate(StaticTileOffsetRange({1, 4, 8}, {4}, {2, 1, 0}))) {
+    SmallVector<int64_t> expected = {0, static_cast<int64_t>(idx) % 4,
+                                     (static_cast<int64_t>(idx) / 4) * 4};
+    EXPECT_EQ(tileOffset, expected);
+  }
+}
+
+TEST(StaticTileOffsetRange, checkIterator3DPermutation) {
+  // Tile <8x4x2> by <4x2x1> with permutation [1, 0, 2]
+  for (auto [idx, tileOffset] : llvm::enumerate(
+           StaticTileOffsetRange({8, 4, 2}, {4, 2, 1}, {1, 0, 2}))) {
+    SmallVector<int64_t> expected = {((static_cast<int64_t>(idx) / 2) % 2) * 4,
+                                     ((static_cast<int64_t>(idx) / 4) % 2) * 2,
+                                     static_cast<int64_t>(idx) % 2};
+    EXPECT_EQ(tileOffset, expected);
+  }
+
+  // Tile <10x20x30> by <5x10x16> with permutation [2, 0, 1]
+  for (auto [idx, tileOffset] : llvm::enumerate(
+           StaticTileOffsetRange({10, 20, 30}, {5, 10, 15}, {2, 0, 1}))) {
+    SmallVector<int64_t> expected = {((static_cast<int64_t>(idx) / 2) % 2) * 5,
+                                     (static_cast<int64_t>(idx) % 2) * 10,
+                                     (static_cast<int64_t>(idx) / 4) % 2 * 15};
+    EXPECT_EQ(tileOffset, expected);
+  }
+}