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

//===- VectorTransferOpTransforms.cpp - transfer op transforms ------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements functions concerned with optimizing transfer_read and
// transfer_write ops.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/LogicalResult.h"

#define DEBUG_TYPE "vector-transfer-opt"

using namespace mlir;

/// Return the ancestor op in the region or nullptr if the region is not
/// an ancestor of the op.
static Operation *findAncestorOpInRegion(Region *region, Operation *op) {
  LDBG() << "    Finding ancestor of " << *op << " in region";
  for (; op != nullptr && op->getParentRegion() != region;
       op = op->getParentOp())
    ;
  if (op) {
    LDBG() << "    -> Ancestor: " << *op;
  } else {
    LDBG() << "    -> Ancestor: nullptr";
  }
  return op;
}

namespace {

class TransferOptimization {
public:
  TransferOptimization(RewriterBase &rewriter, Operation *op)
      : rewriter(rewriter), dominators(op), postDominators(op) {}
  void deadStoreOp(vector::TransferWriteOp);
  void storeToLoadForwarding(vector::TransferReadOp);
  void removeDeadOp() {
    LDBG() << "Removing " << opToErase.size() << " dead operations";
    for (Operation *op : opToErase) {
      LDBG() << "  -> Erasing: " << *op;
      rewriter.eraseOp(op);
    }
    opToErase.clear();
  }

private:
  RewriterBase &rewriter;
  bool isReachable(Operation *start, Operation *dest);
  DominanceInfo dominators;
  PostDominanceInfo postDominators;
  std::vector<Operation *> opToErase;
};

} // namespace
/// Return true if there is a path from start operation to dest operation,
/// otherwise return false. The operations have to be in the same region.
bool TransferOptimization::isReachable(Operation *start, Operation *dest) {
  LDBG() << "    Checking reachability from " << *start << " to " << *dest;
  assert(start->getParentRegion() == dest->getParentRegion() &&
         "This function only works for ops i the same region");
  // Simple case where the start op dominate the destination.
  if (dominators.dominates(start, dest)) {
    LDBG() << "    -> Start dominates dest, reachable";
    return true;
  }
  bool blockReachable = start->getBlock()->isReachable(dest->getBlock());
  LDBG() << "    -> Block reachable: " << blockReachable;
  return blockReachable;
}

/// For transfer_write to overwrite fully another transfer_write must:
/// 1. Access the same memref with the same indices and vector type.
/// 2. Post-dominate the other transfer_write operation.
/// If several candidates are available, one must be post-dominated by all the
/// others since they are all post-dominating the same transfer_write. We only
/// consider the transfer_write post-dominated by all the other candidates as
/// this will be the first transfer_write executed after the potentially dead
/// transfer_write.
/// If we found such an overwriting transfer_write we know that the original
/// transfer_write is dead if all reads that can be reached from the potentially
/// dead transfer_write are dominated by the overwriting transfer_write.
void TransferOptimization::deadStoreOp(vector::TransferWriteOp write) {
  LDBG() << "=== Starting deadStoreOp analysis for: " << *write.getOperation();
  llvm::SmallVector<Operation *, 8> blockingAccesses;
  Operation *firstOverwriteCandidate = nullptr;
  Value source = memref::skipViewLikeOps(cast<MemrefValue>(write.getBase()));
  LDBG() << "Source memref (after skipping view-like ops): " << source;
  llvm::SmallVector<Operation *, 32> users(source.getUsers().begin(),
                                           source.getUsers().end());
  LDBG() << "Found " << users.size() << " users of source memref";
  llvm::SmallDenseSet<Operation *, 32> processed;
  while (!users.empty()) {
    Operation *user = users.pop_back_val();
    LDBG() << "Processing user: " << *user;
    // If the user has already been processed skip.
    if (!processed.insert(user).second) {
      LDBG() << "  -> Already processed, skipping";
      continue;
    }
    if (auto viewLike = dyn_cast<ViewLikeOpInterface>(user)) {
      LDBG() << "  -> View-like operation, following to destination";
      Value viewDest = viewLike.getViewDest();
      users.append(viewDest.getUsers().begin(), viewDest.getUsers().end());
      continue;
    }
    if (isMemoryEffectFree(user)) {
      LDBG() << "  -> Memory effect free, skipping";
      continue;
    }
    if (user == write.getOperation()) {
      LDBG() << "  -> Same as write operation, skipping";
      continue;
    }
    if (auto nextWrite = dyn_cast<vector::TransferWriteOp>(user)) {
      LDBG() << "  -> Found transfer_write candidate: " << *nextWrite;
      // Check candidate that can override the store.
      bool sameView = memref::isSameViewOrTrivialAlias(
          cast<MemrefValue>(nextWrite.getBase()),
          cast<MemrefValue>(write.getBase()));
      bool sameValue = checkSameValueWAW(nextWrite, write);
      bool postDominates = postDominators.postDominates(nextWrite, write);
      LDBG() << "    -> Same view: " << sameView
             << ", Same value: " << sameValue
             << ", Post-dominates: " << postDominates;

      if (sameView && sameValue && postDominates) {
        LDBG() << "    -> Valid overwrite candidate found";
        if (firstOverwriteCandidate == nullptr ||
            postDominators.postDominates(firstOverwriteCandidate, nextWrite)) {
          LDBG() << "    -> New first overwrite candidate: " << *nextWrite;
          firstOverwriteCandidate = nextWrite;
        } else {
          LDBG() << "    -> Keeping existing first overwrite candidate";
          assert(
              postDominators.postDominates(nextWrite, firstOverwriteCandidate));
        }
        continue;
      }
      LDBG() << "    -> Not a valid overwrite candidate";
    }
    if (auto transferOp = dyn_cast<VectorTransferOpInterface>(user)) {
      LDBG() << "  -> Found vector transfer operation: " << *transferOp;
      // Don't need to consider disjoint accesses.
      bool isDisjoint = vector::isDisjointTransferSet(
          cast<VectorTransferOpInterface>(write.getOperation()),
          cast<VectorTransferOpInterface>(transferOp.getOperation()),
          /*testDynamicValueUsingBounds=*/true);
      LDBG() << "    -> Is disjoint: " << isDisjoint;
      if (isDisjoint) {
        LDBG() << "    -> Skipping disjoint access";
        continue;
      }
    }
    LDBG() << "  -> Adding to blocking accesses: " << *user;
    blockingAccesses.push_back(user);
  }
  LDBG() << "Finished processing users. Found " << blockingAccesses.size()
         << " blocking accesses";

  if (firstOverwriteCandidate == nullptr) {
    LDBG() << "No overwrite candidate found, store is not dead";
    return;
  }

  LDBG() << "First overwrite candidate: " << *firstOverwriteCandidate;
  Region *topRegion = firstOverwriteCandidate->getParentRegion();
  Operation *writeAncestor = findAncestorOpInRegion(topRegion, write);
  assert(writeAncestor &&
         "write op should be recursively part of the top region");
  LDBG() << "Write ancestor in top region: " << *writeAncestor;

  LDBG() << "Checking " << blockingAccesses.size()
         << " blocking accesses for reachability";
  for (Operation *access : blockingAccesses) {
    LDBG() << "Checking blocking access: " << *access;
    Operation *accessAncestor = findAncestorOpInRegion(topRegion, access);
    // TODO: if the access and write have the same ancestor we could recurse in
    // the region to know if the access is reachable with more precision.
    if (accessAncestor == nullptr) {
      LDBG() << "  -> No ancestor in top region, skipping";
      continue;
    }

    bool isReachableFromWrite = isReachable(writeAncestor, accessAncestor);
    LDBG() << "  -> Is reachable from write: " << isReachableFromWrite;
    if (!isReachableFromWrite) {
      LDBG() << "  -> Not reachable, skipping";
      continue;
    }

    bool overwriteDominatesAccess =
        dominators.dominates(firstOverwriteCandidate, accessAncestor);
    LDBG() << "  -> Overwrite dominates access: " << overwriteDominatesAccess;
    if (!overwriteDominatesAccess) {
      LDBG() << "Store may not be dead due to op: " << *accessAncestor;
      return;
    }
    LDBG() << "  -> Access is dominated by overwrite, continuing";
  }
  LDBG() << "Found dead store: " << *write.getOperation()
         << " overwritten by: " << *firstOverwriteCandidate;
  opToErase.push_back(write.getOperation());
}

