diff --git a/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp b/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp --- a/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp +++ b/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp @@ -23,6 +23,7 @@ #include "mlir/Target/LLVMIR/TypeFromLLVM.h" #include "mlir/Tools/mlir-translate/Translation.h" +#include "llvm/ADT/PostOrderIterator.h" #include "llvm/ADT/StringSet.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/IR/Attributes.h" @@ -306,6 +307,22 @@ return DataLayoutSpecAttr::get(context, entries); } +/// Get a topologically sorted list of blocks for the given function. +static SetVector +getTopologicallySortedBlocks(llvm::Function *func) { + SetVector blocks; + for (llvm::BasicBlock &bb : *func) { + if (blocks.count(&bb) == 0) { + llvm::ReversePostOrderTraversal traversal(&bb); + blocks.insert(traversal.begin(), traversal.end()); + } + } + assert(blocks.size() == func->getBasicBlockList().size() && + "some blocks are not sorted"); + + return blocks; +} + // Handles importing globals and functions from an LLVM module. namespace { class Importer { @@ -327,6 +344,18 @@ return mlir; } + /// Stores the mapping between an LLVM block and its MLIR counterpart. + void mapBlock(llvm::BasicBlock *llvm, Block *mlir) { + auto result = blockMapping.try_emplace(llvm, mlir); + (void)result; + assert(result.second && "attempting to map a block that is already mapped"); + } + + /// Returns the MLIR block mapped to the given LLVM block. + Block *lookupBlock(llvm::BasicBlock *block) const { + return blockMapping.lookup(block); + } + /// Returns the remapped version of `value` or a placeholder that will be /// remapped later if the defining instruction has not yet been visited. Value processValue(llvm::Value *value); @@ -413,13 +442,10 @@ return std::prev(module.getBody()->end()); } - /// Remapped blocks, for the current function. - DenseMap blocks; - /// Mappings between original and imported values. These are function-local. + /// Function-local mapping between original and imported block. + DenseMap blockMapping; + /// Function-local mapping between original and imported values. DenseMap valueMapping; - /// Instructions that had not been defined when first encountered as a use. - /// Maps to the dummy Operation that was created in processValue(). - DenseMap unknownInstMap; /// Uniquing map of GlobalVariables. DenseMap globals; /// The stateful type translator (contains named structs). @@ -750,16 +776,7 @@ if (it != valueMapping.end()) return it->second; - // We don't expect to see instructions in dominator order. If we haven't seen - // this instruction yet, create an unknown op and remap it later. - if (isa(value)) { - Type type = convertType(value->getType()); - unknownInstMap[value] = - b.create(UnknownLoc::get(context), b.getStringAttr("llvm.unknown"), - /*operands=*/{}, type); - return unknownInstMap[value]->getResult(0); - } - + // Process constants such as immediate arguments that have no mapping. if (auto *c = dyn_cast(value)) return processConstant(c); @@ -829,7 +846,7 @@ SmallVector blockArguments; if (failed(processBranchArgs(brInst, succ, blockArguments))) return failure(); - state.addSuccessors(blocks[succ]); + state.addSuccessors(lookupBlock(succ)); state.addOperands(blockArguments); operandSegmentSizes[i + 1] = blockArguments.size(); } @@ -866,10 +883,10 @@ return failure(); caseOperandRefs[i] = caseOperands[i]; caseValues[i] = caseHandle.getCaseValue()->getSExtValue(); - caseBlocks[i] = blocks[succBB]; + caseBlocks[i] = lookupBlock(succBB); } - b.create(loc, condition, blocks[defaultBB], defaultBlockArgs, + b.create(loc, condition, lookupBlock(defaultBB), defaultBlockArgs, caseValues, caseBlocks, caseOperandRefs); return success(); } @@ -931,12 +948,12 @@ if (llvm::Function *callee = ii->getCalledFunction()) { op = b.create( loc, tys, SymbolRefAttr::get(b.getContext(), callee->getName()), ops, - blocks[ii->getNormalDest()], normalArgs, blocks[ii->getUnwindDest()], - unwindArgs); + lookupBlock(ii->getNormalDest()), normalArgs, + lookupBlock(ii->getUnwindDest()), unwindArgs); } else { ops.insert(ops.begin(), processValue(ii->getCalledOperand())); - op = b.create(loc, tys, ops, blocks[ii->getNormalDest()], - normalArgs, blocks[ii->getUnwindDest()], + op = b.create(loc, tys, ops, lookupBlock(ii->getNormalDest()), + normalArgs, lookupBlock(ii->getUnwindDest()), unwindArgs); } @@ -1001,9 +1018,8 @@ } LogicalResult Importer::processFunction(llvm::Function *f) { - blocks.clear(); + blockMapping.clear(); valueMapping.clear(); - unknownInstMap.clear(); auto functionType = convertType(f->getFunctionType()).dyn_cast(); @@ -1064,34 +1080,28 @@ return success(); // Eagerly create all blocks. - SmallVector blockList; for (llvm::BasicBlock &bb : *f) { - blockList.push_back(b.createBlock(&fop.getBody(), fop.getBody().end())); - blocks[&bb] = blockList.back(); + Block *block = b.createBlock(&fop.getBody(), fop.getBody().end()); + mapBlock(&bb, block); } - currentEntryBlock = blockList[0]; + currentEntryBlock = &fop.getFunctionBody().getBlocks().front(); // Add function arguments to the entry block. - for (const auto &kv : llvm::enumerate(f->args())) { - mapValue(&kv.value(), - blockList[0]->addArgument(functionType.getParamType(kv.index()), - fop.getLoc())); + for (const auto &it : llvm::enumerate(f->args())) { + BlockArgument blockArg = fop.getFunctionBody().addArgument( + functionType.getParamType(it.index()), fop.getLoc()); + mapValue(&it.value(), blockArg); } - for (auto bbs : llvm::zip(*f, blockList)) { - if (failed(processBasicBlock(&std::get<0>(bbs), std::get<1>(bbs)))) + // Process the blocks in topological order. The ordered traversal ensures + // operands defined in a dominating block have a valid mapping to an MLIR + // value once a block is translated. + SetVector blocks = getTopologicallySortedBlocks(f); + for (llvm::BasicBlock *bb : blocks) { + if (failed(processBasicBlock(bb, lookupBlock(bb)))) return failure(); } - // Now that all instructions are guaranteed to have been visited, ensure - // any unknown uses we encountered are remapped. - for (auto &llvmAndUnknown : unknownInstMap) { - assert(valueMapping.count(llvmAndUnknown.first)); - Value newValue = valueMapping[llvmAndUnknown.first]; - Value oldValue = llvmAndUnknown.second->getResult(0); - oldValue.replaceAllUsesWith(newValue); - llvmAndUnknown.second->erase(); - } return success(); } diff --git a/mlir/test/Target/LLVMIR/Import/basic.ll b/mlir/test/Target/LLVMIR/Import/basic.ll --- a/mlir/test/Target/LLVMIR/Import/basic.ll +++ b/mlir/test/Target/LLVMIR/Import/basic.ll @@ -176,49 +176,6 @@ } ; CHECK-DBG: } loc(#[[UNKNOWNLOC]]) -; Test that instructions that dominate can be out of sequential order. -; CHECK-LABEL: llvm.func @f2(%arg0: i64) -> i64 { -; CHECK-DAG: %[[c3:[0-9]+]] = llvm.mlir.constant(3 : i64) : i64 -define i64 @f2(i64 %a) noduplicate { -entry: -; CHECK: llvm.br ^bb2 - br label %next - -; CHECK: ^bb1: -end: -; CHECK: llvm.return %1 - ret i64 %b - -; CHECK: ^bb2: -next: -; CHECK: %1 = llvm.add %arg0, %[[c3]] : i64 - %b = add i64 %a, 3 -; CHECK: llvm.br ^bb1 - br label %end -} - -; Test arguments/phis. -; CHECK-LABEL: llvm.func @f2_phis(%arg0: i64) -> i64 { -; CHECK-DAG: %[[c3:[0-9]+]] = llvm.mlir.constant(3 : i64) : i64 -define i64 @f2_phis(i64 %a) noduplicate { -entry: -; CHECK: llvm.br ^bb2 - br label %next - -; CHECK: ^bb1(%1: i64): -end: - %c = phi i64 [ %b, %next ] -; CHECK: llvm.return %1 - ret i64 %c - -; CHECK: ^bb2: -next: -; CHECK: %2 = llvm.add %arg0, %[[c3]] : i64 - %b = add i64 %a, 3 -; CHECK: llvm.br ^bb1 - br label %end -} - ; CHECK-LABEL: llvm.func @f3() -> !llvm.ptr define i32* @f3() { ; CHECK: %[[c:[0-9]+]] = llvm.mlir.addressof @g2 : !llvm.ptr @@ -342,112 +299,6 @@ ret i32 0 } -; Switch instruction -declare void @g(i32) - -; CHECK-LABEL: llvm.func @simple_switch(%arg0: i32) { -define void @simple_switch(i32 %val) { -; CHECK: %[[C0:.+]] = llvm.mlir.constant(11 : i32) : i32 -; CHECK: %[[C1:.+]] = llvm.mlir.constant(87 : i32) : i32 -; CHECK: %[[C2:.+]] = llvm.mlir.constant(78 : i32) : i32 -; CHECK: %[[C3:.+]] = llvm.mlir.constant(94 : i32) : i32 -; CHECK: %[[C4:.+]] = llvm.mlir.constant(1 : i32) : i32 -; CHECK: llvm.switch %arg0 : i32, ^[[BB5:.+]] [ -; CHECK: 0: ^[[BB1:.+]], -; CHECK: 9: ^[[BB2:.+]], -; CHECK: 994: ^[[BB3:.+]], -; CHECK: 1154: ^[[BB4:.+]] -; CHECK: ] - switch i32 %val, label %def [ - i32 0, label %one - i32 9, label %two - i32 994, label %three - i32 1154, label %four - ] - -; CHECK: ^[[BB1]]: -; CHECK: llvm.call @g(%[[C4]]) : (i32) -> () -; CHECK: llvm.return -one: - call void @g(i32 1) - ret void -; CHECK: ^[[BB2]]: -; CHECK: llvm.call @g(%[[C3]]) : (i32) -> () -; CHECK: llvm.return -two: - call void @g(i32 94) - ret void -; CHECK: ^[[BB3]]: -; CHECK: llvm.call @g(%[[C2]]) : (i32) -> () -; CHECK: llvm.return -three: - call void @g(i32 78) - ret void -; CHECK: ^[[BB4]]: -; CHECK: llvm.call @g(%[[C1]]) : (i32) -> () -; CHECK: llvm.return -four: - call void @g(i32 87) - ret void -; CHECK: ^[[BB5]]: -; CHECK: llvm.call @g(%[[C0]]) : (i32) -> () -; CHECK: llvm.return -def: - call void @g(i32 11) - ret void -} - -; CHECK-LABEL: llvm.func @switch_args(%arg0: i32) { -define void @switch_args(i32 %val) { - ; CHECK: %[[C0:.+]] = llvm.mlir.constant(44 : i32) : i32 - ; CHECK: %[[C1:.+]] = llvm.mlir.constant(34 : i32) : i32 - ; CHECK: %[[C2:.+]] = llvm.mlir.constant(33 : i32) : i32 - %pred = icmp ult i32 %val, 87 - br i1 %pred, label %bbs, label %bb1 - -bb1: - %vx = add i32 %val, 22 - %pred2 = icmp ult i32 %val, 94 - br i1 %pred2, label %bb2, label %bb3 - -bb2: - %vx0 = add i32 %val, 23 - br label %one - -bb3: - br label %def - -; CHECK: %[[V1:.+]] = llvm.add %arg0, %[[C2]] : i32 -; CHECK: %[[V2:.+]] = llvm.add %arg0, %[[C1]] : i32 -; CHECK: %[[V3:.+]] = llvm.add %arg0, %[[C0]] : i32 -; CHECK: llvm.switch %arg0 : i32, ^[[BBD:.+]](%[[V3]] : i32) [ -; CHECK: 0: ^[[BB1:.+]](%[[V1]], %[[V2]] : i32, i32) -; CHECK: ] -bbs: - %vy = add i32 %val, 33 - %vy0 = add i32 %val, 34 - %vz = add i32 %val, 44 - switch i32 %val, label %def [ - i32 0, label %one - ] - -; CHECK: ^[[BB1]](%[[BA0:.+]]: i32, %[[BA1:.+]]: i32): -one: ; pred: bb2, bbs - %v0 = phi i32 [%vx, %bb2], [%vy, %bbs] - %v1 = phi i32 [%vx0, %bb2], [%vy0, %bbs] - ; CHECK: llvm.add %[[BA0]], %[[BA1]] : i32 - %vf = add i32 %v0, %v1 - call void @g(i32 %vf) - ret void - -; CHECK: ^[[BBD]](%[[BA2:.+]]: i32): -def: ; pred: bb3, bbs - %v2 = phi i32 [%vx, %bb3], [%vz, %bbs] - ; CHECK: llvm.call @g(%[[BA2]]) - call void @g(i32 %v2) - ret void -} - ; Varadic function definition %struct.va_list = type { i8* } diff --git a/mlir/test/Target/LLVMIR/Import/control-flow.ll b/mlir/test/Target/LLVMIR/Import/control-flow.ll new file mode 100644 --- /dev/null +++ b/mlir/test/Target/LLVMIR/Import/control-flow.ll @@ -0,0 +1,109 @@ +; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s + +; Verify the import works if the blocks are not topologically sorted. +; CHECK-LABEL: @dominance_order +; CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]] +define i64 @dominance_order(i64 %arg1) { + ; CHECK: llvm.br ^[[BB2:.+]] + br label %bb2 +bb1: + ; CHECK: ^[[BB1:[a-zA-Z0-9]+]]: + ; CHECK: llvm.return %[[VAL1:.+]] : i64 + ret i64 %1 +bb2: + ; CHECK: ^[[BB2]]: + ; CHECK: %[[VAL1]] = llvm.add %[[ARG1]] + %1 = add i64 %arg1, 3 + ; CHECK: llvm.br ^[[BB1]] + br label %bb1 +} + +; // ----- + +; CHECK-LABEL: @block_argument +; CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]] +; CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]] +define i64 @block_argument(i1 %arg1, i64 %arg2) { +entry: + ; CHECK: llvm.cond_br %[[ARG1]] + ; CHECK-SAME: ^[[BB1:.+]](%[[ARG2]] : i64) + ; CHECK-SAME: ^[[BB2:.+]] + br i1 %arg1, label %bb1, label %bb2 +bb1: + ; CHECK: ^[[BB1]](%[[BA1:.+]]: i64): + ; CHECK: llvm.return %[[BA1]] : i64 + %0 = phi i64 [ %arg2, %entry ], [ %1, %bb2 ] + ret i64 %0 +bb2: + ; CHECK: ^[[BB2]]: + ; CHECK: %[[VAL1:.+]] = llvm.add %[[ARG2]] + ; CHECK: llvm.br ^[[BB1]](%[[VAL1]] + %1 = add i64 %arg2, 3 + br label %bb1 +} + +; // ----- + +; CHECK-LABEL: @simple_switch( +; CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]] +define i32 @simple_switch(i32 %arg1) { + ; CHECK: %[[VAL1:.+]] = llvm.add + ; CHECK: %[[VAL2:.+]] = llvm.sub + ; CHECK: %[[VAL3:.+]] = llvm.mul + %1 = add i32 %arg1, 42 + %2 = sub i32 %arg1, 42 + %3 = mul i32 %arg1, 42 + ; CHECK: llvm.switch %arg0 : i32, ^[[BBD:.+]] [ + ; CHECK: 0: ^[[BB1:.+]], + ; CHECK: 9: ^[[BB2:.+]] + ; CHECK: ] + switch i32 %arg1, label %bbd [ + i32 0, label %bb1 + i32 9, label %bb2 + ] +bb1: + ; CHECK: ^[[BB1]]: + ; CHECK: llvm.return %[[VAL1]] + ret i32 %1 +bb2: + ; CHECK: ^[[BB2]]: + ; CHECK: llvm.return %[[VAL2]] + ret i32 %2 +bbd: + ; CHECK: ^[[BBD]]: + ; CHECK: llvm.return %[[VAL3]] + ret i32 %3 +} + +; // ----- + +; CHECK-LABEL: @switch_args +; CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]] +define i32 @switch_args(i32 %arg1) { +entry: + ; CHECK: %[[VAL1:.+]] = llvm.add + ; CHECK: %[[VAL2:.+]] = llvm.sub + ; CHECK: %[[VAL3:.+]] = llvm.mul + %0 = add i32 %arg1, 42 + %1 = sub i32 %arg1, 42 + %2 = mul i32 %arg1, 42 + ; CHECK: llvm.switch %arg0 : i32, ^[[BBD:.+]](%[[VAL3]] : i32) [ + ; CHECK: 0: ^[[BB1:.+]](%[[VAL1]], %[[VAL2]] : i32, i32) + ; CHECK: ] + switch i32 %arg1, label %bbd [ + i32 0, label %bb1 + ] +bb1: + ; CHECK: ^[[BB1]](%[[BA1:.+]]: i32, %[[BA2:.+]]: i32): + ; CHECK: %[[VAL1:.*]] = llvm.add %[[BA1]], %[[BA2]] : i32 + %3 = phi i32 [%0, %entry] + %4 = phi i32 [%1, %entry] + %5 = add i32 %3, %4 + ; CHECK: llvm.br ^[[BBD]](%[[VAL1]] + br label %bbd +bbd: + ; CHECK: ^[[BBD]](%[[BA3:.+]]: i32): + ; CHECK: llvm.return %[[BA3]] + %6 = phi i32 [%2, %entry], [%5, %bb1] + ret i32 %6 +}