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

//===- VectorUtils.cpp - MLIR Utilities for VectorOps   ------------------===//
//
// Part of the MLIR 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
//
//===----------------------------------------------------------------------===//
//
// This file implements utility methods for working with the Vector dialect.
//
//===----------------------------------------------------------------------===//

#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/Arith/Utils/Utils.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/Matchers.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Support/LLVM.h"

#include "llvm/ADT/DenseSet.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/InterleavedRange.h"

#define DEBUG_TYPE "vector-utils"

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 (isa<UnrankedMemRefType, MemRefType>(source.getType()))
    return b.createOrFold<memref::DimOp>(loc, source, dim);
  if (isa<UnrankedTensorType, RankedTensorType>(source.getType()))
    return b.createOrFold<tensor::DimOp>(loc, source, dim);
  llvm_unreachable("Expected MemRefType or TensorType");
}

/// Given the n-D transpose pattern 'transp', return true if 'dim0' and 'dim1'
/// should be transposed with each other within the context of their 2D
/// transposition slice.
///
/// Example 1: dim0 = 0, dim1 = 2, transp = [2, 1, 0]
///   Return true: dim0 and dim1 are transposed within the context of their 2D
///   transposition slice ([1, 0]).
///
/// Example 2: dim0 = 0, dim1 = 1, transp = [2, 1, 0]
///   Return true: dim0 and dim1 are transposed within the context of their 2D
///   transposition slice ([1, 0]). Paradoxically, note how dim1 (1) is *not*
///   transposed within the full context of the transposition.
///
/// Example 3: dim0 = 0, dim1 = 1, transp = [2, 0, 1]
///   Return false: dim0 and dim1 are *not* transposed within the context of
///   their 2D transposition slice ([0, 1]). Paradoxically, note how dim0 (0)
///   and dim1 (1) are transposed within the full context of the of the
///   transposition.
static bool areDimsTransposedIn2DSlice(int64_t dim0, int64_t dim1,
                                       ArrayRef<int64_t> transp) {
  // Perform a linear scan along the dimensions of the transposed pattern. If
  // dim0 is found first, dim0 and dim1 are not transposed within the context of
  // their 2D slice. Otherwise, 'dim1' is found first and they are transposed.
  for (int64_t permDim : transp) {
    if (permDim == dim0)
      return false;
    if (permDim == dim1)
      return true;
  }

  llvm_unreachable("Ill-formed transpose pattern");
}

FailureOr<std::pair<int, int>>
mlir::vector::isTranspose2DSlice(vector::TransposeOp op) {
  VectorType srcType = op.getSourceVectorType();
  SmallVector<int64_t> srcGtOneDims;
  for (auto [index, size] : llvm::enumerate(srcType.getShape()))
    if (size > 1)
      srcGtOneDims.push_back(index);

  if (srcGtOneDims.size() != 2)
    return failure();

  // Check whether the two source vector dimensions that are greater than one
  // must be transposed with each other so that we can apply one of the 2-D
  // transpose patterns. Otherwise, these patterns are not applicable.
  if (!areDimsTransposedIn2DSlice(srcGtOneDims[0], srcGtOneDims[1],
                                  op.getPermutation()))
    return failure();

  return std::pair<int, int>(srcGtOneDims[0], srcGtOneDims[1]);
}

