Merge pull request rust-lang#108 from ptersilie/splitblocksaftercalls

Add pass for splitting blocks after calls.

llvm/include/llvm/InitializePasses.h

@@ -341,6 +341,7 @@ void initializeWasmEHPreparePass(PassRegistry&);
 void initializeWinEHPreparePass(PassRegistry&);
 void initializeWriteBitcodePassPass(PassRegistry&);
 void initializeXRayInstrumentationPass(PassRegistry&);
+void initializeYkSplitBlocksAfterCallsPass(PassRegistry&);
 
 } // end namespace llvm
 

llvm/include/llvm/Transforms/Yk/SplitBlocksAfterCalls.h

@@ -0,0 +1,10 @@
+#ifndef LLVM_TRANSFORMS_YK_SPLITBLOCKSAFTERCALLS_H
+#define LLVM_TRANSFORMS_YK_SPLITBLOCKSAFTERCALLS_H
+
+#include "llvm/Pass.h"
+
+namespace llvm {
+ModulePass *createYkSplitBlocksAfterCallsPass();
+} // namespace llvm
+
+#endif

llvm/lib/CodeGen/CodeGen.cpp

@@ -150,4 +150,5 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeWasmEHPreparePass(Registry);
   initializeWinEHPreparePass(Registry);
   initializeXRayInstrumentationPass(Registry);
+  initializeYkSplitBlocksAfterCallsPass(Registry);
 }

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -52,6 +52,7 @@
 #include "llvm/Transforms/Yk/ControlPoint.h"
 #include "llvm/Transforms/Yk/Linkage.h"
 #include "llvm/Transforms/Yk/ShadowStack.h"
+#include "llvm/Transforms/Yk/SplitBlocksAfterCalls.h"
 #include "llvm/Transforms/Yk/Stackmaps.h"
 #include "llvm/Transforms/Yk/NoCallsInEntryBlocks.h"
 #include <cassert>
@@ -287,6 +288,10 @@ static cl::opt<bool>
     YkNoCallsInEntryBlocks("yk-no-calls-in-entryblocks", cl::init(false), cl::NotHidden,
                       cl::desc("Ensure there are no calls in the entryblock."));
 
+static cl::opt<bool>
+    YkSplitBlocksAfterCalls("yk-split-blocks-after-calls", cl::init(false), cl::NotHidden,
+                      cl::desc("Split blocks after function calls."));
+
 /// Allow standard passes to be disabled by command line options. This supports
 /// simple binary flags that either suppress the pass or do nothing.
 /// i.e. -disable-mypass=false has no effect.
@@ -1162,6 +1167,16 @@ bool TargetPassConfig::addISelPasses() {
     addPass(createYkLinkagePass());
   }
 
+  if (YkSplitBlocksAfterCalls) {
+    if (!YkNoCallsInEntryBlocks) {
+      // YKFIXME: Merge the two passes together. Then modify the control point
+      // pass to make sure we split the block right after the control point, as
+      // we can no longer rely on this pass to do so.
+      report_fatal_error("--yk-split-blocks-after-calls requires --yk-no-calls-in-entryblocks.");
+    }
+    addPass(createYkSplitBlocksAfterCallsPass());
+  }
+
   if (YkInsertStackMaps) {
     addPass(createYkStackmapsPass());
   }

llvm/lib/Transforms/Yk/CMakeLists.txt

