Skip to content

Commit ad1bad2

Browse files
bcardosolopesbcardosolopes
committed
[MLIR][LLVM] Support for indirectbr
Now that LLVM dialect has blockaddress support, introduce import/translation for indirectbr intruction.
1 parent 58b91d1 commit ad1bad2

File tree

10 files changed

+374
-20
lines changed

10 files changed

+374
-20
lines changed

‎mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1705,6 +1705,65 @@ def LLVM_BlockTagOp : LLVM_Op<"blocktag"> {
17051705
let hasVerifier = 0;
17061706
}
17071707

1708+
//===----------------------------------------------------------------------===//
1709+
// IndirectBrOp
1710+
//===----------------------------------------------------------------------===//
1711+
1712+
def LLVM_IndirectBrOp : LLVM_TerminatorOp<"indirectbr",
1713+
[SameVariadicOperandSize, DeclareOpInterfaceMethods<BranchOpInterface>,
1714+
Pure]> {
1715+
let description = [{
1716+
Transfer control flow to address in `$addr`. A list of possible target
1717+
blocks in `$successors` can be provided and maybe used as a hint in LLVM:
1718+
1719+
```mlir
1720+
...
1721+
llvm.func @g(...
1722+
%dest = llvm.blockaddress <function = @g, tag = <id = 0>> : !llvm.ptr
1723+
llvm.indirectbr %dest : !llvm.ptr, [
1724+
^head
1725+
]
1726+
^head:
1727+
llvm.blocktag <id = 0>
1728+
llvm.return %arg0 : i32
1729+
...
1730+
```
1731+
1732+
It also supports a list of operands that can be passed to a target block:
1733+
1734+
```mlir
1735+
llvm.indirectbr %dest : !llvm.ptr, [
1736+
^head(%arg0 : i32),
1737+
^tail(%arg1 : i32)
1738+
]
1739+
^head(%r0 : i32):
1740+
llvm.return %r0 : i32
1741+
^tail(%r1 : i32):
1742+
llvm.return %r1 : i32
1743+
1744+
```
1745+
}];
1746+
let arguments = (ins LLVM_AnyPointer:$addr,
1747+
VariadicOfVariadic<AnyType, "indbr_operand_segments">:$succOperands,
1748+
DenseI32ArrayAttr:$indbr_operand_segments
1749+
);
1750+
let successors = (successor VariadicSuccessor<AnySuccessor>:$successors);
1751+
let assemblyFormat = [{
1752+
$addr `:` type($addr) `,`
1753+
custom<IndirectBrOpSucessors>(ref(type($addr)),
1754+
$successors,
1755+
$succOperands,
1756+
type($succOperands))
1757+
attr-dict
1758+
}];
1759+
1760+
let builders = [
1761+
OpBuilder<(ins "Value":$addr,
1762+
CArg<"BlockRange", "{}">:$successors,
1763+
CArg<"ArrayRef<ValueRange>", "{}">:$succOperands)>
1764+
];
1765+
}
1766+
17081767
def LLVM_ComdatSelectorOp : LLVM_Op<"comdat_selector", [Symbol]> {
17091768
let arguments = (ins
17101769
SymbolNameAttr:$sym_name,

‎mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3879,6 +3879,64 @@ LogicalResult BlockAddressOp::verify() {
38793879
/// attribute.
38803880
OpFoldResult BlockAddressOp::fold(FoldAdaptor) { return getBlockAddr(); }
38813881

3882+
//===----------------------------------------------------------------------===//
3883+
// LLVM::IndirectBrOp
3884+
//===----------------------------------------------------------------------===//
3885+
3886+
SuccessorOperands IndirectBrOp::getSuccessorOperands(unsigned index) {
3887+
assert(index < getNumSuccessors() && "invalid successor index");
3888+
return SuccessorOperands(getSuccOperandsMutable()[index]);
3889+
}
3890+
3891+
static ParseResult parseIndirectBrOpSucessors(
3892+
OpAsmParser &parser, Type &flagType,
3893+
SmallVectorImpl<Block *> &succOperandBlocks,
3894+
SmallVectorImpl<SmallVector<OpAsmParser::UnresolvedOperand>> &succOperands,
3895+
SmallVectorImpl<SmallVector<Type>> &succOperandsTypes) {
3896+
if (failed(parser.parseCommaSeparatedList(
3897+
OpAsmParser::Delimiter::Square,
3898+
[&]() {
3899+
Block *destination = nullptr;
3900+
SmallVector<OpAsmParser::UnresolvedOperand> operands;
3901+
SmallVector<Type> operandTypes;
3902+
3903+
if (parser.parseSuccessor(destination).failed())
3904+
return failure();
3905+
3906+
if (succeeded(parser.parseOptionalLParen())) {
3907+
if (failed(parser.parseOperandList(
3908+
operands, OpAsmParser::Delimiter::None)) ||
3909+
failed(parser.parseColonTypeList(operandTypes)) ||
3910+
failed(parser.parseRParen()))
3911+
return failure();
3912+
}
3913+
succOperandBlocks.push_back(destination);
3914+
succOperands.emplace_back(operands);
3915+
succOperandsTypes.emplace_back(operandTypes);
3916+
return success();
3917+
},
3918+
"successor blocks")))
3919+
return failure();
3920+
return success();
3921+
}
3922+
3923+
static void
3924+
printIndirectBrOpSucessors(OpAsmPrinter &p, IndirectBrOp op, Type flagType,
3925+
SuccessorRange succs, OperandRangeRange succOperands,
3926+
const TypeRangeRange &succOperandsTypes) {
3927+
p << "[";
3928+
llvm::interleave(
3929+
llvm::zip(succs, succOperands),
3930+
[&](auto i) {
3931+
p.printNewline();
3932+
p.printSuccessorAndUseList(std::get<0>(i), std::get<1>(i));
3933+
},
3934+
[&] { p << ','; });
3935+
if (!succOperands.empty())
3936+
p.printNewline();
3937+
p << "]";
3938+
}
3939+
38823940
//===----------------------------------------------------------------------===//
38833941
// AssumeOp (intrinsic)
38843942
//===----------------------------------------------------------------------===//

‎mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -503,6 +503,15 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
503503
moduleTranslation.mapBranch(&opInst, switchInst);
504504
return success();
505505
}
506+
if (auto indBrOp = dyn_cast<LLVM::IndirectBrOp>(opInst)) {
507+
llvm::IndirectBrInst *indBr = builder.CreateIndirectBr(
508+
moduleTranslation.lookupValue(indBrOp.getAddr()),
509+
indBrOp->getNumSuccessors());
510+
for (auto *succ : indBrOp.getSuccessors())
511+
indBr->addDestination(moduleTranslation.lookupBlock(succ));
512+
moduleTranslation.mapBranch(&opInst, indBr);
513+
return success();
514+
}
506515

507516
// Emit addressof. We need to look up the global value referenced by the
508517
// operation and store it in the MLIR-to-LLVM value mapping. This does not

‎mlir/lib/Target/LLVMIR/ModuleImport.cpp‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1988,6 +1988,31 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
19881988
return success();
19891989
}
19901990

1991+
if (inst->getOpcode() == llvm::Instruction::IndirectBr) {
1992+
auto *indBrInst = cast<llvm::IndirectBrInst>(inst);
1993+
1994+
FailureOr<Value> basePtr = convertValue(indBrInst->getAddress());
1995+
if (failed(basePtr))
1996+
return failure();
1997+
1998+
SmallVector<Block *> succBlocks;
1999+
SmallVector<ValueRange> succBlockArgs;
2000+
for (auto i : llvm::seq<unsigned>(0, indBrInst->getNumSuccessors())) {
2001+
llvm::BasicBlock *succ = indBrInst->getSuccessor(i);
2002+
SmallVector<Value> blockArgs;
2003+
if (failed(convertBranchArgs(indBrInst, succ, blockArgs)))
2004+
return failure();
2005+
succBlocks.push_back(lookupBlock(succ));
2006+
succBlockArgs.push_back(blockArgs);
2007+
}
2008+
Location loc = translateLoc(inst->getDebugLoc());
2009+
auto indBrOp = builder.create<LLVM::IndirectBrOp>(
2010+
loc, TypeRange{}, *basePtr, succBlockArgs, succBlocks);
2011+
2012+
mapNoResultOp(inst, indBrOp);
2013+
return success();
2014+
}
2015+
19912016
// Convert all instructions that have an mlirBuilder.
19922017
if (succeeded(convertInstructionImpl(builder, inst, *this, iface)))
19932018
return success();
@@ -1998,8 +2023,8 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
19982023
LogicalResult ModuleImport::processInstruction(llvm::Instruction *inst) {
19992024
// FIXME: Support uses of SubtargetData.
20002025
// FIXME: Add support for call / operand attributes.
2001-
// FIXME: Add support for the indirectbr, cleanupret, catchret, catchswitch,
2002-
// callbr, vaarg, catchpad, cleanuppad instructions.
2026+
// FIXME: Add support for the cleanupret, catchret, catchswitch, callbr,
2027+
// vaarg, catchpad, cleanuppad instructions.
20032028

20042029
// Convert LLVM intrinsics calls to MLIR intrinsics.
20052030
if (auto *intrinsic = dyn_cast<llvm::IntrinsicInst>(inst))

‎mlir/lib/Target/LLVMIR/ModuleTranslation.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -805,6 +805,12 @@ static Value getPHISourceValue(Block *current, Block *pred,
805805
return switchOp.getCaseOperands(i.index())[index];
806806
}
807807

808+
if (auto indBrOp = dyn_cast<LLVM::IndirectBrOp>(terminator)) {
809+
// For indirect branches we take operands for each successor.
810+
for (const auto &i : llvm::enumerate(indBrOp->getSuccessors()))
811+
return indBrOp.getSuccessorOperands(i.index())[index];
812+
}
813+
808814
if (auto invokeOp = dyn_cast<LLVM::InvokeOp>(terminator)) {
809815
return invokeOp.getNormalDest() == current
810816
? invokeOp.getNormalDestOperands()[index]

‎mlir/test/Dialect/LLVMIR/indirectbr.mlir‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
// RUN: mlir-opt -split-input-file --verify-roundtrip %s | FileCheck %s
2+
3+
llvm.func @ib0(%dest : !llvm.ptr, %arg0 : i32, %arg1 : i32) -> i32 {
4+
llvm.indirectbr %dest : !llvm.ptr, [
5+
^head(%arg0 : i32),
6+
^tail(%arg1 : i32)
7+
]
8+
^head(%r0 : i32):
9+
llvm.return %r0 : i32
10+
^tail(%r1 : i32):
11+
llvm.return %r1 : i32
12+
}
13+
14+
// CHECK: llvm.func @ib0(%[[Addr:.*]]: !llvm.ptr, %[[A0:.*]]: i32, %[[A1:.*]]: i32) -> i32 {
15+
// CHECK: llvm.indirectbr %[[Addr]] : !llvm.ptr, [
16+
// CHECK: ^bb1(%[[A0:.*]] : i32)
17+
// CHECK: ^bb2(%[[A1:.*]] : i32)
18+
// CHECK: ]
19+
// CHECK: ^bb1(%[[Op0:.*]]: i32):
20+
// CHECK: llvm.return %[[Op0]] : i32
21+
// CHECK: ^bb2(%[[Op1:.*]]: i32):
22+
// CHECK: llvm.return %[[Op1]] : i32
23+
// CHECK: }
24+
25+
// -----
26+
27+
llvm.func @ib1(%dest : !llvm.ptr) {
28+
llvm.indirectbr %dest : !llvm.ptr, []
29+
}
30+
31+
// CHECK: llvm.func @ib1(%[[Addr:.*]]: !llvm.ptr) {
32+
// CHECK: llvm.indirectbr %[[Addr]] : !llvm.ptr, []
33+
// CHECK: }
34+
35+
// -----
36+
37+
// CHECK: llvm.func @test_indirectbr_phi(
38+
// CHECK-SAME: %arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: i32) -> i32 {
39+
llvm.func @callee(!llvm.ptr, i32, i32) -> i32
40+
llvm.func @test_indirectbr_phi(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: i32) -> i32 {
41+
%0 = llvm.mlir.undef : i1
42+
%1 = llvm.mlir.addressof @test_indirectbr_phi : !llvm.ptr
43+
%2 = llvm.blockaddress <function = @test_indirectbr_phi, tag = <id = 1>> : !llvm.ptr
44+
// CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
45+
%3 = llvm.mlir.constant(1 : i32) : i32
46+
// CHECK: %[[TWO:.*]] = llvm.mlir.constant(2 : i32) : i32
47+
%4 = llvm.mlir.constant(2 : i32) : i32
48+
%5 = llvm.select %0, %2, %arg0 : i1, !llvm.ptr
49+
// CHECK: llvm.indirectbr {{.*}} : !llvm.ptr, [
50+
// CHECK: ^[[HEAD_BB:.*]],
51+
// CHECK: ^[[TAIL_BB:.*]](%[[ONE]] : i32)
52+
// CHECK: ]
53+
llvm.indirectbr %5 : !llvm.ptr, [
54+
^bb1,
55+
^bb2(%3 : i32)
56+
]
57+
^bb1:
58+
// CHECK: ^[[HEAD_BB]]:
59+
// CHECK: llvm.indirectbr {{.*}} : !llvm.ptr, [
60+
// CHECK: ^[[TAIL_BB]](%[[TWO]] : i32),
61+
// CHECK: ^[[END_BB:.*]]
62+
// CHECK: ]
63+
%6 = llvm.select %0, %2, %arg0 : i1, !llvm.ptr
64+
llvm.indirectbr %6 : !llvm.ptr, [
65+
^bb2(%4 : i32),
66+
^bb3
67+
]
68+
^bb2(%7: i32):
69+
// CHECK: ^[[TAIL_BB]](%[[BLOCK_ARG:.*]]: i32):
70+
// CHECK: {{.*}} = llvm.call @callee({{.*}}, %[[BLOCK_ARG]])
71+
// CHECK: llvm.return
72+
%8 = llvm.call @callee(%arg1, %arg3, %7) : (!llvm.ptr, i32, i32) -> i32
73+
llvm.return %8 : i32
74+
^bb3:
75+
// CHECK: ^[[END_BB]]:
76+
// CHECK: llvm.blocktag
77+
// CHECK: llvm.return
78+
// CHECK: }
79+
llvm.blocktag <id = 1>
80+
llvm.return %arg3 : i32
81+
}

‎mlir/test/Target/LLVMIR/Import/import-failure.ll‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,5 @@
11
; RUN: not mlir-translate -import-llvm -emit-expensive-warnings -split-input-file %s 2>&1 -o /dev/null | FileCheck %s
22

3-
; CHECK: <unknown>
4-
; CHECK-SAME: error: unhandled instruction: indirectbr ptr %dst, [label %bb1, label %bb2]
5-
define i32 @unhandled_instruction(ptr %dst) {
6-
indirectbr ptr %dst, [label %bb1, label %bb2]
7-
bb1:
8-
ret i32 0
9-
bb2:
10-
ret i32 1
11-
}
12-
13-
; // -----
14-
153
; Check that debug intrinsics with an unsupported argument are dropped.
164

175
declare void @llvm.dbg.value(metadata, metadata, metadata)

‎mlir/test/Target/LLVMIR/Import/indirectbr.ll‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
; RUN: mlir-translate --import-llvm %s -split-input-file | FileCheck %s
2+
3+
; CHECK: llvm.func @basic(%[[arg0:.*]]: !llvm.ptr)
4+
define i32 @basic(ptr %dst) {
5+
; CHECK: llvm.indirectbr %[[arg0]] : !llvm.ptr, [
6+
; CHECK: ^[[bb1:.*]],
7+
; CHECK: ^[[bb2:.*]]
8+
; CHECK: ]
9+
indirectbr ptr %dst, [label %bb1, label %bb2]
10+
bb1:
11+
; CHECK: ^[[bb1]]:
12+
; CHECK: llvm.return
13+
ret i32 0
14+
bb2:
15+
; CHECK: ^[[bb2]]:
16+
; CHECK: llvm.return
17+
ret i32 1
18+
}
19+
20+
; // -----
21+
22+
; CHECK: llvm.mlir.global external @addr()
23+
@addr = global ptr null
24+
25+
; CHECK-LABEL: llvm.func @test_indirectbr() {
26+
define void @test_indirectbr() {
27+
; CHECK: %[[BA:.*]] = llvm.blockaddress <function = @test_indirectbr, tag = <id = 1>> : !llvm.ptr
28+
; CHECK: {{.*}} = llvm.mlir.addressof @addr : !llvm.ptr
29+
; CHECK: llvm.store %[[BA]], {{.*}} : !llvm.ptr, !llvm.ptr
30+
store ptr blockaddress(@test_indirectbr, %block), ptr @addr
31+
; CHECK: %[[TARGET_ADDR:.*]] = llvm.load {{.*}} : !llvm.ptr -> !llvm.ptr
32+
%val = load ptr, ptr @addr
33+
; CHECK: llvm.indirectbr %[[TARGET_ADDR]] : !llvm.ptr, [
34+
; CHECK: ^[[TARGET_BB:.*]]
35+
; CHECK: ]
36+
indirectbr ptr %val, [label %block]
37+
; CHECK: ^[[TARGET_BB]]:
38+
; CHECK: llvm.blocktag <id = 1>
39+
; CHECK: llvm.return
40+
; CHECK: }
41+
block:
42+
ret void
43+
}
44+
45+
; // -----
46+
47+
; CHECK: llvm.func @callee(!llvm.ptr, i32, i32) -> i32
48+
declare i32 @callee(ptr %a, i32 %v, i32 %p)
49+
50+
; CHECK: llvm.func @test_indirectbr_phi(
51+
; CHECK-SAME: %arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: i32) -> i32 {
52+
define i32 @test_indirectbr_phi(ptr %address, ptr %a, ptr %b, i32 %v) {
53+
entry:
54+
; CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
55+
; CHECK: %[[TWO:.*]] = llvm.mlir.constant(2 : i32) : i32
56+
%dest = select i1 undef, ptr blockaddress(@test_indirectbr_phi, %end), ptr %address
57+
; CHECK: llvm.indirectbr {{.*}} : !llvm.ptr, [
58+
; CHECK: ^[[HEAD_BB:.*]],
59+
; CHECK: ^[[TAIL_BB:.*]](%[[ONE]] : i32)
60+
; CHECK: ]
61+
indirectbr ptr %dest, [label %head, label %tail]
62+
63+
head:
64+
; CHECK: ^[[HEAD_BB]]:
65+
; CHECK: llvm.indirectbr {{.*}} : !llvm.ptr, [
66+
; CHECK: ^[[TAIL_BB]](%[[TWO]] : i32),
67+
; CHECK: ^[[END_BB:.*]]
68+
; CHECK: ]
69+
%dest2 = select i1 undef, ptr blockaddress(@test_indirectbr_phi, %end), ptr %address
70+
indirectbr ptr %dest2, [label %tail, label %end]
71+
72+
tail:
73+
; CHECK: ^[[TAIL_BB]](%[[BLOCK_ARG:.*]]: i32):
74+
; CHECK: {{.*}} = llvm.call @callee({{.*}}, %[[BLOCK_ARG]])
75+
; CHECK: llvm.return
76+
%p = phi i32 [1, %entry], [2, %head]
77+
%r = call i32 @callee(ptr %a, i32 %v, i32 %p)
78+
ret i32 %r
79+
80+
end:
81+
; CHECK: ^[[END_BB]]:
82+
; CHECK: llvm.blocktag
83+
; CHECK: llvm.return
84+
; CHECK: }
85+
ret i32 %v
86+
}

0 commit comments

Comments
 (0)