Add ModuleClone Yk LLVM Pass.

Module cloning allows us to execute multiple versions of the same
program. Cloned functions can be optimized without affecting the
original ones used for tracing. This approach enables us to run
optimised code when tracing is not needed, while still preserving
tracing functionality.

llvm/include/llvm/InitializePasses.h

@@ -344,6 +344,7 @@ void initializeXRayInstrumentationPass(PassRegistry&);
 void initializeYkStackmapsPass(PassRegistry&);
 void initializeYkSplitBlocksAfterCallsPass(PassRegistry&);
 void initializeYkBasicBlockTracerPass(PassRegistry&);
+void initializeYkModuleClonePass(PassRegistry&);
 } // end namespace llvm
 
 #endif // LLVM_INITIALIZEPASSES_H

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

@@ -18,7 +18,7 @@
 #define CP_PPNAME "llvm.experimental.patchpoint.void"
 
 namespace llvm {
-ModulePass *createYkControlPointPass();
+ModulePass *createYkControlPointPass(uint64_t controlPointCount);
 } // namespace llvm
 
 #endif

llvm/include/llvm/Transforms/Yk/ModuleClone.h

@@ -0,0 +1,13 @@
+#ifndef LLVM_TRANSFORMS_YK_MODULE_CLONE_H
+#define LLVM_TRANSFORMS_YK_MODULE_CLONE_H
+
+#include "llvm/Pass.h"
+
+#define YK_CLONE_PREFIX "__yk_clone_"
+#define YK_CLONE_MODULE_CP_COUNT 2
+
+namespace llvm {
+ModulePass *createYkModuleClonePass();
+} // namespace llvm
+
+#endif

llvm/include/llvm/Transforms/Yk/Stackmaps.h

@@ -4,7 +4,7 @@
 #include "llvm/Pass.h"
 
 namespace llvm {
-ModulePass *createYkStackmapsPass();
+ModulePass *createYkStackmapsPass(uint64_t stackmapIDStart);
 } // namespace llvm
 
 #endif

llvm/lib/CodeGen/CodeGen.cpp

@@ -153,4 +153,5 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeYkStackmapsPass(Registry);
   initializeYkSplitBlocksAfterCallsPass(Registry);
   initializeYkBasicBlockTracerPass(Registry);
+  initializeYkModuleClonePass(Registry);
 }

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -56,6 +56,7 @@
 #include "llvm/Transforms/Yk/Stackmaps.h"
 #include "llvm/Transforms/Yk/NoCallsInEntryBlocks.h"
 #include "llvm/Transforms/Yk/BasicBlockTracer.h"
+#include "llvm/Transforms/Yk/ModuleClone.h"
 #include <cassert>
 #include <optional>
 #include <string>
@@ -297,6 +298,9 @@ static cl::opt<bool>
     YkBasicBlockTracer("yk-basicblock-tracer", cl::init(false), cl::NotHidden,
                       cl::desc("Enables YK Software Tracer capability"));
 
+static cl::opt<bool>
+    YkModuleClone("yk-module-clone", cl::init(false), cl::NotHidden,
+                  cl::desc("Enables YK Module Cloning capability"));
 
 /// Allow standard passes to be disabled by command line options. This supports
 /// simple binary flags that either suppress the pass or do nothing.
