Merge pull request #216 from Pavel-Durov/fix-function-reference-yk-mo…

…dule-clone-pass

Make sure that unoptimised functions calls ONLY unoptimised functions.

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

@@ -3,7 +3,7 @@
 
 #include "llvm/Pass.h"
 
-#define YK_CLONE_PREFIX "__yk_clone_"
+#define YK_CLONE_PREFIX "__yk_opt_"
 #define YK_CLONE_MODULE_CP_COUNT 2
 
 namespace llvm {

llvm/lib/Transforms/Yk/ControlPoint.cpp

@@ -134,9 +134,13 @@ class YkControlPoint : public ModulePass {
           DS_Warning));
       return false;
     }
-    assert(ControlPointCalls.size() == controlPointCount &&
-           "Unexpected number of control point calls");
 
+    if (ControlPointCalls.size() != controlPointCount) {
+      Context.emitError("Unexpected number of control point calls: " +
+                        Twine(ControlPointCalls.size()) +
+                        " expected: " + Twine(controlPointCount));
+      return false;
+    }
     unsigned CPStackMapID = 0;
     Function *NF = nullptr;
 

llvm/lib/Transforms/Yk/ModuleClone.cpp

@@ -1,7 +1,8 @@
 //===- ModuleClone.cpp - Yk Module Cloning Pass ---------------------------===//
 //
-// This pass duplicates functions within the module, producing both the original
-// and new versions of functions. the process is as follows:
+// This pass duplicates functions within the module, producing both the
+// original (unoptimised) and cloned (optimised) versions of these
+// functions. The process is as follows:
 //
 // - **Cloning Criteria:**
 //   - Functions **without** their address taken are cloned. This results in two
@@ -11,18 +12,12 @@
 //     and are not cloned.
 //
 // - **Cloned Function Naming:**
-//   - The cloned functions are renamed by adding the prefix `__yk_clone_` to
+//   - The cloned functions are renamed by adding the prefix `__yk_opt_` to
 //     their original names. This distinguishes them from the original
 //     functions.
 //
-// - **Clone Process:**
-//   - 1. The original module is cloned, creating two copies of the module.
-//   - 2. The functions in the cloned module that satisfy the cloning criteria
-//     are renamed.
-//   - 3. The cloned module is linked back into the original module.
-//
 // - **Optimisation Intent:**
-//   - The **cloned functions** (those with the `__yk_clone_` prefix) are
+//   - The **cloned functions** (those with the `__yk_opt_` prefix) are
 //     intended to be the **optimised versions** of the functions.
 //   - The **original functions** remain **unoptimised**.
 //
