337 lines · cpp
1//===- MemRefMemorySlot.cpp - Memory Slot Interfaces ------------*- C++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements Mem2Reg-related interfaces for MemRef dialect10// operations.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"15#include "mlir/Dialect/MemRef/IR/MemRef.h"16#include "mlir/Dialect/UB/IR/UBOps.h"17#include "mlir/IR/BuiltinDialect.h"18#include "mlir/IR/BuiltinTypes.h"19#include "mlir/IR/Matchers.h"20#include "mlir/IR/Value.h"21#include "mlir/Interfaces/MemorySlotInterfaces.h"22#include "llvm/ADT/ArrayRef.h"23#include "llvm/ADT/TypeSwitch.h"24#include "llvm/Support/ErrorHandling.h"25 26using namespace mlir;27 28//===----------------------------------------------------------------------===//29// Utilities30//===----------------------------------------------------------------------===//31 32/// Walks over the indices of the elements of a tensor of a given `shape` by33/// updating `index` in place to the next index. This returns failure if the34/// provided index was the last index.35static LogicalResult nextIndex(ArrayRef<int64_t> shape,36 MutableArrayRef<int64_t> index) {37 for (size_t i = 0; i < shape.size(); ++i) {38 index[i]++;39 if (index[i] < shape[i])40 return success();41 index[i] = 0;42 }43 return failure();44}45 46/// Calls `walker` for each index within a tensor of a given `shape`, providing47/// the index as an array attribute of the coordinates.48template <typename CallableT>49static void walkIndicesAsAttr(MLIRContext *ctx, ArrayRef<int64_t> shape,50 CallableT &&walker) {51 Type indexType = IndexType::get(ctx);52 SmallVector<int64_t> shapeIter(shape.size(), 0);53 do {54 SmallVector<Attribute> indexAsAttr;55 for (int64_t dim : shapeIter)56 indexAsAttr.push_back(IntegerAttr::get(indexType, dim));57 walker(ArrayAttr::get(ctx, indexAsAttr));58 } while (succeeded(nextIndex(shape, shapeIter)));59}60 61//===----------------------------------------------------------------------===//62// Interfaces for AllocaOp63//===----------------------------------------------------------------------===//64 65SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {66 MemRefType type = getType();67 if (!type.hasStaticShape())68 return {};69 // Make sure the memref contains only a single element.70 if (type.getNumElements() != 1)71 return {};72 73 return {MemorySlot{getResult(), type.getElementType()}};74}75 76Value memref::AllocaOp::getDefaultValue(const MemorySlot &slot,77 OpBuilder &builder) {78 return ub::PoisonOp::create(builder, getLoc(), slot.elemType);79}80 81std::optional<PromotableAllocationOpInterface>82memref::AllocaOp::handlePromotionComplete(const MemorySlot &slot,83 Value defaultValue,84 OpBuilder &builder) {85 if (defaultValue.use_empty())86 defaultValue.getDefiningOp()->erase();87 this->erase();88 return std::nullopt;89}90 91void memref::AllocaOp::handleBlockArgument(const MemorySlot &slot,92 BlockArgument argument,93 OpBuilder &builder) {}94 95SmallVector<DestructurableMemorySlot>96memref::AllocaOp::getDestructurableSlots() {97 MemRefType memrefType = getType();98 auto destructurable = llvm::dyn_cast<DestructurableTypeInterface>(memrefType);99 if (!destructurable)100 return {};101 102 std::optional<DenseMap<Attribute, Type>> destructuredType =103 destructurable.getSubelementIndexMap();104 if (!destructuredType)105 return {};106 107 return {108 DestructurableMemorySlot{{getMemref(), memrefType}, *destructuredType}};109}110 111DenseMap<Attribute, MemorySlot> memref::AllocaOp::destructure(112 const DestructurableMemorySlot &slot,113 const SmallPtrSetImpl<Attribute> &usedIndices, OpBuilder &builder,114 SmallVectorImpl<DestructurableAllocationOpInterface> &newAllocators) {115 builder.setInsertionPointAfter(*this);116 117 DenseMap<Attribute, MemorySlot> slotMap;118 119 auto memrefType = llvm::cast<DestructurableTypeInterface>(getType());120 for (Attribute usedIndex : usedIndices) {121 Type elemType = memrefType.getTypeAtIndex(usedIndex);122 MemRefType elemPtr = MemRefType::get({}, elemType);123 auto subAlloca = memref::AllocaOp::create(builder, getLoc(), elemPtr);124 newAllocators.push_back(subAlloca);125 slotMap.try_emplace<MemorySlot>(usedIndex,126 {subAlloca.getResult(), elemType});127 }128 129 return slotMap;130}131 132std::optional<DestructurableAllocationOpInterface>133memref::AllocaOp::handleDestructuringComplete(134 const DestructurableMemorySlot &slot, OpBuilder &builder) {135 assert(slot.ptr == getResult());136 this->erase();137 return std::nullopt;138}139 140//===----------------------------------------------------------------------===//141// Interfaces for LoadOp/StoreOp142//===----------------------------------------------------------------------===//143 144bool memref::LoadOp::loadsFrom(const MemorySlot &slot) {145 return getMemRef() == slot.ptr;146}147 148bool memref::LoadOp::storesTo(const MemorySlot &slot) { return false; }149 150Value memref::LoadOp::getStored(const MemorySlot &slot, OpBuilder &builder,151 Value reachingDef,152 const DataLayout &dataLayout) {153 llvm_unreachable("getStored should not be called on LoadOp");154}155 156bool memref::LoadOp::canUsesBeRemoved(157 const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,158 SmallVectorImpl<OpOperand *> &newBlockingUses,159 const DataLayout &dataLayout) {160 if (blockingUses.size() != 1)161 return false;162 Value blockingUse = (*blockingUses.begin())->get();163 return blockingUse == slot.ptr && getMemRef() == slot.ptr &&164 getResult().getType() == slot.elemType;165}166 167DeletionKind memref::LoadOp::removeBlockingUses(168 const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,169 OpBuilder &builder, Value reachingDefinition,170 const DataLayout &dataLayout) {171 // `canUsesBeRemoved` checked this blocking use must be the loaded slot172 // pointer.173 getResult().replaceAllUsesWith(reachingDefinition);174 return DeletionKind::Delete;175}176 177/// Returns the index of a memref in attribute form, given its indices. Returns178/// a null pointer if whether the indices form a valid index for the provided179/// MemRefType cannot be computed. The indices must come from a valid memref180/// StoreOp or LoadOp.181static Attribute getAttributeIndexFromIndexOperands(MLIRContext *ctx,182 ValueRange indices,183 MemRefType memrefType) {184 SmallVector<Attribute> index;185 for (auto [coord, dimSize] : llvm::zip(indices, memrefType.getShape())) {186 IntegerAttr coordAttr;187 if (!matchPattern(coord, m_Constant<IntegerAttr>(&coordAttr)))188 return {};189 // MemRefType shape dimensions are always positive (checked by verifier).190 std::optional<uint64_t> coordInt = coordAttr.getValue().tryZExtValue();191 if (!coordInt || coordInt.value() >= static_cast<uint64_t>(dimSize))192 return {};193 index.push_back(coordAttr);194 }195 return ArrayAttr::get(ctx, index);196}197 198bool memref::LoadOp::canRewire(const DestructurableMemorySlot &slot,199 SmallPtrSetImpl<Attribute> &usedIndices,200 SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,201 const DataLayout &dataLayout) {202 if (slot.ptr != getMemRef())203 return false;204 Attribute index = getAttributeIndexFromIndexOperands(205 getContext(), getIndices(), getMemRefType());206 if (!index)207 return false;208 usedIndices.insert(index);209 return true;210}211 212DeletionKind memref::LoadOp::rewire(const DestructurableMemorySlot &slot,213 DenseMap<Attribute, MemorySlot> &subslots,214 OpBuilder &builder,215 const DataLayout &dataLayout) {216 Attribute index = getAttributeIndexFromIndexOperands(217 getContext(), getIndices(), getMemRefType());218 const MemorySlot &memorySlot = subslots.at(index);219 setMemRef(memorySlot.ptr);220 getIndicesMutable().clear();221 return DeletionKind::Keep;222}223 224bool memref::StoreOp::loadsFrom(const MemorySlot &slot) { return false; }225 226bool memref::StoreOp::storesTo(const MemorySlot &slot) {227 return getMemRef() == slot.ptr;228}229 230Value memref::StoreOp::getStored(const MemorySlot &slot, OpBuilder &builder,231 Value reachingDef,232 const DataLayout &dataLayout) {233 return getValue();234}235 236bool memref::StoreOp::canUsesBeRemoved(237 const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,238 SmallVectorImpl<OpOperand *> &newBlockingUses,239 const DataLayout &dataLayout) {240 if (blockingUses.size() != 1)241 return false;242 Value blockingUse = (*blockingUses.begin())->get();243 return blockingUse == slot.ptr && getMemRef() == slot.ptr &&244 getValue() != slot.ptr && getValue().getType() == slot.elemType;245}246 247DeletionKind memref::StoreOp::removeBlockingUses(248 const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,249 OpBuilder &builder, Value reachingDefinition,250 const DataLayout &dataLayout) {251 return DeletionKind::Delete;252}253 254bool memref::StoreOp::canRewire(const DestructurableMemorySlot &slot,255 SmallPtrSetImpl<Attribute> &usedIndices,256 SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,257 const DataLayout &dataLayout) {258 if (slot.ptr != getMemRef() || getValue() == slot.ptr)259 return false;260 Attribute index = getAttributeIndexFromIndexOperands(261 getContext(), getIndices(), getMemRefType());262 if (!index || !slot.subelementTypes.contains(index))263 return false;264 usedIndices.insert(index);265 return true;266}267 268DeletionKind memref::StoreOp::rewire(const DestructurableMemorySlot &slot,269 DenseMap<Attribute, MemorySlot> &subslots,270 OpBuilder &builder,271 const DataLayout &dataLayout) {272 Attribute index = getAttributeIndexFromIndexOperands(273 getContext(), getIndices(), getMemRefType());274 const MemorySlot &memorySlot = subslots.at(index);275 setMemRef(memorySlot.ptr);276 getIndicesMutable().clear();277 return DeletionKind::Keep;278}279 280//===----------------------------------------------------------------------===//281// Interfaces for destructurable types282//===----------------------------------------------------------------------===//283 284namespace {285 286struct MemRefDestructurableTypeExternalModel287 : public DestructurableTypeInterface::ExternalModel<288 MemRefDestructurableTypeExternalModel, MemRefType> {289 std::optional<DenseMap<Attribute, Type>>290 getSubelementIndexMap(Type type) const {291 auto memrefType = llvm::cast<MemRefType>(type);292 constexpr int64_t maxMemrefSizeForDestructuring = 16;293 if (!memrefType.hasStaticShape() ||294 memrefType.getNumElements() > maxMemrefSizeForDestructuring ||295 memrefType.getNumElements() == 1)296 return {};297 298 DenseMap<Attribute, Type> destructured;299 walkIndicesAsAttr(300 memrefType.getContext(), memrefType.getShape(), [&](Attribute index) {301 destructured.insert({index, memrefType.getElementType()});302 });303 304 return destructured;305 }306 307 Type getTypeAtIndex(Type type, Attribute index) const {308 auto memrefType = llvm::cast<MemRefType>(type);309 auto coordArrAttr = llvm::dyn_cast<ArrayAttr>(index);310 if (!coordArrAttr || coordArrAttr.size() != memrefType.getShape().size())311 return {};312 313 Type indexType = IndexType::get(memrefType.getContext());314 for (const auto &[coordAttr, dimSize] :315 llvm::zip(coordArrAttr, memrefType.getShape())) {316 auto coord = llvm::dyn_cast<IntegerAttr>(coordAttr);317 if (!coord || coord.getType() != indexType || coord.getInt() < 0 ||318 coord.getInt() >= dimSize)319 return {};320 }321 322 return memrefType.getElementType();323 }324};325 326} // namespace327 328//===----------------------------------------------------------------------===//329// Register external models330//===----------------------------------------------------------------------===//331 332void mlir::memref::registerMemorySlotExternalModels(DialectRegistry ®istry) {333 registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {334 MemRefType::attachInterface<MemRefDestructurableTypeExternalModel>(*ctx);335 });336}337