@@ -1139,6 +1143,15 @@ bool TargetPassConfig::addISelPasses() {
   addIRPasses();
   addCodeGenPrepare();
   addPassesToHandleExceptions();
+
+  // Default number of control points in a module.
+  int numberOfControlPoints = 1;
+
+  if (YkModuleClone) {
+    numberOfControlPoints = YK_CLONE_MODULE_CP_COUNT;
+    addPass(createYkModuleClonePass());
+  }
+
   if (YkBlockDisambiguate)
     addPass(createYkBlockDisambiguatePass());
 
@@ -1166,7 +1179,7 @@ bool TargetPassConfig::addISelPasses() {
   // decomposed into scalar arguments. The JIT runtime relies on the interface
   // *not* being changed.
   if (YkPatchCtrlPoint) {
-    addPass(createYkControlPointPass());
+    addPass(createYkControlPointPass(numberOfControlPoints));
   }
 
   if (YkLinkage) {
@@ -1184,12 +1197,12 @@ bool TargetPassConfig::addISelPasses() {
   }
 
   if (YkInsertStackMaps) {
-    addPass(createYkStackmapsPass());
+    addPass(createYkStackmapsPass(numberOfControlPoints));
   }
 
   if (YkBasicBlockTracer) {
     addPass(createYkBasicBlockTracerPass());
-  }
+  }  
 
   addISelPrepare();
   return addCoreISelPasses();

llvm/lib/Transforms/Yk/BasicBlockTracer.cpp

@@ -8,6 +8,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
+#include "llvm/Transforms/Yk/ModuleClone.h"
 
 #define DEBUG_TYPE "yk-basicblock-tracer-pass"
 
@@ -42,6 +43,9 @@ struct YkBasicBlockTracer : public ModulePass {
     uint32_t FunctionIndex = 0;
     for (auto &F : M) {
       uint32_t BlockIndex = 0;
+      if (F.getName().startswith(YK_CLONE_PREFIX)) {
+        continue;
+      }
       for (auto &BB : F) {
         builder.SetInsertPoint(&*BB.getFirstInsertionPt());
         builder.CreateCall(TraceFunc, {builder.getInt32(FunctionIndex),

llvm/lib/Transforms/Yk/CMakeLists.txt

@@ -8,6 +8,7 @@ add_llvm_component_library(LLVMYkPasses
   NoCallsInEntryBlocks.cpp
   SplitBlocksAfterCalls.cpp
   BasicBlockTracer.cpp
+  ModuleClone.cpp
 
   DEPENDS
   intrinsics_gen
@@ -16,4 +17,6 @@ add_llvm_component_library(LLVMYkPasses
   Analysis
   Core
   Support
+  TransformUtils # Module Clone
+  Linker         # Module Linker
 )

llvm/lib/Transforms/Yk/ControlPoint.cpp

@@ -61,11 +61,6 @@
 #define YK_CONTROL_POINT_ARG_VARS_IDX 2
 #define YK_CONTROL_POINT_NUM_ARGS 3
 
-// Stackmap ID zero is reserved for the control point.
-//
-// This will need to change when we support >1 control point.
-const unsigned CPStackMapID = 0;
-
 // The number of shadow bytes required for the control point's patchpoint.
 //
 // This must be large enough to accommodate the call to patchpoint target
@@ -86,22 +81,28 @@ const unsigned CPShadow = 13;
 
 using namespace llvm;
 
-/// Find the call to the dummy control point that we want to patch.
-/// Returns either a pointer the call instruction, or `nullptr` if the call
-/// could not be found.
-/// YKFIXME: For now assumes there's only one control point.
-CallInst *findControlPointCall(Module &M) {
+/// @brief Locate calls to the dummy control point that we want to patch.
+///
+/// This function searches for all instances of `YK_DUMMY_CONTROL_POINT`
+/// calls within the LLVM module.
+///
+/// @return A vector of pointers to the `YK_DUMMY_CONTROL_POINT` call
+/// instructions, or an empty vector if no calls are found.
+std::vector<CallInst *> findControlPointCalls(Module &M) {
+  std::vector<CallInst *> controlPointCalls;
+
   // Find the declaration of `yk_mt_control_point()`.
   Function *CtrlPoint = M.getFunction(YK_DUMMY_CONTROL_POINT);
   if (CtrlPoint == nullptr)
-    return nullptr;
+    return controlPointCalls;
 
-  // Find the call site of `yk_mt_control_point()`.
-  const Value::user_iterator U = CtrlPoint->user_begin();
-  if (U == CtrlPoint->user_end())
-    return nullptr;
-
-  return cast<CallInst>(*U);
+  // Find all call sites of `yk_mt_control_point()`.
+  for (User *U : CtrlPoint->users()) {
+    if (CallInst *CI = dyn_cast<CallInst>(U)) {
+      controlPointCalls.insert(controlPointCalls.begin(), CI);
+    }
+  }
+  return controlPointCalls;
 }
 
 namespace llvm {
@@ -110,93 +111,109 @@ void initializeYkControlPointPass(PassRegistry &);
 
 namespace {
 class YkControlPoint : public ModulePass {
+private:
+  uint64_t controlPointCount;
+
 public:
   static char ID;
-  YkControlPoint() : ModulePass(ID) {
+  YkControlPoint(uint64_t controlPointCount)
+      : ModulePass(ID), controlPointCount(controlPointCount) {
     initializeYkControlPointPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnModule(Module &M) override {
     LLVMContext &Context = M.getContext();
+    std::vector<CallInst *> ControlPointCalls = findControlPointCalls(M);
+
     // Locate the "dummy" control point provided by the user.
-    CallInst *OldCtrlPointCall = findControlPointCall(M);
-    if (OldCtrlPointCall == nullptr) {
+    if (ControlPointCalls.empty()) {
       // This program doesn't have a control point. We can't do any
       // transformations on it, but we do still want to compile it.
       Context.diagnose(DiagnosticInfoInlineAsm(
           "ykllvm couldn't find the call to `yk_mt_control_point()`",
           DS_Warning));
       return false;
     }
-
-    // Get function containing the control point.
-    Function *Caller = OldCtrlPointCall->getFunction();
-
-    // Check that the control point is inside a loop.
-    DominatorTree DT(*Caller);
-    const LoopInfo Loops(DT);
-    if (!std::any_of(Loops.begin(), Loops.end(), [OldCtrlPointCall](Loop *L) {
-          return L->contains(OldCtrlPointCall);
-        })) {
-      ;
-      Context.emitError("yk_mt_control_point() must be called inside a loop.");
-      return false;
-    }
-
-    // The old control point should be of the form:
-    //    yk_mt_control_point(YkMT*, YkLocation*)
-    assert(OldCtrlPointCall->arg_size() == YK_OLD_CONTROL_POINT_NUM_ARGS);
-    Type *YkMTTy =
-        OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_MT_IDX)->getType();
-    Type *YkLocTy =
-        OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_LOC_IDX)
-            ->getType();
-
-    // Create a call to the "new" (patched) control point, but do so via a
-    // patchpoint so that we can capture the live variables at exactly the
-    // moment before the call.
-    Type *Int64Ty = Type::getInt64Ty(Context);
-    FunctionType *FType = FunctionType::get(Type::getVoidTy(Context),
-                                            {YkMTTy, YkLocTy, Int64Ty}, false);
-    Function *NF = Function::Create(FType, GlobalVariable::ExternalLinkage,
-                                    YK_NEW_CONTROL_POINT, M);
-
-    IRBuilder<> Builder(OldCtrlPointCall);
-
-    const Intrinsic::ID SMFuncID = Function::lookupIntrinsicID(CP_PPNAME);
-    if (SMFuncID == Intrinsic::not_intrinsic) {
-      Context.emitError("can't find stackmap()");
-      return false;
-    }
-    Function *SMFunc = Intrinsic::getDeclaration(&M, SMFuncID);
-    assert(SMFunc != nullptr);
-
-    // Get live variables.
-    LivenessAnalysis LA(Caller);
-    auto Lives = LA.getLiveVarsBefore(OldCtrlPointCall);
-
-    Value *SMID = ConstantInt::get(Type::getInt64Ty(Context), CPStackMapID);
-    Value *Shadow = ConstantInt::get(Type::getInt32Ty(Context), CPShadow);
-    std::vector<Value *> Args = {
-        SMID,
-        Shadow,
-        NF,
-        ConstantInt::get(Type::getInt32Ty(Context), 3),
-        OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_MT_IDX),
-        OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_LOC_IDX),
-        SMID,
-    };
-
-    for (auto *Live : Lives) {
-      Args.push_back(Live);
+    assert(ControlPointCalls.size() == controlPointCount &&
+           "Unexpected number of control point calls");
+
+    unsigned CPStackMapID = 0;
+    Function *NF = nullptr;
+
+    for (CallInst *OldCtrlPointCall : ControlPointCalls) {
+      // Get function containing the control point.
+      Function *Caller = OldCtrlPointCall->getFunction();
+
+      // Check that the control point is inside a loop.
+      DominatorTree DT(*Caller);
+      const LoopInfo Loops(DT);
+      if (!std::any_of(Loops.begin(), Loops.end(), [OldCtrlPointCall](Loop *L) {
+            return L->contains(OldCtrlPointCall);
+          })) {
+        ;
+        Context.emitError(
+            "yk_mt_control_point() must be called inside a loop.");
+        return false;
+      }
+
+      // The old control point should be of the form:
+      //    yk_mt_control_point(YkMT*, YkLocation*)
+      assert(OldCtrlPointCall->arg_size() == YK_OLD_CONTROL_POINT_NUM_ARGS);
+      Type *YkMTTy =
+          OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_MT_IDX)
+              ->getType();
+      Type *YkLocTy =
+          OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_LOC_IDX)
+              ->getType();
+
+      // Create a call to the "new" (patched) control point, but do so via a
+      // patchpoint so that we can capture the live variables at exactly the
+      // moment before the call.
+      if (NF == nullptr) {
+        Type *Int64Ty = Type::getInt64Ty(Context);
+        FunctionType *FType = FunctionType::get(
+            Type::getVoidTy(Context), {YkMTTy, YkLocTy, Int64Ty}, false);
+        NF = Function::Create(FType, GlobalVariable::ExternalLinkage,
+                              YK_NEW_CONTROL_POINT, M);
+      }
+
+      IRBuilder<> Builder(OldCtrlPointCall);
+
+      const Intrinsic::ID SMFuncID = Function::lookupIntrinsicID(CP_PPNAME);
+      if (SMFuncID == Intrinsic::not_intrinsic) {
+        Context.emitError("can't find stackmap()");
+        return false;
+      }
+      Function *SMFunc = Intrinsic::getDeclaration(&M, SMFuncID);
+      assert(SMFunc != nullptr);
+
+      // Get live variables.
+      LivenessAnalysis LA(Caller);
+      auto Lives = LA.getLiveVarsBefore(OldCtrlPointCall);
+
+      Value *SMID = ConstantInt::get(Type::getInt64Ty(Context), CPStackMapID);
+      Value *Shadow = ConstantInt::get(Type::getInt32Ty(Context), CPShadow);
+      std::vector<Value *> Args = {
+          SMID,
+          Shadow,
+          NF,
+          ConstantInt::get(Type::getInt32Ty(Context), 3),
+          OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_MT_IDX),
+          OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_LOC_IDX),
+          SMID,
+      };
+
+      for (auto *Live : Lives) {
+        Args.push_back(Live);
+      }
+
+      Builder.CreateCall(SMFunc->getFunctionType(), SMFunc,
+                         ArrayRef<Value *>(Args));
+
+      // Replace the call to the dummy control point.
+      OldCtrlPointCall->eraseFromParent();
+      ++CPStackMapID;
     }
-
-    Builder.CreateCall(SMFunc->getFunctionType(), SMFunc,
-                       ArrayRef<Value *>(Args));
-
-    // Replace the call to the dummy control point.
-    OldCtrlPointCall->eraseFromParent();
-
 #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.
@@ -213,4 +230,6 @@ class YkControlPoint : public ModulePass {
 char YkControlPoint::ID = 0;
 INITIALIZE_PASS(YkControlPoint, DEBUG_TYPE, "yk control point", false, false)
 
-ModulePass *llvm::createYkControlPointPass() { return new YkControlPoint(); }
+ModulePass *llvm::createYkControlPointPass(uint64_t controlPointCount) {
+  return new YkControlPoint(controlPointCount);
+}