From bdb2a25e8d7ee153d47e5e9328469a703e219e9a Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sat, 14 Sep 2024 15:29:19 +0100
Subject: [PATCH 01/22] Add module clone pass.

---
 llvm/include/llvm/InitializePasses.h          |  1 +
 llvm/include/llvm/Transforms/Yk/ModuleClone.h | 10 +++
 llvm/lib/CodeGen/CodeGen.cpp                  |  1 +
 llvm/lib/CodeGen/TargetPassConfig.cpp         |  7 ++
 llvm/lib/Transforms/Yk/CMakeLists.txt         |  3 +
 llvm/lib/Transforms/Yk/ModuleClone.cpp        | 78 +++++++++++++++++++
 llvm/lib/YkIR/CMakeLists.txt                  |  2 +
 7 files changed, 102 insertions(+)
 create mode 100644 llvm/include/llvm/Transforms/Yk/ModuleClone.h
 create mode 100644 llvm/lib/Transforms/Yk/ModuleClone.cpp

diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index bab49b7675512f9..fc3d06cf6034325 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/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
diff --git a/llvm/include/llvm/Transforms/Yk/ModuleClone.h b/llvm/include/llvm/Transforms/Yk/ModuleClone.h
new file mode 100644
index 000000000000000..4a6bdcf65db2d39
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Yk/ModuleClone.h
@@ -0,0 +1,10 @@
+#ifndef LLVM_TRANSFORMS_YK_MODULE_CLONE_H
+#define LLVM_TRANSFORMS_YK_MODULE_CLONE_H
+
+#include "llvm/Pass.h"
+
+namespace llvm {
+ModulePass *createYkModuleClonePass();
+} // namespace llvm
+
+#endif
diff --git a/llvm/lib/CodeGen/CodeGen.cpp b/llvm/lib/CodeGen/CodeGen.cpp
index 77d7b6fefa7e394..eb51081c6ceaee5 100644
--- a/llvm/lib/CodeGen/CodeGen.cpp
+++ b/llvm/lib/CodeGen/CodeGen.cpp
@@ -153,4 +153,5 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeYkStackmapsPass(Registry);
   initializeYkSplitBlocksAfterCallsPass(Registry);
   initializeYkBasicBlockTracerPass(Registry);
+  initializeYkModuleClonePass(Registry);
 }
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
index 2fcfa3cb59180c9..bca7d4b6316657f 100644
--- a/llvm/lib/CodeGen/TargetPassConfig.cpp
+++ b/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.
@@ -1190,6 +1194,9 @@ bool TargetPassConfig::addISelPasses() {
   if (YkBasicBlockTracer) {
     addPass(createYkBasicBlockTracerPass());
   }
+  if (YkModuleClone){
+    addPass(createYkModuleClonePass());
+  }
 
   addISelPrepare();
   return addCoreISelPasses();
diff --git a/llvm/lib/Transforms/Yk/CMakeLists.txt b/llvm/lib/Transforms/Yk/CMakeLists.txt
index 3e5fba85082c295..7911cad4dfc5925 100644
--- a/llvm/lib/Transforms/Yk/CMakeLists.txt
+++ b/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
 )
diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
new file mode 100644
index 000000000000000..1cac96b7cb9773c
--- /dev/null
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -0,0 +1,78 @@
+#include "llvm/Transforms/Yk/ModuleClone.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"
+
+#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;
+  const std::string ClonePrefix = "__yk_clone_";
+
+  YkModuleClone() : ModulePass(ID) {
+    initializeYkModuleClonePass(*PassRegistry::getPassRegistry());
+  }
+
+  void updateClonedFunctions(Module &M) {
+    for (llvm::Function &F : M) {
+      if (F.hasExternalLinkage() && F.isDeclaration()) {
+        continue;
+      }
+      F.setName(ClonePrefix + F.getName().str());
+    }
+  }
+
+  void updateClonedGlobals(Module &M) {
+    for (llvm::GlobalVariable &GV : M.globals()) {
+      std::string GlobalName = GV.getName().str();
+      if (GlobalName.find('.') == 0) {
+        continue; // This is likely not user-defined. Example: `.L.str`.
+      }
+      GV.setInitializer(nullptr); // Remove the initializer
+      GV.setLinkage(llvm::GlobalValue::ExternalLinkage); // Set external linkage
+    }
+  }
+
+  bool runOnModule(Module &M) override {
+    if (M.getName() == "src/perf/collect.c") {
+      return false;
+    }
+    std::unique_ptr<Module> Cloned = CloneModule(M);
+    if (!Cloned) {
+      errs() << "Error: CloneModule returned null for module: " << M.getName()
+             << "\n";
+      return false;
+    }
+    updateClonedFunctions(*Cloned);
+    updateClonedGlobals(*Cloned);
+
+    llvm::Linker TheLinker(M);
+    if (TheLinker.linkInModule(std::move(Cloned))) {
+      llvm::report_fatal_error("Error linking the modules");
+      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(); }
diff --git a/llvm/lib/YkIR/CMakeLists.txt b/llvm/lib/YkIR/CMakeLists.txt
index 2464ccc11f8e84a..32919c07f8352a0 100644
--- a/llvm/lib/YkIR/CMakeLists.txt
+++ b/llvm/lib/YkIR/CMakeLists.txt
@@ -5,4 +5,6 @@ add_llvm_component_library(LLVMYkIR
   Core
   MC
   Support
+  TransformUtils
+  Linker
   )

From 78a13d19410981f48a0d5e583ce39617cf5bf3e8 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sun, 15 Sep 2024 14:23:24 +0100
Subject: [PATCH 02/22] Add module clone test.

---
 llvm/test/Transforms/Yk/ModuleClone.ll | 88 ++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)
 create mode 100644 llvm/test/Transforms/Yk/ModuleClone.ll

diff --git a/llvm/test/Transforms/Yk/ModuleClone.ll b/llvm/test/Transforms/Yk/ModuleClone.ll
new file mode 100644
index 000000000000000..2eb9cbc0c5975e2
--- /dev/null
+++ b/llvm/test/Transforms/Yk/ModuleClone.ll
@@ -0,0 +1,88 @@
+; RUN: llc -stop-after=yk-module-clone-pass --yk-module-clone < %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 @my_func(i32 %x) {
+entry:
+  %0 = add i32 %x, 1
+  ret i32 %0
+}
+
+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
+}
+
+; ===========================================
+; 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: entry:
+; CHECK: %0 = add i32 %x, 1
+; CHECK: ret i32 %0
+
+; ===========================================
+; Check cloned function: __yk_clone_my_func
+; ===========================================
+; CHECK-LABEL: define dso_local i32 @__yk_clone_my_func(i32 %x)
+; CHECK: entry:
+; CHECK: %0 = add i32 %x, 1
+; CHECK: ret i32 %0
+
+; ===========================================
+; Check original function: main
+; ===========================================
+; CHECK-LABEL: define dso_local i32 @main()
+; CHECK: entry:
+; CHECK: %0 = call i32 @my_func(i32 10)
+; CHECK: %1 = load i32, ptr @my_global
+; CHECK: %2 = call i32 (ptr, ...) @printf
+
+; ===========================================
+; Check cloned function: __yk_clone_main
+; ===========================================
+; CHECK-LABEL: define dso_local i32 @__yk_clone_main()
+; CHECK: entry:
+; CHECK: %0 = call i32 @__yk_clone_my_func(i32 10)
+; CHECK: %1 = load i32, ptr @my_global
+; CHECK: %2 = call i32 (ptr, ...) @printf

From e5e7176c55ddba67d80e14ff615134a11c5fde4a Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sun, 15 Sep 2024 15:27:46 +0100
Subject: [PATCH 03/22] Update header guard in ModuleClone.h

---
 llvm/include/llvm/Transforms/Yk/ModuleClone.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Yk/ModuleClone.h b/llvm/include/llvm/Transforms/Yk/ModuleClone.h
index 4a6bdcf65db2d39..9ea138d53a5e551 100644
--- a/llvm/include/llvm/Transforms/Yk/ModuleClone.h
+++ b/llvm/include/llvm/Transforms/Yk/ModuleClone.h
@@ -1,5 +1,5 @@
-#ifndef LLVM_TRANSFORMS_YK_MODULE_CLONE_H
-#define LLVM_TRANSFORMS_YK_MODULE_CLONE_H
+#ifndef LLVM_TRANSFORMS_YK_MODULECLONE_H
+#define LLVM_TRANSFORMS_YK_MODULECLONE_H
 
 #include "llvm/Pass.h"
 

