Index: mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td =================================================================== --- mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td +++ mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td @@ -462,6 +462,13 @@ let successors = (successor AnySuccessor:$trueDest, AnySuccessor:$falseDest); let assemblyFormat = "$condition `,` successors attr-dict"; } + +def LLVM_IndirectBrOp : LLVM_TerminatorOp<"indirectbr", []> { + let arguments = (ins LLVM_Type:$address); + let successors = (successor VariadicSuccessor:$destinations); + let printer = [{ printIndirectBrOp(p, *this); }]; + let parser = [{ return parseIndirectBrOp(parser, result); }]; +} def LLVM_ReturnOp : LLVM_TerminatorOp<"return", []>, Arguments<(ins Variadic:$args)> { string llvmBuilder = [{ Index: mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp =================================================================== --- mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp +++ mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp @@ -780,6 +780,50 @@ return success(); } +///===---------------------------------------------------------------------===// +/// Printing/parsing for LLVM::IndirectBrOp. +///===---------------------------------------------------------------------===// + +static void printIndirectBrOp(OpAsmPrinter &p, IndirectBrOp op) { + p << op.getOperationName() << ' ' << op.address() << " ["; + interleaveComma(llvm::seq(0, op.getNumSuccessors()), p, + [&](int i) { p.printSuccessorAndUseList(op, i); }); + p << "] : " << op.address().getType(); +} + +/// ::= `llvm.indirect_br` ssa-use `,` +/// `[` bb-id (`[` ssa-use-and-type-list `]`)? +/// (`,` bb-id (`[` ssa-use-and-type-list `]`)?)* `]` +/// attribute-dict? `:` type +static ParseResult parseIndirectBrOp(OpAsmParser &parser, + OperationState &result) { + OpAsmParser::OperandType address; + Type ty; + + if (parser.parseOperand(address) || parser.parseLSquare()) + return failure(); + + do { + Block *dest; + SmallVector oper; + if (parser.parseSuccessorAndUseList(dest, oper)) + return failure(); + result.addSuccessor(dest, oper); + } while (succeeded(parser.parseOptionalComma())); + + // Temporary addressList to resolve address and then insert at the beginning + // of operand list. + SmallVector addressList; + if (parser.parseRSquare() || + parser.parseOptionalAttrDict(result.attributes) || + parser.parseColonType(ty) || + parser.resolveOperand(address, ty, addressList)) + return failure(); + result.operands.insert(result.operands.begin(), addressList[0]); + + return success(); +} + //===----------------------------------------------------------------------===// // Printing/parsing for LLVM::ReturnOp. //===----------------------------------------------------------------------===// Index: mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp =================================================================== --- mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp +++ mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp @@ -645,6 +645,28 @@ b.createOperation(state); return success(); } + case llvm::Instruction::IndirectBr: { + auto *indBr = cast(inst); + + Value address = processValue(indBr->getAddress()); + if (!address) + return failure(); + + SmallVector destinations; + SmallVector destinationOperands; + destinations.reserve(indBr->getNumDestinations()); + destinationOperands.reserve(indBr->getNumOperands()); + + for (llvm::BasicBlock *succ : indBr->successors()) { + SmallVector blockArguments; + if (failed(processBranchArgs(indBr, succ, blockArguments))) + return failure(); + destinations.push_back(blocks[succ]); + destinationOperands.push_back(blockArguments); + } + b.create(loc, address, destinations, destinationOperands); + return success(); + } case llvm::Instruction::PHI: { LLVMType type = processType(inst->getType()); if (!type) Index: mlir/lib/Target/LLVMIR/ModuleTranslation.cpp =================================================================== --- mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -361,6 +361,14 @@ blockMapping[condbrOp.getSuccessor(1)]); return success(); } + if (auto indBrOp = dyn_cast(opInst)) { + unsigned nSucc = indBrOp.getNumSuccessors(); + llvm::IndirectBrInst *ibi = + builder.CreateIndirectBr(valueMapping.lookup(indBrOp.address()), nSucc); + for (unsigned idx = 0; idx < nSucc; idx++) + ibi->addDestination(blockMapping[indBrOp.getSuccessor(idx)]); + return success(); + } // Emit addressof. We need to look up the global value referenced by the // operation and store it in the MLIR-to-LLVM value mapping. This does not Index: mlir/test/Dialect/LLVMIR/roundtrip.mlir =================================================================== --- mlir/test/Dialect/LLVMIR/roundtrip.mlir +++ mlir/test/Dialect/LLVMIR/roundtrip.mlir @@ -282,3 +282,36 @@ ^bb4: llvm.return %0 : !llvm.i32 } + +// CHECK-LABEL: @indirectBr +llvm.func @indirectBr(%arg0: !llvm<"i8*">, %arg1: !llvm<"i32*">) { + + %0 = llvm.mlir.constant(2 : i32) : !llvm.i32 + %1 = llvm.mlir.constant(1 : i32) : !llvm.i32 + %2 = llvm.mlir.constant(0 : i32) : !llvm.i32 + // CHECK: %[[c1:.*]] = constant 0 : i32 + %3 = constant 0 : i32 + // CHECK: %[[c2:.*]] = constant 1 : i32 + %4 = constant 1 : i32 + // CHECK: llvm.indirectbr %arg0 [^[[bb1:.*]], ^[[bb2:.*]]] : !llvm<"i8*"> + llvm.indirectbr %arg0 [^bb1, ^bb2] : !llvm<"i8*"> + +// CHECK: ^[[bb1]]: +^bb1: + llvm.store %2, %arg1 : !llvm<"i32*"> + llvm.br ^bb4 + +// CHECK: ^[[bb2]]: +^bb2: + llvm.store %1, %arg1 : !llvm<"i32*"> + llvm.br ^bb4 + +// CHECK:^[[bb3:.*]](%{{[0-9]+}}: i32, %{{[0-9]}}: i32) +^bb3(%x:i32, %t:i32): + llvm.store %0, %arg1 : !llvm<"i32*"> + llvm.br ^bb4 + +^bb4: + // CHECK: llvm.indirectbr %arg0 [^bb2, ^bb2, ^bb3(%[[c1]], %[[c2]] : i32, i32)] : !llvm<"i8*"> + llvm.indirectbr %arg0 [^bb2, ^bb2, ^bb3(%3, %4 : i32, i32)] : !llvm<"i8*"> +} Index: mlir/test/Target/import.ll =================================================================== --- mlir/test/Target/import.ll +++ mlir/test/Target/import.ll @@ -297,3 +297,30 @@ ; CHECK: llvm.return %{{[0-9]+}} : !llvm.i32 ret i32 0 } + +; CHECK-LABEL: llvm.func @indirectbrTest(%arg0: !llvm<"i8*">, %arg1: !llvm<"i32*">) +define void @indirectbrTest(i8* %address, i32* %sink) #0 { + +entry: + ; CHECK: llvm.indirectbr %arg0 [^[[bb1:.+]], ^[[bb2:.+]]] : !llvm<"i8*"> + indirectbr i8* %address, [label %bb1, label %bb2] + +; CHECK: ^[[bb1]]: +bb1: + store volatile i32 0, i32* %sink + br label %latch + +; CHECK: ^[[bb2]]: +bb2: + store volatile i32 1, i32* %sink + br label %latch + +; CHECK: ^[[bb3:[^:]+]]: +bb3: + store volatile i32 2, i32* %sink + br label %latch + +latch: + ; CHECK: llvm.indirectbr %arg0 [^[[bb2]], ^[[bb3]], ^[[bb3]]] : !llvm<"i8*"> + indirectbr i8* %address, [label %bb2, label %bb3, label %bb3] +} Index: mlir/test/Target/llvmir.mlir =================================================================== --- mlir/test/Target/llvmir.mlir +++ mlir/test/Target/llvmir.mlir @@ -1171,3 +1171,30 @@ ^bb3: // pred: ^bb1 %8 = llvm.invoke @bar(%6) to ^bb2 unwind ^bb1 : (!llvm<"i8*">) -> !llvm<"i8*"> } + +// CHECK-LABEL: @indirectbr(i8* %0, i32* %1) +llvm.func @indirectbr(%arg0: !llvm<"i8*">, %arg1: !llvm<"i32*">) { + %0 = llvm.mlir.constant(2 : i32) : !llvm.i32 + %1 = llvm.mlir.constant(1 : i32) : !llvm.i32 + %2 = llvm.mlir.constant(0 : i32) : !llvm.i32 + // CHECK: indirectbr i8* %0, [label %[[bb1:[0-9]+]], label %[[bb2:[0-9]+]]] + llvm.indirectbr %arg0 [^bb1, ^bb2] : !llvm<"i8*"> + +// CHECK: [[bb1]]: +^bb1: + llvm.store %2, %arg1 : !llvm<"i32*"> + llvm.br ^bb4 + +// CHECK: [[bb2]]: +^bb2: + llvm.store %1, %arg1 : !llvm<"i32*"> + llvm.br ^bb4 + +// CHECK: [[bb3:[0-9]+]]: +^bb3: + llvm.store %0, %arg1 : !llvm<"i32*"> + llvm.br ^bb4 +^bb4: + // CHECK: indirectbr i8* %0, [label %[[bb3]], label %[[bb2]], label %[[bb2]]] + llvm.indirectbr %arg0 [^bb3, ^bb2, ^bb2] : !llvm<"i8*"> +}