Local changes

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/ModuleClone.h

@@ -0,0 +1,12 @@
+#ifndef LLVM_TRANSFORMS_YK_MODULE_CLONE_H
+#define LLVM_TRANSFORMS_YK_MODULE_CLONE_H
+
+#include "llvm/Pass.h"
+
+#define YK_CLONE_PREFIX "__yk_clone_"
+
+namespace llvm {
+ModulePass *createYkModuleClonePass();
+} // 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.
@@ -1187,9 +1191,16 @@ bool TargetPassConfig::addISelPasses() {
     addPass(createYkStackmapsPass());
   }
 
+  // YkModuleClone Needs to run before YkBasicBlockTracerPass cause 
+  // YkBasicBlockTracerPass skips cloned functions.
+  if (YkModuleClone) {
+    addPass(createYkModuleClonePass());
+  }
+
   if (YkBasicBlockTracer) {
     addPass(createYkBasicBlockTracerPass());
   }
+
 
   addISelPrepare();
   return addCoreISelPasses();

llvm/lib/Transforms/Yk/BasicBlockTracer.cpp

@@ -1,4 +1,5 @@
 #include "llvm/Transforms/Yk/BasicBlockTracer.h"
+#include "llvm/Transforms/Yk/ModuleClone.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
@@ -42,6 +43,12 @@ struct YkBasicBlockTracer : public ModulePass {
     uint32_t FunctionIndex = 0;
     for (auto &F : M) {
       uint32_t BlockIndex = 0;
+      // FIXME: This cause the ykjit to crash when executing compiled
+      // trace at: https://github.com/ykjit/yk/blob/master/ykrt/src/mt.rs#L366
+      // Skip cloned functions
+      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/ModuleClone.cpp

@@ -0,0 +1,83 @@
+#include "llvm/Transforms/Yk/ModuleClone.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Linker/Linker.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/IR/Verifier.h"
+
+#define DEBUG_TYPE "yk-module-clone-pass"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeYkModuleClonePass(PassRegistry &);
+} // namespace llvm
+
+namespace {
+struct YkModuleClone : public ModulePass {
+  static char ID;
+
+  YkModuleClone() : ModulePass(ID) {
+    initializeYkModuleClonePass(*PassRegistry::getPassRegistry());
+  }
+  void updateClonedFunctions(Module &M) {
+    for (llvm::Function &F : M) {
+      if (F.hasExternalLinkage() && F.isDeclaration()) {
+        continue;
+      }
+      // Skip functions that are address taken
+      if (!F.hasAddressTaken()) {
+        F.setName(Twine(YK_CLONE_PREFIX) + F.getName());
+      }
+    }
+  }
+
+  // This pass duplicates functions within the module:
+  // the original and the cloned version.
+  // - If a function's address is not taken, it is cloned, resulting in 
+  //   two versions of that function in the module.
+  // - If a function's address is taken, the function remains in its 
+  //   original form.
+  bool runOnModule(Module &M) override {
+    std::unique_ptr<Module> Cloned = CloneModule(M);
+    if (!Cloned) {
+      errs() << "Error: CloneModule returned null for module: " << M.getName()
+             << "\n";
+      return false;
+    }
+    updateClonedFunctions(*Cloned);
+
+    // The `OverrideFromSrc` flag instructs the linker to prioritize
+    // definitions from the source module (the second argument) when
+    // conflicts arise. This means that if two functions 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");
+      return false;
+    }
+
+    if (verifyModule(M, &errs())) {
+      errs() << "Module verification failed!";
+      return false;
+    }
+
+    return true;
+  }
+};
+} // namespace
+
+char YkModuleClone::ID = 0;
+
+INITIALIZE_PASS(YkModuleClone, DEBUG_TYPE, "yk module clone", false, false)
+
+ModulePass *llvm::createYkModuleClonePass() { return new YkModuleClone(); }

llvm/lib/YkIR/CMakeLists.txt

@@ -5,4 +5,6 @@ add_llvm_component_library(LLVMYkIR
   Core
   MC
   Support
+  TransformUtils
+  Linker
   )

llvm/test/Transforms/Yk/ModuleClone.ll

