Merge rust-lang#13

13: Make the control point accept a `ykrt::Location` instead of an integer. r=ltratt a=vext01



Co-authored-by: Edd Barrett <vext01@gmail.com>

llvm/include/llvm/Transforms/Yk/ControlPoint.h

@@ -11,11 +11,7 @@
 #define YK_NEW_CONTROL_POINT "yk_new_control_point"
 
 namespace llvm {
-class YkControlPointPass : public PassInfoMixin<YkControlPointPass> {
-public:
-  explicit YkControlPointPass();
-  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
-};
+ModulePass *createYkControlPointPass();
 } // namespace llvm
 
 #endif

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -49,6 +49,7 @@
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/SymbolRewriter.h"
 #include "llvm/Transforms/Yk/BlockDisambiguate.h"
+#include "llvm/Transforms/Yk/ControlPoint.h"
 #include <cassert>
 #include <string>
 
@@ -243,6 +244,11 @@ static cl::opt<bool> YkBlockDisambiguate(
     "yk-block-disambiguate", cl::init(false), cl::NotHidden,
     cl::desc("Disambiguate blocks for yk"));
 
+static cl::opt<bool>
+YkPatchCtrlPoint("yk-patch-control-point",
+  cl::init(false), cl::NotHidden,
+  cl::desc("Patch yk_control_point()"));
+
 /// 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.
@@ -1067,8 +1073,22 @@ bool TargetPassConfig::addISelPasses() {
   addPassesToHandleExceptions();
   if (YkBlockDisambiguate)
     addPass(createYkBlockDisambiguatePass());
-  addISelPrepare();
 
+  // We insert the yk control point pass as late as possible. It has to run
+  // before instruction selection (or the machine IR won't reflect our
+  // patching), but after other passes which mutate the IR (e.g.
+  // `CodegenPrepare` above).
+  //
+  // By running the pass late, we also ensure that no IR optimisation passes
+  // ever see our patched-in control point function. If optimisations were to
+  // be applied then this would be problematic as the interface to the control
+  // point could be changed, e.g. the control point's struct argument could be
+  // decomposed into scalar arguments. The JIT runtime relies on the interface
+  // *not* being changed.
+  if (YkPatchCtrlPoint)
+    addPass(createYkControlPointPass());
+
+  addISelPrepare();
   return addCoreISelPasses();
 }
 

llvm/lib/LTO/LTOBackend.cpp

@@ -75,11 +75,6 @@ static cl::opt<bool> ThinLTOAssumeMerged(
     cl::desc("Assume the input has already undergone ThinLTO function "
              "importing and the other pre-optimization pipeline changes."));
 
-static cl::opt<bool>
-YkPatchCtrlPoint("yk-patch-control-point",
-  cl::init(false), cl::NotHidden,
-  cl::desc("Patch yk_control_point()"));
-
 LLVM_ATTRIBUTE_NORETURN static void reportOpenError(StringRef Path, Twine Msg) {
   errs() << "failed to open " << Path << ": " << Msg << '\n';
   errs().flush();
@@ -306,13 +301,6 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
     MPM.addPass(PB.buildLTODefaultPipeline(OL, ExportSummary));
   }
 
-  // We add the yk control point pass late in the pipeline (after all
-  // optimisation and just before verification and codegen) so that no IR
-  // optimisation passes have a chance to change the interface to the control
-  // point. The JIT runtime relies on the signature not being changed.
-  if (YkPatchCtrlPoint)
-    MPM.addPass(YkControlPointPass());
-
   if (!Conf.DisableVerify)
     MPM.addPass(VerifierPass());
 

llvm/lib/Transforms/Yk/BlockDisambiguate.cpp

@@ -204,6 +204,7 @@ class YkBlockDisambiguate : public ModulePass {
           if (SuccBB == &BB) {
             BasicBlock *DBB = makeDisambiguationBB(Context, &BB, NewBBs);
             BI->setSuccessor(SuccIdx, DBB);
+            BB.replacePhiUsesWith(&BB, DBB);
           }
         }
       } else if (isa<SwitchInst>(TI)) {
@@ -214,6 +215,7 @@ class YkBlockDisambiguate : public ModulePass {
           if (SuccBB == &BB) {
             BasicBlock *DBB = makeDisambiguationBB(Context, &BB, NewBBs);
             SI->setSuccessor(SuccIdx, DBB);
+            BB.replacePhiUsesWith(&BB, DBB);
           }
         }
       }

