174 lines · cpp
1//===- BufferUtils.cpp - buffer transformation utilities ------------------===//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 utilities for buffer optimization passes.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"14 15#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"16#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"17#include "mlir/Dialect/MemRef/IR/MemRef.h"18#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"19#include "mlir/IR/Operation.h"20#include <optional>21 22using namespace mlir;23using namespace mlir::bufferization;24 25//===----------------------------------------------------------------------===//26// BufferPlacementAllocs27//===----------------------------------------------------------------------===//28 29/// Get the start operation to place the given alloc value withing the30// specified placement block.31Operation *BufferPlacementAllocs::getStartOperation(Value allocValue,32 Block *placementBlock,33 const Liveness &liveness) {34 // We have to ensure that we place the alloc before its first use in this35 // block.36 const LivenessBlockInfo &livenessInfo = *liveness.getLiveness(placementBlock);37 Operation *startOperation = livenessInfo.getStartOperation(allocValue);38 // Check whether the start operation lies in the desired placement block.39 // If not, we will use the terminator as this is the last operation in40 // this block.41 if (startOperation->getBlock() != placementBlock) {42 Operation *opInPlacementBlock =43 placementBlock->findAncestorOpInBlock(*startOperation);44 startOperation = opInPlacementBlock ? opInPlacementBlock45 : placementBlock->getTerminator();46 }47 48 return startOperation;49}50 51/// Initializes the internal list by discovering all supported allocation52/// nodes.53BufferPlacementAllocs::BufferPlacementAllocs(Operation *op) { build(op); }54 55/// Searches for and registers all supported allocation entries.56void BufferPlacementAllocs::build(Operation *op) {57 op->walk([&](MemoryEffectOpInterface opInterface) {58 // Try to find a single allocation result.59 SmallVector<MemoryEffects::EffectInstance, 2> effects;60 opInterface.getEffects(effects);61 62 SmallVector<MemoryEffects::EffectInstance, 2> allocateResultEffects;63 llvm::copy_if(64 effects, std::back_inserter(allocateResultEffects),65 [=](MemoryEffects::EffectInstance &it) {66 Value value = it.getValue();67 return isa<MemoryEffects::Allocate>(it.getEffect()) && value &&68 isa<OpResult>(value) &&69 it.getResource() !=70 SideEffects::AutomaticAllocationScopeResource::get();71 });72 // If there is one result only, we will be able to move the allocation and73 // (possibly existing) deallocation ops.74 if (allocateResultEffects.size() != 1)75 return;76 // Get allocation result.77 Value allocValue = allocateResultEffects[0].getValue();78 // Find the associated dealloc value and register the allocation entry.79 std::optional<Operation *> dealloc = memref::findDealloc(allocValue);80 // If the allocation has > 1 dealloc associated with it, skip handling it.81 if (!dealloc)82 return;83 allocs.push_back(std::make_tuple(allocValue, *dealloc));84 });85}86 87//===----------------------------------------------------------------------===//88// BufferPlacementTransformationBase89//===----------------------------------------------------------------------===//90 91/// Constructs a new transformation base using the given root operation.92BufferPlacementTransformationBase::BufferPlacementTransformationBase(93 Operation *op)94 : aliases(op), allocs(op), liveness(op) {}95 96//===----------------------------------------------------------------------===//97// BufferPlacementTransformationBase98//===----------------------------------------------------------------------===//99 100FailureOr<memref::GlobalOp>101bufferization::getGlobalFor(arith::ConstantOp constantOp,102 SymbolTableCollection &symbolTables,103 uint64_t alignment, Attribute memorySpace) {104 auto type = cast<RankedTensorType>(constantOp.getType());105 auto moduleOp = constantOp->getParentOfType<ModuleOp>();106 if (!moduleOp)107 return failure();108 109 // If we already have a global for this constant value, no need to do110 // anything else.111 for (Operation &op : moduleOp.getRegion().getOps()) {112 auto globalOp = dyn_cast<memref::GlobalOp>(&op);113 if (!globalOp)114 continue;115 if (!globalOp.getInitialValue().has_value())116 continue;117 uint64_t opAlignment = globalOp.getAlignment().value_or(0);118 Attribute initialValue = globalOp.getInitialValue().value();119 if (opAlignment == alignment && initialValue == constantOp.getValue())120 return globalOp;121 }122 123 // Create a builder without an insertion point. We will insert using the124 // symbol table to guarantee unique names.125 OpBuilder globalBuilder(moduleOp.getContext());126 SymbolTable &symbolTable = symbolTables.getSymbolTable(moduleOp);127 128 // Create a pretty name.129 SmallString<64> buf;130 llvm::raw_svector_ostream os(buf);131 interleave(type.getShape(), os, "x");132 os << "x" << type.getElementType();133 134 // Add an optional alignment to the global memref.135 IntegerAttr memrefAlignment =136 alignment > 0 ? IntegerAttr::get(globalBuilder.getI64Type(), alignment)137 : IntegerAttr();138 139 // Memref globals always have an identity layout.140 auto memrefType =141 cast<MemRefType>(getMemRefTypeWithStaticIdentityLayout(type));142 if (memorySpace)143 memrefType = MemRefType::Builder(memrefType).setMemorySpace(memorySpace);144 auto global = memref::GlobalOp::create(145 globalBuilder, constantOp.getLoc(),146 (Twine("__constant_") + os.str()).str(),147 /*sym_visibility=*/globalBuilder.getStringAttr("private"),148 /*type=*/memrefType,149 /*initial_value=*/cast<ElementsAttr>(constantOp.getValue()),150 /*constant=*/true,151 /*alignment=*/memrefAlignment);152 symbolTable.insert(global);153 // The symbol table inserts at the end of the module, but globals are a bit154 // nicer if they are at the beginning.155 global->moveBefore(&moduleOp.front());156 return global;157}158 159namespace mlir::bufferization {160void removeSymbol(Operation *op, BufferizationState &state) {161 SymbolTable &symbolTable = state.getSymbolTables().getSymbolTable(162 op->getParentWithTrait<OpTrait::SymbolTable>());163 164 symbolTable.remove(op);165}166 167void insertSymbol(Operation *op, BufferizationState &state) {168 SymbolTable &symbolTable = state.getSymbolTables().getSymbolTable(169 op->getParentWithTrait<OpTrait::SymbolTable>());170 171 symbolTable.insert(op);172}173} // namespace mlir::bufferization174