@@ -0,0 +1,116 @@
+; RUN: llc -stop-after=yk-basicblock-tracer-pass --yk-module-clone --yk-basicblock-tracer < %s  | FileCheck %s
+
+source_filename = "ModuleClone.c"
+target triple = "x86_64-pc-linux-gnu"
+
+@.str = private unnamed_addr constant [13 x i8] c"Hello, world\00", align 1
+@my_global = global i32 42, align 4
+
+declare i32 @printf(ptr, ...)
+declare dso_local ptr @yk_mt_new(ptr noundef)
+declare dso_local void @yk_mt_hot_threshold_set(ptr noundef, i32 noundef)
+declare dso_local i64 @yk_location_new()
+declare dso_local void @yk_mt_control_point(ptr noundef, ptr noundef)
+declare dso_local i32 @fprintf(ptr noundef, ptr noundef, ...)
+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:
+  %0 = add i32 %x, 1
+  ret i32 %0
+}
+
+define dso_local i32 @my_func(i32 %x) {
+entry:
+  %0 = add i32 %x, 1
+  %func_ptr = alloca ptr, align 8
+  store ptr @func_inc_with_address_taken, ptr %func_ptr, align 8
+  %1 = load ptr, ptr %func_ptr, align 8
+  %2 = call i32 %1(i32 42)
+  ret i32 %2
+}
+
+define dso_local i32 @main() {
+entry:
+  %0 = call i32 @my_func(i32 10)
+  %1 = load i32, ptr @my_global
+  %2 = call i32 (ptr, ...) @printf(ptr @.str, i32 %1)
+  ret i32 0
+}
+
+; ======================================================================
+; Original functions - should have trace calls
+; ======================================================================
+; File header checks
+; CHECK: source_filename = "ModuleClone.c"
+; CHECK: target triple = "x86_64-pc-linux-gnu"
+
+; Global variable and string checks
+; CHECK: @.str = private unnamed_addr constant [13 x i8] c"Hello, world\00", align 1
+; CHECK: @my_global = global i32 42, align 4
+
+; Declaration checks
+; CHECK: declare i32 @printf(ptr, ...)
+; CHECK: declare dso_local ptr @yk_mt_new(ptr noundef)
+; CHECK: declare dso_local void @yk_mt_hot_threshold_set(ptr noundef, i32 noundef)
+; CHECK: declare dso_local i64 @yk_location_new()
+; CHECK: declare dso_local void @yk_mt_control_point(ptr noundef, ptr noundef)
+; CHECK: declare dso_local i32 @fprintf(ptr noundef, ptr noundef, ...)
+; CHECK: declare dso_local void @yk_location_drop(i64)
+; CHECK: declare dso_local void @yk_mt_shutdown(ptr noundef)
+
+; Check original function: my_func
+; CHECK-LABEL: define dso_local i32 @my_func(i32 %x)
+; CHECK-NEXT: entry:
+; CHECK-NEXT: call void @__yk_trace_basicblock({{.*}})
+; CHECK-NEXT: %0 = add i32 %x, 1
+; CHECK-NEXT: %func_ptr = alloca ptr, align 8
+; CHECK-NEXT: store ptr @func_inc_with_address_taken, ptr %func_ptr, align 8
+; CHECK-NEXT: %1 = load ptr, ptr %func_ptr, align 8
+; CHECK-NEXT: %2 = call i32 %1(i32 42)
+; CHECK-NEXT: ret i32 %2
+
+; Check original function: main
+; CHECK-LABEL: define dso_local i32 @main()
+; CHECK-NEXT: entry:
+; CHECK-NEXT: call void @__yk_trace_basicblock({{.*}})
+; 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
+
+; ======================================================================
+; Functions whose addresses are 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
+
+; ======================================================================
+; 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-NEXT: entry:
+; CHECK-NOT: call void @__yk_trace_basicblock({{.*}})
+; CHECK-NEXT: %0 = add i32 %x, 1
+; CHECK-NEXT: %func_ptr = alloca ptr, align 8
+; CHECK-NEXT: store ptr @func_inc_with_address_taken, ptr %func_ptr, align 8
+; CHECK-NEXT: %1 = load ptr, ptr %func_ptr, align 8
+; 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-NEXT: entry:
+; CHECK-NOT: call void @__yk_trace_basicblock({{.*}})
+; CHECK-NEXT: %0 = call i32 @__yk_clone_my_func(i32 10)
+; CHECK-NEXT: %1 = load i32, ptr @my_global
+; CHECK-NEXT: %2 = call i32 (ptr, ...) @printf