/// A transfer_write candidate to storeToLoad forwarding must:
/// 1. Access the same memref with the same indices and vector type as the
/// transfer_read.
/// 2. Dominate the transfer_read operation.
/// If several candidates are available, one must be dominated by all the others
/// since they are all dominating the same transfer_read. We only consider the
/// transfer_write dominated by all the other candidates as this will be the
/// last transfer_write executed before the transfer_read.
/// If we found such a candidate we can do the forwarding if all the other
/// potentially aliasing ops that may reach the transfer_read are post-dominated
/// by the transfer_write.
void TransferOptimization::storeToLoadForwarding(vector::TransferReadOp read) {
  LDBG() << "=== Starting storeToLoadForwarding analysis for: "
         << *read.getOperation();
  if (read.hasOutOfBoundsDim()) {
    LDBG() << "Read has out-of-bounds dimensions, skipping";
    return;
  }
  SmallVector<Operation *, 8> blockingWrites;
  vector::TransferWriteOp lastwrite = nullptr;
  Value source = memref::skipViewLikeOps(cast<MemrefValue>(read.getBase()));
  LDBG() << "Source memref (after skipping view-like ops): " << source;
  llvm::SmallVector<Operation *, 32> users(source.getUsers().begin(),
                                           source.getUsers().end());
  LDBG() << "Found " << users.size() << " users of source memref";
  llvm::SmallDenseSet<Operation *, 32> processed;
  while (!users.empty()) {
    Operation *user = users.pop_back_val();
    LDBG() << "Processing user: " << *user;
    // If the user has already been processed skip.
    if (!processed.insert(user).second) {
      LDBG() << "  -> Already processed, skipping";
      continue;
    }
    if (auto viewLike = dyn_cast<ViewLikeOpInterface>(user)) {
      LDBG() << "  -> View-like operation, following to destination";
      Value viewDest = viewLike.getViewDest();
      users.append(viewDest.getUsers().begin(), viewDest.getUsers().end());
      continue;
    }
    if (isMemoryEffectFree(user) || isa<vector::TransferReadOp>(user)) {
      LDBG() << "  -> Memory effect free or transfer_read, skipping";
      continue;
    }
    if (auto write = dyn_cast<vector::TransferWriteOp>(user)) {
      LDBG() << "  -> Found transfer_write candidate: " << *write;
      // If there is a write, but we can prove that it is disjoint we can ignore
      // the write.
      bool isDisjoint = vector::isDisjointTransferSet(
          cast<VectorTransferOpInterface>(write.getOperation()),
          cast<VectorTransferOpInterface>(read.getOperation()),
          /*testDynamicValueUsingBounds=*/true);
      LDBG() << "    -> Is disjoint: " << isDisjoint;
      if (isDisjoint) {
        LDBG() << "    -> Skipping disjoint write";
        continue;
      }

      bool sameView =
          memref::isSameViewOrTrivialAlias(cast<MemrefValue>(read.getBase()),
                                           cast<MemrefValue>(write.getBase()));
      bool dominates = dominators.dominates(write, read);
      bool sameValue = checkSameValueRAW(write, read);
      LDBG() << "    -> Same view: " << sameView << ", Dominates: " << dominates
             << ", Same value: " << sameValue;

      if (sameView && dominates && sameValue) {
        LDBG() << "    -> Valid forwarding candidate found";
        if (lastwrite == nullptr || dominators.dominates(lastwrite, write)) {
          LDBG() << "    -> New last write candidate: " << *write;
          lastwrite = write;
        } else {
          LDBG() << "    -> Keeping existing last write candidate";
          assert(dominators.dominates(write, lastwrite));
        }
        continue;
      }
      LDBG() << "    -> Not a valid forwarding candidate";
    }
    LDBG() << "  -> Adding to blocking writes: " << *user;
    blockingWrites.push_back(user);
  }
  LDBG() << "Finished processing users. Found " << blockingWrites.size()
         << " blocking writes";

  if (lastwrite == nullptr) {
    LDBG() << "No last write candidate found, cannot forward";
    return;
  }

  LDBG() << "Last write candidate: " << *lastwrite;
  Region *topRegion = lastwrite->getParentRegion();
  Operation *readAncestor = findAncestorOpInRegion(topRegion, read);
  assert(readAncestor &&
         "read op should be recursively part of the top region");
  LDBG() << "Read ancestor in top region: " << *readAncestor;

  LDBG() << "Checking " << blockingWrites.size()
         << " blocking writes for post-dominance";
  for (Operation *write : blockingWrites) {
    LDBG() << "Checking blocking write: " << *write;
    Operation *writeAncestor = findAncestorOpInRegion(topRegion, write);
    if (writeAncestor) {
      LDBG() << "  -> Write ancestor: " << *writeAncestor;
    } else {
      LDBG() << "  -> Write ancestor: nullptr";
    }

    // TODO: if the store and read have the same ancestor we could recurse in
    // the region to know if the read is reachable with more precision.
    if (writeAncestor == nullptr) {
      LDBG() << "  -> No ancestor in top region, skipping";
      continue;
    }

    bool isReachableToRead = isReachable(writeAncestor, readAncestor);
    LDBG() << "  -> Is reachable to read: " << isReachableToRead;
    if (!isReachableToRead) {
      LDBG() << "  -> Not reachable, skipping";
      continue;
    }

    bool lastWritePostDominates =
        postDominators.postDominates(lastwrite, write);
    LDBG() << "  -> Last write post-dominates blocking write: "
           << lastWritePostDominates;
    if (!lastWritePostDominates) {
      LDBG() << "Fail to do write to read forwarding due to op: " << *write;
      return;
    }
    LDBG() << "  -> Blocking write is post-dominated, continuing";
  }

  LDBG() << "Forward value from " << *lastwrite.getOperation()
         << " to: " << *read.getOperation();
  read.replaceAllUsesWith(lastwrite.getVector());
  opToErase.push_back(read.getOperation());
}

