202 lines · cpp
1//===- FuncConversions.cpp - Function conversions -------------------------===//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/Transforms/FuncConversions.h"10#include "mlir/Dialect/Func/IR/FuncOps.h"11#include "mlir/Transforms/DialectConversion.h"12 13using namespace mlir;14using namespace mlir::func;15 16/// Flatten the given value ranges into a single vector of values.17static SmallVector<Value> flattenValues(ArrayRef<ValueRange> values) {18 SmallVector<Value> result;19 for (const auto &vals : values)20 llvm::append_range(result, vals);21 return result;22}23 24namespace {25/// Converts the operand and result types of the CallOp, used together with the26/// FuncOpSignatureConversion.27struct CallOpSignatureConversion : public OpConversionPattern<CallOp> {28 using OpConversionPattern<CallOp>::OpConversionPattern;29 30 /// Hook for derived classes to implement combined matching and rewriting.31 LogicalResult32 matchAndRewrite(CallOp callOp, OneToNOpAdaptor adaptor,33 ConversionPatternRewriter &rewriter) const override {34 // Convert the original function results. Keep track of how many result35 // types an original result type is converted into.36 SmallVector<size_t> numResultsReplacments;37 SmallVector<Type, 1> convertedResults;38 size_t numFlattenedResults = 0;39 for (auto [idx, type] : llvm::enumerate(callOp.getResultTypes())) {40 if (failed(typeConverter->convertTypes(type, convertedResults)))41 return failure();42 numResultsReplacments.push_back(convertedResults.size() -43 numFlattenedResults);44 numFlattenedResults = convertedResults.size();45 }46 47 // Substitute with the new result types from the corresponding FuncType48 // conversion.49 auto newCallOp =50 CallOp::create(rewriter, callOp.getLoc(), callOp.getCallee(),51 convertedResults, flattenValues(adaptor.getOperands()));52 SmallVector<ValueRange> replacements;53 size_t offset = 0;54 for (int i = 0, e = callOp->getNumResults(); i < e; ++i) {55 replacements.push_back(56 newCallOp->getResults().slice(offset, numResultsReplacments[i]));57 offset += numResultsReplacments[i];58 }59 assert(offset == convertedResults.size() &&60 "expected that all converted results are used");61 rewriter.replaceOpWithMultiple(callOp, replacements);62 return success();63 }64};65} // namespace66 67void mlir::populateCallOpTypeConversionPattern(RewritePatternSet &patterns,68 const TypeConverter &converter,69 PatternBenefit benefit) {70 patterns.add<CallOpSignatureConversion>(converter, patterns.getContext(),71 benefit);72}73 74namespace {75/// Only needed to support partial conversion of functions where this pattern76/// ensures that the branch operation arguments matches up with the succesor77/// block arguments.78class BranchOpInterfaceTypeConversion79 : public OpInterfaceConversionPattern<BranchOpInterface> {80public:81 using OpInterfaceConversionPattern<82 BranchOpInterface>::OpInterfaceConversionPattern;83 84 BranchOpInterfaceTypeConversion(85 const TypeConverter &typeConverter, MLIRContext *ctx,86 function_ref<bool(BranchOpInterface, int)> shouldConvertBranchOperand,87 PatternBenefit benefit)88 : OpInterfaceConversionPattern(typeConverter, ctx, benefit),89 shouldConvertBranchOperand(shouldConvertBranchOperand) {}90 91 LogicalResult92 matchAndRewrite(BranchOpInterface op, ArrayRef<Value> operands,93 ConversionPatternRewriter &rewriter) const final {94 // For a branch operation, only some operands go to the target blocks, so95 // only rewrite those.96 SmallVector<Value, 4> newOperands(op->operand_begin(), op->operand_end());97 for (int succIdx = 0, succEnd = op->getBlock()->getNumSuccessors();98 succIdx < succEnd; ++succIdx) {99 OperandRange forwardedOperands =100 op.getSuccessorOperands(succIdx).getForwardedOperands();101 if (forwardedOperands.empty())102 continue;103 104 for (int idx = forwardedOperands.getBeginOperandIndex(),105 eidx = idx + forwardedOperands.size();106 idx < eidx; ++idx) {107 if (!shouldConvertBranchOperand || shouldConvertBranchOperand(op, idx))108 newOperands[idx] = operands[idx];109 }110 }111 rewriter.modifyOpInPlace(112 op, [newOperands, op]() { op->setOperands(newOperands); });113 return success();114 }115 116private:117 function_ref<bool(BranchOpInterface, int)> shouldConvertBranchOperand;118};119} // namespace120 121namespace {122/// Only needed to support partial conversion of functions where this pattern123/// ensures that the branch operation arguments matches up with the succesor124/// block arguments.125class ReturnOpTypeConversion : public OpConversionPattern<ReturnOp> {126public:127 using OpConversionPattern<ReturnOp>::OpConversionPattern;128 129 LogicalResult130 matchAndRewrite(ReturnOp op, OneToNOpAdaptor adaptor,131 ConversionPatternRewriter &rewriter) const final {132 rewriter.replaceOpWithNewOp<ReturnOp>(op,133 flattenValues(adaptor.getOperands()));134 return success();135 }136};137} // namespace138 139void mlir::populateBranchOpInterfaceTypeConversionPattern(140 RewritePatternSet &patterns, const TypeConverter &typeConverter,141 function_ref<bool(BranchOpInterface, int)> shouldConvertBranchOperand,142 PatternBenefit benefit) {143 patterns.add<BranchOpInterfaceTypeConversion>(144 typeConverter, patterns.getContext(), shouldConvertBranchOperand,145 benefit);146}147 148bool mlir::isLegalForBranchOpInterfaceTypeConversionPattern(149 Operation *op, const TypeConverter &converter) {150 // All successor operands of branch like operations must be rewritten.151 if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {152 for (int p = 0, e = op->getBlock()->getNumSuccessors(); p < e; ++p) {153 auto successorOperands = branchOp.getSuccessorOperands(p);154 if (!converter.isLegal(155 successorOperands.getForwardedOperands().getTypes()))156 return false;157 }158 return true;159 }160 161 return false;162}163 164void mlir::populateReturnOpTypeConversionPattern(165 RewritePatternSet &patterns, const TypeConverter &typeConverter,166 PatternBenefit benefit) {167 patterns.add<ReturnOpTypeConversion>(typeConverter, patterns.getContext(),168 benefit);169}170 171bool mlir::isLegalForReturnOpTypeConversionPattern(172 Operation *op, const TypeConverter &converter, bool returnOpAlwaysLegal) {173 // If this is a `return` and the user pass wants to convert/transform across174 // function boundaries, then `converter` is invoked to check whether the175 // `return` op is legal.176 if (isa<ReturnOp>(op) && !returnOpAlwaysLegal)177 return converter.isLegal(op);178 179 // ReturnLike operations have to be legalized with their parent. For180 // return this is handled, for other ops they remain as is.181 return op->hasTrait<OpTrait::ReturnLike>();182}183 184bool mlir::isNotBranchOpInterfaceOrReturnLikeOp(Operation *op) {185 // If it is not a terminator, ignore it.186 if (!op->mightHaveTrait<OpTrait::IsTerminator>())187 return true;188 189 // If it is not the last operation in the block, also ignore it. We do190 // this to handle unknown operations, as well.191 Block *block = op->getBlock();192 if (!block || &block->back() != op)193 return true;194 195 // We don't want to handle terminators in nested regions, assume they are196 // always legal.197 if (!isa_and_nonnull<FuncOp>(op->getParentOp()))198 return true;199 200 return false;201}202