1 //===- ControlFlowToLLVM.cpp - ControlFlow to LLVM dialect 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-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This file implements a pass to convert MLIR standard and builtin dialects 10 // into the LLVM IR dialect. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h" 15 16 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h" 17 #include "mlir/Conversion/LLVMCommon/Pattern.h" 18 #include "mlir/Conversion/LLVMCommon/VectorPattern.h" 19 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" 20 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h" 21 #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 22 #include "mlir/IR/BuiltinOps.h" 23 #include "mlir/IR/PatternMatch.h" 24 #include "mlir/Pass/Pass.h" 25 #include "mlir/Transforms/DialectConversion.h" 26 #include "llvm/ADT/StringRef.h" 27 #include <functional> 28 29 namespace mlir { 30 #define GEN_PASS_DEF_CONVERTCONTROLFLOWTOLLVM 31 #include "mlir/Conversion/Passes.h.inc" 32 } // namespace mlir 33 34 using namespace mlir; 35 36 #define PASS_NAME "convert-cf-to-llvm" 37 38 static std::string generateGlobalMsgSymbolName(ModuleOp moduleOp) { 39 std::string prefix = "assert_msg_"; 40 int counter = 0; 41 while (moduleOp.lookupSymbol(prefix + std::to_string(counter))) 42 ++counter; 43 return prefix + std::to_string(counter); 44 } 45 46 /// Generate IR that prints the given string to stderr. 47 static void createPrintMsg(OpBuilder &builder, Location loc, ModuleOp moduleOp, 48 StringRef msg, LLVMTypeConverter &typeConverter) { 49 auto ip = builder.saveInsertionPoint(); 50 builder.setInsertionPointToStart(moduleOp.getBody()); 51 MLIRContext *ctx = builder.getContext(); 52 53 // Create a zero-terminated byte representation and allocate global symbol. 54 SmallVector<uint8_t> elementVals; 55 elementVals.append(msg.begin(), msg.end()); 56 elementVals.push_back(0); 57 auto dataAttrType = RankedTensorType::get( 58 {static_cast<int64_t>(elementVals.size())}, builder.getI8Type()); 59 auto dataAttr = 60 DenseElementsAttr::get(dataAttrType, llvm::ArrayRef(elementVals)); 61 auto arrayTy = 62 LLVM::LLVMArrayType::get(IntegerType::get(ctx, 8), elementVals.size()); 63 std::string symbolName = generateGlobalMsgSymbolName(moduleOp); 64 auto globalOp = builder.create<LLVM::GlobalOp>( 65 loc, arrayTy, /*constant=*/true, LLVM::Linkage::Private, symbolName, 66 dataAttr); 67 68 // Emit call to `printStr` in runtime library. 69 builder.restoreInsertionPoint(ip); 70 auto msgAddr = builder.create<LLVM::AddressOfOp>( 71 loc, typeConverter.getPointerType(arrayTy), globalOp.getName()); 72 SmallVector<LLVM::GEPArg> indices(1, 0); 73 Value gep = builder.create<LLVM::GEPOp>( 74 loc, typeConverter.getPointerType(builder.getI8Type()), arrayTy, msgAddr, 75 indices); 76 Operation *printer = LLVM::lookupOrCreatePrintStrFn( 77 moduleOp, typeConverter.useOpaquePointers()); 78 builder.create<LLVM::CallOp>(loc, TypeRange(), SymbolRefAttr::get(printer), 79 gep); 80 } 81 82 namespace { 83 /// Lower `cf.assert`. The default lowering calls the `abort` function if the 84 /// assertion is violated and has no effect otherwise. The failure message is 85 /// ignored by the default lowering but should be propagated by any custom 86 /// lowering. 87 struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> { 88 explicit AssertOpLowering(LLVMTypeConverter &typeConverter, 89 bool abortOnFailedAssert = true) 90 : ConvertOpToLLVMPattern<cf::AssertOp>(typeConverter, /*benefit=*/1), 91 abortOnFailedAssert(abortOnFailedAssert) {} 92 93 LogicalResult 94 matchAndRewrite(cf::AssertOp op, OpAdaptor adaptor, 95 ConversionPatternRewriter &rewriter) const override { 96 auto loc = op.getLoc(); 97 auto module = op->getParentOfType<ModuleOp>(); 98 99 // Split block at `assert` operation. 100 Block *opBlock = rewriter.getInsertionBlock(); 101 auto opPosition = rewriter.getInsertionPoint(); 102 Block *continuationBlock = rewriter.splitBlock(opBlock, opPosition); 103 104 // Failed block: Generate IR to print the message and call `abort`. 105 Block *failureBlock = rewriter.createBlock(opBlock->getParent()); 106 createPrintMsg(rewriter, loc, module, op.getMsg(), *getTypeConverter()); 107 if (abortOnFailedAssert) { 108 // Insert the `abort` declaration if necessary. 109 auto abortFunc = module.lookupSymbol<LLVM::LLVMFuncOp>("abort"); 110 if (!abortFunc) { 111 OpBuilder::InsertionGuard guard(rewriter); 112 rewriter.setInsertionPointToStart(module.getBody()); 113 auto abortFuncTy = LLVM::LLVMFunctionType::get(getVoidType(), {}); 114 abortFunc = rewriter.create<LLVM::LLVMFuncOp>(rewriter.getUnknownLoc(), 115 "abort", abortFuncTy); 116 } 117 rewriter.create<LLVM::CallOp>(loc, abortFunc, std::nullopt); 118 rewriter.create<LLVM::UnreachableOp>(loc); 119 } else { 120 rewriter.create<LLVM::BrOp>(loc, ValueRange(), continuationBlock); 121 } 122 123 // Generate assertion test. 124 rewriter.setInsertionPointToEnd(opBlock); 125 rewriter.replaceOpWithNewOp<LLVM::CondBrOp>( 126 op, adaptor.getArg(), continuationBlock, failureBlock); 127 128 return success(); 129 } 130 131 private: 132 /// If set to `false`, messages are printed but program execution continues. 133 /// This is useful for testing asserts. 134 bool abortOnFailedAssert = true; 135 }; 136 137 /// The cf->LLVM lowerings for branching ops require that the blocks they jump 138 /// to first have updated types which should be handled by a pattern operating 139 /// on the parent op. 140 static LogicalResult verifyMatchingValues(ConversionPatternRewriter &rewriter, 141 ValueRange operands, 142 ValueRange blockArgs, Location loc, 143 llvm::StringRef messagePrefix) { 144 for (const auto &idxAndTypes : 145 llvm::enumerate(llvm::zip(blockArgs, operands))) { 146 int64_t i = idxAndTypes.index(); 147 Value argValue = 148 rewriter.getRemappedValue(std::get<0>(idxAndTypes.value())); 149 Type operandType = std::get<1>(idxAndTypes.value()).getType(); 150 // In the case of an invalid jump, the block argument will have been 151 // remapped to an UnrealizedConversionCast. In the case of a valid jump, 152 // there might still be a no-op conversion cast with both types being equal. 153 // Consider both of these details to see if the jump would be invalid. 154 if (auto op = dyn_cast_or_null<UnrealizedConversionCastOp>( 155 argValue.getDefiningOp())) { 156 if (op.getOperandTypes().front() != operandType) { 157 return rewriter.notifyMatchFailure(loc, [&](Diagnostic &diag) { 158 diag << messagePrefix; 159 diag << "mismatched types from operand # " << i << " "; 160 diag << operandType; 161 diag << " not compatible with destination block argument type "; 162 diag << op.getOperandTypes().front(); 163 diag << " which should be converted with the parent op."; 164 }); 165 } 166 } 167 } 168 return success(); 169 } 170 171 /// Ensure that all block types were updated and then create an LLVM::BrOp 172 struct BranchOpLowering : public ConvertOpToLLVMPattern<cf::BranchOp> { 173 using ConvertOpToLLVMPattern<cf::BranchOp>::ConvertOpToLLVMPattern; 174 175 LogicalResult 176 matchAndRewrite(cf::BranchOp op, typename cf::BranchOp::Adaptor adaptor, 177 ConversionPatternRewriter &rewriter) const override { 178 if (failed(verifyMatchingValues(rewriter, adaptor.getDestOperands(), 179 op.getSuccessor()->getArguments(), 180 op.getLoc(), 181 /*messagePrefix=*/""))) 182 return failure(); 183 184 rewriter.replaceOpWithNewOp<LLVM::BrOp>( 185 op, adaptor.getOperands(), op->getSuccessors(), op->getAttrs()); 186 return success(); 187 } 188 }; 189 190 /// Ensure that all block types were updated and then create an LLVM::CondBrOp 191 struct CondBranchOpLowering : public ConvertOpToLLVMPattern<cf::CondBranchOp> { 192 using ConvertOpToLLVMPattern<cf::CondBranchOp>::ConvertOpToLLVMPattern; 193 194 LogicalResult 195 matchAndRewrite(cf::CondBranchOp op, 196 typename cf::CondBranchOp::Adaptor adaptor, 197 ConversionPatternRewriter &rewriter) const override { 198 if (failed(verifyMatchingValues(rewriter, adaptor.getFalseDestOperands(), 199 op.getFalseDest()->getArguments(), 200 op.getLoc(), "in false case branch "))) 201 return failure(); 202 if (failed(verifyMatchingValues(rewriter, adaptor.getTrueDestOperands(), 203 op.getTrueDest()->getArguments(), 204 op.getLoc(), "in true case branch "))) 205 return failure(); 206 207 rewriter.replaceOpWithNewOp<LLVM::CondBrOp>( 208 op, adaptor.getOperands(), op->getSuccessors(), op->getAttrs()); 209 return success(); 210 } 211 }; 212 213 /// Ensure that all block types were updated and then create an LLVM::SwitchOp 214 struct SwitchOpLowering : public ConvertOpToLLVMPattern<cf::SwitchOp> { 215 using ConvertOpToLLVMPattern<cf::SwitchOp>::ConvertOpToLLVMPattern; 216 217 LogicalResult 218 matchAndRewrite(cf::SwitchOp op, typename cf::SwitchOp::Adaptor adaptor, 219 ConversionPatternRewriter &rewriter) const override { 220 if (failed(verifyMatchingValues(rewriter, adaptor.getDefaultOperands(), 221 op.getDefaultDestination()->getArguments(), 222 op.getLoc(), "in switch default case "))) 223 return failure(); 224 225 for (const auto &i : llvm::enumerate( 226 llvm::zip(adaptor.getCaseOperands(), op.getCaseDestinations()))) { 227 if (failed(verifyMatchingValues( 228 rewriter, std::get<0>(i.value()), 229 std::get<1>(i.value())->getArguments(), op.getLoc(), 230 "in switch case " + std::to_string(i.index()) + " "))) { 231 return failure(); 232 } 233 } 234 235 rewriter.replaceOpWithNewOp<LLVM::SwitchOp>( 236 op, adaptor.getOperands(), op->getSuccessors(), op->getAttrs()); 237 return success(); 238 } 239 }; 240 241 } // namespace 242 243 void mlir::cf::populateControlFlowToLLVMConversionPatterns( 244 LLVMTypeConverter &converter, RewritePatternSet &patterns) { 245 // clang-format off 246 patterns.add< 247 AssertOpLowering, 248 BranchOpLowering, 249 CondBranchOpLowering, 250 SwitchOpLowering>(converter); 251 // clang-format on 252 } 253 254 void mlir::cf::populateAssertToLLVMConversionPattern( 255 LLVMTypeConverter &converter, RewritePatternSet &patterns, 256 bool abortOnFailure) { 257 patterns.add<AssertOpLowering>(converter, abortOnFailure); 258 } 259 260 //===----------------------------------------------------------------------===// 261 // Pass Definition 262 //===----------------------------------------------------------------------===// 263 264 namespace { 265 /// A pass converting MLIR operations into the LLVM IR dialect. 266 struct ConvertControlFlowToLLVM 267 : public impl::ConvertControlFlowToLLVMBase<ConvertControlFlowToLLVM> { 268 ConvertControlFlowToLLVM() = default; 269 270 /// Run the dialect converter on the module. 271 void runOnOperation() override { 272 LLVMConversionTarget target(getContext()); 273 RewritePatternSet patterns(&getContext()); 274 275 LowerToLLVMOptions options(&getContext()); 276 if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout) 277 options.overrideIndexBitwidth(indexBitwidth); 278 options.useOpaquePointers = useOpaquePointers; 279 280 LLVMTypeConverter converter(&getContext(), options); 281 mlir::cf::populateControlFlowToLLVMConversionPatterns(converter, patterns); 282 283 if (failed(applyPartialConversion(getOperation(), target, 284 std::move(patterns)))) 285 signalPassFailure(); 286 } 287 }; 288 } // namespace 289 290 std::unique_ptr<Pass> mlir::cf::createConvertControlFlowToLLVMPass() { 291 return std::make_unique<ConvertControlFlowToLLVM>(); 292 } 293