564 lines · cpp
1//===- InliningUtils.cpp ---- Misc utilities for inlining -----------------===//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 miscellaneous inlining utilities.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Transforms/InliningUtils.h"14 15#include "mlir/IR/Builders.h"16#include "mlir/IR/BuiltinOps.h"17#include "mlir/IR/IRMapping.h"18#include "mlir/IR/Operation.h"19#include "mlir/Interfaces/CallInterfaces.h"20#include "llvm/Support/Debug.h"21#include "llvm/Support/DebugLog.h"22#include "llvm/Support/raw_ostream.h"23#include <optional>24 25#define DEBUG_TYPE "inlining"26 27using namespace mlir;28 29/// Combine `callee` location with `caller` location to create a stack that30/// represents the call chain.31/// If `callee` location is a `CallSiteLoc`, indicating an existing stack of32/// locations, the `caller` location is appended to the end of it, extending33/// the chain.34/// Otherwise, a single `CallSiteLoc` is created, representing a direct call35/// from `caller` to `callee`.36static LocationAttr stackLocations(Location callee, Location caller) {37 Location lastCallee = callee;38 SmallVector<CallSiteLoc> calleeInliningStack;39 while (auto nextCallSite = dyn_cast<CallSiteLoc>(lastCallee)) {40 calleeInliningStack.push_back(nextCallSite);41 lastCallee = nextCallSite.getCaller();42 }43 44 CallSiteLoc firstCallSite = CallSiteLoc::get(lastCallee, caller);45 for (CallSiteLoc currentCallSite : reverse(calleeInliningStack))46 firstCallSite =47 CallSiteLoc::get(currentCallSite.getCallee(), firstCallSite);48 49 return firstCallSite;50}51 52/// Remap all locations reachable from the inlined blocks with CallSiteLoc53/// locations with the provided caller location.54static void55remapInlinedLocations(iterator_range<Region::iterator> inlinedBlocks,56 Location callerLoc) {57 DenseMap<Location, LocationAttr> mappedLocations;58 auto remapLoc = [&](Location loc) {59 auto [it, inserted] = mappedLocations.try_emplace(loc);60 // Only query the attribute uniquer once per callsite attribute.61 if (inserted) {62 LocationAttr newLoc = stackLocations(loc, callerLoc);63 it->getSecond() = newLoc;64 }65 return it->second;66 };67 68 AttrTypeReplacer attrReplacer;69 attrReplacer.addReplacement(70 [&](LocationAttr loc) -> std::pair<LocationAttr, WalkResult> {71 return {remapLoc(loc), WalkResult::skip()};72 });73 74 for (Block &block : inlinedBlocks) {75 for (BlockArgument &arg : block.getArguments())76 if (LocationAttr newLoc = remapLoc(arg.getLoc()))77 arg.setLoc(newLoc);78 79 for (Operation &op : block)80 attrReplacer.recursivelyReplaceElementsIn(&op, /*replaceAttrs=*/false,81 /*replaceLocs=*/true);82 }83}84 85static void remapInlinedOperands(iterator_range<Region::iterator> inlinedBlocks,86 IRMapping &mapper) {87 auto remapOperands = [&](Operation *op) {88 for (auto &operand : op->getOpOperands())89 if (auto mappedOp = mapper.lookupOrNull(operand.get()))90 operand.set(mappedOp);91 };92 for (auto &block : inlinedBlocks)93 block.walk(remapOperands);94}95 96//===----------------------------------------------------------------------===//97// InlinerInterface98//===----------------------------------------------------------------------===//99 100bool InlinerInterface::isLegalToInline(Operation *call, Operation *callable,101 bool wouldBeCloned) const {102 if (auto *handler = getInterfaceFor(call))103 return handler->isLegalToInline(call, callable, wouldBeCloned);104 return false;105}106 107bool InlinerInterface::isLegalToInline(Region *dest, Region *src,108 bool wouldBeCloned,109 IRMapping &valueMapping) const {110 if (auto *handler = getInterfaceFor(dest->getParentOp()))111 return handler->isLegalToInline(dest, src, wouldBeCloned, valueMapping);112 return false;113}114 115bool InlinerInterface::isLegalToInline(Operation *op, Region *dest,116 bool wouldBeCloned,117 IRMapping &valueMapping) const {118 if (auto *handler = getInterfaceFor(op))119 return handler->isLegalToInline(op, dest, wouldBeCloned, valueMapping);120 return false;121}122 123bool InlinerInterface::shouldAnalyzeRecursively(Operation *op) const {124 auto *handler = getInterfaceFor(op);125 return handler ? handler->shouldAnalyzeRecursively(op) : true;126}127 128/// Handle the given inlined terminator by replacing it with a new operation129/// as necessary.130void InlinerInterface::handleTerminator(Operation *op, Block *newDest) const {131 auto *handler = getInterfaceFor(op);132 assert(handler && "expected valid dialect handler");133 handler->handleTerminator(op, newDest);134}135 136/// Handle the given inlined terminator by replacing it with a new operation137/// as necessary.138void InlinerInterface::handleTerminator(Operation *op,139 ValueRange valuesToRepl) const {140 auto *handler = getInterfaceFor(op);141 assert(handler && "expected valid dialect handler");142 handler->handleTerminator(op, valuesToRepl);143}144 145/// Returns true if the inliner can assume a fast path of not creating a146/// new block, if there is only one block.147bool InlinerInterface::allowSingleBlockOptimization(148 iterator_range<Region::iterator> inlinedBlocks) const {149 if (inlinedBlocks.empty()) {150 return true;151 }152 auto *handler = getInterfaceFor(inlinedBlocks.begin()->getParentOp());153 assert(handler && "expected valid dialect handler");154 return handler->allowSingleBlockOptimization(inlinedBlocks);155}156 157Value InlinerInterface::handleArgument(OpBuilder &builder, Operation *call,158 Operation *callable, Value argument,159 DictionaryAttr argumentAttrs) const {160 auto *handler = getInterfaceFor(callable);161 assert(handler && "expected valid dialect handler");162 return handler->handleArgument(builder, call, callable, argument,163 argumentAttrs);164}165 166Value InlinerInterface::handleResult(OpBuilder &builder, Operation *call,167 Operation *callable, Value result,168 DictionaryAttr resultAttrs) const {169 auto *handler = getInterfaceFor(callable);170 assert(handler && "expected valid dialect handler");171 return handler->handleResult(builder, call, callable, result, resultAttrs);172}173 174void InlinerInterface::processInlinedCallBlocks(175 Operation *call, iterator_range<Region::iterator> inlinedBlocks) const {176 auto *handler = getInterfaceFor(call);177 assert(handler && "expected valid dialect handler");178 handler->processInlinedCallBlocks(call, inlinedBlocks);179}180 181/// Utility to check that all of the operations within 'src' can be inlined.182static bool isLegalToInline(InlinerInterface &interface, Region *src,183 Region *insertRegion, bool shouldCloneInlinedRegion,184 IRMapping &valueMapping) {185 for (auto &block : *src) {186 for (auto &op : block) {187 // UnrealizedConversionCastOp is inlineable but cannot implement the188 // inliner interface due to layering constraints.189 if (isa<UnrealizedConversionCastOp>(op))190 continue;191 192 // Check this operation.193 if (!interface.isLegalToInline(&op, insertRegion,194 shouldCloneInlinedRegion, valueMapping)) {195 LDBG() << "* Illegal to inline because of op: "196 << OpWithFlags(&op, OpPrintingFlags().skipRegions());197 return false;198 }199 // Check any nested regions.200 if (interface.shouldAnalyzeRecursively(&op) &&201 llvm::any_of(op.getRegions(), [&](Region ®ion) {202 return !isLegalToInline(interface, ®ion, insertRegion,203 shouldCloneInlinedRegion, valueMapping);204 }))205 return false;206 }207 }208 return true;209}210 211//===----------------------------------------------------------------------===//212// Inline Methods213//===----------------------------------------------------------------------===//214 215static void handleArgumentImpl(InlinerInterface &interface, OpBuilder &builder,216 CallOpInterface call,217 CallableOpInterface callable,218 IRMapping &mapper) {219 // Unpack the argument attributes if there are any.220 SmallVector<DictionaryAttr> argAttrs(221 callable.getCallableRegion()->getNumArguments(),222 builder.getDictionaryAttr({}));223 if (ArrayAttr arrayAttr = callable.getArgAttrsAttr()) {224 assert(arrayAttr.size() == argAttrs.size());225 for (auto [idx, attr] : llvm::enumerate(arrayAttr))226 argAttrs[idx] = cast<DictionaryAttr>(attr);227 }228 229 // Run the argument attribute handler for the given argument and attribute.230 for (auto [blockArg, argAttr] :231 llvm::zip(callable.getCallableRegion()->getArguments(), argAttrs)) {232 Value newArgument = interface.handleArgument(233 builder, call, callable, mapper.lookup(blockArg), argAttr);234 assert(newArgument.getType() == mapper.lookup(blockArg).getType() &&235 "expected the argument type to not change");236 237 // Update the mapping to point the new argument returned by the handler.238 mapper.map(blockArg, newArgument);239 }240}241 242static void handleResultImpl(InlinerInterface &interface, OpBuilder &builder,243 CallOpInterface call, CallableOpInterface callable,244 ValueRange results) {245 // Unpack the result attributes if there are any.246 SmallVector<DictionaryAttr> resAttrs(results.size(),247 builder.getDictionaryAttr({}));248 if (ArrayAttr arrayAttr = callable.getResAttrsAttr()) {249 assert(arrayAttr.size() == resAttrs.size());250 for (auto [idx, attr] : llvm::enumerate(arrayAttr))251 resAttrs[idx] = cast<DictionaryAttr>(attr);252 }253 254 // Run the result attribute handler for the given result and attribute.255 for (auto [result, resAttr] : llvm::zip(results, resAttrs)) {256 // Store the original result users before running the handler.257 DenseSet<Operation *> resultUsers(llvm::from_range, result.getUsers());258 259 Value newResult =260 interface.handleResult(builder, call, callable, result, resAttr);261 assert(newResult.getType() == result.getType() &&262 "expected the result type to not change");263 264 // Replace the result uses except for the ones introduce by the handler.265 result.replaceUsesWithIf(newResult, [&](OpOperand &operand) {266 return resultUsers.count(operand.getOwner());267 });268 }269}270 271static LogicalResult inlineRegionImpl(272 InlinerInterface &interface,273 function_ref<InlinerInterface::CloneCallbackSigTy> cloneCallback,274 Region *src, Block *inlineBlock, Block::iterator inlinePoint,275 IRMapping &mapper, ValueRange resultsToReplace, TypeRange regionResultTypes,276 std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion,277 CallOpInterface call = {}) {278 assert(resultsToReplace.size() == regionResultTypes.size());279 // We expect the region to have at least one block.280 if (src->empty())281 return failure();282 283 // Check that all of the region arguments have been mapped.284 auto *srcEntryBlock = &src->front();285 if (llvm::any_of(srcEntryBlock->getArguments(),286 [&](BlockArgument arg) { return !mapper.contains(arg); }))287 return failure();288 289 // Check that the operations within the source region are valid to inline.290 Region *insertRegion = inlineBlock->getParent();291 if (!interface.isLegalToInline(insertRegion, src, shouldCloneInlinedRegion,292 mapper) ||293 !isLegalToInline(interface, src, insertRegion, shouldCloneInlinedRegion,294 mapper))295 return failure();296 297 // Run the argument attribute handler before inlining the callable region.298 OpBuilder builder(inlineBlock, inlinePoint);299 auto callable = dyn_cast<CallableOpInterface>(src->getParentOp());300 if (call && callable)301 handleArgumentImpl(interface, builder, call, callable, mapper);302 303 // Clone the callee's source into the caller.304 Block *postInsertBlock = inlineBlock->splitBlock(inlinePoint);305 cloneCallback(builder, src, inlineBlock, postInsertBlock, mapper,306 shouldCloneInlinedRegion);307 308 // Get the range of newly inserted blocks.309 auto newBlocks = llvm::make_range(std::next(inlineBlock->getIterator()),310 postInsertBlock->getIterator());311 Block *firstNewBlock = &*newBlocks.begin();312 313 // Remap the locations of the inlined operations if a valid source location314 // was provided.315 if (inlineLoc && !llvm::isa<UnknownLoc>(*inlineLoc))316 remapInlinedLocations(newBlocks, *inlineLoc);317 318 // If the blocks were moved in-place, make sure to remap any necessary319 // operands.320 if (!shouldCloneInlinedRegion)321 remapInlinedOperands(newBlocks, mapper);322 323 // Process the newly inlined blocks.324 if (call)325 interface.processInlinedCallBlocks(call, newBlocks);326 interface.processInlinedBlocks(newBlocks);327 328 bool singleBlockFastPath = interface.allowSingleBlockOptimization(newBlocks);329 330 // Handle the case where only a single block was inlined.331 if (singleBlockFastPath && llvm::hasSingleElement(newBlocks)) {332 // Run the result attribute handler on the terminator operands.333 Operation *firstBlockTerminator = firstNewBlock->getTerminator();334 builder.setInsertionPoint(firstBlockTerminator);335 if (call && callable)336 handleResultImpl(interface, builder, call, callable,337 firstBlockTerminator->getOperands());338 339 // Have the interface handle the terminator of this block.340 interface.handleTerminator(firstBlockTerminator, resultsToReplace);341 firstBlockTerminator->erase();342 343 // Merge the post insert block into the cloned entry block.344 firstNewBlock->getOperations().splice(firstNewBlock->end(),345 postInsertBlock->getOperations());346 postInsertBlock->erase();347 } else {348 // Otherwise, there were multiple blocks inlined. Add arguments to the post349 // insertion block to represent the results to replace.350 for (const auto &resultToRepl : llvm::enumerate(resultsToReplace)) {351 resultToRepl.value().replaceAllUsesWith(352 postInsertBlock->addArgument(regionResultTypes[resultToRepl.index()],353 resultToRepl.value().getLoc()));354 }355 356 // Run the result attribute handler on the post insertion block arguments.357 builder.setInsertionPointToStart(postInsertBlock);358 if (call && callable)359 handleResultImpl(interface, builder, call, callable,360 postInsertBlock->getArguments());361 362 /// Handle the terminators for each of the new blocks.363 for (auto &newBlock : newBlocks)364 interface.handleTerminator(newBlock.getTerminator(), postInsertBlock);365 }366 367 // Splice the instructions of the inlined entry block into the insert block.368 inlineBlock->getOperations().splice(inlineBlock->end(),369 firstNewBlock->getOperations());370 firstNewBlock->erase();371 return success();372}373 374static LogicalResult inlineRegionImpl(375 InlinerInterface &interface,376 function_ref<InlinerInterface::CloneCallbackSigTy> cloneCallback,377 Region *src, Block *inlineBlock, Block::iterator inlinePoint,378 ValueRange inlinedOperands, ValueRange resultsToReplace,379 std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion,380 CallOpInterface call = {}) {381 // We expect the region to have at least one block.382 if (src->empty())383 return failure();384 385 auto *entryBlock = &src->front();386 if (inlinedOperands.size() != entryBlock->getNumArguments())387 return failure();388 389 // Map the provided call operands to the arguments of the region.390 IRMapping mapper;391 for (unsigned i = 0, e = inlinedOperands.size(); i != e; ++i) {392 // Verify that the types of the provided values match the function argument393 // types.394 BlockArgument regionArg = entryBlock->getArgument(i);395 if (inlinedOperands[i].getType() != regionArg.getType())396 return failure();397 mapper.map(regionArg, inlinedOperands[i]);398 }399 400 // Call into the main region inliner function.401 return inlineRegionImpl(interface, cloneCallback, src, inlineBlock,402 inlinePoint, mapper, resultsToReplace,403 resultsToReplace.getTypes(), inlineLoc,404 shouldCloneInlinedRegion, call);405}406 407LogicalResult mlir::inlineRegion(408 InlinerInterface &interface,409 function_ref<InlinerInterface::CloneCallbackSigTy> cloneCallback,410 Region *src, Operation *inlinePoint, IRMapping &mapper,411 ValueRange resultsToReplace, TypeRange regionResultTypes,412 std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion) {413 return inlineRegion(interface, cloneCallback, src, inlinePoint->getBlock(),414 ++inlinePoint->getIterator(), mapper, resultsToReplace,415 regionResultTypes, inlineLoc, shouldCloneInlinedRegion);416}417 418LogicalResult mlir::inlineRegion(419 InlinerInterface &interface,420 function_ref<InlinerInterface::CloneCallbackSigTy> cloneCallback,421 Region *src, Block *inlineBlock, Block::iterator inlinePoint,422 IRMapping &mapper, ValueRange resultsToReplace, TypeRange regionResultTypes,423 std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion) {424 return inlineRegionImpl(425 interface, cloneCallback, src, inlineBlock, inlinePoint, mapper,426 resultsToReplace, regionResultTypes, inlineLoc, shouldCloneInlinedRegion);427}428 429LogicalResult mlir::inlineRegion(430 InlinerInterface &interface,431 function_ref<InlinerInterface::CloneCallbackSigTy> cloneCallback,432 Region *src, Operation *inlinePoint, ValueRange inlinedOperands,433 ValueRange resultsToReplace, std::optional<Location> inlineLoc,434 bool shouldCloneInlinedRegion) {435 return inlineRegion(interface, cloneCallback, src, inlinePoint->getBlock(),436 ++inlinePoint->getIterator(), inlinedOperands,437 resultsToReplace, inlineLoc, shouldCloneInlinedRegion);438}439 440LogicalResult mlir::inlineRegion(441 InlinerInterface &interface,442 function_ref<InlinerInterface::CloneCallbackSigTy> cloneCallback,443 Region *src, Block *inlineBlock, Block::iterator inlinePoint,444 ValueRange inlinedOperands, ValueRange resultsToReplace,445 std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion) {446 return inlineRegionImpl(interface, cloneCallback, src, inlineBlock,447 inlinePoint, inlinedOperands, resultsToReplace,448 inlineLoc, shouldCloneInlinedRegion);449}450 451/// Utility function used to generate a cast operation from the given interface,452/// or return nullptr if a cast could not be generated.453static Value materializeConversion(const DialectInlinerInterface *interface,454 SmallVectorImpl<Operation *> &castOps,455 OpBuilder &castBuilder, Value arg, Type type,456 Location conversionLoc) {457 if (!interface)458 return nullptr;459 460 // Check to see if the interface for the call can materialize a conversion.461 Operation *castOp = interface->materializeCallConversion(castBuilder, arg,462 type, conversionLoc);463 if (!castOp)464 return nullptr;465 castOps.push_back(castOp);466 467 // Ensure that the generated cast is correct.468 assert(castOp->getNumOperands() == 1 && castOp->getOperand(0) == arg &&469 castOp->getNumResults() == 1 && *castOp->result_type_begin() == type);470 return castOp->getResult(0);471}472 473/// This function inlines a given region, 'src', of a callable operation,474/// 'callable', into the location defined by the given call operation. This475/// function returns failure if inlining is not possible, success otherwise. On476/// failure, no changes are made to the module. 'shouldCloneInlinedRegion'477/// corresponds to whether the source region should be cloned into the 'call' or478/// spliced directly.479LogicalResult mlir::inlineCall(480 InlinerInterface &interface,481 function_ref<InlinerInterface::CloneCallbackSigTy> cloneCallback,482 CallOpInterface call, CallableOpInterface callable, Region *src,483 bool shouldCloneInlinedRegion) {484 // We expect the region to have at least one block.485 if (src->empty())486 return failure();487 auto *entryBlock = &src->front();488 ArrayRef<Type> callableResultTypes = callable.getResultTypes();489 490 // Make sure that the number of arguments and results matchup between the call491 // and the region.492 SmallVector<Value, 8> callOperands(call.getArgOperands());493 SmallVector<Value, 8> callResults(call->getResults());494 if (callOperands.size() != entryBlock->getNumArguments() ||495 callResults.size() != callableResultTypes.size())496 return failure();497 498 // A set of cast operations generated to matchup the signature of the region499 // with the signature of the call.500 SmallVector<Operation *, 4> castOps;501 castOps.reserve(callOperands.size() + callResults.size());502 503 // Functor used to cleanup generated state on failure.504 auto cleanupState = [&] {505 for (auto *op : castOps) {506 op->getResult(0).replaceAllUsesWith(op->getOperand(0));507 op->erase();508 }509 return failure();510 };511 512 // Builder used for any conversion operations that need to be materialized.513 OpBuilder castBuilder(call);514 Location castLoc = call.getLoc();515 const auto *callInterface = interface.getInterfaceFor(call->getDialect());516 517 // Map the provided call operands to the arguments of the region.518 IRMapping mapper;519 for (unsigned i = 0, e = callOperands.size(); i != e; ++i) {520 BlockArgument regionArg = entryBlock->getArgument(i);521 Value operand = callOperands[i];522 523 // If the call operand doesn't match the expected region argument, try to524 // generate a cast.525 Type regionArgType = regionArg.getType();526 if (operand.getType() != regionArgType) {527 if (!(operand = materializeConversion(callInterface, castOps, castBuilder,528 operand, regionArgType, castLoc)))529 return cleanupState();530 }531 mapper.map(regionArg, operand);532 }533 534 // Ensure that the resultant values of the call match the callable.535 castBuilder.setInsertionPointAfter(call);536 for (unsigned i = 0, e = callResults.size(); i != e; ++i) {537 Value callResult = callResults[i];538 if (callResult.getType() == callableResultTypes[i])539 continue;540 541 // Generate a conversion that will produce the original type, so that the IR542 // is still valid after the original call gets replaced.543 Value castResult =544 materializeConversion(callInterface, castOps, castBuilder, callResult,545 callResult.getType(), castLoc);546 if (!castResult)547 return cleanupState();548 callResult.replaceAllUsesWith(castResult);549 castResult.getDefiningOp()->replaceUsesOfWith(castResult, callResult);550 }551 552 // Check that it is legal to inline the callable into the call.553 if (!interface.isLegalToInline(call, callable, shouldCloneInlinedRegion))554 return cleanupState();555 556 // Attempt to inline the call.557 if (failed(inlineRegionImpl(interface, cloneCallback, src, call->getBlock(),558 ++call->getIterator(), mapper, callResults,559 callableResultTypes, call.getLoc(),560 shouldCloneInlinedRegion, call)))561 return cleanupState();562 return success();563}564