Merge pull request #202 from Pavel-Durov/clone-llvm-module-pass

Module clone pass.

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);
+}