Skip to content

Commit 05ab891

Browse files
committed
Add pass for splitting blocks after calls.
This effectively turns call instructions into terminators, making it easier to process traces.
1 parent 9813084 commit 05ab891

File tree

7 files changed

+184
-0
lines changed

7 files changed

+184
-0
lines changed

llvm/include/llvm/InitializePasses.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,7 @@ void initializeWasmEHPreparePass(PassRegistry&);
341341
void initializeWinEHPreparePass(PassRegistry&);
342342
void initializeWriteBitcodePassPass(PassRegistry&);
343343
void initializeXRayInstrumentationPass(PassRegistry&);
344+
void initializeYkSplitBlocksAfterCallsPass(PassRegistry&);
344345

345346
} // end namespace llvm
346347

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
#ifndef LLVM_TRANSFORMS_YK_SPLITBLOCKSAFTERCALLS_H
2+
#define LLVM_TRANSFORMS_YK_SPLITBLOCKSAFTERCALLS_H
3+
4+
#include "llvm/Pass.h"
5+
6+
namespace llvm {
7+
ModulePass *createYkSplitBlocksAfterCallsPass();
8+
} // namespace llvm
9+
10+
#endif

llvm/lib/CodeGen/CodeGen.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,4 +150,5 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
150150
initializeWasmEHPreparePass(Registry);
151151
initializeWinEHPreparePass(Registry);
152152
initializeXRayInstrumentationPass(Registry);
153+
initializeYkSplitBlocksAfterCallsPass(Registry);
153154
}

llvm/lib/CodeGen/TargetPassConfig.cpp

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
#include "llvm/Transforms/Yk/ControlPoint.h"
5353
#include "llvm/Transforms/Yk/Linkage.h"
5454
#include "llvm/Transforms/Yk/ShadowStack.h"
55+
#include "llvm/Transforms/Yk/SplitBlocksAfterCalls.h"
5556
#include "llvm/Transforms/Yk/Stackmaps.h"
5657
#include "llvm/Transforms/Yk/NoCallsInEntryBlocks.h"
5758
#include <cassert>
@@ -287,6 +288,10 @@ static cl::opt<bool>
287288
YkNoCallsInEntryBlocks("yk-no-calls-in-entryblocks", cl::init(false), cl::NotHidden,
288289
cl::desc("Ensure there are no calls in the entryblock."));
289290

291+
static cl::opt<bool>
292+
YkSplitBlocksAfterCalls("yk-split-blocks-after-calls", cl::init(false), cl::NotHidden,
293+
cl::desc("Split blocks after function calls."));
294+
290295
/// Allow standard passes to be disabled by command line options. This supports
291296
/// simple binary flags that either suppress the pass or do nothing.
292297
/// i.e. -disable-mypass=false has no effect.
@@ -1162,6 +1167,16 @@ bool TargetPassConfig::addISelPasses() {
11621167
addPass(createYkLinkagePass());
11631168
}
11641169

1170+
if (YkSplitBlocksAfterCalls) {
1171+
if (!YkNoCallsInEntryBlocks) {
1172+
// YKFIXME: Merge the two passes together. Then modify the control point
1173+
// pass to make sure we split the block right after the control point, as
1174+
// we can no longer rely on this pass to do so.
1175+
report_fatal_error("--yk-split-blocks-after-calls requires --yk-no-calls-in-entryblocks.");
1176+
}
1177+
addPass(createYkSplitBlocksAfterCallsPass());
1178+
}
1179+
11651180
if (YkInsertStackMaps) {
11661181
addPass(createYkStackmapsPass());
11671182
}