@@ -6,6 +6,7 @@ add_llvm_component_library(LLVMYkPasses
   StackMaps.cpp
   ShadowStack.cpp
   NoCallsInEntryBlocks.cpp
+  SplitBlocksAfterCalls.cpp
 
   DEPENDS
   intrinsics_gen

llvm/lib/Transforms/Yk/SplitBlocksAfterCalls.cpp

@@ -0,0 +1,120 @@
+//===- SplitBlocksAfterCalls.cpp -===//
+//
+// Makes function calls effectively terminators by splitting blocks after each
+// call. This ensures that there can only be at most one call per block. This
+// is used in order to detect recursion and external function calls within a
+// trace.
+
+#include "llvm/Transforms/Yk/SplitBlocksAfterCalls.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Verifier.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/Yk/LivenessAnalysis.h"
+
+#include <map>
+
+#define DEBUG_TYPE "yk-splitblocksaftercalls"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeYkSplitBlocksAfterCallsPass(PassRegistry &);
+} // namespace llvm
+
+namespace {
+
+class YkSplitBlocksAfterCalls : public ModulePass {
+public:
+  static char ID;
+  YkSplitBlocksAfterCalls() : ModulePass(ID) {
+    initializeYkSplitBlocksAfterCallsPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnModule(Module &M) override {
+    LLVMContext &Context = M.getContext();
+
+    for (Function &F : M) {
+      if (F.empty()) // skip declarations.
+        continue;
+      // As we will be modifying the blocks of this function inplace, we
+      // require a work list to process all existing and newly inserted blocks
+      // in order to not miss any.
+      std::vector<BasicBlock *> Todo;
+      std::set<BasicBlock *> Seen;
+      BasicBlock &Entry = F.getEntryBlock();
+
+      // This pass requires the `NoCallsInEntryBlocksPass` to have run first,
+      // which in turn needs to run before the shadowstack pass. Otherwise,
+      // this pass would split the block after the shadowstack malloc, which
+      // results in allocas outside of the entry block which breaks stackmaps.
+      Instruction *T = Entry.getTerminator();
+      for (size_t I = 0; I < T->getNumSuccessors(); I++) {
+        Todo.push_back(T->getSuccessor(I));
+      }
+      Seen.insert(&Entry);
+      while (!Todo.empty()) {
+        BasicBlock *Next = Todo.back();
+        Todo.pop_back();
+        if (Seen.count(Next) > 0) {
+          continue;
+        }
+        Seen.insert(Next);
+
+        for (Instruction &I : *Next) {
+          if (I.isDebugOrPseudoInst()) {
+            continue;
+          }
+          if (isa<CallInst>(I)) {
+            // YKFIXME: Can we determine at compile time if inline asm contains
+            // calls or jumps, e.g. via `getAsmString`, and then not split the
+            // block after them?
+            CallInst *CI = cast<CallInst>(&I);
+            Function *F = CI->getCalledFunction();
+            if (F && F->getName() == "llvm.frameaddress.p0") {
+              // This call is always inlined so we don't need to split the
+              // block here.
+              continue;
+            }
+            // YKFIXME: If the next instruction is an unconditional branch, we
+            // don't need to split the block here.
+
+            // Since `splitBasicBlock` splits before the given instruction,
+            // pass the instruction following this call instead.
+            Next->splitBasicBlock(I.getNextNode());
+            break;
+          }
+        }
+
+        // Add successors to todo list.
+        Instruction *T = Next->getTerminator();
+        for (size_t I = 0; I < T->getNumSuccessors(); I++) {
+          Todo.insert(Todo.begin(), T->getSuccessor(I));
+        }
+      }
+    }
+
+#ifndef NDEBUG
+    // Our pass runs after LLVM normally does its verify pass. In debug builds
+    // we run it again to check that our pass is generating valid IR.
+    if (verifyModule(M, &errs())) {
+      Context.emitError("Stackmap insertion pass generated invalid IR!");
+      return false;
+    }
+#endif
+    return true;
+  }
+};
+} // namespace
+
+char YkSplitBlocksAfterCalls::ID = 0;
+INITIALIZE_PASS(YkSplitBlocksAfterCalls, DEBUG_TYPE, "yk stackmaps", false,
+                false)
+
+ModulePass *llvm::createYkSplitBlocksAfterCallsPass() {
+  return new YkSplitBlocksAfterCalls();
+}

llvm/test/CodeGen/X86/yk-split-blocks-after-calls.ll

@@ -0,0 +1,36 @@
+; RUN: llc -stop-after yk-splitblocksaftercalls --yk-split-blocks-after-calls --yk-no-calls-in-entryblocks < %s  | FileCheck %s
+
+; Check that the block in the main function containing two calls to foo, has
+; been split after each of the calls.
+
+; CHECK-LABEL: define dso_local i32 @main
+; CHECK-NEXT: %2 = add nsw i32 %0, 1
+; CHECK-NEXT: br label %3
+; CHECK-LABEL: 3:
+; CHECK-NEXT: %4 = call i32 @foo(i32 noundef %2)
+; CHECK-NEXT: br label %5
+; CHECK-LABEL: 5:
+; CHECK-NEXT: %6 = add nsw i32 %0, 1
+; CHECK-NEXT: %7 = call i32 @foo(i32 noundef %4)
+; CHECK-NEXT: br label %8
+
+@.str = private unnamed_addr constant [13 x i8] c"%d %d %d %d\0A\00", align 1
+
+define dso_local i32 @foo(i32 noundef %0) #0 {
+  %2 = add nsw i32 %0, 10
+  ret i32 %2
+}
+
+declare i32 @printf(ptr noundef, ...) #2
+
+define dso_local i32 @main(i32 noundef %0) #0 {
+  %2 = add nsw i32 %0, 1
+  %3 = call i32 @foo(i32 noundef %2)
+  %4 = add nsw i32 %0, 1
+  %5 = call i32 @foo(i32 noundef %3)
+  ret i32 0
+}
+
+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" }
+attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+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" }