MemorySlotInterfacesImpl.cpp

Go to the documentation of this file.

//===- MemorySlotInterfacesImpl.cpp - Mem2Reg for Vector ops --------------===//
//
// 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 memory slot interface models for
// `vector.transfer_read`, `vector.transfer_write`, and `vector.mask`. They
// implement `PromotableMemOpInterface` to replace vector transfers with uses
// and definitions of a vector SSA value. The models for `memref.subview`
// live in the MemRef dialect.
//
// A read uses the slot's current value; a write defines its next value.
// Static subviews use a vector matching the subview's shape: reads extract that
// sub-vector from the parent, and writes insert it back into the parent.
// Supported transfers through dynamic subviews use a vector type matching the
// whole static parent, and the subview's slot holds that whole vector. Reads
// pad lanes outside the runtime extent, and writes preserve the previous value
// of those lanes. `arith.select` implements this behavior and also handles
// lanes disabled by a transfer's mask operand.
//
// `vector.mask` owns its wrapped transfer's blocking use and is promoted as a
// whole. Its region cannot accept extra operations, so `arith.select` is built
// outside it.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/Vector/Transforms/MemorySlotInterfacesImpl.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/Transforms/MemorySlotInterfacesImpl.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Interfaces/MemorySlotInterfaces.h"
 
using namespace mlir;
using namespace mlir::vector;
 
//===----------------------------------------------------------------------===//
//  Utilities
//===----------------------------------------------------------------------===//
 
/// Returns whether `xferOp` can be promoted to use or define `slot`'s vector
/// value. The transfer's memref base must be the slot pointer, its vector type
/// must equal `slot.valueType`, its indices must be zero, and its permutation
/// map must be the identity. Out-of-bounds dimensions are allowed only for a
/// supported dynamic subview.
///
/// More general transfers may be supportable, but their value depends on
/// motivating use cases, implementation complexity, and generated code quality.
static bool
isPromotableTransfer(VectorTransferOpInterface xferOp, const MemorySlot &slot,
                     const SmallPtrSetImpl<OpOperand *> &blockingUses) {
  // The transfer's memref base must be the slot pointer.
  if (blockingUses.size() != 1)
    return false;
  Value blockingUse = (*blockingUses.begin())->get();
  if (blockingUse != slot.ptr || xferOp.getBase() != slot.ptr)
    return false;
 
  // Reject the tensor form (already implied, since slot pointers are memrefs).
  if (!isa<MemRefType>(xferOp.getBase().getType()))
    return false;
 
  // Exact type match pins rank/extents/element type/scalable dims.
  if (xferOp.getVectorType() != slot.valueType)
    return false;
 
  // Access must start at the buffer origin in every dimension.
  for (Value index : xferOp.getIndices()) {
    std::optional<int64_t> constIndex = getConstantIntValue(index);
    if (!constIndex || *constIndex != 0)
      return false;
  }
 
  // Identity map: no broadcast or transpose.
  if (!xferOp.getPermutationMap().isIdentity())
    return false;
 
  // Out-of-bounds is allowed only for a dynamic view.
  if (xferOp.hasOutOfBoundsDim() && !memref::isDynamicSubViewSlot(slot.ptr))
    return false;
 
  return true;
}
 
//===----------------------------------------------------------------------===//
//  PromotableMemOpInterface models
//===----------------------------------------------------------------------===//
 
namespace {
/// `PromotableMemOpInterface` model for `vector.transfer_read`.
struct TransferReadOpMemOpModel
    : public PromotableMemOpInterface::ExternalModel<TransferReadOpMemOpModel,
                                                     vector::TransferReadOp> {
  bool loadsFrom(Operation *op, const MemorySlot &slot) const {
    return cast<vector::TransferReadOp>(op).getBase() == slot.ptr;
  }
 
  bool storesTo(Operation *op, const MemorySlot &slot) const { return false; }
 
  Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
                  Value reachingDef, const DataLayout &dataLayout) const {
    llvm_unreachable("getStored should not be called on TransferReadOp");
  }
 
  bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
                        const SmallPtrSetImpl<OpOperand *> &blockingUses,
                        SmallVectorImpl<OpOperand *> &newBlockingUses,
                        const DataLayout &dataLayout) const {
    return isPromotableTransfer(cast<VectorTransferOpInterface>(op), slot,
                                blockingUses);
  }
 
  // Replace the read with the reaching definition, selecting padding for lanes
  // outside the subview's extent or disabled by the transfer's mask.
  // PromotableAliaserInterface::projectSlotValueToAliasValue supplies the whole
  // parent vector for a dynamic subview. Apply the extent mask here because
  // this hook has access to the transfer's padding operand. If the transfer
  // also has a mask operand, combine the two masks with arith.andi.
  DeletionKind
  removeBlockingUses(Operation *op, const MemorySlot &slot,
                     const SmallPtrSetImpl<OpOperand *> &blockingUses,
                     OpBuilder &builder, Value reachingDefinition,
                     const DataLayout &dataLayout) const {
    auto readOp = cast<vector::TransferReadOp>(op);
    Location loc = op->getLoc();
    Value mask = memref::buildDynamicSubViewMask(builder, loc, slot.ptr);
    if (Value opMask = readOp.getMask())
      mask = mask
                 ? arith::AndIOp::create(builder, loc, mask, opMask).getResult()
                 : opMask;
 
    Value result = reachingDefinition;
    if (mask) {
      Value padSplat = vector::BroadcastOp::create(
          builder, loc, readOp.getVectorType(), readOp.getPadding());
      result = arith::SelectOp::create(builder, loc, mask, reachingDefinition,
                                       padSplat);
    }
    readOp.getResult().replaceAllUsesWith(result);
    return DeletionKind::Delete;
  }
};
 
/// `PromotableMemOpInterface` model for `vector.transfer_write`.
struct TransferWriteOpMemOpModel
    : public PromotableMemOpInterface::ExternalModel<TransferWriteOpMemOpModel,
                                                     vector::TransferWriteOp> {
  bool loadsFrom(Operation *op, const MemorySlot &slot) const { return false; }
 
  bool storesTo(Operation *op, const MemorySlot &slot) const {
    return cast<vector::TransferWriteOp>(op).getBase() == slot.ptr;
  }
 
  // The stored value is the transfer's vector, with inactive lanes preserved
  // from the reaching definition when the transfer has a mask.
  Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
                  Value reachingDef, const DataLayout &dataLayout) const {
    auto writeOp = cast<vector::TransferWriteOp>(op);
    Value stored = writeOp.getValueToStore();
    // PromotableAliaserInterface::projectAliasValueToSlotValue separately
    // preserves parent lanes outside a dynamic subview's extent.
    if (Value mask = writeOp.getMask())
      stored = arith::SelectOp::create(builder, op->getLoc(), mask, stored,
                                       reachingDef);
    return stored;
  }
 
  bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
                        const SmallPtrSetImpl<OpOperand *> &blockingUses,
                        SmallVectorImpl<OpOperand *> &newBlockingUses,
                        const DataLayout &dataLayout) const {
    return isPromotableTransfer(cast<VectorTransferOpInterface>(op), slot,
                                blockingUses);
  }
 
  // `getStored` already provided the value for later uses, so the write can
  // simply be erased.
  DeletionKind
  removeBlockingUses(Operation *op, const MemorySlot &slot,
                     const SmallPtrSetImpl<OpOperand *> &blockingUses,
                     OpBuilder &builder, Value reachingDefinition,
                     const DataLayout &dataLayout) const {
    return DeletionKind::Delete;
  }
};
 
