Move the YkControlPoint pass later in the compilation pipeline.

This is required because one of the codegen passes, CodegenPrepare, can
mutate the IR and if we run before this pass, then our view of live
variables isn't true.

Author: vext01

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/ControlPoint.cpp

@@ -79,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>
@@ -266,87 +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");
+    // 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;
+    }
 
-  // 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++;
-  }
+    // 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;
+    }
 
-  // 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++;
-  }
+    // 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++;
+    }
+
+    // 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 the call to the dummy control point.
+    OldCtrlPointCall->eraseFromParent();
 
-  // Replace the call to the dummy control point.
-  OldCtrlPointCall->eraseFromParent();
+    // Generate new control point logic.
+    createControlPoint(M, NF, LiveVals, CtrlPointReturnTy, YkLocTy);
+    return true;
+  }
+};
+} // namespace
 
-  // Generate new control point logic.
-  createControlPoint(M, NF, LiveVals, CtrlPointReturnTy, YkLocTy);
+char YkControlPoint::ID = 0;
+INITIALIZE_PASS(YkControlPoint, DEBUG_TYPE, "yk control point", false, false)
 
-  return PreservedAnalyses::none();
-}
+ModulePass *llvm::createYkControlPointPass() { return new YkControlPoint(); }