/// 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
/// at-most-one MemRef index that is invariant along said loop. This index is
/// guaranteed to be at most one by construction: otherwise the MemRef is not
/// vectorizable.
/// If this invariant index is found, it is added to the permutation_map at the
/// proper 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.
///
/// Returns an empty AffineMap if `enclosingLoopToVectorDim` is empty,
/// signalling that no permutation map can be constructed given
/// `enclosingLoopToVectorDim`.
///
/// Examples can be found in the documentation of `makePermutationMap`, in 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> perm(enclosingLoopToVectorDim.size(),
                               getAffineConstantExpr(0, context));

  for (auto kvp : enclosingLoopToVectorDim) {
    assert(kvp.second < perm.size());
    auto invariants = affine::getInvariantAccesses(
        cast<affine::AffineForOp>(kvp.first).getInductionVar(), indices);
    unsigned numIndices = indices.size();
    unsigned countInvariantIndices = 0;
    for (unsigned dim = 0; dim < numIndices; ++dim) {
      if (!invariants.count(indices[dim])) {
        assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
               "permutationMap already has an entry along dim");
        perm[kvp.second] = getAffineDimExpr(dim, context);
      } else {
        ++countInvariantIndices;
      }
    }
    assert((countInvariantIndices == numIndices ||
            countInvariantIndices == numIndices - 1) &&
           "Vectorization prerequisite violated: at most 1 index may be "
           "invariant wrt a vectorized loop");
    (void)countInvariantIndices;
  }
  return AffineMap::get(indices.size(), 0, perm, context);
}

/// Implementation detail that walks up the parents and records the ones with
/// the specified type.
/// TODO: could also be implemented as a collect parents followed by a
/// filter and made available outside this file.
template <typename T>
static SetVector<Operation *> getParentsOfType(Block *block) {
  SetVector<Operation *> res;
  auto *current = block->getParentOp();
  while (current) {
    if ([[maybe_unused]] auto typedParent = dyn_cast<T>(current)) {
      assert(res.count(current) == 0 && "Already inserted");
      res.insert(current);
    }
    current = current->getParentOp();
  }
  return res;
}

/// Returns the enclosing AffineForOp, from closest to farthest.
static SetVector<Operation *> getEnclosingforOps(Block *block) {
  return getParentsOfType<affine::AffineForOp>(block);
}

AffineMap mlir::makePermutationMap(
    Block *insertPoint, ArrayRef<Value> indices,
    const DenseMap<Operation *, unsigned> &loopToVectorDim) {
  DenseMap<Operation *, unsigned> enclosingLoopToVectorDim;
  auto enclosingLoops = getEnclosingforOps(insertPoint);
  for (auto *forInst : enclosingLoops) {
    auto it = loopToVectorDim.find(forInst);
    if (it != loopToVectorDim.end()) {
      enclosingLoopToVectorDim.insert(*it);
    }
  }
  return ::makePermutationMap(indices, enclosingLoopToVectorDim);
}

AffineMap mlir::makePermutationMap(
    Operation *op, ArrayRef<Value> indices,
    const DenseMap<Operation *, unsigned> &loopToVectorDim) {
  return makePermutationMap(op->getBlock(), indices, loopToVectorDim);
}