llvm/lib/Transforms/Yk/ControlPoint.cpp

@@ -46,10 +46,11 @@
 // YKFIXME: The control point cannot yet be used in an interpreter using
 // threaded dispatch.
 //
-// YKFIXME: The tracing logic is currently over-simplified:
+// YKFIXME: The tracing logic is currently over-simplified. The following items
+// need to be fixed:
 //
-//  - A JIT location is assumed to be a simple integer program counter. It
-//    should be a `ykrt::Location`.
+//  - The address of `YkLocation` instances are used for identity, but they are
+//    intended to be freely moved by the user.
 //
 //  - Tracing starts when we encounter a location for which we have no machine
 //    code. A hot counter should be used instead.
@@ -78,6 +79,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include <llvm/IR/Dominators.h>
 #include <llvm/IR/IRBuilder.h>
@@ -123,7 +125,7 @@ void createJITStatePrint(IRBuilder<> &Builder, Module *Mod, std::string Str) {
 /// Generates the new control point, which includes all logic to start/stop
 /// tracing and to compile/execute traces.
 void createControlPoint(Module &Mod, Function *F, std::vector<Value *> LiveVars,
-                        StructType *YkCtrlPointStruct) {
+                        StructType *YkCtrlPointStruct, Type *YkLocTy) {
   auto &Context = Mod.getContext();
 
   // Create control point blocks and setup the IRBuilder.
@@ -166,9 +168,8 @@ void createControlPoint(Module &Mod, Function *F, std::vector<Value *> LiveVars,
       PtNull, "compiled_trace", (GlobalVariable *)nullptr);
 
   GlobalVariable *GVStartLoc = new GlobalVariable(
-      Mod, Type::getInt32Ty(Context), false, GlobalVariable::InternalLinkage,
-      ConstantInt::get(Context, APInt(32, -1)), "start_loc",
-      (GlobalVariable *)nullptr);
+      Mod, YkLocTy, false, GlobalVariable::InternalLinkage,
+      Constant::getNullValue(YkLocTy), "start_loc", (GlobalVariable *)nullptr);
 
   // Create control point entry block. Checks if we are currently tracing.
   Value *GVTracingVal = Builder.CreateLoad(Type::getInt8Ty(Context), GVTracing);
@@ -196,8 +197,7 @@ void createControlPoint(Module &Mod, Function *F, std::vector<Value *> LiveVars,
   // Create block that checks if we've reached the same location again so we
   // can execute a compiled trace.
   Builder.SetInsertPoint(BBHasTrace);
-  Value *ValStartLoc =
-      Builder.CreateLoad(Type::getInt32Ty(Context), GVStartLoc);
+  Value *ValStartLoc = Builder.CreateLoad(YkLocTy, GVStartLoc);
   Value *ExecTraceCond = Builder.CreateICmp(CmpInst::Predicate::ICMP_EQ,
                                             ValStartLoc, F->getArg(0));
   Builder.CreateCondBr(ExecTraceCond, BBExecuteTrace, BBReturn);
@@ -220,8 +220,7 @@ void createControlPoint(Module &Mod, Function *F, std::vector<Value *> LiveVars,
 
   // Create block that decides when to stop tracing.
   Builder.SetInsertPoint(BBTracing);
-  Value *ValStartLoc2 =
-      Builder.CreateLoad(Type::getInt32Ty(Context), GVStartLoc);
+  Value *ValStartLoc2 = Builder.CreateLoad(YkLocTy, GVStartLoc);
   Value *StopTracingCond = Builder.CreateICmp(CmpInst::Predicate::ICMP_EQ,
                                               ValStartLoc2, F->getArg(0));
   Builder.CreateCondBr(StopTracingCond, BBStopTracing, BBReturn);
@@ -268,86 +267,103 @@ std::vector<Value *> getLiveVars(DominatorTree &DT, CallInst *OldCtrlPoint) {
   return Vec;
 }
 
-YkControlPointPass::YkControlPointPass() {}
-
-PreservedAnalyses YkControlPointPass::run(Module &M,
-                                          ModuleAnalysisManager &AM) {
-  LLVMContext &Context = M.getContext();
+namespace llvm {
+void initializeYkControlPointPass(PassRegistry &);
+}
 
-  // Locate the "dummy" control point provided by the user.
-  CallInst *OldCtrlPointCall = findControlPointCall(M);
-  if (OldCtrlPointCall == nullptr) {
-    Context.emitError("ykllvm couldn't find the call to `yk_control_point()`");
-    return PreservedAnalyses::all();
+namespace {
+class YkControlPoint : public ModulePass {
+public:
+  static char ID;
+  YkControlPoint() : ModulePass(ID) {
+    initializeYkControlPointPass(*PassRegistry::getPassRegistry());
   }
 
-  // Replace old control point call.
-  IRBuilder<> Builder(OldCtrlPointCall);
-
-  // Get function containing the control point.
-  Function *Caller = OldCtrlPointCall->getFunction();
-
-  // Find all live variables just before the call to the control point.
-  DominatorTree DT(*Caller);
-  std::vector<Value *> LiveVals = getLiveVars(DT, OldCtrlPointCall);
-  if (LiveVals.size() == 0) {
-    Context.emitError(
-        "The interpreter loop has no live variables!\n"
-        "ykllvm doesn't support this scenario, as such an interpreter would "
-        "make little sense.");
-    return PreservedAnalyses::all();
-  }
+  bool runOnModule(Module &M) override {
+    LLVMContext &Context = M.getContext();
 
-  // Generate the YkCtrlPointVars struct. This struct is used to package up a
-  // copy of all LLVM variables that are live just before the call to the
-  // control point. These are passed in to the patched control point so that
-  // they can be used as inputs and outputs to JITted trace code. The control
-  // point returns a new YkCtrlPointVars whose members may have been mutated
-  // by JITted trace code (if a trace was executed).
-  std::vector<Type *> TypeParams;
-  for (Value *V : LiveVals) {
-    TypeParams.push_back(V->getType());
-  }
-  StructType *CtrlPointReturnTy =
-      StructType::create(TypeParams, "YkCtrlPointVars");
-
-  // Create the new control point.
-  FunctionType *FType = FunctionType::get(
-      CtrlPointReturnTy, {Type::getInt32Ty(Context), CtrlPointReturnTy}, false);
-  Function *NF = Function::Create(FType, GlobalVariable::ExternalLinkage,
-                                  YK_NEW_CONTROL_POINT, M);
-
-  // Instantiate the YkCtrlPointStruct to pass in to the control point.
-  Value *InputStruct = cast<Value>(Constant::getNullValue(CtrlPointReturnTy));
-  unsigned LvIdx = 0;
-  for (Value *LV : LiveVals) {
-    InputStruct = Builder.CreateInsertValue(InputStruct, LV, LvIdx);
-    assert(LvIdx != UINT_MAX);
-    LvIdx++;
-  }
+    // Locate the "dummy" control point provided by the user.
+    CallInst *OldCtrlPointCall = findControlPointCall(M);
+    if (OldCtrlPointCall == nullptr) {
+      Context.emitError(
+          "ykllvm couldn't find the call to `yk_control_point()`");
+      return false;
+    }
 
-  // Insert call to the new control point.
-  CallInst *CtrlPointRet =
-      Builder.CreateCall(NF, {OldCtrlPointCall->getArgOperand(0), InputStruct});
-
-  // Once the control point returns we need to extract the (potentially
-  // mutated) values from the returned YkCtrlPointStruct and reassign them to
-  // their corresponding live variables. In LLVM IR we can do this by simply
-  // replacing all future references with the new values.
-  LvIdx = 0;
-  for (Value *LV : LiveVals) {
-    Value *New = Builder.CreateExtractValue(cast<Value>(CtrlPointRet), LvIdx);
-    LV->replaceUsesWithIf(
-        New, [&](Use &U) { return DT.dominates(CtrlPointRet, U); });
-    assert(LvIdx != UINT_MAX);
-    LvIdx++;
-  }
+    // Replace old control point call.
+    IRBuilder<> Builder(OldCtrlPointCall);
+
+    // Get function containing the control point.
+    Function *Caller = OldCtrlPointCall->getFunction();
+
+    // Find all live variables just before the call to the control point.
+    DominatorTree DT(*Caller);
+    std::vector<Value *> LiveVals = getLiveVars(DT, OldCtrlPointCall);
+    if (LiveVals.size() == 0) {
+      Context.emitError(
+          "The interpreter loop has no live variables!\n"
+          "ykllvm doesn't support this scenario, as such an interpreter would "
+          "make little sense.");
+      return false;
+    }
 
-  // Replace the call to the dummy control point.
-  OldCtrlPointCall->eraseFromParent();
+    // Generate the YkCtrlPointVars struct. This struct is used to package up a
+    // copy of all LLVM variables that are live just before the call to the
+    // control point. These are passed in to the patched control point so that
+    // they can be used as inputs and outputs to JITted trace code. The control
+    // point returns a new YkCtrlPointVars whose members may have been mutated
+    // by JITted trace code (if a trace was executed).
+    std::vector<Type *> TypeParams;
+    for (Value *V : LiveVals) {
+      TypeParams.push_back(V->getType());
+    }
+    StructType *CtrlPointReturnTy =
+        StructType::create(TypeParams, "YkCtrlPointVars");
+
+    // Create the new control point.
+    Type *YkLocTy = OldCtrlPointCall->getArgOperand(0)->getType();
+    FunctionType *FType = FunctionType::get(
+        CtrlPointReturnTy, {YkLocTy, CtrlPointReturnTy}, false);
+    Function *NF = Function::Create(FType, GlobalVariable::ExternalLinkage,
+                                    YK_NEW_CONTROL_POINT, M);
+
+    // Instantiate the YkCtrlPointStruct to pass in to the control point.
+    Value *InputStruct = cast<Value>(Constant::getNullValue(CtrlPointReturnTy));
+    unsigned LvIdx = 0;
+    for (Value *LV : LiveVals) {
+      InputStruct = Builder.CreateInsertValue(InputStruct, LV, LvIdx);
+      assert(LvIdx != UINT_MAX);
+      LvIdx++;
+    }
 
-  // Generate new control point logic.
-  createControlPoint(M, NF, LiveVals, CtrlPointReturnTy);
+    // Insert call to the new control point.
+    CallInst *CtrlPointRet = Builder.CreateCall(
+        NF, {OldCtrlPointCall->getArgOperand(0), InputStruct});
+
+    // Once the control point returns we need to extract the (potentially
+    // mutated) values from the returned YkCtrlPointStruct and reassign them to
+    // their corresponding live variables. In LLVM IR we can do this by simply
+    // replacing all future references with the new values.
+    LvIdx = 0;
+    for (Value *LV : LiveVals) {
+      Value *New = Builder.CreateExtractValue(cast<Value>(CtrlPointRet), LvIdx);
+      LV->replaceUsesWithIf(
+          New, [&](Use &U) { return DT.dominates(CtrlPointRet, U); });
+      assert(LvIdx != UINT_MAX);
+      LvIdx++;
+    }
 
-  return PreservedAnalyses::none();
-}
+    // Replace the call to the dummy control point.
+    OldCtrlPointCall->eraseFromParent();
+
+    // Generate new control point logic.
+    createControlPoint(M, NF, LiveVals, CtrlPointReturnTy, YkLocTy);
+    return true;
+  }
+};
+} // namespace
+
+char YkControlPoint::ID = 0;
+INITIALIZE_PASS(YkControlPoint, DEBUG_TYPE, "yk control point", false, false)
+
+ModulePass *llvm::createYkControlPointPass() { return new YkControlPoint(); }