415 lines · cpp
1//===- FuncTransformOps.cpp - Implementation of CF transform ops ----------===//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/Func/TransformOps/FuncTransformOps.h"10 11#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"12#include "mlir/Conversion/LLVMCommon/TypeConverter.h"13#include "mlir/Dialect/Func/IR/FuncOps.h"14#include "mlir/Dialect/Func/Utils/Utils.h"15#include "mlir/Dialect/LLVMIR/LLVMDialect.h"16#include "mlir/Dialect/Transform/IR/TransformDialect.h"17#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"18#include "mlir/IR/PatternMatch.h"19#include "mlir/Transforms/DialectConversion.h"20#include "llvm/ADT/STLExtras.h"21 22using namespace mlir;23 24//===----------------------------------------------------------------------===//25// Apply...ConversionPatternsOp26//===----------------------------------------------------------------------===//27 28void transform::ApplyFuncToLLVMConversionPatternsOp::populatePatterns(29 TypeConverter &typeConverter, RewritePatternSet &patterns) {30 populateFuncToLLVMConversionPatterns(31 static_cast<LLVMTypeConverter &>(typeConverter), patterns);32}33 34LogicalResult35transform::ApplyFuncToLLVMConversionPatternsOp::verifyTypeConverter(36 transform::TypeConverterBuilderOpInterface builder) {37 if (builder.getTypeConverterType() != "LLVMTypeConverter")38 return emitOpError("expected LLVMTypeConverter");39 return success();40}41 42//===----------------------------------------------------------------------===//43// CastAndCallOp44//===----------------------------------------------------------------------===//45 46DiagnosedSilenceableFailure47transform::CastAndCallOp::apply(transform::TransformRewriter &rewriter,48 transform::TransformResults &results,49 transform::TransformState &state) {50 SmallVector<Value> inputs;51 if (getInputs())52 llvm::append_range(inputs, state.getPayloadValues(getInputs()));53 54 SetVector<Value> outputs;55 if (getOutputs()) {56 outputs.insert_range(state.getPayloadValues(getOutputs()));57 58 // Verify that the set of output values to be replaced is unique.59 if (outputs.size() !=60 llvm::range_size(state.getPayloadValues(getOutputs()))) {61 return emitSilenceableFailure(getLoc())62 << "cast and call output values must be unique";63 }64 }65 66 // Get the insertion point for the call.67 auto insertionOps = state.getPayloadOps(getInsertionPoint());68 if (!llvm::hasSingleElement(insertionOps)) {69 return emitSilenceableFailure(getLoc())70 << "Only one op can be specified as an insertion point";71 }72 bool insertAfter = getInsertAfter();73 Operation *insertionPoint = *insertionOps.begin();74 75 // Check that all inputs dominate the insertion point, and the insertion76 // point dominates all users of the outputs.77 DominanceInfo dom(insertionPoint);78 for (Value output : outputs) {79 for (Operation *user : output.getUsers()) {80 // If we are inserting after the insertion point operation, the81 // insertion point operation must properly dominate the user. Otherwise82 // basic dominance is enough.83 bool doesDominate = insertAfter84 ? dom.properlyDominates(insertionPoint, user)85 : dom.dominates(insertionPoint, user);86 if (!doesDominate) {87 return emitDefiniteFailure()88 << "User " << user << " is not dominated by insertion point "89 << insertionPoint;90 }91 }92 }93 94 for (Value input : inputs) {95 // If we are inserting before the insertion point operation, the96 // input must properly dominate the insertion point operation. Otherwise97 // basic dominance is enough.98 bool doesDominate = insertAfter99 ? dom.dominates(input, insertionPoint)100 : dom.properlyDominates(input, insertionPoint);101 if (!doesDominate) {102 return emitDefiniteFailure()103 << "input " << input << " does not dominate insertion point "104 << insertionPoint;105 }106 }107 108 // Get the function to call. This can either be specified by symbol or as a109 // transform handle.110 func::FuncOp targetFunction = nullptr;111 if (getFunctionName()) {112 targetFunction = SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(113 insertionPoint, *getFunctionName());114 if (!targetFunction) {115 return emitDefiniteFailure()116 << "unresolved symbol " << *getFunctionName();117 }118 } else if (getFunction()) {119 auto payloadOps = state.getPayloadOps(getFunction());120 if (!llvm::hasSingleElement(payloadOps)) {121 return emitDefiniteFailure() << "requires a single function to call";122 }123 targetFunction = dyn_cast<func::FuncOp>(*payloadOps.begin());124 if (!targetFunction) {125 return emitDefiniteFailure() << "invalid non-function callee";126 }127 } else {128 llvm_unreachable("Invalid CastAndCall op without a function to call");129 return emitDefiniteFailure();130 }131 132 // Verify that the function argument and result lengths match the inputs and133 // outputs given to this op.134 if (targetFunction.getNumArguments() != inputs.size()) {135 return emitSilenceableFailure(targetFunction.getLoc())136 << "mismatch between number of function arguments "137 << targetFunction.getNumArguments() << " and number of inputs "138 << inputs.size();139 }140 if (targetFunction.getNumResults() != outputs.size()) {141 return emitSilenceableFailure(targetFunction.getLoc())142 << "mismatch between number of function results "143 << targetFunction->getNumResults() << " and number of outputs "144 << outputs.size();145 }146 147 // Gather all specified converters.148 mlir::TypeConverter converter;149 if (!getRegion().empty()) {150 for (Operation &op : getRegion().front()) {151 cast<transform::TypeConverterBuilderOpInterface>(&op)152 .populateTypeMaterializations(converter);153 }154 }155 156 if (insertAfter)157 rewriter.setInsertionPointAfter(insertionPoint);158 else159 rewriter.setInsertionPoint(insertionPoint);160 161 for (auto [input, type] :162 llvm::zip_equal(inputs, targetFunction.getArgumentTypes())) {163 if (input.getType() != type) {164 Value newInput = converter.materializeSourceConversion(165 rewriter, input.getLoc(), type, input);166 if (!newInput) {167 return emitDefiniteFailure() << "Failed to materialize conversion of "168 << input << " to type " << type;169 }170 input = newInput;171 }172 }173 174 auto callOp = func::CallOp::create(rewriter, insertionPoint->getLoc(),175 targetFunction, inputs);176 177 // Cast the call results back to the expected types. If any conversions fail178 // this is a definite failure as the call has been constructed at this point.179 for (auto [output, newOutput] :180 llvm::zip_equal(outputs, callOp.getResults())) {181 Value convertedOutput = newOutput;182 if (output.getType() != newOutput.getType()) {183 convertedOutput = converter.materializeTargetConversion(184 rewriter, output.getLoc(), output.getType(), newOutput);185 if (!convertedOutput) {186 return emitDefiniteFailure()187 << "Failed to materialize conversion of " << newOutput188 << " to type " << output.getType();189 }190 }191 rewriter.replaceAllUsesExcept(output, convertedOutput, callOp);192 }193 results.set(cast<OpResult>(getResult()), {callOp});194 return DiagnosedSilenceableFailure::success();195}196 197LogicalResult transform::CastAndCallOp::verify() {198 if (!getRegion().empty()) {199 for (Operation &op : getRegion().front()) {200 if (!isa<transform::TypeConverterBuilderOpInterface>(&op)) {201 InFlightDiagnostic diag = emitOpError()202 << "expected children ops to implement "203 "TypeConverterBuilderOpInterface";204 diag.attachNote(op.getLoc()) << "op without interface";205 return diag;206 }207 }208 }209 if (!getFunction() && !getFunctionName()) {210 return emitOpError() << "expected a function handle or name to call";211 }212 if (getFunction() && getFunctionName()) {213 return emitOpError() << "function handle and name are mutually exclusive";214 }215 return success();216}217 218void transform::CastAndCallOp::getEffects(219 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {220 transform::onlyReadsHandle(getInsertionPointMutable(), effects);221 if (getInputs())222 transform::onlyReadsHandle(getInputsMutable(), effects);223 if (getOutputs())224 transform::onlyReadsHandle(getOutputsMutable(), effects);225 if (getFunction())226 transform::onlyReadsHandle(getFunctionMutable(), effects);227 transform::producesHandle(getOperation()->getOpResults(), effects);228 transform::modifiesPayload(effects);229}230 231//===----------------------------------------------------------------------===//232// ReplaceFuncSignatureOp233//===----------------------------------------------------------------------===//234 235DiagnosedSilenceableFailure236transform::ReplaceFuncSignatureOp::apply(transform::TransformRewriter &rewriter,237 transform::TransformResults &results,238 transform::TransformState &state) {239 auto payloadOps = state.getPayloadOps(getModule());240 if (!llvm::hasSingleElement(payloadOps))241 return emitDefiniteFailure() << "requires a single module to operate on";242 243 auto targetModuleOp = dyn_cast<ModuleOp>(*payloadOps.begin());244 if (!targetModuleOp)245 return emitSilenceableFailure(getLoc())246 << "target is expected to be module operation";247 248 func::FuncOp funcOp =249 targetModuleOp.lookupSymbol<func::FuncOp>(getFunctionName());250 if (!funcOp)251 return emitSilenceableFailure(getLoc())252 << "function with name '" << getFunctionName() << "' not found";253 254 unsigned numArgs = funcOp.getNumArguments();255 unsigned numResults = funcOp.getNumResults();256 // Check that the number of arguments and results matches the257 // interchange sizes.258 if (numArgs != getArgsInterchange().size())259 return emitSilenceableFailure(getLoc())260 << "function with name '" << getFunctionName() << "' has " << numArgs261 << " arguments, but " << getArgsInterchange().size()262 << " args interchange were given";263 264 if (numResults != getResultsInterchange().size())265 return emitSilenceableFailure(getLoc())266 << "function with name '" << getFunctionName() << "' has "267 << numResults << " results, but " << getResultsInterchange().size()268 << " results interchange were given";269 270 // Check that the args and results interchanges are unique.271 SetVector<unsigned> argsInterchange, resultsInterchange;272 argsInterchange.insert_range(getArgsInterchange());273 resultsInterchange.insert_range(getResultsInterchange());274 if (argsInterchange.size() != getArgsInterchange().size())275 return emitSilenceableFailure(getLoc())276 << "args interchange must be unique";277 278 if (resultsInterchange.size() != getResultsInterchange().size())279 return emitSilenceableFailure(getLoc())280 << "results interchange must be unique";281 282 // Check that the args and results interchange indices are in bounds.283 for (unsigned index : argsInterchange) {284 if (index >= numArgs) {285 return emitSilenceableFailure(getLoc())286 << "args interchange index " << index287 << " is out of bounds for function with name '"288 << getFunctionName() << "' with " << numArgs << " arguments";289 }290 }291 for (unsigned index : resultsInterchange) {292 if (index >= numResults) {293 return emitSilenceableFailure(getLoc())294 << "results interchange index " << index295 << " is out of bounds for function with name '"296 << getFunctionName() << "' with " << numResults << " results";297 }298 }299 300 llvm::SmallVector<int> oldArgToNewArg(argsInterchange.size());301 for (auto [newArgIdx, oldArgIdx] : llvm::enumerate(argsInterchange))302 oldArgToNewArg[oldArgIdx] = newArgIdx;303 304 llvm::SmallVector<int> oldResToNewRes(resultsInterchange.size());305 for (auto [newResIdx, oldResIdx] : llvm::enumerate(resultsInterchange))306 oldResToNewRes[oldResIdx] = newResIdx;307 308 FailureOr<func::FuncOp> newFuncOpOrFailure = func::replaceFuncWithNewMapping(309 rewriter, funcOp, oldArgToNewArg, oldResToNewRes);310 if (failed(newFuncOpOrFailure))311 return emitSilenceableFailure(getLoc())312 << "failed to replace function signature '" << getFunctionName()313 << "' with new order";314 315 if (getAdjustFuncCalls()) {316 SmallVector<func::CallOp> callOps;317 targetModuleOp.walk([&](func::CallOp callOp) {318 if (callOp.getCallee() == getFunctionName().getRootReference().getValue())319 callOps.push_back(callOp);320 });321 322 for (func::CallOp callOp : callOps)323 func::replaceCallOpWithNewMapping(rewriter, callOp, oldArgToNewArg,324 oldResToNewRes);325 }326 327 results.set(cast<OpResult>(getTransformedModule()), {targetModuleOp});328 results.set(cast<OpResult>(getTransformedFunction()), {*newFuncOpOrFailure});329 330 return DiagnosedSilenceableFailure::success();331}332 333void transform::ReplaceFuncSignatureOp::getEffects(334 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {335 transform::consumesHandle(getModuleMutable(), effects);336 transform::producesHandle(getOperation()->getOpResults(), effects);337 transform::modifiesPayload(effects);338}339 340//===----------------------------------------------------------------------===//341// DeduplicateFuncArgsOp342//===----------------------------------------------------------------------===//343 344DiagnosedSilenceableFailure345transform::DeduplicateFuncArgsOp::apply(transform::TransformRewriter &rewriter,346 transform::TransformResults &results,347 transform::TransformState &state) {348 auto payloadOps = state.getPayloadOps(getModule());349 if (!llvm::hasSingleElement(payloadOps))350 return emitDefiniteFailure() << "requires a single module to operate on";351 352 auto targetModuleOp = dyn_cast<ModuleOp>(*payloadOps.begin());353 if (!targetModuleOp)354 return emitSilenceableFailure(getLoc())355 << "target is expected to be module operation";356 357 func::FuncOp funcOp =358 targetModuleOp.lookupSymbol<func::FuncOp>(getFunctionName());359 if (!funcOp)360 return emitSilenceableFailure(getLoc())361 << "function with name '" << getFunctionName() << "' is not found";362 363 auto transformationResult =364 func::deduplicateArgsOfFuncOp(rewriter, funcOp, targetModuleOp);365 if (failed(transformationResult))366 return emitSilenceableFailure(getLoc())367 << "failed to deduplicate function arguments of function "368 << funcOp.getName();369 370 auto [newFuncOp, newCallOp] = *transformationResult;371 372 results.set(cast<OpResult>(getTransformedModule()), {targetModuleOp});373 results.set(cast<OpResult>(getTransformedFunction()), {newFuncOp});374 375 return DiagnosedSilenceableFailure::success();376}377 378void transform::DeduplicateFuncArgsOp::getEffects(379 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {380 transform::consumesHandle(getModuleMutable(), effects);381 transform::producesHandle(getOperation()->getOpResults(), effects);382 transform::modifiesPayload(effects);383}384 385//===----------------------------------------------------------------------===//386// Transform op registration387//===----------------------------------------------------------------------===//388 389namespace {390class FuncTransformDialectExtension391 : public transform::TransformDialectExtension<392 FuncTransformDialectExtension> {393public:394 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FuncTransformDialectExtension)395 396 using Base::Base;397 398 void init() {399 declareGeneratedDialect<LLVM::LLVMDialect>();400 401 registerTransformOps<402#define GET_OP_LIST403#include "mlir/Dialect/Func/TransformOps/FuncTransformOps.cpp.inc"404 >();405 }406};407} // namespace408 409#define GET_OP_CLASSES410#include "mlir/Dialect/Func/TransformOps/FuncTransformOps.cpp.inc"411 412void mlir::func::registerTransformDialectExtension(DialectRegistry ®istry) {413 registry.addExtensions<FuncTransformDialectExtension>();414}415