From d223b645c1eb821c53ff09123ad8e7d7eaa0aeae Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Mon, 16 Sep 2024 17:30:45 +0100
Subject: [PATCH 04/22] Remove src/perf/collect.c check.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index 1cac96b7cb9773c..92ada3a8a50b2a5 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -49,9 +49,6 @@ struct YkModuleClone : public ModulePass {
   }
 
   bool runOnModule(Module &M) override {
-    if (M.getName() == "src/perf/collect.c") {
-      return false;
-    }
     std::unique_ptr<Module> Cloned = CloneModule(M);
     if (!Cloned) {
       errs() << "Error: CloneModule returned null for module: " << M.getName()

From f47b1cd6edf31686218a8f522e417010bdafca75 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Mon, 16 Sep 2024 17:32:19 +0100
Subject: [PATCH 05/22] Revert "Update header guard in ModuleClone.h"

This reverts commit e5e7176c55ddba67d80e14ff615134a11c5fde4a.
---
 llvm/include/llvm/Transforms/Yk/ModuleClone.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Yk/ModuleClone.h b/llvm/include/llvm/Transforms/Yk/ModuleClone.h
index 9ea138d53a5e551..4a6bdcf65db2d39 100644
--- a/llvm/include/llvm/Transforms/Yk/ModuleClone.h
+++ b/llvm/include/llvm/Transforms/Yk/ModuleClone.h
@@ -1,5 +1,5 @@
-#ifndef LLVM_TRANSFORMS_YK_MODULECLONE_H
-#define LLVM_TRANSFORMS_YK_MODULECLONE_H
+#ifndef LLVM_TRANSFORMS_YK_MODULE_CLONE_H
+#define LLVM_TRANSFORMS_YK_MODULE_CLONE_H
 
 #include "llvm/Pass.h"
 

From 4a77e9548af93922b53d6e8b8a5725d99dd8e418 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Mon, 16 Sep 2024 17:49:43 +0100
Subject: [PATCH 06/22] Use static linkModules function over linkInModule.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index 92ada3a8a50b2a5..06cd6ab49d7ebd5 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -58,8 +58,7 @@ struct YkModuleClone : public ModulePass {
     updateClonedFunctions(*Cloned);
     updateClonedGlobals(*Cloned);
 
-    llvm::Linker TheLinker(M);
-    if (TheLinker.linkInModule(std::move(Cloned))) {
+    if (Linker::linkModules(M, std::move(Cloned))) {
       llvm::report_fatal_error("Error linking the modules");
       return false;
     }

From f6d9c60847b7e2157da5e7c0b1d9358257cbc40a Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Mon, 16 Sep 2024 17:54:33 +0100
Subject: [PATCH 07/22] Add comments.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index 06cd6ab49d7ebd5..9b5ca8dfd19f200 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -33,18 +33,20 @@ struct YkModuleClone : public ModulePass {
       if (F.hasExternalLinkage() && F.isDeclaration()) {
         continue;
       }
+      // Rename the function to avoid name conflicts
       F.setName(ClonePrefix + F.getName().str());
     }
   }
 
   void updateClonedGlobals(Module &M) {
     for (llvm::GlobalVariable &GV : M.globals()) {
-      std::string GlobalName = GV.getName().str();
-      if (GlobalName.find('.') == 0) {
+      if (GV.getName().str().find('.') == 0) {
         continue; // This is likely not user-defined. Example: `.L.str`.
       }
-      GV.setInitializer(nullptr); // Remove the initializer
-      GV.setLinkage(llvm::GlobalValue::ExternalLinkage); // Set external linkage
+      // Remove the initializer
+      GV.setInitializer(nullptr);
+      // Set external linkage
+      GV.setLinkage(llvm::GlobalValue::ExternalLinkage);
     }
   }
 

From d9444b8ae140cd1a70d8f79a3d4004a733da597b Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Thu, 19 Sep 2024 19:09:39 +0100
Subject: [PATCH 08/22] Skip cloned functions in BasicBlockTracer pass.

---
 llvm/include/llvm/Transforms/Yk/ModuleClone.h |  2 ++
 llvm/lib/CodeGen/TargetPassConfig.cpp         | 10 +++++++---
 llvm/lib/Transforms/Yk/BasicBlockTracer.cpp   |  5 +++++
 llvm/lib/Transforms/Yk/ModuleClone.cpp        |  4 ++--
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Yk/ModuleClone.h b/llvm/include/llvm/Transforms/Yk/ModuleClone.h
index 4a6bdcf65db2d39..4353fa8b97965ee 100644
--- a/llvm/include/llvm/Transforms/Yk/ModuleClone.h
+++ b/llvm/include/llvm/Transforms/Yk/ModuleClone.h
@@ -3,6 +3,8 @@
 
 #include "llvm/Pass.h"
 
+#define YK_CLONE_PREFIX "__yk_clone_"
+
 namespace llvm {
 ModulePass *createYkModuleClonePass();
 } // namespace llvm
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
index bca7d4b6316657f..d59fa9c06382f43 100644
--- a/llvm/lib/CodeGen/TargetPassConfig.cpp
+++ b/llvm/lib/CodeGen/TargetPassConfig.cpp
@@ -1191,12 +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());
   }
-  if (YkModuleClone){
-    addPass(createYkModuleClonePass());
-  }
+  
 
   addISelPrepare();
   return addCoreISelPasses();
diff --git a/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp b/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp
index 24851bb8ef65469..35710f085e29e72 100644
--- a/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp
+++ b/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,10 @@ struct YkBasicBlockTracer : public ModulePass {
     uint32_t FunctionIndex = 0;
     for (auto &F : M) {
       uint32_t BlockIndex = 0;
+      // 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),
diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index 9b5ca8dfd19f200..1302137cced9a93 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -1,4 +1,5 @@
 #include "llvm/Transforms/Yk/ModuleClone.h"
+#include "llvm/Transforms/Yk/BasicBlockTracer.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
@@ -22,7 +23,6 @@ void initializeYkModuleClonePass(PassRegistry &);
 namespace {
 struct YkModuleClone : public ModulePass {
   static char ID;
-  const std::string ClonePrefix = "__yk_clone_";
 
   YkModuleClone() : ModulePass(ID) {
     initializeYkModuleClonePass(*PassRegistry::getPassRegistry());
@@ -34,7 +34,7 @@ struct YkModuleClone : public ModulePass {
         continue;
       }
       // Rename the function to avoid name conflicts
-      F.setName(ClonePrefix + F.getName().str());
+      F.setName(YK_CLONE_PREFIX + F.getName().str());
     }
   }
 

From efe57dfd699f3b2c2e3aa1f5a10718b968b5c25b Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 12:33:57 +0100
Subject: [PATCH 09/22] Add fixme comment.

---
 llvm/lib/Transforms/Yk/BasicBlockTracer.cpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp b/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp
index 35710f085e29e72..f3c7083dbdc9f4f 100644
--- a/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp
+++ b/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp
@@ -43,10 +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;
-      }
+      // if (F.getName().startswith(YK_CLONE_PREFIX)) {
+      //   continue;
+      // }
       for (auto &BB : F) {
         builder.SetInsertPoint(&*BB.getFirstInsertionPt());
         builder.CreateCall(TraceFunc, {builder.getInt32(FunctionIndex),

From 54bc3b9cfa2db6798c234db36937d9d256c620c4 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 13:04:52 +0100
Subject: [PATCH 10/22] Clone only funcitons that have no refrence taken to
 them.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 28 ++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index 1302137cced9a93..1f1cc2e2a14c128 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -29,12 +29,17 @@ struct YkModuleClone : public ModulePass {
   }
 
   void updateClonedFunctions(Module &M) {
+    std::vector<Function*> functionsToRemove;
     for (llvm::Function &F : M) {
       if (F.hasExternalLinkage() && F.isDeclaration()) {
         continue;
       }
-      // Rename the function to avoid name conflicts
-      F.setName(YK_CLONE_PREFIX + F.getName().str());
+      // Skip functions that are address taken
+      if (!F.hasAddressTaken()) {
+        // Rename the function. This will lead to having multiple 
+        // verisons of the same function in the linked module.
+        F.setName(YK_CLONE_PREFIX + F.getName().str());
+      }
     }
   }
 
@@ -50,6 +55,17 @@ struct YkModuleClone : public ModulePass {
     }
   }
 
+
+
+  // This pass creates two distinct versions of the module functions. 
+  // It clones the module, updates functions and globals, and links the 
+  // cloned module to the original module.
+  // Functions:
+  // -  If function address is not taken, it is duplicated - resulting in 
+  //    two versions of that function in a module
+  // -  If functions address is taken, they remain as is.
+  // Globals:
+  // -  Global variables are not cloned.
   bool runOnModule(Module &M) override {
     std::unique_ptr<Module> Cloned = CloneModule(M);
     if (!Cloned) {
@@ -60,11 +76,15 @@ struct YkModuleClone : public ModulePass {
     updateClonedFunctions(*Cloned);
     updateClonedGlobals(*Cloned);
 
-    if (Linker::linkModules(M, std::move(Cloned))) {
+    // The `OverrideFromSrc` flag tells the linker to 
+    // prefer definitions from the source module (the 2nd argument) when
+    // there are conflicts. This means that if there are two functions 
+    // with the same name, the one from the Cloned module will be used.
+    if (Linker::linkModules(M, std::move(Cloned), Linker::Flags::OverrideFromSrc)) {
       llvm::report_fatal_error("Error linking the modules");
       return false;
     }
-    return true;
+    return false;
   }
 };
 } // namespace

From 2d482bea97871e4aa0d83b58a47718570943f6f2 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 13:05:42 +0100
Subject: [PATCH 11/22] Update returned value.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index 1f1cc2e2a14c128..ebbbefe08f47529 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -84,7 +84,7 @@ struct YkModuleClone : public ModulePass {
       llvm::report_fatal_error("Error linking the modules");
       return false;
     }
-    return false;
+    return true;
   }
 };
 } // namespace

From bbf8d79fc875c3d926d0e8f063ac824c4e0cb4d6 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 13:41:11 +0100
Subject: [PATCH 12/22] Fix .ll test.

---
 llvm/test/Transforms/Yk/ModuleClone.ll | 39 +++++++++++++-------------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/llvm/test/Transforms/Yk/ModuleClone.ll b/llvm/test/Transforms/Yk/ModuleClone.ll
index 2eb9cbc0c5975e2..a7ec1493760e33e 100644
--- a/llvm/test/Transforms/Yk/ModuleClone.ll
+++ b/llvm/test/Transforms/Yk/ModuleClone.ll
@@ -1,4 +1,4 @@
-; RUN: llc -stop-after=yk-module-clone-pass --yk-module-clone < %s | FileCheck %s
+; RUN: llc -stop-after=yk-module-clone-pass --yk-module-clone < %s  | FileCheck %s
 
 source_filename = "ModuleClone.c"
 target triple = "x86_64-pc-linux-gnu"
@@ -57,32 +57,33 @@ entry:
 ; Check original function: my_func
 ; ===========================================
 ; CHECK-LABEL: define dso_local i32 @my_func(i32 %x)
-; CHECK: entry:
-; CHECK: %0 = add i32 %x, 1
-; CHECK: ret i32 %0
+; CHECK-NEXT: entry:
+; CHECK-NEXT: %0 = add i32 %x, 1
+; CHECK-NEXT: ret i32 %0
 
 ; ===========================================
-; Check cloned function: __yk_clone_my_func
+; Check original function: main
 ; ===========================================
-; CHECK-LABEL: define dso_local i32 @__yk_clone_my_func(i32 %x)
-; CHECK: entry:
-; CHECK: %0 = add i32 %x, 1
-; CHECK: ret i32 %0
+; CHECK-LABEL: define dso_local i32 @main()
+; CHECK-NEXT: 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 original function: main
+; Check cloned function: __yk_clone_my_func
 ; ===========================================
-; CHECK-LABEL: define dso_local i32 @main()
-; CHECK: entry:
-; CHECK: %0 = call i32 @my_func(i32 10)
-; CHECK: %1 = load i32, ptr @my_global
-; CHECK: %2 = call i32 (ptr, ...) @printf
+; CHECK-LABEL: define dso_local i32 @__yk_clone_my_func(i32 %x)
+; CHECK-NEXT: entry:
+; CHECK-NEXT: %0 = add i32 %x, 1
+; CHECK-NEXT: ret i32 %0
 
 ; ===========================================
 ; Check cloned function: __yk_clone_main
 ; ===========================================
 ; CHECK-LABEL: define dso_local i32 @__yk_clone_main()
-; CHECK: entry:
-; CHECK: %0 = call i32 @__yk_clone_my_func(i32 10)
-; CHECK: %1 = load i32, ptr @my_global
-; CHECK: %2 = call i32 (ptr, ...) @printf
+; CHECK-NEXT: entry:
+; 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
+

From 6930afa62f29080b198ec4a32f699c390420cbf4 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 15:14:41 +0100
Subject: [PATCH 13/22] Add test for referenced functions in a module.

---
 llvm/test/Transforms/Yk/ModuleClone.ll | 58 ++++++++++++++++++--------
 1 file changed, 41 insertions(+), 17 deletions(-)

diff --git a/llvm/test/Transforms/Yk/ModuleClone.ll b/llvm/test/Transforms/Yk/ModuleClone.ll
index a7ec1493760e33e..f7f1753a9512cf7 100644
--- a/llvm/test/Transforms/Yk/ModuleClone.ll
+++ b/llvm/test/Transforms/Yk/ModuleClone.ll
@@ -15,12 +15,22 @@ 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 @my_func(i32 %x) {
+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)
@@ -29,21 +39,19 @@ entry:
   ret i32 0
 }
 
-; ===========================================
+; ======================================================================
+; Original module 
+; ======================================================================
 ; 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)
@@ -53,34 +61,50 @@ entry:
 ; 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: %0 = add i32 %x, 1
-; CHECK-NEXT: ret i32 %0
+; 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: %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: %0 = add i32 %x, 1
+; CHECK-NEXT: ret i32 %0
+
+; ======================================================================
+; Functions with taken address should be not cloned
+; ======================================================================
+; Check that function was not cloned
+; CHECK-NOT: define {{.*}} @_yk_clone_func_inc_with_address_taken
+
+; ======================================================================
+; Cloned functions
+; ======================================================================
+
 ; Check cloned function: __yk_clone_my_func
-; ===========================================
 ; CHECK-LABEL: define dso_local i32 @__yk_clone_my_func(i32 %x)
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT: %0 = add i32 %x, 1
-; CHECK-NEXT: ret i32 %0
+; 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-NEXT: %0 = call i32 @__yk_clone_my_func(i32 10)

From 53b518951c3b38e4649b3de1825c460ced2679e5 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 15:22:53 +0100
Subject: [PATCH 14/22] Remove globals update.

Updating globals is no longer required since the use of the Linker:Flags
flags. Llvm cloning mechanism will make sure to override the globabl
from source.
---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 42 ++++++++------------------
 1 file changed, 12 insertions(+), 30 deletions(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index ebbbefe08f47529..cbc04353f76f1c5 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -1,5 +1,4 @@
 #include "llvm/Transforms/Yk/ModuleClone.h"
-#include "llvm/Transforms/Yk/BasicBlockTracer.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
@@ -11,6 +10,7 @@
 #include "llvm/Linker/Linker.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Yk/BasicBlockTracer.h"
 
 #define DEBUG_TYPE "yk-module-clone-pass"
 
@@ -29,43 +29,25 @@ struct YkModuleClone : public ModulePass {
   }
 
   void updateClonedFunctions(Module &M) {
-    std::vector<Function*> functionsToRemove;
+    std::vector<Function *> functionsToRemove;
     for (llvm::Function &F : M) {
       if (F.hasExternalLinkage() && F.isDeclaration()) {
         continue;
       }
       // Skip functions that are address taken
       if (!F.hasAddressTaken()) {
-        // Rename the function. This will lead to having multiple 
+        // Rename the function. This will lead to having multiple
         // verisons of the same function in the linked module.
         F.setName(YK_CLONE_PREFIX + F.getName().str());
       }
     }
   }
 
-  void updateClonedGlobals(Module &M) {
-    for (llvm::GlobalVariable &GV : M.globals()) {
-      if (GV.getName().str().find('.') == 0) {
-        continue; // This is likely not user-defined. Example: `.L.str`.
-      }
-      // Remove the initializer
-      GV.setInitializer(nullptr);
-      // Set external linkage
-      GV.setLinkage(llvm::GlobalValue::ExternalLinkage);
-    }
-  }
-
-
-
-  // This pass creates two distinct versions of the module functions. 
-  // It clones the module, updates functions and globals, and links the 
-  // cloned module to the original module.
-  // Functions:
-  // -  If function address is not taken, it is duplicated - resulting in 
-  //    two versions of that function in a module
-  // -  If functions address is taken, they remain as is.
-  // Globals:
-  // -  Global variables are not cloned.
+  // This pass creates two versions for module functions.
+  // original.
+  // If function address is not taken, it is duplicated - resulting in
+  // two versions of that function in a module. If function address is
+  // taken, it remain in its original form.
   bool runOnModule(Module &M) override {
     std::unique_ptr<Module> Cloned = CloneModule(M);
     if (!Cloned) {
@@ -74,13 +56,13 @@ struct YkModuleClone : public ModulePass {
       return false;
     }
     updateClonedFunctions(*Cloned);
-    updateClonedGlobals(*Cloned);
 
-    // The `OverrideFromSrc` flag tells the linker to 
+    // The `OverrideFromSrc` flag tells the linker to
     // prefer definitions from the source module (the 2nd argument) when
-    // there are conflicts. This means that if there are two functions 
+    // there are conflicts. This means that if there are two functions
     // with the same name, the one from the Cloned module will be used.
-    if (Linker::linkModules(M, std::move(Cloned), Linker::Flags::OverrideFromSrc)) {
+    if (Linker::linkModules(M, std::move(Cloned),
+                            Linker::Flags::OverrideFromSrc)) {
       llvm::report_fatal_error("Error linking the modules");
       return false;
     }

From 192a245688d31941056a2c7fb35504c9ebf3d8ac Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 15:43:03 +0100
Subject: [PATCH 15/22] Remove unused variable.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index cbc04353f76f1c5..ad72bf86033750d 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -29,7 +29,6 @@ struct YkModuleClone : public ModulePass {
   }
 
   void updateClonedFunctions(Module &M) {
-    std::vector<Function *> functionsToRemove;
     for (llvm::Function &F : M) {
       if (F.hasExternalLinkage() && F.isDeclaration()) {
         continue;

From 8042923b149031eb7bafc49e88d97dedf82d1fd5 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 15:44:30 +0100
Subject: [PATCH 16/22] Remove extra line.

---
 llvm/test/Transforms/Yk/ModuleClone.ll | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/test/Transforms/Yk/ModuleClone.ll b/llvm/test/Transforms/Yk/ModuleClone.ll
index f7f1753a9512cf7..b8f68697f6ab6b2 100644
--- a/llvm/test/Transforms/Yk/ModuleClone.ll
+++ b/llvm/test/Transforms/Yk/ModuleClone.ll
@@ -110,4 +110,3 @@ entry:
 ; 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
-

From 06795ada82e7a6e0eb906d6f3d96e652a583407a Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 20 Sep 2024 15:47:30 +0100
Subject: [PATCH 17/22] Remove unused import.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index ad72bf86033750d..c48644623c15441 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -10,7 +10,6 @@
 #include "llvm/Linker/Linker.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Utils/Cloning.h"
-#include "llvm/Transforms/Yk/BasicBlockTracer.h"
 
 #define DEBUG_TYPE "yk-module-clone-pass"
 

From b745abf2c7b59d7ec8533ddfc6f0a5edf6d70256 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sun, 22 Sep 2024 12:07:49 +0100
Subject: [PATCH 18/22] Use Twine for string concatination.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index c48644623c15441..dff82112d5a5fa2 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -1,4 +1,5 @@
 #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"
@@ -36,7 +37,7 @@ struct YkModuleClone : public ModulePass {
       if (!F.hasAddressTaken()) {
         // Rename the function. This will lead to having multiple
         // verisons of the same function in the linked module.
-        F.setName(YK_CLONE_PREFIX + F.getName().str());
+        F.setName(Twine(YK_CLONE_PREFIX) + F.getName());
       }
     }
   }

From 3a56536e04ff795a1fdde063d3a49849626a70b9 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sun, 22 Sep 2024 12:11:09 +0100
Subject: [PATCH 19/22] Add verifyModule after module linking.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index dff82112d5a5fa2..bfe625a6ad1810f 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -11,6 +11,7 @@
 #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"
 
@@ -65,6 +66,12 @@ struct YkModuleClone : public ModulePass {
       llvm::report_fatal_error("Error linking the modules");
       return false;
     }
+
+    if (verifyModule(M, &errs())) {
+      errs() << "Module verification failed!";
+      return false;
+    }
+
     return true;
   }
 };

From 904ea2d44c9c480563a3bdbe39ea85f23531acbe Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sun, 22 Sep 2024 12:15:31 +0100
Subject: [PATCH 20/22] Update comments.

---
 llvm/lib/Transforms/Yk/ModuleClone.cpp | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp
index bfe625a6ad1810f..cd0464e93eb22b6 100644
--- a/llvm/lib/Transforms/Yk/ModuleClone.cpp
+++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp
@@ -28,7 +28,6 @@ struct YkModuleClone : public ModulePass {
   YkModuleClone() : ModulePass(ID) {
     initializeYkModuleClonePass(*PassRegistry::getPassRegistry());
   }
-
   void updateClonedFunctions(Module &M) {
     for (llvm::Function &F : M) {
       if (F.hasExternalLinkage() && F.isDeclaration()) {
@@ -36,18 +35,17 @@ struct YkModuleClone : public ModulePass {
       }
       // Skip functions that are address taken
       if (!F.hasAddressTaken()) {
-        // Rename the function. This will lead to having multiple
-        // verisons of the same function in the linked module.
         F.setName(Twine(YK_CLONE_PREFIX) + F.getName());
       }
     }
   }
 
-  // This pass creates two versions for module functions.
-  // original.
-  // If function address is not taken, it is duplicated - resulting in
-  // two versions of that function in a module. If function address is
-  // taken, it remain in its original form.
+  // 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) {
@@ -57,10 +55,11 @@ struct YkModuleClone : public ModulePass {
     }
     updateClonedFunctions(*Cloned);
 
-    // The `OverrideFromSrc` flag tells the linker to
-    // prefer definitions from the source module (the 2nd argument) when
-    // there are conflicts. This means that if there are two functions
-    // with the same name, the one from the Cloned module will be used.
+    // 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");

From 65e56f868d98f998544adc9c05e33d4767ceb44d Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sun, 22 Sep 2024 12:19:48 +0100
Subject: [PATCH 21/22] Update test.

---
 llvm/test/Transforms/Yk/ModuleClone.ll | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/llvm/test/Transforms/Yk/ModuleClone.ll b/llvm/test/Transforms/Yk/ModuleClone.ll
index b8f68697f6ab6b2..6473777f621a6c2 100644
--- a/llvm/test/Transforms/Yk/ModuleClone.ll
+++ b/llvm/test/Transforms/Yk/ModuleClone.ll
@@ -40,10 +40,9 @@ entry:
 }
 
 ; ======================================================================
-; Original module 
+; Original functions
 ; ======================================================================
 ; File header checks
-
 ; CHECK: source_filename = "ModuleClone.c"
 ; CHECK: target triple = "x86_64-pc-linux-gnu"
 
@@ -85,15 +84,15 @@ entry:
 ; CHECK-NEXT: ret i32 %0
 
 ; ======================================================================
-; Functions with taken address should be not cloned
+; Functions whose addresses are taken should not be cloned
 ; ======================================================================
-; Check that function was not cloned
-; CHECK-NOT: define {{.*}} @_yk_clone_func_inc_with_address_taken
+; 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
 ; ======================================================================
-
 ; Check cloned function: __yk_clone_my_func
 ; CHECK-LABEL: define dso_local i32 @__yk_clone_my_func(i32 %x)
 ; CHECK-NEXT: entry:

From 43507e01f0736ad632f613cb23b3ba0c2f25a578 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sun, 22 Sep 2024 13:00:09 +0100
Subject: [PATCH 22/22] Add tests for ModuleClone and BasicBlockTracer passes.

Assert on tracing calls in cloned functions. Assert on not having
tracing calls in original functions.
---
 llvm/lib/Transforms/Yk/BasicBlockTracer.cpp |  6 +++---
 llvm/test/Transforms/Yk/ModuleClone.ll      | 11 ++++++++---
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp b/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp
index f3c7083dbdc9f4f..3098a6df835be2a 100644
--- a/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp
+++ b/llvm/lib/Transforms/Yk/BasicBlockTracer.cpp
@@ -46,9 +46,9 @@ struct YkBasicBlockTracer : public ModulePass {
       // 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;
-      // }
+      if (F.getName().startswith(YK_CLONE_PREFIX)) {
+        continue;
+      }
       for (auto &BB : F) {
         builder.SetInsertPoint(&*BB.getFirstInsertionPt());
         builder.CreateCall(TraceFunc, {builder.getInt32(FunctionIndex),
diff --git a/llvm/test/Transforms/Yk/ModuleClone.ll b/llvm/test/Transforms/Yk/ModuleClone.ll
index 6473777f621a6c2..5ea725ff9414f1d 100644
--- a/llvm/test/Transforms/Yk/ModuleClone.ll
+++ b/llvm/test/Transforms/Yk/ModuleClone.ll
@@ -1,4 +1,4 @@
-; RUN: llc -stop-after=yk-module-clone-pass --yk-module-clone < %s  | FileCheck %s
+; 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"
@@ -40,7 +40,7 @@ entry:
 }
 
 ; ======================================================================
-; Original functions
+; Original functions - should have trace calls
 ; ======================================================================
 ; File header checks
 ; CHECK: source_filename = "ModuleClone.c"
@@ -63,6 +63,7 @@ entry:
 ; 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
@@ -73,6 +74,7 @@ entry:
 ; 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
@@ -80,6 +82,7 @@ entry:
 ; 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
 
@@ -91,11 +94,12 @@ entry:
 ; CHECK-NOT: define dso_local i32 @__yk_clone_func_inc_with_address_taken
 
 ; ======================================================================
-; Cloned functions
+; 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
@@ -106,6 +110,7 @@ entry:
 ; 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