345 lines · cpp
1//===- BufferResultsToOutParams.cpp - Calling convention conversion -------===//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#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"10#include "mlir/Dialect/Bufferization/Transforms/Passes.h"11 12#include "mlir/Dialect/Func/IR/FuncOps.h"13#include "mlir/Dialect/MemRef/IR/MemRef.h"14#include "mlir/IR/Operation.h"15 16namespace mlir {17namespace bufferization {18#define GEN_PASS_DEF_BUFFERRESULTSTOOUTPARAMSPASS19#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"20} // namespace bufferization21} // namespace mlir22 23using namespace mlir;24using AllocationFn = bufferization::BufferResultsToOutParamsOpts::AllocationFn;25using MemCpyFn = bufferization::BufferResultsToOutParamsOpts::MemCpyFn;26using AllocDynamicSizesMap =27 llvm::DenseMap<func::FuncOp, SmallVector<SmallVector<Value>>>;28 29/// Return `true` if the given MemRef type has a fully dynamic layout.30static bool hasFullyDynamicLayoutMap(MemRefType type) {31 int64_t offset;32 SmallVector<int64_t, 4> strides;33 if (failed(type.getStridesAndOffset(strides, offset)))34 return false;35 if (!llvm::all_of(strides, ShapedType::isDynamic))36 return false;37 if (ShapedType::isStatic(offset))38 return false;39 return true;40}41 42/// Return `true` if the given MemRef type has a static identity layout (i.e.,43/// no layout).44static bool hasStaticIdentityLayout(MemRefType type) {45 return type.getLayout().isIdentity();46}47 48/// Return the dynamic shapes of the `memref` based on the defining op. If the49/// complete dynamic shape fails to be captured, return an empty value.50/// Currently, only function block arguments are supported for capturing.51static SmallVector<Value> getDynamicSize(Value memref, func::FuncOp funcOp) {52 Operation *defOp = memref.getDefiningOp();53 if (!defOp)54 return {};55 auto operands = defOp->getOperands();56 SmallVector<Value> dynamicSizes;57 for (Value size : operands) {58 if (!isa<IndexType>(size.getType()))59 continue;60 61 BlockArgument sizeSrc = dyn_cast<BlockArgument>(size);62 if (!sizeSrc)63 return {};64 auto arguments = funcOp.getArguments();65 auto iter = llvm::find(arguments, sizeSrc);66 if (iter == arguments.end())67 return {};68 dynamicSizes.push_back(*iter);69 }70 return dynamicSizes;71}72 73/// Returns the dynamic sizes at the callee, through the call relationship74/// between the caller and callee.75static SmallVector<Value> mapDynamicSizeAtCaller(func::CallOp call,76 func::FuncOp callee,77 ValueRange dynamicSizes) {78 SmallVector<Value> mappedDynamicSizes;79 for (Value size : dynamicSizes) {80 for (auto [src, dst] :81 llvm::zip_first(call.getOperands(), callee.getArguments())) {82 if (size != dst)83 continue;84 mappedDynamicSizes.push_back(src);85 }86 }87 assert(mappedDynamicSizes.size() == dynamicSizes.size() &&88 "could not find all dynamic sizes");89 return mappedDynamicSizes;90}91 92// Updates the func op and entry block.93//94// Any args appended to the entry block are added to `appendedEntryArgs`.95// If `addResultAttribute` is true, adds the unit attribute `bufferize.result`96// to each newly created function argument.97static LogicalResult98updateFuncOp(func::FuncOp func,99 SmallVectorImpl<BlockArgument> &appendedEntryArgs,100 bool addResultAttribute) {101 auto functionType = func.getFunctionType();102 103 // Collect information about the results will become appended arguments.104 SmallVector<Type, 6> erasedResultTypes;105 BitVector erasedResultIndices(functionType.getNumResults());106 for (const auto &resultType : llvm::enumerate(functionType.getResults())) {107 if (auto memrefType = dyn_cast<MemRefType>(resultType.value())) {108 if (!hasStaticIdentityLayout(memrefType) &&109 !hasFullyDynamicLayoutMap(memrefType)) {110 // Only buffers with static identity layout can be allocated. These can111 // be casted to memrefs with fully dynamic layout map. Other layout maps112 // are not supported.113 return func->emitError()114 << "cannot create out param for result with unsupported layout";115 }116 erasedResultIndices.set(resultType.index());117 erasedResultTypes.push_back(memrefType);118 }119 }120 121 // Add the new arguments to the function type.122 auto newArgTypes = llvm::to_vector<6>(123 llvm::concat<const Type>(functionType.getInputs(), erasedResultTypes));124 auto newFunctionType = FunctionType::get(func.getContext(), newArgTypes,125 functionType.getResults());126 func.setType(newFunctionType);127 128 // Transfer the result attributes to arg attributes.129 auto erasedIndicesIt = erasedResultIndices.set_bits_begin();130 for (int i = 0, e = erasedResultTypes.size(); i < e; ++i, ++erasedIndicesIt) {131 func.setArgAttrs(functionType.getNumInputs() + i,132 func.getResultAttrs(*erasedIndicesIt));133 if (addResultAttribute)134 func.setArgAttr(functionType.getNumInputs() + i,135 StringAttr::get(func.getContext(), "bufferize.result"),136 UnitAttr::get(func.getContext()));137 }138 139 // Erase the results.140 if (failed(func.eraseResults(erasedResultIndices)))141 return failure();142 143 // Add the new arguments to the entry block if the function is not external.144 if (func.isExternal())145 return success();146 Location loc = func.getLoc();147 for (Type type : erasedResultTypes)148 appendedEntryArgs.push_back(func.front().addArgument(type, loc));149 150 return success();151}152 153// Updates all ReturnOps in the scope of the given func::FuncOp by either154// keeping them as return values or copying the associated buffer contents into155// the given out-params.156static LogicalResult157updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,158 AllocDynamicSizesMap &map,159 const bufferization::BufferResultsToOutParamsOpts &options) {160 auto res = func.walk([&](func::ReturnOp op) {161 SmallVector<Value, 6> copyIntoOutParams;162 SmallVector<Value, 6> keepAsReturnOperands;163 for (Value operand : op.getOperands()) {164 if (isa<MemRefType>(operand.getType()))165 copyIntoOutParams.push_back(operand);166 else167 keepAsReturnOperands.push_back(operand);168 }169 OpBuilder builder(op);170 SmallVector<SmallVector<Value>> dynamicSizes;171 for (auto [orig, arg] : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {172 bool hoistStaticAllocs =173 options.hoistStaticAllocs &&174 cast<MemRefType>(orig.getType()).hasStaticShape();175 bool hoistDynamicAllocs =176 options.hoistDynamicAllocs &&177 !cast<MemRefType>(orig.getType()).hasStaticShape();178 if ((hoistStaticAllocs || hoistDynamicAllocs) &&179 isa_and_nonnull<bufferization::AllocationOpInterface>(180 orig.getDefiningOp())) {181 orig.replaceAllUsesWith(arg);182 if (hoistDynamicAllocs) {183 SmallVector<Value> dynamicSize = getDynamicSize(orig, func);184 dynamicSizes.push_back(dynamicSize);185 }186 orig.getDefiningOp()->erase();187 } else {188 if (failed(options.memCpyFn(builder, op.getLoc(), orig, arg)))189 return WalkResult::interrupt();190 }191 }192 func::ReturnOp::create(builder, op.getLoc(), keepAsReturnOperands);193 op.erase();194 auto dynamicSizePair =195 std::pair<func::FuncOp, SmallVector<SmallVector<Value>>>(func,196 dynamicSizes);197 map.insert(dynamicSizePair);198 return WalkResult::advance();199 });200 return failure(res.wasInterrupted());201}202 203// Updates all CallOps in the scope of the given ModuleOp by allocating204// temporary buffers for newly introduced out params.205static LogicalResult206updateCalls(ModuleOp module, const AllocDynamicSizesMap &map,207 const bufferization::BufferResultsToOutParamsOpts &options) {208 bool didFail = false;209 SymbolTable symtab(module);210 module.walk([&](func::CallOp op) {211 auto callee = symtab.lookup<func::FuncOp>(op.getCallee());212 if (!callee) {213 op.emitError() << "cannot find callee '" << op.getCallee() << "' in "214 << "symbol table";215 didFail = true;216 return;217 }218 if (!options.filterFn(&callee))219 return;220 if (callee.isPublic() && !options.modifyPublicFunctions)221 return;222 if (callee.isExternal())223 return;224 225 SmallVector<Value, 6> replaceWithNewCallResults;226 SmallVector<Value, 6> replaceWithOutParams;227 for (OpResult result : op.getResults()) {228 if (isa<MemRefType>(result.getType()))229 replaceWithOutParams.push_back(result);230 else231 replaceWithNewCallResults.push_back(result);232 }233 SmallVector<Value, 6> outParams;234 OpBuilder builder(op);235 SmallVector<SmallVector<Value>> dynamicSizes = map.lookup(callee);236 size_t dynamicSizesIndex = 0;237 for (Value memref : replaceWithOutParams) {238 SmallVector<Value> dynamicSize = dynamicSizes.size() > dynamicSizesIndex239 ? dynamicSizes[dynamicSizesIndex]240 : SmallVector<Value>();241 bool memrefStaticShape =242 cast<MemRefType>(memref.getType()).hasStaticShape();243 if (!memrefStaticShape && dynamicSize.empty()) {244 op.emitError()245 << "cannot create out param for dynamically shaped result";246 didFail = true;247 return;248 }249 auto memrefType = cast<MemRefType>(memref.getType());250 auto allocType =251 MemRefType::get(memrefType.getShape(), memrefType.getElementType(),252 AffineMap(), memrefType.getMemorySpace());253 254 if (memrefStaticShape) {255 dynamicSize = {};256 } else {257 ++dynamicSizesIndex;258 dynamicSize = mapDynamicSizeAtCaller(op, callee, dynamicSize);259 }260 auto maybeOutParam =261 options.allocationFn(builder, op.getLoc(), allocType, dynamicSize);262 if (failed(maybeOutParam)) {263 op.emitError() << "failed to create allocation op";264 didFail = true;265 return;266 }267 Value outParam = maybeOutParam.value();268 if (!hasStaticIdentityLayout(memrefType)) {269 // Layout maps are already checked in `updateFuncOp`.270 assert(hasFullyDynamicLayoutMap(memrefType) &&271 "layout map not supported");272 outParam =273 memref::CastOp::create(builder, op.getLoc(), memrefType, outParam);274 }275 memref.replaceAllUsesWith(outParam);276 outParams.push_back(outParam);277 }278 279 auto newOperands = llvm::to_vector<6>(op.getOperands());280 newOperands.append(outParams.begin(), outParams.end());281 auto newResultTypes = llvm::to_vector<6>(llvm::map_range(282 replaceWithNewCallResults, [](Value v) { return v.getType(); }));283 auto newCall = func::CallOp::create(284 builder, op.getLoc(), op.getCalleeAttr(), newResultTypes, newOperands);285 for (auto t : llvm::zip(replaceWithNewCallResults, newCall.getResults()))286 std::get<0>(t).replaceAllUsesWith(std::get<1>(t));287 op.erase();288 });289 290 return failure(didFail);291}292 293LogicalResult mlir::bufferization::promoteBufferResultsToOutParams(294 ModuleOp module,295 const bufferization::BufferResultsToOutParamsOpts &options) {296 // It maps the shape source of the dynamic shape memref returned by each297 // function.298 AllocDynamicSizesMap map;299 for (auto func : module.getOps<func::FuncOp>()) {300 if (func.isPublic() && !options.modifyPublicFunctions)301 continue;302 if (func.isExternal())303 continue;304 if (!options.filterFn(&func))305 continue;306 SmallVector<BlockArgument, 6> appendedEntryArgs;307 if (failed(308 updateFuncOp(func, appendedEntryArgs, options.addResultAttribute)))309 return failure();310 if (failed(updateReturnOps(func, appendedEntryArgs, map, options))) {311 return failure();312 }313 }314 if (failed(updateCalls(module, map, options)))315 return failure();316 return success();317}318 319namespace {320struct BufferResultsToOutParamsPass321 : bufferization::impl::BufferResultsToOutParamsPassBase<322 BufferResultsToOutParamsPass> {323 using Base::Base;324 325 void runOnOperation() override {326 // Convert from pass options in tablegen to BufferResultsToOutParamsOpts.327 if (addResultAttribute)328 options.addResultAttribute = true;329 if (hoistStaticAllocs)330 options.hoistStaticAllocs = true;331 if (hoistDynamicAllocs)332 options.hoistDynamicAllocs = true;333 if (modifyPublicFunctions)334 options.modifyPublicFunctions = true;335 336 if (failed(bufferization::promoteBufferResultsToOutParams(getOperation(),337 options)))338 return signalPassFailure();339 }340 341private:342 bufferization::BufferResultsToOutParamsOpts options;343};344} // namespace345