/// Converts OpFoldResults to int64_t shape without unit dims.
static SmallVector<int64_t> getReducedShape(ArrayRef<OpFoldResult> mixedSizes) {
  SmallVector<int64_t> reducedShape;
  for (const auto size : mixedSizes) {
    if (llvm::dyn_cast_if_present<Value>(size)) {
      reducedShape.push_back(ShapedType::kDynamic);
      continue;
    }

    auto value = cast<IntegerAttr>(cast<Attribute>(size)).getValue();
    if (value == 1)
      continue;
    reducedShape.push_back(value.getSExtValue());
  }
  return reducedShape;
}

/// Drops unit dimensions from the input MemRefType.
static MemRefType dropUnitDims(MemRefType inputType,
                               ArrayRef<OpFoldResult> offsets,
                               ArrayRef<OpFoldResult> sizes,
                               ArrayRef<OpFoldResult> strides) {
  auto targetShape = getReducedShape(sizes);
  MemRefType rankReducedType = memref::SubViewOp::inferRankReducedResultType(
      targetShape, inputType, offsets, sizes, strides);
  return rankReducedType.canonicalizeStridedLayout();
}

/// Creates a rank-reducing memref.subview op that drops unit dims from its
/// input. Or just returns the input if it was already without unit dims.
static Value rankReducingSubviewDroppingUnitDims(PatternRewriter &rewriter,
                                                 mlir::Location loc,
                                                 Value input) {
  MemRefType inputType = cast<MemRefType>(input.getType());
  SmallVector<OpFoldResult> offsets(inputType.getRank(),
                                    rewriter.getIndexAttr(0));
  SmallVector<OpFoldResult> sizes = memref::getMixedSizes(rewriter, loc, input);
  SmallVector<OpFoldResult> strides(inputType.getRank(),
                                    rewriter.getIndexAttr(1));
  MemRefType resultType = dropUnitDims(inputType, offsets, sizes, strides);

  if (resultType.canonicalizeStridedLayout() ==
      inputType.canonicalizeStridedLayout())
    return input;
  return memref::SubViewOp::create(rewriter, loc, resultType, input, offsets,
                                   sizes, strides);
}

/// Returns the number of dims that aren't unit dims.
static int getReducedRank(ArrayRef<int64_t> shape) {
  return llvm::count_if(shape, [](int64_t dimSize) { return dimSize != 1; });
}

/// Trims non-scalable one dimensions from `oldType` and returns the result
/// type.
static VectorType trimNonScalableUnitDims(VectorType oldType) {
  SmallVector<int64_t> newShape;
  SmallVector<bool> newScalableDims;
  for (auto [dimIdx, dimSize] : llvm::enumerate(oldType.getShape())) {
    if (dimSize == 1 && !oldType.getScalableDims()[dimIdx])
      continue;
    newShape.push_back(dimSize);
    newScalableDims.push_back(oldType.getScalableDims()[dimIdx]);
  }
  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
}

static bool isUnitDimMask(Value maskDimSize) {
  return matchPattern(maskDimSize, m_One());
}
static bool isUnitDimMask(int64_t maskDimSize) { return maskDimSize == 1; }

/// Rewrites a mask op to drop non-scalable unit dimensions.
/// Supports vector.create_mask and vector.constant_mask.
template <typename MaskOp>
static FailureOr<Value> maskDropNonScalableUnitDims(PatternRewriter &rewriter,
                                                    Location loc, MaskOp op) {
  auto type = op.getType();
  VectorType reducedType = trimNonScalableUnitDims(type);
  if (reducedType.getRank() == type.getRank())
    return failure();

  using ElemType = std::decay_t<decltype(*op.getMaskDimSizes().begin())>;
  SmallVector<ElemType> reduced;
  for (auto [dim, dimIsScalable, elem] : llvm::zip_equal(
           type.getShape(), type.getScalableDims(), op.getMaskDimSizes())) {
    if (dim == 1 && !dimIsScalable) {
      if (!isUnitDimMask(elem))
        return failure();
      continue;
    }
    reduced.push_back(elem);
  }
  return MaskOp::create(rewriter, loc, reducedType, reduced).getResult();
}

