[NFC][mlir] VectorUtils / IndexingUtils simplifications and cleanups

[nicolasvasilache]

Summary

This revision refactors and cleans up a bunch of infra related to vector, shapes and indexing into more reusable APIs.

[nicolasvasilache]

Drop spurious filter in assert, update commit message.

[nicolasvasilache]

Cleanup variable name.

[springerm]

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

35

dup

46

backtick missing

50

typo

53

this one does not make sense to me. expected None because 3/2 does not divide evenly. but even then, the other two values do also not make sense to me.

[nicolasvasilache]

Address comments.

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

 class ArrayAttr;
 
 /// Computes and returns the linearized index of 'offsets' w.r.t. 'basis'.
 int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis);
 
 /// Given the strides together with a linear index in the dimension
 /// space, returns the vector-space offsets in each dimension for a
 /// de-linearized index.
-SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> strides,
-                                    int64_t linearIndex);
+SmallVector<int64_t> delinearize(ArrayRef<int64_t> strides,
 
+/// Given a set of sizes, compute and return the strides (i.e. the number of
+/// linear incides to skip along the (k-1) most minor dimensions to get the next
+/// k-slice). This is also the basis that one can use to linearize an n-D offset
+/// confined to `[0 .. sizes]`.
+SmallVector<int64_t> computeStrides(ArrayRef<int64_t> sizes);
+
+/// Return a vector containing llvm::zip of v1 and v2 multiplied elementwise.
+SmallVector<int64_t> computeElementwiseMul(ArrayRef<int64_t> v1,
+                                           ArrayRef<int64_t> v2);
+
+/// Compute and return the multi-dimensional integral ratio of `subShape` to
+/// the trailing dimensions of `shape`. This represents how many times
+/// `subShape` fits within `shape`.
+/// If integral division is not possible, return None.
+/// The trailing `subShape.size()` entries of both shapes are assumed (and
+/// enforced) to only contain noonnegative values.
+///
+/// Examples:
+///   - shapeRatio({3, 5, 8}, {2, 5, 2}) returns {3, 2, 1}.
+///   - shapeRatio({3, 8}, {2, 5, 2}) returns None (subshape has higher rank).
+///   - shapeRatio({42, 2, 10, 32}, {2, 5, 2}) returns {42, 1, 2, 16} which is
+///     derived as {42(leading shape dim), 2/2, 10/5, 32/2}.
+///   - shapeRatio({42, 2, 11, 32}, {2, 5, 2}) returns None  which is
+///     derived as {42(leading shape dim), 2/2, 11/5(not divisible), 32/2}.
+Optional<SmallVector<int64_t>> computeShapeRatio(ArrayRef<int64_t> shape,
+                                                 ArrayRef<int64_t> subShape);
+
+/// Return the number of elements of basis (i.e. the max linear index).
+/// Return `0` if `basis` is empty.
+int64_t computeMaxLinearIndex(ArrayRef<int64_t> basis);
+
 /// 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;
 }
 
 /// Helper that returns a subset of `arrayAttr` as a vector of int64_t.
-SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr,
-                                       unsigned dropFront = 0,
-                                       unsigned dropBack = 0);
+SmallVector<int64_t> getI64SubArray(ArrayAttr arrayAttr, unsigned dropFront = 0,
 
 /// Computes and returns linearized affine expression w.r.t. `basis`.
 mlir::AffineExpr getLinearAffineExpr(ArrayRef<int64_t> basis, mlir::Builder &b);
 
 /// Given the strides in the dimension space, returns the affine expressions for
 /// vector-space offsets in each dimension for a de-linearized index.
-SmallVector<mlir::AffineExpr, 4>
+SmallVector<mlir::AffineExpr>
 getDelinearizedAffineExpr(ArrayRef<int64_t> strides, mlir::Builder &b);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_INDEXINGUTILS_H

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

struct UnrollVectorOptions {
   /// attempted on the operation.
   FilterConstraintFnType filterConstraint = nullptr;
   UnrollVectorOptions &setFilterConstraint(FilterConstraintFnType constraint) {
     filterConstraint = std::move(constraint);
     return *this;
   }
 
   using NativeShapeFnType =
-      std::function<Optional<SmallVector<int64_t, 4>>(Operation *op)>;
+      std::function<Optional<SmallVector<int64_t>>(Operation *op)>;
   /// Function that returns the shape of the vector to unroll to for a given
   /// operation. The unrolling is aborted if the function returns `llvm::None`.
   NativeShapeFnType nativeShape = nullptr;
   UnrollVectorOptions &setNativeShapeFn(NativeShapeFnType fn) {
     nativeShape = std::move(fn);
     return *this;
   }
 
   /// Set the native shape to use for unrolling.
   UnrollVectorOptions &setNativeShape(ArrayRef<int64_t> shape) {
-    SmallVector<int64_t, 4> tsShape(shape.begin(), shape.end());
-    nativeShape = [=](Operation *) -> Optional<SmallVector<int64_t, 4>> {
+    SmallVector<int64_t> tsShape(shape.begin(), shape.end());
+    nativeShape = [=](Operation *) -> Optional<SmallVector<int64_t>> {
       return tsShape;
     };
     return *this;
   }
 
   /// Function that returns the traversal order (in terms of "for loop order",
   /// i.e. slowest varying dimension to fastest varying dimension) that shoudl
   /// be used when unrolling the given operation into units of the native vector

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

 class VectorTransferOpInterface;
 
 namespace vector {
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
 } // namespace vector
 
-/// Return the number of elements of basis, `0` if empty.
-int64_t computeMaxLinearIndex(ArrayRef<int64_t> basis);
-
-/// Given the shape and sizes of a vector, returns the corresponding
-/// strides for each dimension.
-/// TODO: needs better doc of how it is used.
-SmallVector<int64_t, 4> computeStrides(ArrayRef<int64_t> shape,
-                                       ArrayRef<int64_t> sizes);
-
-/// Given the target sizes of a vector, together with vector-space offsets,
-/// returns the element-space offsets for each dimension.
-SmallVector<int64_t, 4>
-computeElementOffsetsFromVectorSliceOffsets(ArrayRef<int64_t> sizes,
-                                            ArrayRef<int64_t> vectorOffsets);
-
-/// Computes and returns the multi-dimensional ratio of `superShape` to
-/// `subShape`. This is calculated by performing a traversal from minor to major
-/// dimensions (i.e. in reverse shape order). If integral division is not
-/// possible, returns None.
-/// The ArrayRefs are assumed (and enforced) to only contain > 1 values.
-/// This constraint comes from the fact that they are meant to be used with
-/// VectorTypes, for which the property holds by construction.
-///
-/// Examples:
-///   - shapeRatio({3, 4, 5, 8}, {2, 5, 2}) returns {3, 2, 1, 4}
-///   - shapeRatio({3, 4, 4, 8}, {2, 5, 2}) returns None
-///   - shapeRatio({1, 2, 10, 32}, {2, 5, 2}) returns {1, 1, 2, 16}
-Optional<SmallVector<int64_t, 4>> shapeRatio(ArrayRef<int64_t> superShape,
-                                             ArrayRef<int64_t> subShape);
-
-/// Computes and returns the multi-dimensional ratio of the shapes of
-/// `superVector` to `subVector`. If integral division is not possible, returns
-/// None.
-/// Assumes and enforces that the VectorTypes have the same elemental type.
-Optional<SmallVector<int64_t, 4>> shapeRatio(VectorType superVectorType,
-                                             VectorType subVectorType);
-
 /// Constructs a permutation map of invariant memref indices to vector
 /// dimension.
 ///
 /// If no index is found to be invariant, 0 is added to the permutation_map and
 /// corresponds to a vector broadcast along that dimension.
 ///
 /// The implementation uses the knowledge of the mapping of loops to
 /// vector dimension. `loopToVectorDim` carries this information as a map with:

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

VecOpToScalarOp<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const {
   if (!vecType.hasRank())
     return rewriter.notifyMatchFailure(op, "unknown vector rank");
   ArrayRef<int64_t> shape = vecType.getShape();
   int64_t numElements = vecType.getNumElements();
 
   Value result = rewriter.create<arith::ConstantOp>(
       loc, DenseElementsAttr::get(
                vecType, IntegerAttr::get(vecType.getElementType(), 0)));
-  SmallVector<int64_t> ones(shape.size(), 1);
-  SmallVector<int64_t> strides = computeStrides(shape, ones);
+  SmallVector<int64_t> strides = computeStrides(shape);
   for (int64_t linearIndex = 0; linearIndex < numElements; ++linearIndex) {
     SmallVector<int64_t> positions = delinearize(strides, linearIndex);
     SmallVector<Value> operands;
     for (Value input : op->getOperands())
       operands.push_back(
           rewriter.create<vector::ExtractOp>(loc, input, positions));
     Value scalarOp =
         rewriter.create<Op>(loc, vecType.getElementType(), operands);

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

VecOpToScalarOp<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const {
   if (!vecType.hasRank())
     return failure();
   auto shape = vecType.getShape();
   int64_t numElements = vecType.getNumElements();
 
   Value result = rewriter.create<arith::ConstantOp>(
       loc, DenseElementsAttr::get(
                vecType, FloatAttr::get(vecType.getElementType(), 0.0)));
-  SmallVector<int64_t> ones(shape.size(), 1);
-  SmallVector<int64_t> strides = computeStrides(shape, ones);
+  SmallVector<int64_t> strides = computeStrides(shape);
   for (auto linearIndex = 0; linearIndex < numElements; ++linearIndex) {
     SmallVector<int64_t> positions = delinearize(strides, linearIndex);
     SmallVector<Value> operands;
     for (auto input : op->getOperands())
       operands.push_back(
           rewriter.create<vector::ExtractOp>(loc, input, positions));
     Value scalarOp =
         rewriter.create<Op>(loc, vecType.getElementType(), operands);

mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp

for (unsigned i = 0; i < operands.size(); ++i) {
       auto expandedType = VectorType::get(expandedShape, eltType);
       expandedOperands[i] =
           builder.create<vector::ShapeCastOp>(expandedType, operand);
     }
   }
 
   // Iterate over all outer dimensions of the compute shape vector type.
   auto iterationDims = ArrayRef<int64_t>(expandedShape).drop_back();
-  int64_t maxLinearIndex = computeMaxLinearIndex(iterationDims);
-
-  SmallVector<int64_t> ones(iterationDims.size(), 1);
-  auto strides = computeStrides(iterationDims, ones);
+  int64_t maxIndex = computeMaxLinearIndex(iterationDims);
+  auto strides = computeStrides(iterationDims);
 
   // Compute results for each one dimensional vector.
-  SmallVector<Value> results(maxLinearIndex);
+  SmallVector<Value> results(maxIndex);
 
-  for (int64_t i = 0; i < maxLinearIndex; ++i) {
+  for (int64_t i = 0; i < maxIndex; ++i) {
     auto offsets = delinearize(strides, i);
 
     SmallVector<Value> extracted(expandedOperands.size());
     for (const auto &tuple : llvm::enumerate(expandedOperands))
       extracted[tuple.index()] =
           builder.create<vector::ExtractOp>(tuple.value(), offsets);
 
     results[i] = compute(extracted);
   }
 
   // Stitch results together into one large vector.
   Type resultEltType = results[0].getType().cast<VectorType>().getElementType();
   Type resultExpandedType = VectorType::get(expandedShape, resultEltType);
   Value result = builder.create<arith::ConstantOp>(
       resultExpandedType, builder.getZeroAttr(resultExpandedType));
 
-  for (int64_t i = 0; i < maxLinearIndex; ++i)
+  for (int64_t i = 0; i < maxIndex; ++i)
     result = builder.create<vector::InsertOp>(results[i], result,
                                               delinearize(strides, i));
 
   // Reshape back to the original vector shape.
   return builder.create<vector::ShapeCastOp>(
       VectorType::get(inputShape, resultEltType), result);
 }
 

mlir/lib/Dialect/Utils/IndexingUtils.cpp

 //===- IndexingUtils.cpp - Helpers related to index computations ----------===//
 //
 // 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 "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 
+#include <numeric>
+
+using namespace mlir;
+
+SmallVector<int64_t> mlir::computeStrides(ArrayRef<int64_t> sizes) {
+  SmallVector<int64_t> strides(sizes.size(), 1);
+  for (int64_t r = strides.size() - 2; r >= 0; --r)
+    strides[r] = strides[r + 1] * sizes[r + 1];
+  return strides;
+}
+
+SmallVector<int64_t> mlir::computeElementwiseMul(ArrayRef<int64_t> v1,
+                                                 ArrayRef<int64_t> v2) {
+  SmallVector<int64_t> result;
+  for (auto it : llvm::zip(v1, v2))
+    result.push_back(std::get<0>(it) * std::get<1>(it));
+  return result;
+}
+
+Optional<SmallVector<int64_t>>
+mlir::computeShapeRatio(ArrayRef<int64_t> shape, ArrayRef<int64_t> subShape) {
+  if (shape.size() < subShape.size())
+    return None;
+  assert(llvm::all_of(shape, [](int64_t s) { return s > 0; }) &&
+         "shape must be nonnegative");
+  assert(llvm::all_of(subShape, [](int64_t s) { return s > 0; }) &&
+         "subShape must be nonnegative");
+
+  // Starting from the end, compute the integer divisors.
+  std::vector<int64_t> result;
+  result.reserve(shape.size());
+  for (auto [size, subSize] :
+       llvm::zip(llvm::reverse(shape), llvm::reverse(subShape))) {
+    // If integral division does not occur, return and let the caller decide.
+    if (size % subSize != 0)
+      return None;
+    result.push_back(size / subSize);
+  }
+  // At this point we computed the ratio (in reverse) for the common size.
+  // Fill with the remaining entries from the shape (still in reverse).
+  int commonSize = subShape.size();
+  std::copy(shape.rbegin() + commonSize, shape.rend(),
+            std::back_inserter(result));
+  // Reverse again to get it back in the proper order and return.
+  return SmallVector<int64_t>{result.rbegin(), result.rend()};
+}
+
 int64_t mlir::linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis) {
   assert(offsets.size() == basis.size());
   int64_t linearIndex = 0;
   for (unsigned idx = 0, e = basis.size(); idx < e; ++idx)
     linearIndex += offsets[idx] * basis[idx];
   return linearIndex;
 }
 
-llvm::SmallVector<int64_t, 4> mlir::delinearize(ArrayRef<int64_t> sliceStrides,
-                                                int64_t index) {
+llvm::SmallVector<int64_t> mlir::delinearize(ArrayRef<int64_t> sliceStrides,
   int64_t rank = sliceStrides.size();
-  SmallVector<int64_t, 4> vectorOffsets(rank);
+  SmallVector<int64_t> vectorOffsets(rank);
   for (int64_t r = 0; r < rank; ++r) {
     assert(sliceStrides[r] > 0);
     vectorOffsets[r] = index / sliceStrides[r];
     index %= sliceStrides[r];
   }
   return vectorOffsets;
 }
 
-llvm::SmallVector<int64_t, 4> mlir::getI64SubArray(ArrayAttr arrayAttr,
-                                                   unsigned dropFront,
-                                                   unsigned dropBack) {
+int64_t mlir::computeMaxLinearIndex(ArrayRef<int64_t> basis) {
+  if (basis.empty())
+    return 0;
+  return std::accumulate(basis.begin(), basis.end(), 1,
+                         std::multiplies<int64_t>());
+}
+
+llvm::SmallVector<int64_t> mlir::getI64SubArray(ArrayAttr arrayAttr,
   assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds");
   auto range = arrayAttr.getAsRange<IntegerAttr>();
-  SmallVector<int64_t, 4> res;
+  SmallVector<int64_t> res;
   res.reserve(arrayAttr.size() - dropFront - dropBack);
   for (auto it = range.begin() + dropFront, eit = range.end() - dropBack;
        it != eit; ++it)
     res.push_back((*it).getValue().getSExtValue());
   return res;
 }
 
 mlir::AffineExpr mlir::getLinearAffineExpr(ArrayRef<int64_t> basis,
                                            mlir::Builder &b) {
   AffineExpr resultExpr = b.getAffineDimExpr(0);
   resultExpr = resultExpr * basis[0];
   for (unsigned i = 1; i < basis.size(); i++)
     resultExpr = resultExpr + b.getAffineDimExpr(i) * basis[i];
   return resultExpr;
 }
 
-llvm::SmallVector<mlir::AffineExpr, 4>
+llvm::SmallVector<mlir::AffineExpr>
 mlir::getDelinearizedAffineExpr(mlir::ArrayRef<int64_t> strides, Builder &b) {
   AffineExpr resultExpr = b.getAffineDimExpr(0);
   int64_t rank = strides.size();
-  SmallVector<AffineExpr, 4> vectorOffsets(rank);
+  SmallVector<AffineExpr> vectorOffsets(rank);
   vectorOffsets[0] = resultExpr.floorDiv(strides[0]);
   resultExpr = resultExpr % strides[0];
   for (unsigned i = 1; i < rank; i++) {
     vectorOffsets[i] = resultExpr.floorDiv(strides[i]);
     resultExpr = resultExpr % strides[i];
   }
   return vectorOffsets;
 }

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

for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
     int64_t idx = map.getDimPosition(i);
     if (idx == index)
       return i;
   }
   return None;
 }
 
 // Helper to construct iterator types with one index removed.
-static SmallVector<Attribute, 4> adjustIter(ArrayAttr iteratorTypes,
-                                            int64_t index) {
-  SmallVector<Attribute, 4> results;
+static SmallVector<Attribute> adjustIter(ArrayAttr iteratorTypes,
+  SmallVector<Attribute> results;
   for (const auto &it : llvm::enumerate(iteratorTypes)) {
     int64_t idx = it.index();
     if (idx == index)
       continue;
     results.push_back(it.value());
   }
   return results;
 }
 
 // Helper to construct an affine map with one index removed.
 static AffineMap adjustMap(AffineMap map, int64_t index,
                            PatternRewriter &rewriter) {
   auto *ctx = rewriter.getContext();
-  SmallVector<AffineExpr, 4> results;
+  SmallVector<AffineExpr> results;
   for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
     int64_t idx = map.getDimPosition(i);
     if (idx == index)
       continue;
     // Re-insert remaining indices, but renamed when occurring
     // after the removed index.
     auto targetExpr = getAffineDimExpr(idx < index ? idx : idx - 1, ctx);
     results.push_back(targetExpr);
for (int64_t d = 0, e = type.getDimSize(0); d < e; d++) {
     Value ins = rewriter.create<vector::ExtractOp>(loc, insType, val, posAttr);
     Value sto = reshapeStore(loc, ins, ext, vType, index - 1, pos, rewriter);
     result = rewriter.create<vector::InsertOp>(loc, type, sto, result, posAttr);
   }
   return result;
 }
 
 template <typename IntType>
-static SmallVector<IntType, 4> extractVector(ArrayAttr arrayAttr) {
+static SmallVector<IntType> extractVector(ArrayAttr arrayAttr) {
   return llvm::to_vector<4>(llvm::map_range(
       arrayAttr.getAsRange<IntegerAttr>(),
       [](IntegerAttr attr) { return static_cast<IntType>(attr.getInt()); }));
 }
 
 /// Helper to create arithmetic operation associated with a kind of contraction.
 static Optional<Value> createContractArithOp(Location loc, Value x, Value y,
                                              Value acc,
LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     Value input = op.getVector();
     VectorType inputType = op.getVectorType();
     VectorType resType = op.getResultType();
 
     // Set up convenience transposition table.
-    SmallVector<int64_t, 4> transp;
+    SmallVector<int64_t> transp;
     for (auto attr : op.getTransp())
       transp.push_back(attr.cast<IntegerAttr>().getInt());
 
     if (vectorTransformOptions.vectorTransposeLowering ==
             vector::VectorTransposeLowering::Shuffle &&
         resType.getRank() == 2 && transp[0] == 1 && transp[1] == 0)
       return rewriter.notifyMatchFailure(
           op, "Options specifies lowering to shuffle");
if (vectorTransformOptions.vectorTransposeLowering ==
       return success();
     }
 
     // Generate unrolled extract/insert ops. We do not unroll the rightmost
     // (i.e., highest-order) dimensions that are not transposed and leave them
     // in vector form to improve performance. Therefore, we prune those
     // dimensions from the shape/transpose data structures used to generate the
     // extract/insert ops.
-    SmallVector<int64_t, 4> prunedTransp;
+    SmallVector<int64_t> prunedTransp;
     pruneNonTransposedDims(transp, prunedTransp);
     size_t numPrunedDims = transp.size() - prunedTransp.size();
     auto prunedInShape = inputType.getShape().drop_back(numPrunedDims);
-    SmallVector<int64_t, 4> ones(prunedInShape.size(), 1);
-    auto prunedInStrides = computeStrides(prunedInShape, ones);
+    auto prunedInStrides = computeStrides(prunedInShape);
 
     // Generates the extract/insert operations for every scalar/vector element
     // of the leftmost transposed dimensions. We traverse every transpose
     // element using a linearized index that we delinearize to generate the
     // appropriate indices for the extract/insert operations.
     Value result = rewriter.create<arith::ConstantOp>(
         loc, resType, rewriter.getZeroAttr(resType));
     int64_t numTransposedElements = ShapedType::getNumElements(prunedInShape);
 
     for (int64_t linearIdx = 0; linearIdx < numTransposedElements;
          ++linearIdx) {
       auto extractIdxs = delinearize(prunedInStrides, linearIdx);
-      SmallVector<int64_t, 4> insertIdxs(extractIdxs);
+      SmallVector<int64_t> insertIdxs(extractIdxs);
       applyPermutationToVector(insertIdxs, prunedTransp);
       Value extractOp =
           rewriter.create<vector::ExtractOp>(loc, input, extractIdxs);
       result =
           rewriter.create<vector::InsertOp>(loc, extractOp, result, insertIdxs);
     }
 
     rewriter.replaceOp(op, result);
public:
   LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     VectorType srcType = op.getVectorType();
     if (srcType.getRank() != 2)
       return rewriter.notifyMatchFailure(op, "Not a 2D transpose");
 
-    SmallVector<int64_t, 4> transp;
+    SmallVector<int64_t> transp;
     for (auto attr : op.getTransp())
       transp.push_back(attr.cast<IntegerAttr>().getInt());
     if (transp[0] != 1 && transp[1] != 0)
       return rewriter.notifyMatchFailure(op, "Not a 2D transpose permutation");
 
     if (vectorTransformOptions.vectorTransposeLowering !=
         VectorTransposeLowering::Shuffle)
       return rewriter.notifyMatchFailure(op, "Options do not ask for Shuffle");
if (!rhsDims.empty()) {
       rhsDims.append(rhsShape.begin(), rhsShape.end());
       auto expandedType =
           VectorType::get(rhsDims, contractOp.getRhsType().getElementType());
       newRhs = rewriter.create<vector::BroadcastOp>(loc, expandedType, newRhs);
     }
     bool isInt = contractOp.getLhsType().getElementType().isIntOrIndex();
     newLhs = rewriter.create<vector::TransposeOp>(loc, newLhs, lhsTranspose);
     newRhs = rewriter.create<vector::TransposeOp>(loc, newRhs, rhsTranspose);
-    SmallVector<int64_t, 4> lhsOffsets(lhsReductionDims.size(), 0);
-    SmallVector<int64_t, 4> rhsOffsets(rhsReductionDims.size(), 0);
+    SmallVector<int64_t> lhsOffsets(lhsReductionDims.size(), 0);
+    SmallVector<int64_t> rhsOffsets(rhsReductionDims.size(), 0);
     newLhs = rewriter.create<vector::ExtractOp>(
         loc, newLhs, rewriter.getI64ArrayAttr(lhsOffsets));
     newRhs = rewriter.create<vector::ExtractOp>(
         loc, newRhs, rewriter.getI64ArrayAttr(rhsOffsets));
     Optional<Value> result =
         createContractArithOp(loc, newLhs, newRhs, contractOp.getAcc(),
                               contractOp.getKind(), rewriter, isInt);
     rewriter.replaceOp(contractOp, {*result});
LogicalResult matchAndRewrite(vector::ConstantMaskOp op,
     }
 
     int64_t trueDim = std::min(dstType.getDimSize(0),
                                dimSizes[0].cast<IntegerAttr>().getInt());
 
     if (rank == 1) {
       // Express constant 1-D case in explicit vector form:
       //   [T,..,T,F,..,F].
-      SmallVector<bool, 4> values(dstType.getDimSize(0));
+      SmallVector<bool> values(dstType.getDimSize(0));
       for (int64_t d = 0; d < trueDim; d++)
         values[d] = true;
       rewriter.replaceOpWithNewOp<arith::ConstantOp>(
           op, dstType, rewriter.getBoolVectorAttr(values));
       return success();
     }
 
     VectorType lowType =
         VectorType::get(dstType.getShape().drop_front(), eltType);
-    SmallVector<int64_t, 4> newDimSizes;
+    SmallVector<int64_t> newDimSizes;
     for (int64_t r = 1; r < rank; r++)
       newDimSizes.push_back(dimSizes[r].cast<IntegerAttr>().getInt());
     Value trueVal = rewriter.create<vector::ConstantMaskOp>(
         loc, lowType, rewriter.getI64ArrayAttr(newDimSizes));
     Value result = rewriter.create<arith::ConstantOp>(
         loc, dstType, rewriter.getZeroAttr(dstType));
     for (int64_t d = 0; d < trueDim; d++) {
       auto pos = rewriter.getI64ArrayAttr(d);
LogicalResult matchAndRewrite(vector::ShapeCastOp op,
     for (int64_t r = 0; r < srcRank; r++)
       numElts *= sourceVectorType.getDimSize(r);
     // Replace with data movement operations:
     //    x[0,0,0] = y[0,0]
     //    x[0,0,1] = y[0,1]
     //    x[0,1,0] = y[0,2]
     // etc., incrementing the two index vectors "row-major"
     // within the source and result shape.
-    SmallVector<int64_t, 4> srcIdx(srcRank);
-    SmallVector<int64_t, 4> resIdx(resRank);
+    SmallVector<int64_t> srcIdx(srcRank);
+    SmallVector<int64_t> resIdx(resRank);
     Value result = rewriter.create<arith::ConstantOp>(
         loc, resultVectorType, rewriter.getZeroAttr(resultVectorType));
     for (int64_t i = 0; i < numElts; i++) {
       if (i != 0) {
         incIdx(srcIdx, sourceVectorType, srcRank - 1);
         incIdx(resIdx, resultVectorType, resRank - 1);
       }
       Value e = rewriter.create<vector::ExtractOp>(loc, op.getSource(), srcIdx);
       result = rewriter.create<vector::InsertOp>(loc, e, result, resIdx);
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 
 private:
-  static void incIdx(SmallVector<int64_t, 4> &idx, VectorType tp, int64_t r) {
+  static void incIdx(SmallVector<int64_t> &idx, VectorType tp, int64_t r) {
     assert(0 <= r && r < tp.getRank());
     if (++idx[r] == tp.getDimSize(r)) {
       idx[r] = 0;
       incIdx(idx, tp, r - 1);
     }
   }
 };
 

 ///    : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
 ///  ```
 struct CombineContractABTranspose final
     : public OpRewritePattern<vector::ContractionOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                 PatternRewriter &rewriter) const override {
-    SmallVector<AffineMap, 4> maps =
+    SmallVector<AffineMap> maps =
         llvm::to_vector<4>(contractOp.getIndexingMapsArray());
     Value lhs = contractOp.getLhs();
     Value rhs = contractOp.getRhs();
     size_t index = 0;
     bool changed = false;
     for (Value *operand : {&lhs, &rhs}) {
       AffineMap &map = maps[index++];
       auto transposeOp = operand->getDefiningOp<vector::TransposeOp>();

 ///    : vector<32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
 ///  ```
 struct CombineContractBroadcast
     : public OpRewritePattern<vector::ContractionOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                 PatternRewriter &rewriter) const override {
-    SmallVector<AffineMap, 4> maps =
+    SmallVector<AffineMap> maps =
         llvm::to_vector<4>(contractOp.getIndexingMapsArray());
     Value lhs = contractOp.getLhs();
     Value rhs = contractOp.getRhs();
     size_t index = 0;
     bool changed = false;
     for (Value *operand : {&lhs, &rhs}) {
       AffineMap &map = maps[index++];
       auto broadcast = operand->getDefiningOp<vector::BroadcastOp>();
LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
 
     // Determine which dims are usused, now that the maps have been composed
     // with the broadcast maps.
     llvm::SmallBitVector unusedDimsBitVector = getUnusedDimsBitVector(maps);
     // Compress unused dims.
     for (auto &m : maps)
       m = compressDims(m, unusedDimsBitVector);
     // Compute the combined iterators.
-    SmallVector<Attribute, 4> iterators;
+    SmallVector<Attribute> iterators;
     for (unsigned i = 0; i < unusedDimsBitVector.size(); ++i) {
       if (!unusedDimsBitVector.test(i))
         iterators.push_back(contractOp.getIteratorTypes().getValue()[i]);
     }
     // Check that compressing unused dims isn't removing all reduction dimension
     // pairs. For example, if the vector.contract had only one reduction
     // iterator and that was a unit-dimension created by a broadcast,
     // then we should bail here, otherwise we would create a contract without
struct ReorderElementwiseOpsOnTranspose final
   using OpTraitRewritePattern::OpTraitRewritePattern;
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     if (op->getNumResults() != 1 || op->getNumRegions() != 0)
       return failure();
 
     // Make sure all operands are transpose/constant ops and collect their
     // transposition maps.
-    SmallVector<ArrayAttr, 4> transposeMaps;
+    SmallVector<ArrayAttr> transposeMaps;
     transposeMaps.reserve(op->getNumOperands());
     // Record the initial type before transposition. We'll use its shape later.
     // Any type will do here as we will check all transpose maps are the same.
     VectorType srcType;
     for (Value operand : op->getOperands()) {
       auto transposeOp = operand.getDefiningOp<vector::TransposeOp>();
       if (transposeOp) {
         transposeMaps.push_back(transposeOp.getTransp());
         srcType = transposeOp.getVectorType();
       } else if (!matchPattern(operand, m_Constant())) {
         return failure();
       }
     }
     if (transposeMaps.empty())
       return failure();
     // This is an elementwise op, so all transposed operands should have the
     // same type. We need to additionally check that all transposes uses the
     // same map.
     if (!llvm::all_equal(transposeMaps))
       return rewriter.notifyMatchFailure(op, "different transpose map");
 
-    SmallVector<Value, 4> srcValues;
+    SmallVector<Value> srcValues;
     srcValues.reserve(op->getNumOperands());
 
     // If there are constant operands, we need to insert inverse transposes for
     // them. Calculate the inverse order first.
     auto order = extractVector<unsigned>(transposeMaps.front());
     SmallVector<int64_t> invOrder(order.size());
     for (int i = 0, e = order.size(); i < e; ++i)
       invOrder[order[i]] = i;
ContractionOpToDotLowering::matchAndRewrite(vector::ContractionOp op,
   static constexpr std::array<int64_t, 2> perm = {1, 0};
   Location loc = op.getLoc();
   Value lhs = op.getLhs(), rhs = op.getRhs();
 
   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
   auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
   AffineExpr m, n, k;
   bindDims(rewriter.getContext(), m, n, k);
-  SmallVector<AffineMap, 4> maps = op.getIndexingMapsArray();
+  SmallVector<AffineMap> maps = op.getIndexingMapsArray();
   //
   // In the following we wish to make the reduction dimension innermost so we
   // can load vectors and just fmul + reduce into a scalar.
   //
   if (isParallelIterator(iteratorTypes[0]) &&
       isParallelIterator(iteratorTypes[1]) &&
       isReductionIterator(iteratorTypes[2])) {
     //

 FailureOr<Value>
 ContractionOpLowering::lowerParallel(vector::ContractionOp op, int64_t lhsIndex,
                                      int64_t rhsIndex,
                                      PatternRewriter &rewriter) const {
   VectorType lhsType = op.getLhsType();
   VectorType rhsType = op.getRhsType();
   VectorType resType = op.getResultType().cast<VectorType>();
   // Find the iterator type index and result index.
-  SmallVector<AffineMap, 4> iMap = op.getIndexingMapsArray();
+  SmallVector<AffineMap> iMap = op.getIndexingMapsArray();
   int64_t iterIndex = -1;
   int64_t dimSize = -1;
   if (lhsIndex >= 0) {
     iterIndex = iMap[0].getDimPosition(lhsIndex);
     if (rhsIndex >= 0 && iterIndex != iMap[1].getDimPosition(rhsIndex))
       return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
         diag << "expected lhsIndex=" << lhsIndex << " and rhsIndex=" << rhsIndex
              << " to map to the same dimension";
ContractionOpLowering::lowerReduction(vector::ContractionOp op,
   VectorType rhsType = op.getRhsType();
   Type resType = op.getResultType();
   if (resType.isa<VectorType>())
     return rewriter.notifyMatchFailure(op,
                                        "did not expect a VectorType result");
   bool isInt = resType.isa<IntegerType>();
   // Use iterator index 0.
   int64_t iterIndex = 0;
-  SmallVector<AffineMap, 4> iMap = op.getIndexingMapsArray();
+  SmallVector<AffineMap> iMap = op.getIndexingMapsArray();
   Optional<int64_t> lookupLhs = getResultIndex(iMap[0], iterIndex);
   Optional<int64_t> lookupRhs = getResultIndex(iMap[1], iterIndex);
   if (!lookupLhs.has_value())
     return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
       diag << "expected iterIndex=" << iterIndex << "to map to a LHS dimension";
     });
   if (!lookupRhs.has_value())
     return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
TransferReadToVectorLoadLowering(MLIRContext *context,
       : OpRewritePattern<vector::TransferReadOp>(context, benefit),
         maxTransferRank(maxRank) {}
 
   LogicalResult matchAndRewrite(vector::TransferReadOp read,
                                 PatternRewriter &rewriter) const override {
     if (maxTransferRank && read.getVectorType().getRank() > *maxTransferRank)
       return failure();
 
-    SmallVector<unsigned, 4> broadcastedDims;
+    SmallVector<unsigned> broadcastedDims;
     // Permutations are handled by VectorToSCF or
     // populateVectorTransferPermutationMapLoweringPatterns.
     // We let the 0-d corner case pass-through as it is supported.
     if (!read.getPermutationMap().isMinorIdentityWithBroadcasting(
             &broadcastedDims))
       return failure();
 
     auto memRefType = read.getShapedType().dyn_cast<MemRefType>();
     if (!memRefType)
       return failure();
 
     // Non-unit strides are handled by VectorToSCF.
     if (!vector::isLastMemrefDimUnitStride(memRefType))
       return failure();
 
     // If there is broadcasting involved then we first load the unbroadcasted
     // vector, and then broadcast it with `vector.broadcast`.
     ArrayRef<int64_t> vectorShape = read.getVectorType().getShape();
-    SmallVector<int64_t, 4> unbroadcastedVectorShape(vectorShape.begin(),
-                                                     vectorShape.end());
+    SmallVector<int64_t> unbroadcastedVectorShape(vectorShape.begin(),
     for (unsigned i : broadcastedDims)
       unbroadcastedVectorShape[i] = 1;
     VectorType unbroadcastedVectorType = VectorType::get(
         unbroadcastedVectorShape, read.getVectorType().getElementType());
 
     // `vector.load` supports vector types as memref's elements only when the
     // resulting vector type is the same as the element type.
     auto memrefElTy = memRefType.getElementType();
LogicalResult matchAndRewrite(vector::TransferWriteOp write,
     }
     return success();
   }
 
   llvm::Optional<unsigned> maxTransferRank;
 };
 
 // Returns the values in `arrayAttr` as an integer vector.
-static SmallVector<int64_t, 4> getIntValueVector(ArrayAttr arrayAttr) {
+static SmallVector<int64_t> getIntValueVector(ArrayAttr arrayAttr) {
   return llvm::to_vector<4>(
       llvm::map_range(arrayAttr.getAsRange<IntegerAttr>(),
                       [](IntegerAttr attr) { return attr.getInt(); }));
 }
 
 // Shuffles vector.bitcast op after vector.extract op.
 //
 // This transforms IR like:
LogicalResult matchAndRewrite(vector::ExtractStridedSliceOp extractOp,
     int64_t expandRatio = castDstLastDim / castSrcLastDim;
 
     // If we have a less number of offsets than the rank, then implicitly we
     // are selecting the full range for the last bitcasted dimension; other
     // dimensions aren't affected. Otherwise, we need to scale down the last
     // dimension's offset given we are extracting from less elements now.
     ArrayAttr newOffsets = extractOp.getOffsets();
     if (newOffsets.size() == rank) {
-      SmallVector<int64_t, 4> offsets = getIntValueVector(newOffsets);
+      SmallVector<int64_t> offsets = getIntValueVector(newOffsets);
       if (offsets.back() % expandRatio != 0)
         return failure();
       offsets.back() = offsets.back() / expandRatio;
       newOffsets = rewriter.getI64ArrayAttr(offsets);
     }
 
     // Similarly for sizes.
     ArrayAttr newSizes = extractOp.getSizes();
     if (newSizes.size() == rank) {
-      SmallVector<int64_t, 4> sizes = getIntValueVector(newSizes);
+      SmallVector<int64_t> sizes = getIntValueVector(newSizes);
       if (sizes.back() % expandRatio != 0)
         return failure();
       sizes.back() = sizes.back() / expandRatio;
       newSizes = rewriter.getI64ArrayAttr(sizes);
     }
 
-    SmallVector<int64_t, 4> dims =
+    SmallVector<int64_t> dims =
         llvm::to_vector<4>(extractOp.getType().cast<VectorType>().getShape());
     dims.back() = dims.back() / expandRatio;
     VectorType newExtractType =
         VectorType::get(dims, castSrcType.getElementType());
 
     auto newExtractOp = rewriter.create<vector::ExtractStridedSliceOp>(
         extractOp.getLoc(), newExtractType, castOp.getSource(), newOffsets,
         newSizes, extractOp.getStrides());
LogicalResult matchAndRewrite(vector::BitCastOp bitcastOp,
     unsigned destinationWidth =
         castDstType.getElementType().getIntOrFloatBitWidth();
     unsigned numElements = destinationWidth / sourceWidth;
     if (insertOp.getSourceVectorType().getNumElements() % numElements != 0)
       return failure();
 
     ArrayAttr newOffsets = insertOp.getOffsets();
     assert(newOffsets.size() == rank);
-    SmallVector<int64_t, 4> offsets = getIntValueVector(newOffsets);
+    SmallVector<int64_t> offsets = getIntValueVector(newOffsets);
     if (offsets.back() % shrinkRatio != 0)
       return failure();
     offsets.back() = offsets.back() / shrinkRatio;
     newOffsets = rewriter.getI64ArrayAttr(offsets);
 
-    SmallVector<int64_t, 4> srcDims =
+    SmallVector<int64_t> srcDims =
         llvm::to_vector<4>(insertOp.getSourceVectorType().getShape());
     srcDims.back() = srcDims.back() / shrinkRatio;
     VectorType newCastSrcType =
         VectorType::get(srcDims, castDstType.getElementType());
 
     auto newCastSrcOp = rewriter.create<vector::BitCastOp>(
         bitcastOp.getLoc(), newCastSrcType, insertOp.getSource());
 
-    SmallVector<int64_t, 4> dstDims =
+    SmallVector<int64_t> dstDims =
         llvm::to_vector<4>(insertOp.getDestVectorType().getShape());
     dstDims.back() = dstDims.back() / shrinkRatio;
     VectorType newCastDstType =
         VectorType::get(dstDims, castDstType.getElementType());
 
     auto newCastDstOp = rewriter.create<vector::BitCastOp>(
         bitcastOp.getLoc(), newCastDstType, insertOp.getDest());
 

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

 
 #define DEBUG_TYPE "vector-unrolling"
 
 using namespace mlir;
 using namespace mlir::vector;
 
 /// During unrolling from `originalShape` to `targetShape` return the offset for
 /// the slice `index`.
-static SmallVector<int64_t, 4> getVectorOffset(ArrayRef<int64_t> originalShape,
-                                               ArrayRef<int64_t> targetShape,
-                                               int64_t index) {
-  SmallVector<int64_t, 4> dstSliceStrides =
-      computeStrides(originalShape, targetShape);
-  SmallVector<int64_t, 4> vectorOffsets = delinearize(dstSliceStrides, index);
-  SmallVector<int64_t, 4> elementOffsets =
-      computeElementOffsetsFromVectorSliceOffsets(targetShape, vectorOffsets);
-  return elementOffsets;
+static SmallVector<int64_t> getVectorOffset(ArrayRef<int64_t> ratioStrides,
+                                            int64_t index,
+                                            ArrayRef<int64_t> targetShape) {
+  return computeElementwiseMul(delinearize(ratioStrides, index), 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).
+/// 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, 4> vectorShape;
+  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() == targetShape.size());
+    assert(originalShape.size() >= targetShape.size());
+    assert(loopOrder.size() == originalShape.size());
 
     // Compute the count for each dimension.
-    SmallVector<int64_t> sliceDimCounts(originalShape.size());
-    for (unsigned r = 0; r < originalShape.size(); ++r) {
-      sliceDimCounts[r] = ceilDiv(originalShape[r], targetShape[r]);
-      maxIndexVal *= sliceDimCounts[r];
-    }
+    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];
     }
public:
 
   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 =
-        computeElementOffsetsFromVectorSliceOffsets(vectorShape, vectorOffsets);
+        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,
static Operation *cloneOpWithOperandsAndTypes(OpBuilder &builder, Location loc,
                                               ArrayRef<Value> operands,
                                               ArrayRef<Type> resultTypes) {
   return builder.create(loc, op->getName().getIdentifier(), operands,
                         resultTypes, op->getAttrs());
 }
 
 /// Return the target shape for unrolling for the given `op`. Return llvm::None
 /// if the op shouldn't be or cannot be unrolled.
-static Optional<SmallVector<int64_t, 4>>
+static Optional<SmallVector<int64_t>>
 getTargetShape(const vector::UnrollVectorOptions &options, Operation *op) {
   if (options.filterConstraint && failed(options.filterConstraint(op)))
     return llvm::None;
   assert(options.nativeShape &&
          "vector unrolling expects the native shape or native"
          "shape call back function to be set");
   auto unrollableVectorOp = dyn_cast<VectorUnrollOpInterface>(op);
   if (!unrollableVectorOp)
     return llvm::None;
   auto maybeUnrollShape = unrollableVectorOp.getShapeForUnroll();
   if (!maybeUnrollShape)
     return llvm::None;
-  Optional<SmallVector<int64_t, 4>> targetShape = options.nativeShape(op);
+  Optional<SmallVector<int64_t>> targetShape = options.nativeShape(op);
   if (!targetShape)
     return llvm::None;
-  auto maybeShapeRatio = shapeRatio(*maybeUnrollShape, *targetShape);
+  auto maybeShapeRatio = computeShapeRatio(*maybeUnrollShape, *targetShape);
   if (!maybeShapeRatio ||
       llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; }))
     return llvm::None;
   return targetShape;
 }
 
 static SmallVector<int64_t>
 getUnrollOrder(unsigned numLoops, Operation *op,
LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
     if (readOp.getTransferRank() == 0)
       return failure();
     if (readOp.getMask())
       return failure();
     auto targetShape = getTargetShape(options, readOp);
     if (!targetShape)
       return failure();
     auto sourceVectorType = readOp.getVectorType();
-    SmallVector<int64_t, 4> strides(targetShape->size(), 1);
+    SmallVector<int64_t> strides(targetShape->size(), 1);
     Location loc = readOp.getLoc();
     ArrayRef<int64_t> originalSize = readOp.getVectorType().getShape();
 
     // Prepare the result vector;
     Value result = rewriter.create<arith::ConstantOp>(
         loc, sourceVectorType, rewriter.getZeroAttr(sourceVectorType));
     auto targetType =
         VectorType::get(*targetShape, sourceVectorType.getElementType());
-    SmallVector<Value, 4> originalIndices(readOp.getIndices().begin(),
-                                          readOp.getIndices().end());
+    SmallVector<Value> originalIndices(readOp.getIndices().begin(),
 
     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, 4> elementOffsets =
-          indexToOffsets.getVectorOffset(i);
-      SmallVector<Value, 4> indices =
+      SmallVector<int64_t> elementOffsets = indexToOffsets.getVectorOffset(i);
+      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());
 
       result = rewriter.create<vector::InsertStridedSliceOp>(
if (writeOp.getTransferRank() == 0)
       return failure();
 
     if (writeOp.getMask())
       return failure();
     auto targetShape = getTargetShape(options, writeOp);
     if (!targetShape)
       return failure();
     auto sourceVectorType = writeOp.getVectorType();
-    SmallVector<int64_t, 4> strides(targetShape->size(), 1);
+    SmallVector<int64_t> strides(targetShape->size(), 1);
     Location loc = writeOp.getLoc();
     ArrayRef<int64_t> originalSize = sourceVectorType.getShape();
-    SmallVector<Value, 4> originalIndices(writeOp.getIndices().begin(),
-                                          writeOp.getIndices().end());
+    SmallVector<Value> originalIndices(writeOp.getIndices().begin(),
 
     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, 4> elementOffsets =
-          indexToOffsets.getVectorOffset(i);
+      SmallVector<int64_t> elementOffsets = indexToOffsets.getVectorOffset(i);
       Value slicedVector = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, writeOp.getVector(), elementOffsets, *targetShape, strides);
-      SmallVector<Value, 4> indices =
+      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());
       // For the tensor case update the destination for the next transfer write.
       if (!slicedWrite->getResults().empty())
         resultTensor = slicedWrite->getResult(0);
UnrollContractionPattern(MLIRContext *context,
         options(options) {}
 
   LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                 PatternRewriter &rewriter) const override {
     auto targetShape = getTargetShape(options, contractOp);
     if (!targetShape)
       return failure();
     auto dstVecType = contractOp.getResultType().cast<VectorType>();
-    SmallVector<int64_t, 4> originalSize = *contractOp.getShapeForUnroll();
+    SmallVector<int64_t> originalSize = *contractOp.getShapeForUnroll();
 
     Location loc = contractOp.getLoc();
     unsigned accIndex = vector::ContractionOp::getAccOperandIndex();
     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, 4> offsets = indexToOffsets.getVectorOffset(i);
-      SmallVector<Value, 4> slicesOperands(contractOp.getNumOperands());
+      SmallVector<int64_t> offsets = indexToOffsets.getVectorOffset(i);
+      SmallVector<Value> slicesOperands(contractOp.getNumOperands());
 
-      // Helper to coompute the new shape of each operand and extract the slice.
+      // 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));
-        SmallVector<int64_t, 4> operandStrides(operandOffets.size(), 1);
+        SmallVector<int64_t> operandStrides(operandOffets.size(), 1);
         slicesOperands[index] = rewriter.create<vector::ExtractStridedSliceOp>(
             loc, operand, operandOffets, operandShape, operandStrides);
       };
 
       // Extract the new lhs operand.
       AffineMap lhsPermutationMap = contractOp.getIndexingMapsArray()[0];
       SmallVector<int64_t> lhsOffets =
           applyPermutationMap(lhsPermutationMap, ArrayRef<int64_t>(offsets));
struct UnrollMultiReductionPattern
   UnrollMultiReductionPattern(MLIRContext *context,
                               const vector::UnrollVectorOptions &options,
                               PatternBenefit benefit = 1)
       : OpRewritePattern<vector::MultiDimReductionOp>(context, benefit),
         options(options) {}
 
   LogicalResult matchAndRewrite(vector::MultiDimReductionOp reductionOp,
                                 PatternRewriter &rewriter) const override {
-    Optional<SmallVector<int64_t, 4>> targetShape =
+    Optional<SmallVector<int64_t>> targetShape =
         getTargetShape(options, reductionOp);
     if (!targetShape)
       return failure();
-    SmallVector<int64_t, 4> originalSize = *reductionOp.getShapeForUnroll();
-    SmallVector<int64_t, 4> ratio = *shapeRatio(originalSize, *targetShape);
+    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, 4> offsets =
-          getVectorOffset(originalSize, *targetShape, i);
+      SmallVector<int64_t> offsets =
+          getVectorOffset(ratioStrides, i, *targetShape);
 
       SmallVector<Value> operands;
-      SmallVector<int64_t, 4> operandStrides(offsets.size(), 1);
+      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())) {
         if (!reductionOp.isReducedDim(i)) {
           destOffset.push_back(offsets[i]);
           dstShape.push_back((*targetShape)[i]);
         }
       }
       Value acc;
-      SmallVector<int64_t, 4> accStrides(destOffset.size(), 1);
+      SmallVector<int64_t> accStrides(destOffset.size(), 1);
       // If a version of the accumulator has already been computed, use it
       // otherwise extract the first version from the original operand.
       auto accIt = accCache.find(destOffset);
       if (accIt != accCache.end())
         acc = accIt->second;
       else
         acc = rewriter.create<vector::ExtractStridedSliceOp>(
             loc, reductionOp.getAcc(), destOffset, dstShape, accStrides);
struct UnrollElementwisePattern : public RewritePattern {
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
       return failure();
     auto targetShape = getTargetShape(options, op);
     if (!targetShape)
       return failure();
     auto dstVecType = op->getResult(0).getType().cast<VectorType>();
-    SmallVector<int64_t, 4> originalSize =
+    SmallVector<int64_t> originalSize =
         *cast<VectorUnrollOpInterface>(op).getShapeForUnroll();
-    SmallVector<int64_t, 4> ratio = *shapeRatio(originalSize, *targetShape);
+    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, 4> strides(targetShape->size(), 1);
+    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, 4> offsets =
-          getVectorOffset(originalSize, *targetShape, i);
-      SmallVector<Value, 4> extractOperands;
+      SmallVector<int64_t> offsets =
+          getVectorOffset(ratioStrides, i, *targetShape);
+      SmallVector<Value> extractOperands;
       for (OpOperand &operand : op->getOpOperands()) {
         auto vecType = operand.get().getType().template dyn_cast<VectorType>();
         if (!vecType) {
           extractOperands.push_back(operand.get());
           continue;
         }
         extractOperands.push_back(
             rewriter.create<vector::ExtractStridedSliceOp>(
struct UnrollReductionPattern : public OpRewritePattern<vector::ReductionOp> {
   UnrollReductionPattern(MLIRContext *context,
                          const vector::UnrollVectorOptions &options,
                          PatternBenefit benefit = 1)
       : OpRewritePattern<vector::ReductionOp>(context, benefit),
         options(options) {}
 
   LogicalResult matchAndRewrite(vector::ReductionOp reductionOp,
                                 PatternRewriter &rewriter) const override {
-    Optional<SmallVector<int64_t, 4>> targetShape =
+    Optional<SmallVector<int64_t>> targetShape =
         getTargetShape(options, reductionOp);
     if (!targetShape)
       return failure();
     SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
-    int64_t ratio = (*shapeRatio(originalSize, *targetShape))[0];
+    auto ratio = *computeShapeRatio(originalSize, *targetShape);
+    int64_t sliceCount = ratio[0];
 
     // Create unrolled vector reduction.
     Location loc = reductionOp.getLoc();
     Value accumulator = nullptr;
-    for (int64_t i = 0; i < ratio; ++i) {
+
+    // 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(originalSize, *targetShape, i);
+          getVectorOffset(ratioStrides, i, *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) {
struct UnrollTranposePattern : public OpRewritePattern<vector::TransposeOp> {
   LogicalResult matchAndRewrite(vector::TransposeOp tranposeOp,
                                 PatternRewriter &rewriter) const override {
     if (tranposeOp.getResultType().getRank() == 0)
       return failure();
     auto targetShape = getTargetShape(options, tranposeOp);
     if (!targetShape)
       return failure();
     auto originalVectorType = tranposeOp.getResultType();
-    SmallVector<int64_t, 4> strides(targetShape->size(), 1);
+    SmallVector<int64_t> strides(targetShape->size(), 1);
     Location loc = tranposeOp.getLoc();
     ArrayRef<int64_t> originalSize = originalVectorType.getShape();
-    SmallVector<int64_t, 4> ratio = *shapeRatio(originalSize, *targetShape);
+    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;
     tranposeOp.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, 4> elementOffsets =
-          getVectorOffset(originalSize, *targetShape, i);
-      SmallVector<int64_t, 4> permutedOffsets(elementOffsets.size());
-      SmallVector<int64_t, 4> permutedShape(elementOffsets.size());
+      SmallVector<int64_t> elementOffsets =
+          getVectorOffset(ratioStrides, i, *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>(
           loc, tranposeOp.getVector(), permutedOffsets, permutedShape, strides);
       Value tranposedSlice =

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

 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 
 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
-#include <numeric>
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value mlir::vector::createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
                                       int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);
   llvm_unreachable("Expected MemRefType or TensorType");
 }
 
-/// Return the number of elements of basis, `0` if empty.
-int64_t mlir::computeMaxLinearIndex(ArrayRef<int64_t> basis) {
-  if (basis.empty())
-    return 0;
-  return std::accumulate(basis.begin(), basis.end(), 1,
-                         std::multiplies<int64_t>());
-}
-
-SmallVector<int64_t, 4> mlir::computeStrides(ArrayRef<int64_t> shape,
-                                             ArrayRef<int64_t> sizes) {
-  int64_t rank = shape.size();
-  // Compute the count for each dimension.
-  SmallVector<int64_t, 4> sliceDimCounts(rank);
-  for (int64_t r = 0; r < rank; ++r)
-    sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
-  // Use that to compute the slice stride for each dimension.
-  SmallVector<int64_t, 4> sliceStrides(rank);
-  sliceStrides[rank - 1] = 1;
-  for (int64_t r = rank - 2; r >= 0; --r)
-    sliceStrides[r] = sliceStrides[r + 1] * sliceDimCounts[r + 1];
-  return sliceStrides;
-}
-
-SmallVector<int64_t, 4> mlir::computeElementOffsetsFromVectorSliceOffsets(
-    ArrayRef<int64_t> sizes, ArrayRef<int64_t> vectorOffsets) {
-  SmallVector<int64_t, 4> result;
-  for (auto it : llvm::zip(vectorOffsets, sizes))
-    result.push_back(std::get<0>(it) * std::get<1>(it));
-  return result;
-}
-
-Optional<SmallVector<int64_t, 4>> mlir::shapeRatio(ArrayRef<int64_t> superShape,
-                                                   ArrayRef<int64_t> subShape) {
-  if (superShape.size() < subShape.size()) {
-    return None;
-  }
-
-  // Starting from the end, compute the integer divisors.
-  std::vector<int64_t> result;
-  result.reserve(superShape.size());
-  for (auto [superSize, subSize] :
-       llvm::zip(llvm::reverse(superShape), llvm::reverse(subShape))) {
-    assert(superSize > 0 && "superSize must be > 0");
-    assert(subSize > 0 && "subSize must be > 0");
-
-    // If integral division does not occur, return and let the caller decide.
-    if (superSize % subSize != 0)
-      return None;
-    result.push_back(superSize / subSize);
-  }
-
-  // At this point we computed the ratio (in reverse) for the common
-  // size. Fill with the remaining entries from the super-vector shape (still in
-  // reverse).
-  int commonSize = subShape.size();
-  std::copy(superShape.rbegin() + commonSize, superShape.rend(),
-            std::back_inserter(result));
-
-  assert(result.size() == superShape.size() &&
-         "super to sub shape ratio is not of the same size as the super rank");
-
-  // Reverse again to get it back in the proper order and return.
-  return SmallVector<int64_t, 4>{result.rbegin(), result.rend()};
-}
-
-Optional<SmallVector<int64_t, 4>> mlir::shapeRatio(VectorType superVectorType,
-                                                   VectorType subVectorType) {
-  assert(superVectorType.getElementType() == subVectorType.getElementType() &&
-         "vector types must be of the same elemental type");
-  return shapeRatio(superVectorType.getShape(), subVectorType.getShape());
-}
-
 /// Constructs a permutation map from memref indices to vector dimension.
 ///
 /// The implementation uses the knowledge of the mapping of enclosing loop to
 /// vector dimension. `enclosingLoopToVectorDim` carries this information as a
 /// map with:
 ///   - keys representing "vectorized enclosing loops";
 ///   - values representing the corresponding vector dimension.
 /// The algorithm traverses "vectorized enclosing loops" and extracts the

 /// header file.
 static AffineMap makePermutationMap(
     ArrayRef<Value> indices,
     const DenseMap<Operation *, unsigned> &enclosingLoopToVectorDim) {
   if (enclosingLoopToVectorDim.empty())
     return AffineMap();
   MLIRContext *context =
       enclosingLoopToVectorDim.begin()->getFirst()->getContext();
-  SmallVector<AffineExpr, 4> perm(enclosingLoopToVectorDim.size(),
-                                  getAffineConstantExpr(0, context));
+  SmallVector<AffineExpr> perm(enclosingLoopToVectorDim.size(),
 
   for (auto kvp : enclosingLoopToVectorDim) {
     assert(kvp.second < perm.size());
     auto invariants = getInvariantAccesses(
         cast<AffineForOp>(kvp.first).getInductionVar(), indices);
     unsigned numIndices = indices.size();
     unsigned countInvariantIndices = 0;
     for (unsigned dim = 0; dim < numIndices; ++dim) {
bool matcher::operatesOnSuperVectorsOf(Operation &op,
   } else {
     // Not a vector.transfer and has more than 1 result, fail hard for now to
     // wake us up when something changes.
     op.emitError("NYI: operation has more than 1 result");
     return false;
   }
 
   // Get the ratio.
-  auto ratio = shapeRatio(superVectorType, subVectorType);
+  auto ratio =
+      computeShapeRatio(superVectorType.getShape(), subVectorType.getShape());
 
   // Sanity check.
   assert((ratio || !mustDivide) &&
          "vector.transfer operation in which super-vector size is not an"
          " integer multiple of sub-vector size");
 
   // This catches cases that are not strictly necessary to have multiplicity but
   // still aren't divisible by the sub-vector shape.
   // This could be useful information if we wanted to reshape at the level of
   // the vector type (but we would have to look at the compute and distinguish
   // between parallel, reduction and possibly other cases.
   return ratio.has_value();
 }

mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp

 
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/NestedMatcher.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
 #include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
void VectorizerTestPass::testVectorShapeRatio(llvm::raw_ostream &outs) {
   pat.match(f, &matches);
   for (auto m : matches) {
     auto *opInst = m.getMatchedOperation();
     // This is a unit test that only checks and prints shape ratio.
     // As a consequence we write only Ops with a single return type for the
     // purpose of this test. If we need to test more intricate behavior in the
     // future we can always extend.
     auto superVectorType = opInst->getResult(0).getType().cast<VectorType>();
-    auto ratio = shapeRatio(superVectorType, subVectorType);
+    auto ratio =
+        computeShapeRatio(superVectorType.getShape(), subVectorType.getShape());
     if (!ratio) {
       opInst->emitRemark("NOT MATCHED");
     } else {
       outs << "\nmatched: " << *opInst << " with shape ratio: ";
       llvm::interleaveComma(MutableArrayRef<int64_t>(*ratio), outs);
     }
   }
 }

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

void runOnOperation() override {
     populateVectorToVectorCanonicalizationPatterns(patterns);
     populateBubbleVectorBitCastOpPatterns(patterns);
     populateCastAwayVectorLeadingOneDimPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 
 private:
   // Return the target shape based on op type.
-  static Optional<SmallVector<int64_t, 4>> getShape(Operation *op) {
+  static Optional<SmallVector<int64_t>> getShape(Operation *op) {
     if (isa<arith::AddFOp, arith::SelectOp, arith::CmpFOp>(op))
-      return SmallVector<int64_t, 4>(2, 2);
+      return SmallVector<int64_t>(2, 2);
     if (isa<vector::ContractionOp>(op))
-      return SmallVector<int64_t, 4>(3, 2);
+      return SmallVector<int64_t>(3, 2);
     // For transfer ops, just propagate the shape coming from
     // InsertStridedSlices/ExtractStridedSlices.
     if (auto readOp = dyn_cast<vector::TransferReadOp>(op)) {
       VectorType dstVec;
       for (Operation *users : readOp->getUsers()) {
         auto extract = dyn_cast<ExtractStridedSliceOp>(users);
         if (!extract)
           return llvm::None;
         auto vecType = extract.getResult().getType().cast<VectorType>();
         if (dstVec && dstVec != vecType)
           return llvm::None;
         dstVec = vecType;
       }
-      return SmallVector<int64_t, 4>(dstVec.getShape().begin(),
-                                     dstVec.getShape().end());
+      return SmallVector<int64_t>(dstVec.getShape().begin(),
     }
     if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {
       auto insert = writeOp.getVector().getDefiningOp<InsertStridedSliceOp>();
       if (!insert)
         return llvm::None;
       ArrayRef<int64_t> shape = insert.getSourceVectorType().getShape();
-      return SmallVector<int64_t, 4>(shape.begin(), shape.end());
+      return SmallVector<int64_t>(shape.begin(), shape.end());
     }
     return llvm::None;
   }
 
   static LogicalResult filter(Operation *op) {
     return success(isa<arith::AddFOp, arith::SelectOp, arith::CmpFOp,
                        ContractionOp, TransferReadOp, TransferWriteOp>(op));
   }
populateVectorUnrollPatterns(
         patterns, UnrollVectorOptions()
                       .setNativeShape(ArrayRef<int64_t>{1, 3, 4, 2})
                       .setFilterConstraint([](Operation *op) {
                         return success(isa<vector::TransposeOp>(op));
                       }));
 
     if (unrollBasedOnType) {
       UnrollVectorOptions::NativeShapeFnType nativeShapeFn =
-          [](Operation *op) -> Optional<SmallVector<int64_t, 4>> {
+          [](Operation *op) -> Optional<SmallVector<int64_t>> {
         vector::ContractionOp contractOp = cast<vector::ContractionOp>(op);
-        SmallVector<int64_t, 4> nativeShape(
-            contractOp.getIteratorTypes().size(), 4);
+        SmallVector<int64_t> nativeShape(contractOp.getIteratorTypes().size(),
+                                         4);
         Type lhsType = contractOp.getLhsType().getElementType();
         nativeShape[nativeShape.size() - 1] = lhsType.isF16() ? 4 : 2;
         return nativeShape;
       };
 
       UnrollVectorOptions opts;
       opts.setNativeShapeFn(nativeShapeFn)
           .setFilterConstraint(
               [](Operation *op) { return success(isa<ContractionOp>(op)); });
 
       if (!unrollOrder.empty()) {
         opts.setUnrollTraversalOrderFn([this](Operation *op)
                                            -> Optional<SmallVector<int64_t>> {
           vector::ContractionOp contractOp = cast<vector::ContractionOp>(op);
           if (contractOp.getIteratorTypes().size() == unrollOrder.size())
             return SmallVector<int64_t>(unrollOrder.begin(), unrollOrder.end());
           return None;
         });
       }
       populateVectorUnrollPatterns(patterns, opts);
     } else {
-      auto nativeShapeFn =
-          [](Operation *op) -> Optional<SmallVector<int64_t, 4>> {
+      auto nativeShapeFn = [](Operation *op) -> Optional<SmallVector<int64_t>> {
         auto contractOp = dyn_cast<ContractionOp>(op);
         if (!contractOp)
           return None;
-        return SmallVector<int64_t, 4>(contractOp.getIteratorTypes().size(), 2);
+        return SmallVector<int64_t>(contractOp.getIteratorTypes().size(), 2);
       };
       populateVectorUnrollPatterns(patterns,
                                    UnrollVectorOptions()
                                        .setNativeShapeFn(nativeShapeFn)
                                        .setFilterConstraint([](Operation *op) {
                                          return success(isa<ContractionOp>(op));
                                        }));
     }

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterface.cpp

explicit LinalgTransformationFilter(
       const FilterFunction &f, ArrayRef<StringAttr> matchDisjunction = {},
       Optional<StringAttr> replacement = None);
 
   LinalgTransformationFilter(LinalgTransformationFilter &&) = default;
   LinalgTransformationFilter(const LinalgTransformationFilter &) = default;
   LogicalResult checkAndNotify(PatternRewriter &rewriter, Operation *op) const;
   void replaceLinalgTransformationFilter(PatternRewriter &rewriter,
                                          Operation *op) const;
-  bool hasReplacementFilter(Operation *op) const;
 
   LinalgTransformationFilter &addFilter(const FilterFunction &f) {
     if (f)
       filters.push_back(f);
     return *this;
   }
 
   template <typename... OpTypes>
private:
   bool matchByDefault;
 };
 
 LinalgTransformationFilter::LinalgTransformationFilter(
     ArrayRef<StringAttr> matchDisjunction, Optional<StringAttr> replacement)
     : matchDisjunction(matchDisjunction.begin(), matchDisjunction.end()),
       replacement(replacement), matchByDefault(false) {}
 
-LinalgTransformationFilter::LinalgTransformationFilter(
-    const FilterFunction &f, ArrayRef<StringAttr> matchDisjunction,
-    Optional<StringAttr> replacement)
-    : matchDisjunction(matchDisjunction.begin(), matchDisjunction.end()),
-      replacement(replacement), matchByDefault(false) {
-  if (f)
-    filters.push_back(f);
-}
-
 LogicalResult
 LinalgTransformationFilter::checkAndNotify(PatternRewriter &rewriter,
                                            Operation *op) const {
   if (llvm::any_of(filters,
                    [&](const FilterFunction &f) { return failed(f(op)); }))
     return failure();
 
   auto attr = op->template getAttrOfType<StringAttr>(kLinalgTransformMarker);

 void LinalgTransformationFilter::replaceLinalgTransformationFilter(
     PatternRewriter &rewriter, Operation *op) const {
   if (replacement.has_value())
     op->setAttr(kLinalgTransformMarker, replacement.value());
   else
     op->removeAttr(rewriter.getStringAttr(kLinalgTransformMarker));
 }
 
-bool LinalgTransformationFilter::hasReplacementFilter(Operation *op) const {
-  if (!replacement)
-    return false;
-  auto attr = op->getAttr(kLinalgTransformMarker).dyn_cast<StringAttr>();
-  return attr && attr == *replacement;
-}
-
 /// Pattern for testing `TileUsingSCFForOp` pattern (that tiles operations using
 /// the `TilingInterface` with `scf.for` ops for iterating over the tiles) while
 /// using a `filter` to avoid recursive application.
 struct TestTileUsingSCFForOp
     : public OpInterfaceRewritePattern<TilingInterface> {
   TestTileUsingSCFForOp(
       MLIRContext *context, scf::SCFTilingOptions options,
       LinalgTransformationFilter filter = LinalgTransformationFilter(),