llvm/lib/Transforms/Yk/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ add_llvm_component_library(LLVMYkPasses
66
StackMaps.cpp
77
ShadowStack.cpp
88
NoCallsInEntryBlocks.cpp
9+
SplitBlocksAfterCalls.cpp
910

1011
DEPENDS
1112
intrinsics_gen
Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
//===- SplitBlocksAfterCalls.cpp -===//
2+
//
3+
// Makes function calls effectively terminators by splitting blocks after each
4+
// call. This ensures that there can only be at most one call per block. This
5+
// is used in order to detect recursion and external function calls within a
6+
// trace.
7+
8+
#include "llvm/Transforms/Yk/SplitBlocksAfterCalls.h"
9+
#include "llvm/IR/BasicBlock.h"
10+
#include "llvm/IR/Function.h"
11+
#include "llvm/IR/IRBuilder.h"
12+
#include "llvm/IR/Instructions.h"
13+
#include "llvm/IR/Module.h"
14+
#include "llvm/IR/Verifier.h"
15+
#include "llvm/InitializePasses.h"
16+
#include "llvm/Pass.h"
17+
#include "llvm/Transforms/Yk/LivenessAnalysis.h"
18+
19+
#include <map>
20+
21+
#define DEBUG_TYPE "yk-splitblocksaftercalls"
22+
23+
using namespace llvm;
24+
25+
namespace llvm {
26+
void initializeYkSplitBlocksAfterCallsPass(PassRegistry &);
27+
} // namespace llvm
28+
29+
namespace {
30+
31+
class YkSplitBlocksAfterCalls : public ModulePass {
32+
public:
33+
static char ID;
34+
YkSplitBlocksAfterCalls() : ModulePass(ID) {
35+
initializeYkSplitBlocksAfterCallsPass(*PassRegistry::getPassRegistry());
36+
}
37+
38+
bool runOnModule(Module &M) override {
39+
LLVMContext &Context = M.getContext();
40+
41+
for (Function &F : M) {
42+
if (F.empty()) // skip declarations.
43+
continue;
44+
// As we will be modifying the blocks of this function inplace, we
45+
// require a work list to process all existing and newly inserted blocks
46+
// in order to not miss any.
47+
std::vector<BasicBlock *> Todo;
48+
std::set<BasicBlock *> Seen;
49+
BasicBlock &Entry = F.getEntryBlock();
50+
51+
// This pass requires the `NoCallsInEntryBlocksPass` to have run first,
52+
// which in turn needs to run before the shadowstack pass. Otherwise,
53+
// this pass would split the block after the shadowstack malloc, which
54+
// results in allocas outside of the entry block which breaks stackmaps.
55+
Instruction *T = Entry.getTerminator();
56+
for (size_t I = 0; I < T->getNumSuccessors(); I++) {
57+
Todo.push_back(T->getSuccessor(I));
58+
}
59+
Seen.insert(&Entry);
60+
while (!Todo.empty()) {
61+
BasicBlock *Next = Todo.back();
62+
Todo.pop_back();
63+
if (Seen.count(Next) > 0) {
64+
continue;
65+
}
66+
Seen.insert(Next);
67+
68+
for (Instruction &I : *Next) {
69+
if (I.isDebugOrPseudoInst()) {
70+
continue;
71+
}
72+
if (isa<CallInst>(I)) {
73+
// YKFIXME: Can we determine at compile time if inline asm contains
74+
// calls or jumps, e.g. via `getAsmString`, and then not split the
75+
// block after them?
76+
CallInst *CI = cast<CallInst>(&I);
77+
Function *F = CI->getCalledFunction();
78+
if (F && F->getName() == "llvm.frameaddress.p0") {
79+
// This call is always inlined so we don't need to split the
80+
// block here.
81+
continue;
82+
}
83+
// YKFIXME: If the next instruction is an unconditional branch, we
84+
// don't need to split the block here.
85+
86+
// Since `splitBasicBlock` splits before the given instruction,
87+
// pass the instruction following this call instead.
88+
Next->splitBasicBlock(I.getNextNode());
89+
break;
90+
}
91+
}
92+
93+
// Add successors to todo list.
94+
Instruction *T = Next->getTerminator();
95+
for (size_t I = 0; I < T->getNumSuccessors(); I++) {
96+
Todo.insert(Todo.begin(), T->getSuccessor(I));
97+
}
98+
}
99+
}
100+
101+
#ifndef NDEBUG
102+
// Our pass runs after LLVM normally does its verify pass. In debug builds
103+
// we run it again to check that our pass is generating valid IR.
104+
if (verifyModule(M, &errs())) {
105+
Context.emitError("Stackmap insertion pass generated invalid IR!");
106+
return false;
107+
}
108+
#endif
109+
return true;
110+
}
111+
};
112+
} // namespace
113+
114+
char YkSplitBlocksAfterCalls::ID = 0;
115+
INITIALIZE_PASS(YkSplitBlocksAfterCalls, DEBUG_TYPE, "yk stackmaps", false,
116+
false)
117+
118+
ModulePass *llvm::createYkSplitBlocksAfterCallsPass() {
119+
return new YkSplitBlocksAfterCalls();
120+
}
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
; RUN: llc -stop-after yk-splitblocksaftercalls --yk-split-blocks-after-calls --yk-no-calls-in-entryblocks < %s | FileCheck %s
2+
3+
; Check that the block in the main function containing two calls to foo, has
4+
; been split after each of the calls.
5+
6+
; CHECK-LABEL: define dso_local i32 @main
7+
; CHECK-NEXT: %2 = add nsw i32 %0, 1
8+
; CHECK-NEXT: br label %3
9+
; CHECK-LABEL: 3:
10+
; CHECK-NEXT: %4 = call i32 @foo(i32 noundef %2)
11+
; CHECK-NEXT: br label %5
12+
; CHECK-LABEL: 5:
13+
; CHECK-NEXT: %6 = add nsw i32 %0, 1
14+
; CHECK-NEXT: %7 = call i32 @foo(i32 noundef %4)
15+
; CHECK-NEXT: br label %8
16+
17+
@.str = private unnamed_addr constant [13 x i8] c"%d %d %d %d\0A\00", align 1
18+
19+
define dso_local i32 @foo(i32 noundef %0) #0 {
20+
%2 = add nsw i32 %0, 10
21+
ret i32 %2
22+
}
23+
24+
declare i32 @printf(ptr noundef, ...) #2
25+
26+
define dso_local i32 @main(i32 noundef %0) #0 {
27+
%2 = add nsw i32 %0, 1
28+
%3 = call i32 @foo(i32 noundef %2)
29+
%4 = add nsw i32 %0, 1
30+
%5 = call i32 @foo(i32 noundef %3)
31+
ret i32 0
32+
}
33+
34+
attributes #0 = { noinline nounwind optnone uwtable "frame-pointer"="all" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
35+
attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
36+
attributes #2 = { "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }

0 commit comments

Comments
 (0)