namespace {

/// Rewrites `vector.transfer_read` ops where the source has unit dims, by
/// inserting a memref.subview dropping those unit dims. The vector shapes are
/// also reduced accordingly.
class TransferReadDropUnitDimsPattern
    : public vector::MaskableOpRewritePattern<vector::TransferReadOp> {
  using MaskableOpRewritePattern::MaskableOpRewritePattern;

  FailureOr<Value>
  matchAndRewriteMaskableOp(vector::TransferReadOp transferReadOp,
                            vector::MaskingOpInterface maskingOp,
                            PatternRewriter &rewriter) const override {
    LDBG() << "=== TransferReadDropUnitDimsPattern: Analyzing "
           << *transferReadOp;
    auto loc = transferReadOp.getLoc();
    Value vector = transferReadOp.getVector();
    VectorType vectorType = cast<VectorType>(vector.getType());
    Value source = transferReadOp.getBase();
    MemRefType sourceType = dyn_cast<MemRefType>(source.getType());
    // TODO: support tensor types.
    if (!sourceType) {
      LDBG() << "  -> Not a MemRefType, skipping";
      return failure();
    }
    // TODO: generalize this pattern, relax the requirements here.
    if (transferReadOp.hasOutOfBoundsDim()) {
      LDBG() << "  -> Has out-of-bounds dimensions, skipping";
      return failure();
    }
    if (!transferReadOp.getPermutationMap().isMinorIdentity()) {
      LDBG() << "  -> Not minor identity permutation map, skipping";
      return failure();
    }
    // Check if the source shape can be further reduced.
    int reducedRank = getReducedRank(sourceType.getShape());
    LDBG() << "  -> Source rank: " << sourceType.getRank()
           << ", Reduced rank: " << reducedRank;
    if (reducedRank == sourceType.getRank()) {
      LDBG() << "  -> No unit dimensions to drop, skipping";
      return failure();
    }
    // TODO: Extend vector.mask to support 0-d vectors. In the meantime, bail
    // out.
    if (reducedRank == 0 && maskingOp) {
      LDBG() << "  -> 0-d vector with masking not supported, skipping";
      return failure();
    }
    // Check if the reduced vector shape matches the reduced source shape.
    // Otherwise, this case is not supported yet.
    VectorType reducedVectorType = trimNonScalableUnitDims(vectorType);
    LDBG() << "  -> Vector type: " << vectorType
           << ", Reduced vector type: " << reducedVectorType;
    if (reducedRank != reducedVectorType.getRank()) {
      LDBG() << "  -> Reduced ranks don't match, skipping";
      return failure();
    }
    if (llvm::any_of(transferReadOp.getIndices(), [](Value v) {
          return getConstantIntValue(v) != static_cast<int64_t>(0);
        })) {
      LDBG() << "  -> Non-zero indices found, skipping";
      return failure();
    }

    Value maskOp = transferReadOp.getMask();
    if (maskOp) {
      LDBG() << "  -> Processing mask operation";
      auto maskVectorType = cast<VectorType>(maskOp.getType());
      FailureOr<Value> rankReducedMaskOp = failure();
      if (auto createMaskOp = maskOp.getDefiningOp<vector::CreateMaskOp>())
        rankReducedMaskOp =
            maskDropNonScalableUnitDims(rewriter, loc, createMaskOp);
      else if (auto constantMaskOp =
                   maskOp.getDefiningOp<vector::ConstantMaskOp>())
        rankReducedMaskOp =
            maskDropNonScalableUnitDims(rewriter, loc, constantMaskOp);
      else
        return rewriter.notifyMatchFailure(
            transferReadOp,
            "unsupported mask op, only 'vector.create_mask' and "
            "'vector.constant_mask' are currently supported");

      if (succeeded(rankReducedMaskOp)) {
        maskOp = *rankReducedMaskOp;
        LDBG() << "  -> Successfully reduced mask dimensions";
      } else if (maskVectorType.getRank() != reducedVectorType.getRank()) {
        return rewriter.notifyMatchFailure(
            transferReadOp, "Mask reduction required, but failed");
      }
    }

    LDBG() << "  -> Creating rank-reduced subview and new transfer_read";
    Value reducedShapeSource =
        rankReducingSubviewDroppingUnitDims(rewriter, loc, source);
    Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
    Repeated<Value> zeros(reducedRank, c0);
    auto identityMap = rewriter.getMultiDimIdentityMap(reducedRank);
    SmallVector<bool> inBounds(reducedVectorType.getRank(), true);
    Operation *newTransferReadOp = vector::TransferReadOp::create(
        rewriter, loc, reducedVectorType, reducedShapeSource, zeros,
        identityMap, transferReadOp.getPadding(), maskOp,
        rewriter.getBoolArrayAttr(inBounds));
    LDBG() << "  -> Created new transfer_read: " << *newTransferReadOp;

    if (maskingOp) {
      LDBG() << "  -> Applying masking operation";
      auto shapeCastMask = rewriter.createOrFold<vector::ShapeCastOp>(
          loc, reducedVectorType.cloneWith(std::nullopt, rewriter.getI1Type()),
          maskingOp.getMask());
      newTransferReadOp = mlir::vector::maskOperation(
          rewriter, newTransferReadOp, shapeCastMask);
    }

    auto shapeCast = rewriter.createOrFold<vector::ShapeCastOp>(
        loc, vectorType, newTransferReadOp->getResults()[0]);
    LDBG() << "  -> Created shape cast: " << *shapeCast.getDefiningOp();
    LDBG() << "  -> Pattern match successful, returning result";

    return shapeCast;
  }
};

/// Rewrites `vector.transfer_write` ops where the "source" (i.e. destination)
/// has unit dims, by inserting a `memref.subview` dropping those unit dims. The
/// vector shapes are also reduced accordingly.
class TransferWriteDropUnitDimsPattern
    : public vector::MaskableOpRewritePattern<vector::TransferWriteOp> {
  using MaskableOpRewritePattern::MaskableOpRewritePattern;

  FailureOr<Value>
  matchAndRewriteMaskableOp(vector::TransferWriteOp transferWriteOp,
                            vector::MaskingOpInterface maskingOp,
                            PatternRewriter &rewriter) const override {
    LDBG() << "=== TransferWriteDropUnitDimsPattern: Analyzing "
           << *transferWriteOp;
    auto loc = transferWriteOp.getLoc();
    Value vector = transferWriteOp.getVector();
    VectorType vectorType = cast<VectorType>(vector.getType());
    Value source = transferWriteOp.getBase();
    MemRefType sourceType = dyn_cast<MemRefType>(source.getType());
    // TODO: support tensor type.
    if (!sourceType) {
      LDBG() << "  -> Not a MemRefType, skipping";
      return failure();
    }
    // TODO: generalize this pattern, relax the requirements here.
    if (transferWriteOp.hasOutOfBoundsDim()) {
      LDBG() << "  -> Has out-of-bounds dimensions, skipping";
      return failure();
    }
    if (!transferWriteOp.getPermutationMap().isMinorIdentity()) {
      LDBG() << "  -> Not minor identity permutation map, skipping";
      return failure();
    }
    // Check if the destination shape can be further reduced.
    int reducedRank = getReducedRank(sourceType.getShape());
    LDBG() << "  -> Source rank: " << sourceType.getRank()
           << ", Reduced rank: " << reducedRank;
    if (reducedRank == sourceType.getRank()) {
      LDBG() << "  -> No unit dimensions to drop, skipping";
      return failure();
    }
    // TODO: Extend vector.mask to support 0-d vectors. In the meantime, bail
    // out.
    if (reducedRank == 0 && maskingOp) {
      LDBG() << "  -> 0-d vector with masking not supported, skipping";
      return failure();
    }
    // Check if the reduced vector shape matches the reduced destination shape.
    // Otherwise, this case is not supported yet.
    VectorType reducedVectorType = trimNonScalableUnitDims(vectorType);
    LDBG() << "  -> Vector type: " << vectorType
           << ", Reduced vector type: " << reducedVectorType;
    if (reducedRank != reducedVectorType.getRank()) {
      LDBG() << "  -> Reduced ranks don't match, skipping";
      return failure();
    }
    if (llvm::any_of(transferWriteOp.getIndices(), [](Value v) {
          return getConstantIntValue(v) != static_cast<int64_t>(0);
        })) {
      LDBG() << "  -> Non-zero indices found, skipping";
      return failure();
    }

    Value maskOp = transferWriteOp.getMask();
    if (maskOp) {
      LDBG() << "  -> Processing mask operation";
      auto maskVectorType = cast<VectorType>(maskOp.getType());
      FailureOr<Value> rankReducedMask = failure();
      if (auto createMaskOp = maskOp.getDefiningOp<vector::CreateMaskOp>())
        rankReducedMask =
            maskDropNonScalableUnitDims(rewriter, loc, createMaskOp);
      else if (auto constantMaskOp =
                   maskOp.getDefiningOp<vector::ConstantMaskOp>())
        rankReducedMask =
            maskDropNonScalableUnitDims(rewriter, loc, constantMaskOp);
      else
        return rewriter.notifyMatchFailure(
            transferWriteOp,
            "unsupported mask op, only 'vector.create_mask' and "
            "'vector.constant_mask' are currently supported");

      if (succeeded(rankReducedMask)) {
        maskOp = *rankReducedMask;
        LDBG() << "  -> Successfully reduced mask dimensions";
      } else if (maskVectorType.getRank() != reducedVectorType.getRank()) {
        return rewriter.notifyMatchFailure(
            transferWriteOp, "Mask reduction required, but failed");
      }
    }
    LDBG() << "  -> Creating rank-reduced subview and new transfer_write";
    Value reducedShapeSource =
        rankReducingSubviewDroppingUnitDims(rewriter, loc, source);
    Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
    Repeated<Value> zeros(reducedRank, c0);
    auto identityMap = rewriter.getMultiDimIdentityMap(reducedRank);
    SmallVector<bool> inBounds(reducedVectorType.getRank(), true);
    auto shapeCastSrc = rewriter.createOrFold<vector::ShapeCastOp>(
        loc, reducedVectorType, vector);
    Operation *newXferWrite = vector::TransferWriteOp::create(
        rewriter, loc, Type(), shapeCastSrc, reducedShapeSource, zeros,
        identityMap, maskOp, rewriter.getBoolArrayAttr(inBounds));
    LDBG() << "  -> Created new transfer_write: " << *newXferWrite;

    if (maskingOp) {
      LDBG() << "  -> Applying masking operation";
      auto shapeCastMask = rewriter.createOrFold<vector::ShapeCastOp>(
          loc, reducedVectorType.cloneWith(std::nullopt, rewriter.getI1Type()),
          maskingOp.getMask());
      newXferWrite =
          mlir::vector::maskOperation(rewriter, newXferWrite, shapeCastMask);
    }

    if (transferWriteOp.hasPureTensorSemantics()) {
      LDBG() << "  -> Pattern match successful (tensor semantics), returning "
                "result";
      return newXferWrite->getResults()[0];
    }

    // With Memref semantics, there's no return value. Use empty value to signal
    // success.
    LDBG() << "  -> Pattern match successful (memref semantics)";
    return Value();
  }
};

} // namespace

