Skip to content

Commit c179847

Browse files
bcardosolopesbcardosolopes
authored
Reapply [MLIR][LLVM] Support for indirectbr (#136378)
Fix msan issue that caused revert in #135695 ### Original message Now that LLVM dialect has `blockaddress` support, introduce import/translation for `indirectbr` instruction.
1 parent b95ec24 commit c179847

File tree

11 files changed

+411
-40
lines changed

11 files changed

+411
-40
lines changed

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

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1690,8 +1690,7 @@ def LLVM_BlockTagOp : LLVM_Op<"blocktag"> {
16901690
the target address.
16911691

16921692
A given function should have at most one `llvm.blocktag` operation with a
1693-
given `tag`. This operation cannot be used as a terminator but can be
1694-
placed everywhere else in a block.
1693+
given `tag`. This operation cannot be used as a terminator.
16951694

16961695
Example:
16971696

@@ -1711,6 +1710,66 @@ def LLVM_BlockTagOp : LLVM_Op<"blocktag"> {
17111710
let hasVerifier = 0;
17121711
}
17131712

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

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

Lines changed: 86 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2240,24 +2240,21 @@ static LogicalResult verifyComdat(Operation *op,
22402240

22412241
static LogicalResult verifyBlockTags(LLVMFuncOp funcOp) {
22422242
llvm::DenseSet<BlockTagAttr> blockTags;
2243-
BlockTagOp badBlockTagOp;
2244-
if (funcOp
2245-
.walk([&](BlockTagOp blockTagOp) {
2246-
if (blockTags.contains(blockTagOp.getTag())) {
2247-
badBlockTagOp = blockTagOp;
2248-
return WalkResult::interrupt();
2249-
}
2250-
blockTags.insert(blockTagOp.getTag());
2251-
return WalkResult::advance();
2252-
})
2253-
.wasInterrupted()) {
2254-
badBlockTagOp.emitError()
2255-
<< "duplicate block tag '" << badBlockTagOp.getTag().getId()
2256-
<< "' in the same function: ";
2257-
return failure();
2258-
}
2243+
// Note that presence of `BlockTagOp`s currently can't prevent an unrecheable
2244+
// block to be removed by canonicalizer's region simplify pass, which needs to
2245+
// be dialect aware to allow extra constraints to be described.
2246+
WalkResult res = funcOp.walk([&](BlockTagOp blockTagOp) {
2247+
if (blockTags.contains(blockTagOp.getTag())) {
2248+
blockTagOp.emitError()
2249+
<< "duplicate block tag '" << blockTagOp.getTag().getId()
2250+
<< "' in the same function: ";
2251+
return WalkResult::interrupt();
2252+
}
2253+
blockTags.insert(blockTagOp.getTag());
2254+
return WalkResult::advance();
2255+
});
22592256

2260-
return success();
2257+
return failure(res.wasInterrupted());
22612258
}
22622259

22632260
/// Parse common attributes that might show up in the same order in both
@@ -3818,6 +3815,78 @@ LogicalResult BlockAddressOp::verify() {
38183815
/// attribute.
38193816
OpFoldResult BlockAddressOp::fold(FoldAdaptor) { return getBlockAddr(); }
38203817

3818+
//===----------------------------------------------------------------------===//
3819+
// LLVM::IndirectBrOp
3820+
//===----------------------------------------------------------------------===//
3821+
3822+
SuccessorOperands IndirectBrOp::getSuccessorOperands(unsigned index) {
3823+
assert(index < getNumSuccessors() && "invalid successor index");
3824+
return SuccessorOperands(getSuccOperandsMutable()[index]);
3825+
}
3826+
3827+
void IndirectBrOp::build(OpBuilder &odsBuilder, OperationState &odsState,
3828+
Value addr, ArrayRef<ValueRange> succOperands,
3829+
BlockRange successors) {
3830+
odsState.addOperands(addr);
3831+
for (ValueRange range : succOperands)
3832+
odsState.addOperands(range);
3833+
SmallVector<int32_t> rangeSegments;
3834+
for (ValueRange range : succOperands)
3835+
rangeSegments.push_back(range.size());
3836+
odsState.getOrAddProperties<Properties>().indbr_operand_segments =
3837+
odsBuilder.getDenseI32ArrayAttr(rangeSegments);
3838+
odsState.addSuccessors(successors);
3839+
}
3840+
3841+
static ParseResult parseIndirectBrOpSucessors(
3842+
OpAsmParser &parser, Type &flagType,
3843+
SmallVectorImpl<Block *> &succOperandBlocks,
3844+
SmallVectorImpl<SmallVector<OpAsmParser::UnresolvedOperand>> &succOperands,
3845+
SmallVectorImpl<SmallVector<Type>> &succOperandsTypes) {
3846+
if (failed(parser.parseCommaSeparatedList(
3847+
OpAsmParser::Delimiter::Square,
3848+
[&]() {
3849+
Block *destination = nullptr;
3850+
SmallVector<OpAsmParser::UnresolvedOperand> operands;
3851+
SmallVector<Type> operandTypes;
3852+
3853+
if (parser.parseSuccessor(destination).failed())
3854+
return failure();
3855+
3856+
if (succeeded(parser.parseOptionalLParen())) {
3857+
if (failed(parser.parseOperandList(
3858+
operands, OpAsmParser::Delimiter::None)) ||
3859+
failed(parser.parseColonTypeList(operandTypes)) ||
3860+
failed(parser.parseRParen()))
3861+
return failure();
3862+
}
3863+
succOperandBlocks.push_back(destination);
3864+
succOperands.emplace_back(operands);
3865+
succOperandsTypes.emplace_back(operandTypes);
3866+
return success();
3867+
},
3868+
"successor blocks")))
3869+
return failure();
3870+
return success();
3871+
}
3872+
3873+
static void
3874+
printIndirectBrOpSucessors(OpAsmPrinter &p, IndirectBrOp op, Type flagType,
3875+
SuccessorRange succs, OperandRangeRange succOperands,
3876+
const TypeRangeRange &succOperandsTypes) {
3877+
p << "[";
3878+
llvm::interleave(
3879+
llvm::zip(succs, succOperands),
3880+
[&](auto i) {
3881+
p.printNewline();
3882+
p.printSuccessorAndUseList(std::get<0>(i), std::get<1>(i));
3883+
},
3884+
[&] { p << ','; });
3885+
if (!succOperands.empty())
3886+
p.printNewline();
3887+
p << "]";
3888+
}
3889+
38213890
//===----------------------------------------------------------------------===//
38223891
// AssumeOp (intrinsic)
38233892
//===----------------------------------------------------------------------===//

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

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

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

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

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1988,6 +1988,33 @@ 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<SmallVector<Value>> 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+
SmallVector<ValueRange> succBlockArgsRange =
2009+
llvm::to_vector_of<ValueRange>(succBlockArgs);
2010+
Location loc = translateLoc(inst->getDebugLoc());
2011+
auto indBrOp = builder.create<LLVM::IndirectBrOp>(
2012+
loc, *basePtr, succBlockArgsRange, succBlocks);
2013+
2014+
mapNoResultOp(inst, indBrOp);
2015+
return success();
2016+
}
2017+
19912018
// Convert all instructions that have an mlirBuilder.
19922019
if (succeeded(convertInstructionImpl(builder, inst, *this, iface)))
19932020
return success();
@@ -1998,8 +2025,8 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
19982025
LogicalResult ModuleImport::processInstruction(llvm::Instruction *inst) {
19992026
// FIXME: Support uses of SubtargetData.
20002027
// FIXME: Add support for call / operand attributes.
2001-
// FIXME: Add support for the indirectbr, cleanupret, catchret, catchswitch,
2002-
// callbr, vaarg, catchpad, cleanuppad instructions.
2028+
// FIXME: Add support for the cleanupret, catchret, catchswitch, callbr,
2029+
// vaarg, catchpad, cleanuppad instructions.
20032030

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

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -805,6 +805,14 @@ 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+
if (indBrOp->getSuccessor(i.index()) == current)
812+
return indBrOp.getSuccessorOperands(i.index())[index];
813+
}
814+
}
815+
808816
if (auto invokeOp = dyn_cast<LLVM::InvokeOp>(terminator)) {
809817
return invokeOp.getNormalDest() == current
810818
? invokeOp.getNormalDestOperands()[index]

‎mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// RUN: mlir-opt %s -pass-pipeline='builtin.module(llvm.func(canonicalize{region-simplify=aggressive}))' -split-input-file | FileCheck %s
1+
// RUN: mlir-opt %s -pass-pipeline='builtin.module(llvm.func(canonicalize{region-simplify=aggressive}))' -verify-diagnostics -split-input-file | FileCheck %s
22

33
llvm.mlir.global private @x() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
44
%0 = llvm.blockaddress <function = @ba, tag = <id = 2>> : !llvm.ptr

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

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

‎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)

0 commit comments

Comments
 (0)