bool matcher::operatesOnSuperVectorsOf(Operation &op,
                                       VectorType subVectorType) {
  // First, extract the vector type and distinguish between:
  //   a. ops that *must* lower a super-vector (i.e. vector.transfer_read,
  //      vector.transfer_write); and
  //   b. ops that *may* lower a super-vector (all other ops).
  // The ops that *may* lower a super-vector only do so if the super-vector to
  // sub-vector ratio exists. The ops that *must* lower a super-vector are
  // explicitly checked for this property.
  /// TODO: there should be a single function for all ops to do this so we
  /// do not have to special case. Maybe a trait, or just a method, unclear atm.
  VectorType superVectorType;
  if (auto transfer = dyn_cast<VectorTransferOpInterface>(op)) {
    superVectorType = transfer.getVectorType();
  } else if (op.getNumResults() == 0) {
    if (!isa<func::ReturnOp>(op)) {
      op.emitError("NYI: assuming only return operations can have 0 "
                   " results at this point");
    }
    return false;
  } else if (op.getNumResults() == 1) {
    if (auto v = dyn_cast<VectorType>(op.getResult(0).getType())) {
      superVectorType = v;
    } else {
      // Not a vector type.
      return false;
    }
  } 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. If the shapes are incompatible (e.g., different ranks or
  // non-integer divisibility), the operation does not operate on a super-vector
  // of the given sub-vector type.
  auto ratio =
      computeShapeRatio(superVectorType.getShape(), subVectorType.getShape());
  return ratio.has_value();
}

bool vector::isContiguousSlice(MemRefType memrefType, VectorType vectorType) {
  if (vectorType.isScalable())
    return false;

  // Ignore a leading sequence of adjacent unit dimensions in the vector.
  ArrayRef<int64_t> vectorShape =
      vectorType.getShape().drop_while([](auto v) { return v == 1; });
  auto vecRank = vectorShape.size();

  // A single element is always contiguous.
  if (vecRank == 0)
    return true;

  if (!memrefType.areTrailingDimsContiguous(vecRank))
    return false;

  // Extract the trailing dims of the input memref
  auto memrefShape = memrefType.getShape().take_back(vecRank);

  // Compare the dims of `vectorType` against `memrefType`.
  // All of the dimensions, except the first must match.
  return llvm::equal(vectorShape.drop_front(), memrefShape.drop_front());
}

std::optional<StaticTileOffsetRange>
vector::createUnrollIterator(VectorType vType, int64_t targetRank) {
  if (vType.getRank() <= targetRank)
    return {};
  // Attempt to unroll until targetRank or the first scalable dimension (which
  // cannot be unrolled).
  auto shapeToUnroll = vType.getShape().drop_back(targetRank);
  auto inputScalableVecDimsToUnroll =
      vType.getScalableDims().drop_back(targetRank);
  const auto *it = llvm::find(inputScalableVecDimsToUnroll, true);
  auto firstScalableDim = it - inputScalableVecDimsToUnroll.begin();
  if (firstScalableDim == 0)
    return {};
  // All scalable dimensions should be removed now.
  inputScalableVecDimsToUnroll =
      inputScalableVecDimsToUnroll.slice(0, firstScalableDim);
  assert(!llvm::is_contained(inputScalableVecDimsToUnroll, true) &&
         "unexpected leading scalable dimension");
  // Create an unroll iterator for leading dimensions.
  shapeToUnroll = shapeToUnroll.slice(0, firstScalableDim);
  return StaticTileOffsetRange(shapeToUnroll, /*unrollStep=*/1);
}

SmallVector<OpFoldResult> vector::getMixedSizesXfer(bool hasTensorSemantics,
                                                    Operation *xfer,
                                                    RewriterBase &rewriter) {
  auto loc = xfer->getLoc();

  Value base =
      TypeSwitch<Operation *, Value>(xfer)
          .Case([&](vector::TransferReadOp readOp) { return readOp.getBase(); })
          .Case([&](vector::TransferWriteOp writeOp) {
            return writeOp.getOperand(1);
          });

  SmallVector<OpFoldResult> mixedSourceDims =
      hasTensorSemantics ? tensor::getMixedSizes(rewriter, loc, base)
                         : memref::getMixedSizes(rewriter, loc, base);
  return mixedSourceDims;
}

bool vector::isLinearizableVector(VectorType type) {
  return (type.getRank() > 1) && (type.getNumScalableDims() <= 1);
}

/// Determines whether a mask for xfer_read/write is trivially "all true"
///
/// Given all the inputs required to generate a mask (mask sizes and shapes),
/// and an xfer_read/write operation (indices and the source/destination tensor
/// shape), determines whether the corresponding mask would be trivially
/// foldable (i.e., trivially "all true").
///
/// Use this method to avoid generating spurious masks and relying on
/// vectorization post-processing to remove them.
///
/// Pre-conditions for a mask to be trivially foldable:
///   * All involved shapes (mask + destination tensor) are static.
///   * All indices are constant.
///   * All mask sizes are constant (including `arith.constant`).
///
/// If the pre-conditions are met, the method checks for each destination
/// dimension `d`:
///   (1) maskShape[d] <= destDimSize[rankDiff + d]
///   (2) index[d] + maskSize[d] <= destDimSize[rankDiff + d]
///   (3) maskShape[d] <= maskSize[d]
///
/// rankDiff = rank(dest) - rank(mask).
///
/// This method takes a conservative view: it may return false even if the mask
/// is technically foldable.
///
/// EXAMPLE 1 (trivially foldable, all shapes match, mask sizes match the shape
/// of the dest tensor):
///   %c0 = arith.constant 0 : index
///   %mask = vector.create_mask 5, 1
///   vector.mask %mask {
///     vector.transfer_write %vecToStore_1, %dest{[%c0, %c0]
///       {in_bounds = [true, true]}
///     : vector<5x1xi32>, tensor<5x1xi32>
///   }
///
/// EXAMPLE 2 (not trivially foldable - vector shape exceeds the tensor shape,
/// mask is required to avoid out-of-bounds write):
///   %c0 = arith.constant 0 : index
///   %mask = vector.create_mask 5, 1
///   vector.mask %mask {
///     vector.transfer_write %vecToStore_2, %dest[%c0, %c0]
///      {in_bounds = [true, true]}
///     : vector<8x1xi32>, tensor<5x1xi32>
///   }
static bool isMaskTriviallyFoldable(SmallVector<OpFoldResult> &maskSizes,
                                    SmallVector<Value> &indices,
                                    ArrayRef<int64_t> baseShape,
                                    ArrayRef<int64_t> maskShape) {
  // Masking is unavoidable in the case of dynamic tensors.
  if (ShapedType::isDynamicShape(baseShape))
    return false;

  // Collect all constant mask sizes.
  SmallVector<int64_t, 4> cstMaskSizes;
  for (auto [i, dimSize] : llvm::enumerate(maskSizes)) {
    if (auto intSize = getConstantIntValue(dimSize)) {
      cstMaskSizes.push_back(*intSize);
    }
  }

  // If any of the mask sizes is non-constant, bail out.
  if (cstMaskSizes.size() != maskShape.size())
    return false;

  // Collect all constant indices.
  SmallVector<int64_t, 4> cstIndices;
  for (auto [i, idx] : llvm::enumerate(indices)) {
    APSInt intVal;
    if (matchPattern(idx, m_ConstantInt(&intVal))) {
      cstIndices.push_back(intVal.getSExtValue());
    }
  }

  // If any of the indices is non-constant, bail out.
  if (cstIndices.size() != baseShape.size())
    return false;

  // Go over all destination dims and check (1), (2) and (3). Take into account
  // that:
  //  * The number of mask sizes will match the rank of the vector to
  //    load/store. This could be lower than the rank of the destination tensor.
  //  * Mask sizes could be larger than the corresponding mask shape (hence
  //    `clamp`).
  // TODO: The 2nd item should be rejected by the verifier.
  int64_t rankDiff = baseShape.size() - cstMaskSizes.size();
  for (auto [i, idx] : llvm::enumerate(cstMaskSizes)) {
    if (/*(1)*/ maskShape[i] > baseShape[rankDiff + i] ||
        /*(2)*/ baseShape[rankDiff + i] <
            (std::clamp(cstMaskSizes[i], int64_t(0), maskShape[i]) +
             cstIndices[i]) ||
        /*(3)*/ maskShape[i] > cstMaskSizes[i])
      return false;
  }

  return true;
}

Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
                                     Value source,
                                     ArrayRef<int64_t> inputVectorSizes,
                                     std::optional<Value> padValue,
                                     bool useInBoundsInsteadOfMasking,
                                     ArrayRef<bool> inputScalableVecDims) {
  VectorType vecToReadTy = VectorType::get(
      inputVectorSizes, cast<ShapedType>(source.getType()).getElementType(),
      inputScalableVecDims);

  return createReadOrMaskedRead(builder, loc, source, vecToReadTy, padValue,
                                useInBoundsInsteadOfMasking);
}