/// Creates a memref.collapse_shape collapsing all inner dimensions of the
/// input starting at `firstDimToCollapse`.
static Value collapseInnerDims(PatternRewriter &rewriter, mlir::Location loc,
                               Value input, int64_t firstDimToCollapse) {
  ShapedType inputType = cast<ShapedType>(input.getType());
  if (inputType.getRank() == 1)
    return input;
  SmallVector<ReassociationIndices> reassociation;
  for (int64_t i = 0; i < firstDimToCollapse; ++i)
    reassociation.push_back(ReassociationIndices{i});
  ReassociationIndices collapsedIndices;
  for (int64_t i = firstDimToCollapse; i < inputType.getRank(); ++i)
    collapsedIndices.push_back(i);
  reassociation.push_back(collapsedIndices);
  return memref::CollapseShapeOp::create(rewriter, loc, input, reassociation);
}

/// Returns the new indices that collapses the inner dimensions starting from
/// the `firstDimToCollapse` dimension.
static SmallVector<Value> getCollapsedIndices(RewriterBase &rewriter,
                                              Location loc,
                                              ArrayRef<int64_t> shape,
                                              ValueRange indices,
                                              int64_t firstDimToCollapse) {
  assert(firstDimToCollapse < static_cast<int64_t>(indices.size()));

  // If all the collapsed indices are zero then no extra logic is needed.
  // Otherwise, a new offset/index has to be computed.
  SmallVector<Value> indicesAfterCollapsing(
      indices.begin(), indices.begin() + firstDimToCollapse);
  SmallVector<Value> indicesToCollapse(indices.begin() + firstDimToCollapse,
                                       indices.end());
  if (llvm::all_of(indicesToCollapse, isZeroInteger)) {
    indicesAfterCollapsing.push_back(indicesToCollapse[0]);
    return indicesAfterCollapsing;
  }

  // Compute the remaining trailing index/offset required for reading from
  // the collapsed memref:
  //
  //    offset = 0
  //    for (i = firstDimToCollapse; i < outputRank; ++i)
  //      offset += sourceType.getDimSize(i) * transferReadOp.indices[i]
  //
  // For this example:
  //   %2 = vector.transfer_read/write %arg4[%c0, %arg0, %c0] (...) :
  //      memref<1x43x2xi32>, vector<1x2xi32>
  // which would be collapsed to:
  //   %1 = vector.transfer_read/write %collapse_shape[%c0, %offset] (...) :
  //      memref<1x86xi32>, vector<2xi32>
  // one would get the following offset:
  //    %offset = %arg0 * 43
  OpFoldResult collapsedOffset =
      arith::ConstantIndexOp::create(rewriter, loc, 0).getResult();

  auto collapsedStrides = computeSuffixProduct(
      ArrayRef<int64_t>(shape.begin() + firstDimToCollapse, shape.end()));

  // Compute the collapsed offset.
  auto &&[collapsedExpr, collapsedVals] =
      computeLinearIndex(collapsedOffset, collapsedStrides, indicesToCollapse);
  collapsedOffset = affine::makeComposedFoldedAffineApply(
      rewriter, loc, collapsedExpr, collapsedVals);

  if (auto value = dyn_cast<Value>(collapsedOffset)) {
    indicesAfterCollapsing.push_back(value);
  } else {
    indicesAfterCollapsing.push_back(arith::ConstantIndexOp::create(
        rewriter, loc, *getConstantIntValue(collapsedOffset)));
  }

  return indicesAfterCollapsing;
}

