Merge pull request rust-lang#105 from ptersilie/nomoreentryblockcalls

Add pass disallowing calls in entry blocks.

Author: ltratt

llvm/include/llvm/Transforms/Yk/NoCallsInEntryBlocks.h

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

llvm/lib/CodeGen/StackMaps.cpp

@@ -305,8 +305,22 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
     // each one of them. Note, that the stackmap may track either of %rbx or
     // %rcx, resulting in different ways below to retrieve the mappings.
     int ExtraReg = 0;
+    Register R = MOI->getReg();
     if (MOI->isReg()) {
-      Register R = MOI->getReg();
+      if (MOI->isKill()) {
+        // There's no point in tracking a killed register. Instead, try and
+        // find a copy of the register and use that to increase the likelyhood
+        // of us tracking that value.
+        // YKFIXME: This is likely only a temporary fix. In the future we might
+        // want to allow stackmaps to track arbitrarily many locations per live
+        // variable. Currently, we can only track two.
+        for (auto I : SpillOffsets) {
+          if (I.second == R) {
+            R = I.first;
+            break;
+          }
+        }
+      }
       if (SpillOffsets.count(R) > 0) {
         RHS = SpillOffsets[R];
         assert(SpillOffsets[R] != 0);
@@ -333,9 +347,9 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
       }
     }
 
-    unsigned DwarfRegNum = getDwarfRegNum(MOI->getReg(), TRI);
+    unsigned DwarfRegNum = getDwarfRegNum(R, TRI);
     unsigned LLVMRegNum = *TRI->getLLVMRegNum(DwarfRegNum, false);
-    unsigned SubRegIdx = TRI->getSubRegIndex(LLVMRegNum, MOI->getReg());
+    unsigned SubRegIdx = TRI->getSubRegIndex(LLVMRegNum, R);
     if (SubRegIdx)
       Offset = TRI->getSubRegIdxOffset(SubRegIdx);
 

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -53,6 +53,7 @@
 #include "llvm/Transforms/Yk/Linkage.h"
 #include "llvm/Transforms/Yk/ShadowStack.h"
 #include "llvm/Transforms/Yk/Stackmaps.h"
+#include "llvm/Transforms/Yk/NoCallsInEntryBlocks.h"
 #include <cassert>
 #include <optional>
 #include <string>
@@ -282,6 +283,10 @@ static cl::opt<bool>
     YkInsertStackMaps("yk-insert-stackmaps", cl::init(false), cl::NotHidden,
                       cl::desc("Insert stackmaps for JIT deoptimisation"));
 
+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."));
+
 /// 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.