/// Compute the in_bounds attribute for a transfer op given its permutation map
/// and the source being accessed. Dimension i is in-bounds when the map result
/// is an AffineDimExpr pointing to a static source dimension divisible by the
/// vector size, or an AffineConstantExpr (broadcast).
static SmallVector<bool> computeInBoundsFromPermutationMap(
    AffineMap permutationMap, VectorType vectorType, ShapedType sourceType) {
  SmallVector<bool> inBounds(vectorType.getRank(), false);
  for (unsigned i = 0; i < (unsigned)vectorType.getRank(); ++i) {
    AffineExpr expr = permutationMap.getResult(i);
    if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
      unsigned memDim = dimExpr.getPosition();
      if (!sourceType.isDynamicDim(memDim) &&
          sourceType.getDimSize(memDim) % vectorType.getDimSize(i) == 0)
        inBounds[i] = true;
    } else if (isa<AffineConstantExpr>(expr)) {
      inBounds[i] = true;
    }
  }
  return inBounds;
}

Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
                                     Value source,
                                     const VectorType &vecToReadTy,
                                     std::optional<Value> padValue,
                                     bool useInBoundsInsteadOfMasking,
                                     ArrayRef<Value> customIndices,
                                     AffineMap permutationMap) {
  assert(!llvm::is_contained(vecToReadTy.getScalableDims(),
                             ShapedType::kDynamic) &&
         "invalid input vector sizes");
  auto sourceShapedType = cast<ShapedType>(source.getType());
  auto sourceShape = sourceShapedType.getShape();

  int64_t vecToReadRank = vecToReadTy.getRank();
  auto vecToReadShape = vecToReadTy.getShape();

  // The permutation map maps the source's index space to the vector's, so its
  // dims must match the source rank and its results the vector rank. Without a
  // map, a minor identity is implied, requiring the two ranks to match.
  assert(sourceShape.size() == (permutationMap
                                    ? permutationMap.getNumDims()
                                    : static_cast<size_t>(vecToReadRank)) &&
         "expected source rank to match permutation map dims or vector rank.");
  assert((!permutationMap || permutationMap.getNumResults() ==
                                 static_cast<size_t>(vecToReadRank)) &&
         "expected permutation map results to match vector rank.");
  assert((!padValue.has_value() ||
          padValue.value().getType() == sourceShapedType.getElementType()) &&
         "expected same pad element type to match source element type");

  SmallVector<bool> inBoundsVal(vecToReadRank, true);

  if (useInBoundsInsteadOfMasking) {
    if (permutationMap) {
      // Update the inBounds attribute.
      // FIXME: This computation is too weak - it ignores the read indices.
      inBoundsVal = computeInBoundsFromPermutationMap(
          permutationMap, vecToReadTy, cast<ShapedType>(source.getType()));
    } else {
      // Update the inBounds attribute.
      // FIXME: This computation is too weak - it ignores the read indices.
      for (unsigned i = 0; i < vecToReadRank; i++)
        inBoundsVal[i] = (sourceShape[i] == vecToReadShape[i]) &&
                         ShapedType::isStatic(sourceShape[i]);
    }
  }
  // The transfer op expects one index per source dimension.
  assert(
      (customIndices.empty() || customIndices.size() == sourceShape.size()) &&
      "expected as many custom indices as source dims.");
  SmallVector<Value> indices;
  customIndices.empty()
      ? indices.assign(sourceShape.size(),
                       arith::ConstantIndexOp::create(builder, loc, 0))
      : indices.assign(customIndices.begin(), customIndices.end());

  // A null permutation map means the builder defaults to a minor identity map.
  auto transferReadOp =
      vector::TransferReadOp::create(builder, loc, /*vectorType=*/vecToReadTy,
                                     /*source=*/source,
                                     /*indices=*/indices,
                                     /*padding=*/padValue,
                                     /*permutationMap=*/permutationMap,
                                     /*inBounds=*/inBoundsVal);

  if (useInBoundsInsteadOfMasking)
    return transferReadOp;

  SmallVector<OpFoldResult> mixedSourceDims =
      isa<MemRefType>(source.getType())
          ? memref::getMixedSizes(builder, loc, source)
          : tensor::getMixedSizes(builder, loc, source);

  if (isMaskTriviallyFoldable(mixedSourceDims, indices, sourceShape,
                              vecToReadShape))
    return transferReadOp;

  auto maskType = vecToReadTy.cloneWith(/*shape=*/{}, builder.getI1Type());
  Value mask =
      vector::CreateMaskOp::create(builder, loc, maskType, mixedSourceDims);
  return mlir::vector::maskOperation(builder, transferReadOp, mask)
      ->getResult(0);
}

Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
                                            Value vecToStore, Value dest,
                                            SmallVector<Value> writeIndices,
                                            bool useInBoundsInsteadOfMasking,
                                            AffineMap permutationMap) {

  ShapedType destType = cast<ShapedType>(dest.getType());
  int64_t destRank = destType.getRank();
  auto destShape = destType.getShape();

  VectorType vecToStoreType = cast<VectorType>(vecToStore.getType());
  int64_t vecToStoreRank = vecToStoreType.getRank();
  auto vecToStoreShape = vecToStoreType.getShape();

  // Compute the in_bounds attribute
  SmallVector<bool> inBoundsVal(vecToStoreRank, true);
  if (useInBoundsInsteadOfMasking) {
    if (permutationMap) {
      // Update the inBounds attribute.
      // FIXME: This computation is too weak - it ignores the write indices.
      inBoundsVal = computeInBoundsFromPermutationMap(
          permutationMap, vecToStoreType, cast<ShapedType>(dest.getType()));
    } else {
      // Update the inBounds attribute.
      // FIXME: This computation is too weak - it ignores the write indices.
      for (unsigned i = 0; i < vecToStoreRank; i++)
        inBoundsVal[i] =
            (destShape[destRank - vecToStoreRank + i] >= vecToStoreShape[i]) &&
            ShapedType::isStatic(destShape[destRank - vecToStoreRank + i]);
    }
  }

  // If missing, initialize the write indices to 0.
  bool useDefaultWriteIdxs = writeIndices.empty();
  assert((useDefaultWriteIdxs ||
          writeIndices.size() == static_cast<size_t>(destRank)) &&
         "Invalid number of write indices!");
  if (useDefaultWriteIdxs) {
    auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
    writeIndices.assign(destRank, zero);
  }

  // Generate the xfer_write Op. A null permutation map means the builder
  // defaults to a minor identity map.
  Operation *write =
      vector::TransferWriteOp::create(builder, loc,
                                      /*vector=*/vecToStore,
                                      /*dest=*/dest,
                                      /*indices=*/writeIndices,
                                      /*permutationMap=*/permutationMap,
                                      /*inBounds=*/inBoundsVal);

  // If masking is disabled, exit.
  if (useInBoundsInsteadOfMasking)
    return write;

  // The dest dims written to are those the permutation map selects. A null map
  // means a minor identity: the trailing dims.
  AffineMap map = permutationMap
                      ? permutationMap
                      : AffineMap::getMinorIdentityMap(destRank, vecToStoreRank,
                                                       builder.getContext());

  // Check if masking is needed. If not, exit.
  if (llvm::equal(vecToStoreShape, applyPermutationMap(map, destShape)))
    return write;

  // Compute the mask and mask the write Op.
  VectorType writeMaskType = inferTransferOpMaskType(vecToStoreType, map);

  SmallVector<OpFoldResult> destSizes =
      isa<MemRefType>(dest.getType())
          ? memref::getMixedSizes(builder, loc, dest)
          : tensor::getMixedSizes(builder, loc, dest);

  // Compute sizes for write-mask: `dim - index` for each dest dim written to,
  // in increasing order, as the mask dims are.
  SmallVector<int64_t> writtenDims;
  for (AffineExpr expr : map.getResults())
    writtenDims.push_back(cast<AffineDimExpr>(expr).getPosition());
  llvm::sort(writtenDims);
  SmallVector<OpFoldResult> maskSizes;
  for (int64_t dim : writtenDims) {
    Value size = getValueOrCreateConstantIndexOp(builder, loc, destSizes[dim]);
    maskSizes.push_back(
        builder.createOrFold<arith::SubIOp>(loc, size, writeIndices[dim]));
  }

  // TODO: isMaskTriviallyFoldable checks the trailing dims against the leading
  // indices. Make it use the dims that the map selects and their indices.
  if (!permutationMap && isMaskTriviallyFoldable(maskSizes, writeIndices,
                                                 destShape, vecToStoreShape))
    return write;

  Value maskForWrite =
      builder.createOrFold<vector::CreateMaskOp>(loc, writeMaskType, maskSizes);
  return mlir::vector::maskOperation(builder, write, maskForWrite);
}