namespace {
/// Rewrites contiguous row-major vector.transfer_read ops by inserting
/// memref.collapse_shape on the source so that the resulting
/// vector.transfer_read has a 1D source. Requires the source shape to be
/// already reduced i.e. without unit dims.
///
/// If `targetVectorBitwidth` is provided, the flattening will only happen if
/// the trailing dimension of the vector read is smaller than the provided
/// bitwidth.
class FlattenContiguousRowMajorTransferReadPattern
    : public OpRewritePattern<vector::TransferReadOp> {
public:
  FlattenContiguousRowMajorTransferReadPattern(MLIRContext *context,
                                               unsigned vectorBitwidth,
                                               PatternBenefit benefit)
      : OpRewritePattern<vector::TransferReadOp>(context, benefit),
        targetVectorBitwidth(vectorBitwidth) {}

  LogicalResult matchAndRewrite(vector::TransferReadOp transferReadOp,
                                PatternRewriter &rewriter) const override {
    LDBG() << "=== FlattenContiguousRowMajorTransferReadPattern: Analyzing "
           << *transferReadOp;
    auto loc = transferReadOp.getLoc();
    Value vector = transferReadOp.getVector();
    VectorType vectorType = cast<VectorType>(vector.getType());
    auto source = transferReadOp.getBase();
    MemRefType sourceType = dyn_cast<MemRefType>(source.getType());

    // 0. Check pre-conditions
    // Contiguity check is valid on tensors only.
    if (!sourceType) {
      LDBG() << "  -> Not a MemRefType, skipping";
      return failure();
    }
    // If this is already 0D/1D, there's nothing to do.
    if (vectorType.getRank() <= 1) {
      LDBG() << "  -> Already 0D/1D, skipping";
      return failure();
    }
    if (!vectorType.getElementType().isSignlessIntOrFloat()) {
      LDBG() << "  -> Not signless int or float, skipping";
      return failure();
    }
    unsigned trailingVectorDimBitwidth =
        vectorType.getShape().back() * vectorType.getElementTypeBitWidth();
    LDBG() << "  -> Trailing vector dim bitwidth: " << trailingVectorDimBitwidth
           << ", target: " << targetVectorBitwidth;
    if (trailingVectorDimBitwidth >= targetVectorBitwidth) {
      LDBG() << "  -> Trailing dim bitwidth >= target, skipping";
      return failure();
    }
    if (!vector::isContiguousSlice(sourceType, vectorType)) {
      LDBG() << "  -> Not contiguous slice, skipping";
      return failure();
    }
    // TODO: generalize this pattern, relax the requirements here.
    if (transferReadOp.hasOutOfBoundsDim()) {
      LDBG() << "  -> Has out-of-bounds dimensions, skipping";
      return failure();
    }
    if (!transferReadOp.getPermutationMap().isMinorIdentity()) {
      LDBG() << "  -> Not minor identity permutation map, skipping";
      return failure();
    }
    if (transferReadOp.getMask()) {
      LDBG() << "  -> Has mask, skipping";
      return failure();
    }

    // Determine vector dimensions to collapse.
    // Ignore a leading sequence of adjacent unit dimensions in the vector.
    ArrayRef<int64_t> collapsedVectorShape =
        vectorType.getShape().drop_while([](auto v) { return v == 1; });
    size_t collapsedVecRank = collapsedVectorShape.size();
    // Limit the collapse of multi-dimensional unit vectors (e.g. <1x1x1xf32>)
    // to a 1D single-element vector.
    if (collapsedVecRank == 0)
      collapsedVecRank = 1;

    // Determine the first memref dimension to collapse - just enough so we can
    // read a flattened vector.
    int64_t firstDimToCollapse = sourceType.getRank() - collapsedVecRank;
    LDBG() << "  -> First dimension to collapse: " << firstDimToCollapse;

    // 1. Collapse the source memref
    LDBG() << "  -> Collapsing source memref";
    Value collapsedSource =
        collapseInnerDims(rewriter, loc, source, firstDimToCollapse);
    MemRefType collapsedSourceType =
        cast<MemRefType>(collapsedSource.getType());
    int64_t collapsedRank = collapsedSourceType.getRank();
    assert(collapsedRank == firstDimToCollapse + 1);
    LDBG() << "  -> Collapsed source type: " << collapsedSourceType;

    // 2. Generate input args for a new vector.transfer_read that will read
    // from the collapsed memref.
    // 2.1. New dim exprs + affine map
    SmallVector<AffineExpr, 1> dimExprs{
        getAffineDimExpr(firstDimToCollapse, rewriter.getContext())};
    auto collapsedMap =
        AffineMap::get(collapsedRank, 0, dimExprs, rewriter.getContext());

    // 2.2 New indices
    SmallVector<Value> collapsedIndices =
        getCollapsedIndices(rewriter, loc, sourceType.getShape(),
                            transferReadOp.getIndices(), firstDimToCollapse);

    // 3. Create new vector.transfer_read that reads from the collapsed memref
    VectorType flatVectorType = VectorType::get({vectorType.getNumElements()},
                                                vectorType.getElementType());
    LDBG() << "  -> Creating flattened vector type: " << flatVectorType;
    vector::TransferReadOp flatRead = vector::TransferReadOp::create(
        rewriter, loc, flatVectorType, collapsedSource, collapsedIndices,
        transferReadOp.getPadding(), collapsedMap);
    flatRead.setInBoundsAttr(rewriter.getBoolArrayAttr({true}));
    LDBG() << "  -> Created flat transfer_read: " << *flatRead;

    // 4. Replace the old transfer_read with the new one reading from the
    // collapsed shape
    LDBG() << "  -> Replacing with shape cast";
    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
        transferReadOp, cast<VectorType>(vector.getType()), flatRead);
    LDBG() << "  -> Pattern match successful";
    return success();
  }

private:
  // Minimum bitwidth that the trailing vector dimension should have after
  // flattening.
  unsigned targetVectorBitwidth;
};