/// Returns the operation wrapped by the `vector.mask` `op`.
PromotableMemOpInterface getWrappedMemOp(Operation *op) {
  Operation *maskableOp = cast<vector::MaskOp>(op).getMaskableOp();
  if (!maskableOp)
    return {};
  if (!isa<vector::TransferReadOp, vector::TransferWriteOp>(maskableOp))
    return {};
  return cast<PromotableMemOpInterface>(maskableOp);
}
 
/// `PromotableMemOpInterface` model for `vector.mask` wrapping a promotable
/// memory operation. `vector.mask` owns the wrapped operation's slot use,
/// delegates promotion to it, and applies its mask to the result.
struct MaskOpMemOpModel
    : public PromotableMemOpInterface::ExternalModel<MaskOpMemOpModel,
                                                     vector::MaskOp> {
  bool ownsWrappedBlockingUses(Operation *op) const {
    return static_cast<bool>(getWrappedMemOp(op));
  }
 
  bool loadsFrom(Operation *op, const MemorySlot &slot) const {
    PromotableMemOpInterface memOp = getWrappedMemOp(op);
    return memOp && memOp.loadsFrom(slot);
  }
 
  bool storesTo(Operation *op, const MemorySlot &slot) const {
    PromotableMemOpInterface memOp = getWrappedMemOp(op);
    return memOp && memOp.storesTo(slot);
  }
 
  Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
                  Value reachingDef, const DataLayout &dataLayout) const {
    Value stored =
        getWrappedMemOp(op).getStored(slot, builder, reachingDef, dataLayout);
    return arith::SelectOp::create(builder, op->getLoc(),
                                   cast<vector::MaskOp>(op).getMask(), stored,
                                   reachingDef);
  }
 
  bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
                        const SmallPtrSetImpl<OpOperand *> &blockingUses,
                        SmallVectorImpl<OpOperand *> &newBlockingUses,
                        const DataLayout &dataLayout) const {
    PromotableMemOpInterface memOp = getWrappedMemOp(op);
    // Transfers with their own mask operand are unsupported.
    if (!memOp || cast<VectorTransferOpInterface>(*memOp).getMask())
      return false;
    return memOp.canUsesBeRemoved(slot, blockingUses, newBlockingUses,
                                  dataLayout);
  }
 
  // Let the wrapped memory op promote itself, then apply the wrapper's mask to
  // the value it now yields.
  DeletionKind
  removeBlockingUses(Operation *op, const MemorySlot &slot,
                     const SmallPtrSetImpl<OpOperand *> &blockingUses,
                     OpBuilder &builder, Value reachingDefinition,
                     const DataLayout &dataLayout) const {
    auto maskOp = cast<vector::MaskOp>(op);
    PromotableMemOpInterface memOp = getWrappedMemOp(op);
    memOp.removeBlockingUses(slot, blockingUses, builder, reachingDefinition,
                             dataLayout);
    auto readOp = dyn_cast<vector::TransferReadOp>(*memOp);
    if (!readOp)
      return DeletionKind::Delete;
 
    Location loc = op->getLoc();
    Value enabled = maskOp.getMaskBlock()->getTerminator()->getOperand(0);
    Value disabled =
        maskOp.hasPassthru()
            ? maskOp.getPassthru()
            : vector::BroadcastOp::create(builder, loc, readOp.getVectorType(),
                                          readOp.getPadding());
    Value result = arith::SelectOp::create(builder, loc, maskOp.getMask(),
                                           enabled, disabled);
    maskOp->getResult(0).replaceAllUsesWith(result);
    return DeletionKind::Delete;
  }
};
} // namespace
 
//===----------------------------------------------------------------------===//
//  Register external models
//===----------------------------------------------------------------------===//
 
void mlir::vector::registerMemorySlotInterfacesExternalModels(
    DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
    TransferReadOp::attachInterface<TransferReadOpMemOpModel>(*ctx);
    TransferWriteOp::attachInterface<TransferWriteOpMemOpModel>(*ctx);
    MaskOp::attachInterface<MaskOpMemOpModel>(*ctx);
  });
}