LogicalResult
vector::isValidMaskedInputVector(ArrayRef<int64_t> shape,
                                 ArrayRef<int64_t> inputVectorSizes) {
  LDBG() << "Iteration space static sizes:" << llvm::interleaved(shape);

  if (inputVectorSizes.size() != shape.size()) {
    LDBG() << "Input vector sizes don't match the number of loops";
    return failure();
  }
  if (ShapedType::isDynamicShape(inputVectorSizes)) {
    LDBG() << "Input vector sizes can't have dynamic dimensions";
    return failure();
  }
  if (!llvm::all_of(llvm::zip(shape, inputVectorSizes),
                    [](std::tuple<int64_t, int64_t> sizePair) {
                      int64_t staticSize = std::get<0>(sizePair);
                      int64_t inputSize = std::get<1>(sizePair);
                      return ShapedType::isDynamic(staticSize) ||
                             staticSize <= inputSize;
                    })) {
    LDBG() << "Input vector sizes must be greater than or equal to iteration "
              "space static sizes";
    return failure();
  }
  return success();
}

/// Takes a 2+ dimensional vector as an input
/// returns n vector values produced by n vector.extract operations.
/// I.e. calling unrollVectorValue([[%v]], rewriter) such that
///
///   %v : vector<nxaxb...>
///
/// will produce the following IR changes
///
///   %v0 = vector.extract %v[0] : vector<axbx...> from vector<nxaxb...>
///   %v1 = vector.extract %v[1] : vector<axbx...> from vector<nxaxb...>
///   ...
///   %vnminusone = vector.extract %v[n-1] : vector<axbx...> from ...
///
/// and returns SmallVector<Value> r = {[[%v0]], [[%v1]], ..., [[%vnminusone]]}
FailureOr<SmallVector<Value>>
vector::unrollVectorValue(TypedValue<VectorType> vector,
                          RewriterBase &rewriter) {
  SmallVector<Value> subvectors;
  VectorType ty = cast<VectorType>(vector.getType());
  Location loc = vector.getLoc();
  if (ty.getRank() < 2)
    return rewriter.notifyMatchFailure(loc, "already 1-D");

  // Unrolling doesn't take vscale into account. Pattern is disabled for
  // vectors with leading scalable dim(s).
  if (ty.getScalableDims().front())
    return rewriter.notifyMatchFailure(loc, "cannot unroll scalable dim");

  for (int64_t i = 0, e = ty.getShape().front(); i < e; ++i) {
    subvectors.push_back(vector::ExtractOp::create(rewriter, loc, vector, i));
  }

  return subvectors;
}

LogicalResult vector::unrollVectorOp(Operation *op, PatternRewriter &rewriter,
                                     vector::UnrollVectorOpFn unrollFn) {
  assert(op->getNumResults() == 1 && "expected single result");
  assert(isa<VectorType>(op->getResult(0).getType()) && "expected vector type");
  VectorType resultTy = cast<VectorType>(op->getResult(0).getType());
  if (resultTy.getRank() < 2)
    return rewriter.notifyMatchFailure(op, "already 1-D");

  // Unrolling doesn't take vscale into account. Pattern is disabled for
  // vectors with leading scalable dim(s).
  if (resultTy.getScalableDims().front())
    return rewriter.notifyMatchFailure(op, "cannot unroll scalable dim");

  Location loc = op->getLoc();
  Value result = ub::PoisonOp::create(rewriter, loc, resultTy);
  VectorType subTy = VectorType::Builder(resultTy).dropDim(0);

  for (int64_t i = 0, e = resultTy.getShape().front(); i < e; ++i) {
    Value subVector = unrollFn(rewriter, loc, subTy, i);
    result = vector::InsertOp::create(rewriter, loc, subVector, result, i);
  }

  rewriter.replaceOp(op, result);
  return success();
}