/// Rewrites contiguous row-major vector.transfer_write ops by inserting
/// memref.collapse_shape on the source so that the resulting
/// vector.transfer_write has a 1D source. Requires the source shape to be
/// already reduced i.e. without unit dims.
///
/// If `targetVectorBitwidth` is provided, the flattening will only happen if
/// the trailing dimension of the vector read is smaller than the provided
/// bitwidth.
class FlattenContiguousRowMajorTransferWritePattern
    : public OpRewritePattern<vector::TransferWriteOp> {
public:
  FlattenContiguousRowMajorTransferWritePattern(MLIRContext *context,
                                                unsigned vectorBitwidth,
                                                PatternBenefit benefit)
      : OpRewritePattern<vector::TransferWriteOp>(context, benefit),
        targetVectorBitwidth(vectorBitwidth) {}

  LogicalResult matchAndRewrite(vector::TransferWriteOp transferWriteOp,
                                PatternRewriter &rewriter) const override {
    auto loc = transferWriteOp.getLoc();
    Value vector = transferWriteOp.getVector();
    VectorType vectorType = cast<VectorType>(vector.getType());
    Value source = transferWriteOp.getBase();
    MemRefType sourceType = dyn_cast<MemRefType>(source.getType());

    // 0. Check pre-conditions
    // Contiguity check is valid on tensors only.
    if (!sourceType)
      return failure();
    // If this is already 0D/1D, there's nothing to do.
    if (vectorType.getRank() <= 1)
      // Already 0D/1D, nothing to do.
      return failure();
    if (!vectorType.getElementType().isSignlessIntOrFloat())
      return failure();
    unsigned trailingVectorDimBitwidth =
        vectorType.getShape().back() * vectorType.getElementTypeBitWidth();
    if (trailingVectorDimBitwidth >= targetVectorBitwidth)
      return failure();
    if (!vector::isContiguousSlice(sourceType, vectorType))
      return failure();
    // TODO: generalize this pattern, relax the requirements here.
    if (transferWriteOp.hasOutOfBoundsDim())
      return failure();
    if (!transferWriteOp.getPermutationMap().isMinorIdentity())
      return failure();
    if (transferWriteOp.getMask())
      return failure();

    // Determine vector dimensions to collapse.
    // Ignore a leading sequence of adjacent unit dimensions in the vector.
    ArrayRef<int64_t> collapsedVectorShape =
        vectorType.getShape().drop_while([](auto v) { return v == 1; });
    size_t collapsedVecRank = collapsedVectorShape.size();
    // Limit the collapse of multi-dimensional unit vectors (e.g. <1x1x1xf32>)
    // to a 1D single-element vector.
    if (collapsedVecRank == 0)
      collapsedVecRank = 1;

    // Determine the first memref dimension to collapse - just enough so we can
    // read a flattened vector.
    int64_t firstDimToCollapse = sourceType.getRank() - collapsedVecRank;

    // 1. Collapse the source memref
    Value collapsedSource =
        collapseInnerDims(rewriter, loc, source, firstDimToCollapse);
    MemRefType collapsedSourceType =
        cast<MemRefType>(collapsedSource.getType());
    int64_t collapsedRank = collapsedSourceType.getRank();
    assert(collapsedRank == firstDimToCollapse + 1);

    // 2. Generate input args for a new vector.transfer_read that will read
    // from the collapsed memref.
    // 2.1. New dim exprs + affine map
    SmallVector<AffineExpr, 1> dimExprs{
        getAffineDimExpr(firstDimToCollapse, rewriter.getContext())};
    auto collapsedMap =
        AffineMap::get(collapsedRank, 0, dimExprs, rewriter.getContext());

    // 2.2 New indices
    SmallVector<Value> collapsedIndices =
        getCollapsedIndices(rewriter, loc, sourceType.getShape(),
                            transferWriteOp.getIndices(), firstDimToCollapse);

    // 3. Create new vector.transfer_write that writes to the collapsed memref
    VectorType flatVectorType = VectorType::get({vectorType.getNumElements()},
                                                vectorType.getElementType());
    Value flatVector =
        vector::ShapeCastOp::create(rewriter, loc, flatVectorType, vector);
    vector::TransferWriteOp flatWrite = vector::TransferWriteOp::create(
        rewriter, loc, flatVector, collapsedSource, collapsedIndices,
        collapsedMap);
    flatWrite.setInBoundsAttr(rewriter.getBoolArrayAttr({true}));

    // 4. Replace the old transfer_write with the new one writing the
    // collapsed shape
    rewriter.eraseOp(transferWriteOp);
    return success();
  }

private:
  // Minimum bitwidth that the trailing vector dimension should have after
  // flattening.
  unsigned targetVectorBitwidth;
};