@@ -1126,6 +1131,15 @@ bool TargetPassConfig::addISelPasses() {
   if (YkBlockDisambiguate)
     addPass(createYkBlockDisambiguatePass());
 
+  if (YkNoCallsInEntryBlocks) {
+    // Make sure this pass runs before the shadowstack pass, so that we don't
+    // split the entry block after that pass inserted a `malloc` into `main`.
+    // This would otherwise result in allocas being moved outside the entry
+    // block which makes them dynamic allocas, and results in stackmaps not
+    // being able to record the functions stack size.
+    addPass(createYkNoCallsInEntryBlocksPass());
+  }
+
   if (YkShadowStack) {
     addPass(createYkShadowStackPass());
   }

llvm/lib/Target/X86/X86AsmPrinter.cpp

@@ -65,7 +65,21 @@ void clearRhs(Register Reg, std::map<Register, int64_t> &SpillMap) {
   auto I = SpillMap.begin();
   while (I != SpillMap.end()) {
     if (I->second == Reg) {
-      I = SpillMap.erase(I);
+      // If there's a mapping A => B, where B has been reassigned and has a
+      // mapping B => C, then transitively apply the mapping so that A => C.
+      // This makes sure we don't remove mappings (especially to stack slots)
+      // when a register is reassigned. Example:
+      //   store $r13 in [$rbp - 8]
+      //   $rcx = $r13
+      //   $r13 = $rbx
+      // Upon assigning `$r13 = $rbx`, we apply `$r13`'s previous mapping to
+      // `$rcx` so that `SpillMap[$rcx] = -8`.
+      if(SpillMap.count(Reg) > 0) {
+        SpillMap[I->first] = SpillMap[Reg];
+        ++I;
+      } else {
+        I = SpillMap.erase(I);
+      }
     } else {
       ++I;
     }
@@ -99,13 +113,14 @@ void processInstructions(
       const MachineOperand Rhs = Instr.getOperand(1);
       assert(Lhs.isReg() && "Is register.");
       assert(Rhs.isReg() && "Is register.");
+      // Reassigning a new value to LHS means any mappings to Lhs are now void
+      // and need to be removed. We need to do this before updating the
+      // mapping, so the transitive property of the SpillMap isn't violated
+      // (see `clearRhs` for more info).
+      clearRhs(Lhs.getReg(), SpillMap);
       SpillMap[Lhs.getReg()] = Rhs.getReg();
       // YKFIXME: If the `mov` instruction has a killed-flag, remove the
       // register from the map.
-
-      // Reassigning a new value to Lhs means any mappings to Lhs are now void
-      // and need to be removed.
-      clearRhs(Lhs.getReg(), SpillMap);
       continue;
     }
 

llvm/lib/Transforms/Yk/CMakeLists.txt

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

llvm/lib/Transforms/Yk/NoCallsInEntryBlocks.cpp

@@ -0,0 +1,107 @@
+// NoCallsInEntryBlocks.cpp - Ensure there are no calls in the first block.
+//
+// This pass splits up the entry block so that there are no calls inside it
+// (one exception being the malloc call inserted by the shadowstack pass which
+// will be applied later). This allows us to detect call backs from external
+// functions more easily, which in turn simplifies our JITModBuilder
+// implementation and fixes the soundness issue that existed in the old
+// StackAdjust approach.
+//
+// As a simplified example, imagine this function:
+//
+// void foo(int argc) {
+//    if (argc == 0) {
+//      extfunc(foo)
+//    }
+// }
+//
+// where `extfunc` may or may not call back to `foo`. Previously, this might
+// result in a trace like
+//
+// foo.0
+// extfunc
+// foo.0
+// ...
+//
+// Upon seeing the second `foo.0` we are unable to determine whether we are
+// observing a callback to `foo` or are returning from `extfunc` which hasn't
+// called foo. When running this pass we force `extfunc` not to be in the entry
+// block, resulting in a trace like
+//
+// foo.0
+// foo.1
+// extfunc
+// foo.1
+// ...
+//
+// Here it is clear that `extfunc` has not called `foo` as we are seeing block
+// `foo.1` next. If `extfunc` had called `foo`, the next block would be
+// `foo.0`.
+
+#include "llvm/Transforms/Yk/NoCallsInEntryBlocks.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"
+
+#define DEBUG_TYPE "yk-no-more-entryblock-calls"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeYkNoCallsInEntryBlocksPass(PassRegistry &);
+} // namespace llvm
+
+namespace {
+
+class YkNoCallsInEntryBlocks : public ModulePass {
+public:
+  static char ID;
+  YkNoCallsInEntryBlocks() : ModulePass(ID) {
+    initializeYkNoCallsInEntryBlocksPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnModule(Module &M) override {
+    LLVMContext &Context = M.getContext();
+
+    // split block at first function call
+    std::set<Instruction *> Calls;
+    for (Function &F : M) {
+      if (F.empty()) // skip declarations.
+        continue;
+      BasicBlock &BB = F.getEntryBlock();
+      auto CurrFuncName = F.getName();
+      for (Instruction &I : BB) {
+        if (isa<CallInst>(I)) {
+          Function *F = cast<CallInst>(I).getCalledFunction();
+          BB.splitBasicBlock(&I);
+          break;
+        }
+      }
+    }
+
+#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 YkNoCallsInEntryBlocks::ID = 0;
+INITIALIZE_PASS(YkNoCallsInEntryBlocks, DEBUG_TYPE, "no more entry block calls",
+                false, false)
+
+ModulePass *llvm::createYkNoCallsInEntryBlocksPass() {
+  return new YkNoCallsInEntryBlocks();
+}