@@ -58,39 +53,96 @@ struct YkModuleClone : public ModulePass {
   YkModuleClone() : ModulePass(ID) {
     initializeYkModuleClonePass(*PassRegistry::getPassRegistry());
   }
-  void updateClonedFunctions(Module &M) {
+
+  /**
+   * @brief Clones eligible functions within the given module.
+   *
+   * This function iterates over all functions in the provided LLVM module `M`
+   * and clones those that meet the following criteria:
+   *
+   * - The function does **not** have external linkage and is **not** a
+   * declaration.
+   * - The function's address is **not** taken.
+   *
+   * @param M The LLVM module containing the functions to be cloned.
+   * @return A map where the keys are the original function names and the
+   *         values are pointers to the cloned `Function` objects. Returns
+   *         a map if cloning succeeds, or `nullopt` if cloning fails.
+   */
+  std::optional<std::map<std::string, Function *>>
+  cloneFunctionsInModule(Module &M) {
+    LLVMContext &Context = M.getContext();
+    std::map<std::string, Function *> ClonedFuncs;
     for (llvm::Function &F : M) {
+      // Skip external declarations.
       if (F.hasExternalLinkage() && F.isDeclaration()) {
         continue;
       }
-      // Skip functions that are address taken
-      if (!F.hasAddressTaken()) {
-        F.setName(Twine(YK_CLONE_PREFIX) + F.getName());
+      // Skip already cloned functions or functions with address taken.
+      if (F.hasAddressTaken() || F.getName().startswith(YK_CLONE_PREFIX)) {
+        continue;
+      }
+      ValueToValueMapTy VMap;
+      Function *ClonedFunc = CloneFunction(&F, VMap);
+      if (ClonedFunc == nullptr) {
+        Context.emitError("Failed to clone function: " + F.getName());
+        return std::nullopt;
+      }
+      // Copy arguments
+      auto DestArgIt = ClonedFunc->arg_begin();
+      for (const Argument &OrigArg : F.args()) {
+        DestArgIt->setName(OrigArg.getName());
+        VMap[&OrigArg] = &*DestArgIt++;
       }
+      // Rename function
+      auto originalName = F.getName().str();
+      auto cloneName = Twine(YK_CLONE_PREFIX) + originalName;
+      ClonedFunc->setName(cloneName);
+      ClonedFuncs[originalName] = ClonedFunc;
     }
+    return ClonedFuncs;
   }
 
-  bool runOnModule(Module &M) override {
-    std::unique_ptr<Module> Cloned = CloneModule(M);
-    if (!Cloned) {
-      llvm::report_fatal_error("Error cloning the module");
-      return false;
+  /**
+   * @brief Updates call instructions in cloned functions to reference
+   * other cloned functions.
+   *
+   * @param M The LLVM module containing the functions.
+   * @param ClonedFuncs A map of cloned function names to functions.
+   */
+  void
+  updateClonedFunctionCalls(Module &M,
+                            std::map<std::string, Function *> &ClonedFuncs) {
+    for (auto &Entry : ClonedFuncs) {
+      Function *ClonedFunc = Entry.second;
+      for (BasicBlock &BB : *ClonedFunc) {
+        for (Instruction &I : BB) {
+          if (auto *Call = dyn_cast<CallInst>(&I)) {
+            Function *CalledFunc = Call->getCalledFunction();
+            if (CalledFunc && !CalledFunc->isIntrinsic()) {
+              std::string CalledName = CalledFunc->getName().str();
+              auto It = ClonedFuncs.find(CalledName);
+              if (It != ClonedFuncs.end()) {
+                Call->setCalledFunction(It->second);
+              }
+            }
+          }
+        }
+      }
     }
-    updateClonedFunctions(*Cloned);
+  }
 
-    // The `OverrideFromSrc` flag instructs the linker to prioritise
-    // definitions from the source module (the second argument) when
-    // conflicts arise. This means that if two global variables, functions,
-    // or constants have the same name in both the original and cloned modules,
-    // the version from the cloned module will overwrite the original.
-    if (Linker::linkModules(M, std::move(Cloned),
-                            Linker::Flags::OverrideFromSrc)) {
-      llvm::report_fatal_error("Error linking the modules");
+  bool runOnModule(Module &M) override {
+    LLVMContext &Context = M.getContext();
+    auto clonedFunctions = cloneFunctionsInModule(M);
+    if (!clonedFunctions) {
+      Context.emitError("Failed to clone functions in module");
       return false;
     }
+    updateClonedFunctionCalls(M, *clonedFunctions);
 
     if (verifyModule(M, &errs())) {
-      errs() << "Module verification failed!";
+      Context.emitError("Module verification failed!");
       return false;
     }
 

llvm/test/Transforms/Yk/ModuleClone.ll

@@ -16,6 +16,12 @@ declare dso_local void @yk_location_drop(i64)
 declare dso_local void @yk_mt_shutdown(ptr noundef)
 
 define dso_local i32 @func_inc_with_address_taken(i32 %x) {
+entry:
+  %call = call i32 @inc(i32 %x)
+  ret i32 %call
+}
+
+define dso_local i32 @inc(i32 %x) {
 entry:
   %0 = add i32 %x, 1
   ret i32 %0
@@ -60,6 +66,20 @@ entry:
 ; CHECK: declare dso_local void @yk_location_drop(i64)
 ; CHECK: declare dso_local void @yk_mt_shutdown(ptr noundef)
 
+; Check that func_inc_with_address_taken is present in its original form
+; CHECK-LABEL: define dso_local i32 @func_inc_with_address_taken(i32 %x)
+; CHECK-NEXT: entry:
+; CHECK-NEXT: call void @__yk_trace_basicblock({{.*}})
+; CHECK-NEXT: %call = call i32 @inc(i32 %x)
+; CHECK-NEXT: ret i32 %call
+
+; Check original function: inc
+; CHECK-LABEL: define dso_local i32 @inc(i32 %x)
+; CHECK-NEXT: entry:
+; CHECK-NEXT: call void @__yk_trace_basicblock({{.*}})
+; CHECK-NEXT: %0 = add i32 %x, 1
+; CHECK-NEXT: ret i32 %0
+
 ; Check original function: my_func
 ; CHECK-LABEL: define dso_local i32 @my_func(i32 %x)
 ; CHECK-NEXT: entry:
@@ -78,26 +98,27 @@ entry:
 ; CHECK-NEXT: %0 = call i32 @my_func(i32 10)
 ; CHECK-NEXT: %1 = load i32, ptr @my_global
 ; CHECK-NEXT: %2 = call i32 (ptr, ...) @printf
-
-; Check that func_inc_with_address_taken is present in its original form
-; CHECK-LABEL: define dso_local i32 @func_inc_with_address_taken(i32 %x)
-; CHECK-NEXT: entry:
-; CHECK-NEXT: call void @__yk_trace_basicblock({{.*}})
-; CHECK-NEXT: %0 = add i32 %x, 1
-; CHECK-NEXT: ret i32 %0
+; CHECK-NEXT: ret i32 0
 
 ; ======================================================================
 ; Functions whose addresses are taken should not be cloned
 ; ======================================================================
-; Functions with their addresses taken should not be cloned. 
+; Functions with their addresses taken should not be cloned.
 ; `func_inc_with_address_taken` is used by pointer and thus remains unaltered.
-; CHECK-NOT: define dso_local i32 @__yk_clone_func_inc_with_address_taken
+; CHECK-NOT: define dso_local i32 @__yk_opt_func_inc_with_address_taken
 
 ; ======================================================================
 ; Cloned functions - should have no trace calls
 ; ======================================================================
-; Check cloned function: __yk_clone_my_func
-; CHECK-LABEL: define dso_local i32 @__yk_clone_my_func(i32 %x)
+; Check cloned function: __yk_opt_inc
+; CHECK-LABEL: define dso_local i32 @__yk_opt_inc(i32 %x)
+; CHECK-NEXT: entry:
+; CHECK-NOT: call void @yk_trace_basicblock({{.*}})
+; CHECK-NEXT: %0 = add i32 %x, 1
+; CHECK-NEXT: ret i32 %0
+
+; Check cloned function: __yk_opt_my_func
+; CHECK-LABEL: define dso_local i32 @__yk_opt_my_func(i32 %x)
 ; CHECK-NEXT: entry:
 ; CHECK-NOT: call void @__yk_trace_basicblock({{.*}})
 ; CHECK-NEXT: %0 = add i32 %x, 1
@@ -107,10 +128,10 @@ entry:
 ; CHECK-NEXT: %2 = call i32 %1(i32 42)
 ; CHECK-NEXT: ret i32 %2
 
-; Check cloned function: __yk_clone_main
-; CHECK-LABEL: define dso_local i32 @__yk_clone_main()
+; Check cloned function: __yk_opt_main
+; CHECK-LABEL: define dso_local i32 @__yk_opt_main()
 ; CHECK-NEXT: entry:
 ; CHECK-NOT: call void @__yk_trace_basicblock({{.*}})
-; CHECK-NEXT: %0 = call i32 @__yk_clone_my_func(i32 10)
+; CHECK-NEXT: %0 = call i32 @__yk_opt_my_func(i32 10)
 ; CHECK-NEXT: %1 = load i32, ptr @my_global
 ; CHECK-NEXT: %2 = call i32 (ptr, ...) @printf