/// Rewrite `vector.extract(vector.transfer_read)` to `memref.load`.
///
/// All the users of the transfer op must be `vector.extract` ops. If
/// `allowMultipleUses` is set to true, rewrite transfer ops with any number of
/// users. Otherwise, rewrite only if the extract op is the single user of the
/// transfer op. Rewriting a single vector load with multiple scalar loads may
/// negatively affect performance.
class RewriteScalarExtractOfTransferRead
    : public OpRewritePattern<vector::ExtractOp> {
public:
  RewriteScalarExtractOfTransferRead(MLIRContext *context,
                                     PatternBenefit benefit,
                                     bool allowMultipleUses)
      : OpRewritePattern(context, benefit),
        allowMultipleUses(allowMultipleUses) {}

  LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
                                PatternRewriter &rewriter) const override {
    // Match phase.
    auto xferOp = extractOp.getSource().getDefiningOp<vector::TransferReadOp>();
    if (!xferOp)
      return failure();
    // Check that we are extracting a scalar and not a sub-vector.
    if (isa<VectorType>(extractOp.getResult().getType()))
      return failure();
    // If multiple uses are not allowed, check if xfer has a single use.
    if (!allowMultipleUses && !xferOp.getResult().hasOneUse())
      return failure();
    // If multiple uses are allowed, check if all the xfer uses are extract ops.
    if (allowMultipleUses &&
        !llvm::all_of(xferOp->getUses(), [](OpOperand &use) {
          return isa<vector::ExtractOp>(use.getOwner());
        }))
      return failure();
    // Mask not supported.
    if (xferOp.getMask())
      return failure();
    // Map not supported.
    if (!xferOp.getPermutationMap().isMinorIdentity())
      return failure();
    // Cannot rewrite if the indices may be out of bounds.
    if (xferOp.hasOutOfBoundsDim())
      return failure();

    // Rewrite phase: construct scalar load.
    SmallVector<Value> newIndices(xferOp.getIndices().begin(),
                                  xferOp.getIndices().end());
    for (auto [i, pos] : llvm::enumerate(extractOp.getMixedPosition())) {
      int64_t idx = newIndices.size() - extractOp.getNumIndices() + i;

      // Compute affine expression `newIndices[idx] + pos` where `pos` can be
      // either a constant or a value.
      OpFoldResult composedIdx;
      if (auto attr = dyn_cast<Attribute>(pos)) {
        int64_t offset = cast<IntegerAttr>(attr).getInt();
        composedIdx = affine::makeComposedFoldedAffineApply(
            rewriter, extractOp.getLoc(),
            rewriter.getAffineSymbolExpr(0) + offset, {newIndices[idx]});
      } else {
        Value dynamicOffset = cast<Value>(pos);
        AffineExpr sym0, sym1;
        bindSymbols(rewriter.getContext(), sym0, sym1);
        composedIdx = affine::makeComposedFoldedAffineApply(
            rewriter, extractOp.getLoc(), sym0 + sym1,
            {newIndices[idx], dynamicOffset});
      }

      // Update the corresponding index with the folded result.
      if (auto value = dyn_cast<Value>(composedIdx)) {
        newIndices[idx] = value;
      } else {
        newIndices[idx] = arith::ConstantIndexOp::create(
            rewriter, extractOp.getLoc(), *getConstantIntValue(composedIdx));
      }
    }
    if (isa<MemRefType>(xferOp.getBase().getType())) {
      rewriter.replaceOpWithNewOp<memref::LoadOp>(extractOp, xferOp.getBase(),
                                                  newIndices);
    } else {
      rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
          extractOp, xferOp.getBase(), newIndices);
    }

    return success();
  }

private:
  bool allowMultipleUses;
};

/// Rewrite transfer_writes of vectors of size 1 (e.g., vector<1x1xf32>)
/// to memref.store.
class RewriteScalarWrite : public OpRewritePattern<vector::TransferWriteOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp,
                                PatternRewriter &rewriter) const override {
    // Must be a scalar write.
    auto vecType = xferOp.getVectorType();
    if (!llvm::all_of(vecType.getShape(), [](int64_t sz) { return sz == 1; }))
      return failure();
    // Mask not supported.
    if (xferOp.getMask())
      return failure();
    // Map not supported.
    if (!xferOp.getPermutationMap().isMinorIdentity())
      return failure();
    // Only float and integer element types are supported.
    Value scalar = vector::ExtractOp::create(rewriter, xferOp.getLoc(),
                                             xferOp.getVector());
    // Construct a scalar store.
    if (isa<MemRefType>(xferOp.getBase().getType())) {
      rewriter.replaceOpWithNewOp<memref::StoreOp>(
          xferOp, scalar, xferOp.getBase(), xferOp.getIndices());
    } else {
      rewriter.replaceOpWithNewOp<tensor::InsertOp>(
          xferOp, scalar, xferOp.getBase(), xferOp.getIndices());
    }
    return success();
  }
};

} // namespace

void mlir::vector::transferOpflowOpt(RewriterBase &rewriter,
                                     Operation *rootOp) {
  LDBG() << "=== Starting transferOpflowOpt on root operation: "
         << OpWithFlags(rootOp, OpPrintingFlags().skipRegions());
  TransferOptimization opt(rewriter, rootOp);

  // Run store to load forwarding first since it can expose more dead store
  // opportunity.
  LDBG() << "Phase 1: Store-to-load forwarding";
  int readCount = 0;
  rootOp->walk([&](vector::TransferReadOp read) {
    if (isa<MemRefType>(read.getShapedType())) {
      LDBG() << "Processing transfer_read #" << ++readCount << ": " << *read;
      opt.storeToLoadForwarding(read);
    }
  });
  LDBG() << "Phase 1 complete. Removing dead operations from forwarding";
  opt.removeDeadOp();

  LDBG() << "Phase 2: Dead store elimination";
  int writeCount = 0;
  rootOp->walk([&](vector::TransferWriteOp write) {
    if (isa<MemRefType>(write.getShapedType())) {
      LDBG() << "Processing transfer_write #" << ++writeCount << ": " << *write;
      opt.deadStoreOp(write);
    }
  });
  LDBG() << "Phase 2 complete. Removing dead operations from dead store "
            "elimination";
  opt.removeDeadOp();
  LDBG() << "=== transferOpflowOpt complete";
}

void mlir::vector::populateScalarVectorTransferLoweringPatterns(
    RewritePatternSet &patterns, PatternBenefit benefit,
    bool allowMultipleUses) {
  patterns.add<RewriteScalarExtractOfTransferRead>(patterns.getContext(),
                                                   benefit, allowMultipleUses);
  patterns.add<RewriteScalarWrite>(patterns.getContext(), benefit);
}

void mlir::vector::populateVectorTransferDropUnitDimsPatterns(
    RewritePatternSet &patterns, PatternBenefit benefit) {
  patterns
      .add<TransferReadDropUnitDimsPattern, TransferWriteDropUnitDimsPattern>(
          patterns.getContext(), benefit);
}

void mlir::vector::populateFlattenVectorTransferPatterns(
    RewritePatternSet &patterns, unsigned targetVectorBitwidth,
    PatternBenefit benefit) {
  patterns.add<FlattenContiguousRowMajorTransferReadPattern,
               FlattenContiguousRowMajorTransferWritePattern>(
      patterns.getContext(), targetVectorBitwidth, benefit);
  populateDropUnitDimWithShapeCastPatterns(patterns, benefit);
}