Module clone pass. #202

Pavel-Durov · 2024-09-14T14:51:08Z

YkModuleClone LLVM pass, which creates duplicate versions of functions within a module.

Key changes

Function Cloning

Functions that do not have their addresses taken are cloned. This results in two versions of the same function within the module: the original and the cloned version. Cloned functions are renamed by appending a YK_CLONE_PREFIX to their original names. It is assumed that each function has unique name within the module.

Example of original and cloned functions in a module

func func_example(%arg0: i32) -> i32 {
  bb0:
....
func __yk_clone_func_example(%arg0: i32) -> i32 {
  bb0:
...

Functions whose addresses are taken are left unchanged and are not cloned, ensuring their original version remains in the module.

Module Linking

After cloning, the pass uses llvm::Linker::linkModules() to link the cloned module back into the original. It uses the OverrideFromSrc flag to prefer function definitions from the cloned module if there are conflicts.

ltratt · 2024-09-14T14:53:55Z

We will need some (at least basic) tests before we can think about merging this: without those, it's actually a bit hard to know what's going on!

ltratt · 2024-09-14T14:55:46Z

Do we need to close this? Can't we just add more commits?

Pavel-Durov · 2024-09-14T14:56:16Z

Do we need to close this? Can't we just add more commits?

no... I clicked on the the close button by mistake...

Pavel-Durov · 2024-09-14T14:57:42Z

We will need some (at least basic) tests before we can think about merging this: without those, it's actually a bit hard to know what's going on!

Ok, will add similar tests to llvm-lit llvm/test/Transforms/Yk/BasicBlockTracer.ll

Pavel-Durov · 2024-09-15T13:23:51Z

Added ModuleClone.ll test 👉 78a13d1

llvm-lit ./llvm/test/Transforms/Yk/*
              0   0%    0.00kB/s    0:00:00 (xfr#0, to-chk=0/150768)   
-- Testing: 3 tests, 3 workers --
PASS: LLVM :: Transforms/Yk/ModuleClone.ll (1 of 3)
PASS: LLVM :: Transforms/Yk/BasicBlockTracer.ll (2 of 3)
PASS: LLVM :: Transforms/Yk/Linkage/YkLinkage.ll (3 of 3)

ltratt · 2024-09-15T14:47:52Z

Is there any way in LLVM to generate "fresh" names (maybe using name mangling?) rather than __yk_clone? Or making the name an argument rather than hard coded? [Maybe this is hard, I dunno.]

vext01 · 2024-09-16T08:34:45Z

I think it will do it automatically. If I recall it will start suffixing things .0, .1 etc.

From the top of my head it already does this kind of thing by itself during optimisation passes that require symbols to be copied, so I'd stick with a "well known yk-specific prefix" to avoid accidentally calling things we don't mean to.

llvm/lib/Transforms/Yk/ModuleClone.cpp

vext01 · 2024-09-16T08:39:36Z

llvm/lib/Transforms/Yk/ModuleClone.cpp

+  }
+
+  bool runOnModule(Module &M) override {
+    if (M.getName() == "src/perf/collect.c") {


What's the story with this?

I thought it was an issue, but I can't reproduce it anymore
removed 👉 d223b64

vext01 · 2024-09-16T08:40:45Z

llvm/lib/Transforms/Yk/ModuleClone.cpp

+    updateClonedGlobals(*Cloned);
+
+    llvm::Linker TheLinker(M);
+    if (TheLinker.linkInModule(std::move(Cloned))) {


Ah hah! So you found a bit to do the merging for you. What exactly does this call do?

This call attempts to link/merge the both modules std::move(Cloned) moves ownership of the cloned module to the linker, meaning that after this call the cloned module is no longer valid in the original context.
If the modules are successfully linked, the function returns false (curious choise). If there's an error (such as conflicting function or global names that cannot be resolved), it returns true.

https://llvm.org/doxygen/classllvm_1_1Linker.html#a5a864cc4d71e0234e9f8d786f8716804

I changed it to be using the static linkModules function, they do the same thing but linkModules feels more intuitive to me :)
👉 4a77e95

I changed the linkInModule function to use Linker::Flags::OverrideFromSrc it helps with resolving conflicts, for example when we don't want to clone the functions since its address has been taken.
Additionally, it removed the need to handle globals as they will be overridden based on source definition.

In which commit?

This one 👉 54bc3b9

Tests updated here 👉 6930afa

vext01 · 2024-09-16T08:42:24Z

I'd stick with a "well known yk-specific prefix"

I imagine this will also aid debugging: knowing if we cloned it, or something else did.

Pavel-Durov · 2024-09-21T11:58:10Z

I think it will do it automatically. If I recall it will start suffixing things .0, .1 etc.

From the top of my head it already does this kind of thing by itself during optimisation passes that require symbols to be copied, so I'd stick with a "well known yk-specific prefix" to avoid accidentally calling things we don't mean to.

Yes, LLVM automatically appends suffixes to functions with the same name, such as in cases of overloading. This ensures that function names within an LLVM module are unique. Therefore, cloning and renaming functions with a "static" yk suffix should be safe as well.

Pavel-Durov · 2024-09-26T19:18:10Z

@vext01 as discussed, I removed the stackmaps from the cloned module (this commit)@vext01 as discussed, I removed the stackmaps from the cloned module (this commit 👉 5c5740b).

Example of the original basicblock:

20:                                               ; preds = %18
  %21 = load i32, ptr %3, align 4, !dbg !60
  call void @llvm.dbg.value(metadata i32 %21, metadata !48, metadata !DIExpression()), !dbg !27
  %22 = getelementptr i8, ptr %3, i32 16
  store ptr %22, ptr @shadowstack_0, align 8, !dbg !61
  %23 = call i32 @decrement(i32 noundef %21), !dbg !61
  call void (i64, i32, ...) @llvm.experimental.stackmap(i64 2, i32 0, ptr %3, ptr %4, ptr %5, i32 %21), !dbg !27
  br label %24, !dbg !27

Example of the cloned basicblock:

20:                                               ; preds = %18
  %21 = load i32, ptr %3, align 4, !dbg !100
  call void @llvm.dbg.value(metadata i32 %21, metadata !90, metadata !DIExpression()), !dbg !76
  %22 = getelementptr i8, ptr %3, i32 16
  store ptr %22, ptr @shadowstack_0, align 8, !dbg !101
  %23 = call i32 @__yk_clone_decrement(i32 noundef %21), !dbg !101
  br label %24, !dbg !76

As expected I get assertion error in YkIRWriter.cpp:

llvm/lib/YkIR/YkIRWriter.cpp:585: void {anonymous}::YkIRWriter::serialiseStackmapCall(llvm::CallInst*, {anonymous}::FuncLowerCtxt&): Assertion `I' failed.

vext01 · 2024-09-27T08:21:19Z

I think this is because the serialiser expects to find stackmaps for branches and calls, which for the cloned module, there are none.

The quick fix is to not serialise IR for the cloned modules.

ptersilie · 2024-09-27T10:25:49Z

llvm/lib/CodeGen/TargetPassConfig.cpp

+  // `YkModuleClone` needs to run before `YkBasicBlockTracerPass` and
+  // `YkInsertStackMaps` because they both skip cloned functions.
+  if (YkModuleClone) {
+    addPass(createYkModuleClonePass());


This is not a bad spot for the cloning pass. But I think we might want to run it even earlier before the block-splitting and no calls in entry blocks pass (and maybe also before the blockdisambiguate pass). Though that means we need to manually add the control point and a single stackmap for that control point into the cloned module. We need those in order to jump between the two modules.

Got it.
Made changes here 👉 c1225df.
I think it actually works by reusing the same YkStackmaps pass with different configurations for stackmapIDStart - 1 by default, 2 in case of YkModuleClone is active.

Excellent. Does this still insert stackmaps into the cloned module? The cloned module only needs a single stackmap at the control point. Though, while important, it's kind of an optimisation, so you may choose to do this in another PR.

yes, the the patchpoint emits stackmaps.
For the cloned module there's a patchpoint in the original main function

18: ; preds = %15 call void @__yk_trace_basicblock(i32 2, i32 8), !dbg !57 %19 = load ptr, ptr %4, align 8, !dbg !57 call void @llvm.dbg.value(metadata ptr %19, metadata !32, metadata !DIExpression()), !dbg !27 call void (i64, i32, ptr, i32, ...) @llvm.experimental.patchpoint.void(i64 0, i32 13, ptr @__ykrt_control_point, i32 3, ptr %19, ptr %3, i64 0, ptr %3, ptr %4, ptr %5, ptr %19), !dbg !59 br label %20, !dbg !60

and cloned main function:

18: ; preds = %15 %19 = load ptr, ptr %4, align 8, !dbg !97 call void @llvm.dbg.value(metadata ptr %19, metadata !80, metadata !DIExpression()), !dbg !76 call void (i64, i32, ptr, i32, ...) @llvm.experimental.patchpoint.void(i64 1, i32 13, ptr @__ykrt_control_point, i32 3, ptr %19, ptr %3, i64 1, ptr %3, ptr %4, ptr %5, ptr %19), !dbg !99 br label %20, !dbg !100

Pavel-Durov · 2024-09-28T14:04:33Z

I think this is because the serialiser expects to find stackmaps for branches and calls, which for the cloned module, there are none.

The quick fix is to not serialise IR for the cloned modules.

Done here 👉 d73b87c, (and later patched with e5c2e43, 2e837df)

Pavel-Durov · 2024-09-28T14:20:45Z

llvm/lib/Transforms/Yk/ControlPoint.cpp

-CallInst *findControlPointCall(Module &M) {
+/// Find the calls to the dummy control point that we want to patch.
+/// Returns either a pointer the a vector of instructions
+std::vector<CallInst *> findControlPointCalls(Module &M) {


Now we have two control points in both original and the cloned main functions.

main (stack map id: 0):

21: ; preds = %18 call void @__yk_trace_basicblock(i32 0, i32 9), !dbg !70 %22 = load ptr, ptr %4, align 8, !dbg !70 call void @llvm.dbg.value(metadata ptr %22, metadata !41, metadata !DIExpression()), !dbg !36 call void (i64, i32, ptr, i32, ...) @llvm.experimental.patchpoint.void(i64 0, i32 13, ptr @__ykrt_control_point, i32 3, ptr %22, ptr %3, i64 0, ptr %3, ptr %4, ptr %5, ptr %6, ptr %22), !dbg !72 br label %23, !dbg !73

__yk_clone_main (stack map id: 1):

21: ; preds = %18 %22 = load ptr, ptr %4, align 8, !dbg !113 call void @llvm.dbg.value(metadata ptr %22, metadata !93, metadata !DIExpression()), !dbg !89 call void (i64, i32, ptr, i32, ...) @llvm.experimental.patchpoint.void(i64 1, i32 13, ptr @__ykrt_control_point, i32 3, ptr %22, ptr %3, i64 1, ptr %3, ptr %4, ptr %5, ptr %6, ptr %22), !dbg !115 br label %23, !dbg !116 ****

Pavel-Durov · 2024-09-29T13:13:52Z

I think this is because the serialiser expects to find stackmaps for branches and calls, which for the cloned module, there are none.

I had to update the header as well 👉 934af15

ptersilie · 2024-09-30T08:28:47Z

llvm/lib/YkIR/YkIRWriter.cpp

@@ -1742,7 +1743,10 @@ class YkIRWriter {
    OutStreamer.emitSizeT(M.size());
    // funcs:
    for (llvm::Function &F : M) {
-      serialiseFunc(F);
+      // Skip cloned functions


Is it enough to just skip the cloned functions? What about the rest of the module, e.g. M.size() above. Is this still serialised somewhere? Do we not need to skip this in the de-serialiser?

I see that in a later commit you updated how M.size() is calculated. But now I'm confused. Do the cloned functions live in the same module as the original? Or do we have two separate modules?

Maybe I got confused here, but I thought that we don't need the cloned functions in the IR?
#202 (comment)

In the current version of this PR both function versions (original and cloned, excluding the functions that have reference taken to them) live in the same module, but we serialize only the non-cloned version in YkIRWriter.

Ah I see. I think that answers my question. Though now I'm curious how we are going to optimise one version of the function while not optimising the other. That's something we need to think about soon.

Pavel-Durov · 2024-10-08T19:47:22Z

Finally, I figured out that the serialisation/deserialisation issue was caused by a misalignment of function indexes (actually noted in the FIXME comment) 👉 bd74adf

With this fix all tests are passing in the yk using SWT and the Clone Pass 💪

Pavel-Durov · 2024-10-10T12:08:59Z

This PR is ready for another review.
Summary of changes:

Clone module functions that have no address taken to them
Insert control point + stackmap in the clone - when this pass is enabled we will have two control points.
Skip cloned functions serialisation in YkIRWriter.

vext01 · 2024-10-11T12:46:02Z

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

+
+#include "llvm/Pass.h"
+
+#define YK_CLONE_PREFIX "__yk_clone_"


Can we document somewhere (maybe in a doc comment at the top of the pass cpp file?) what the cloned versions are? Are they intended to be the optimised versions, or the unoptimisaed versions?

sounds good, will add.

Added doc comment 👉 71f1973

llvm/lib/CodeGen/TargetPassConfig.cpp

llvm/lib/Transforms/Yk/BasicBlockTracer.cpp

llvm/lib/Transforms/Yk/ControlPoint.cpp

vext01 · 2024-10-11T13:17:23Z

llvm/lib/Transforms/Yk/ModuleClone.cpp

+    // 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.


Does this mean that the "deleting the initialiser" trick is no longer needed?

Yep, we don't need it, linkModules resolves these conflicts for us. #202 (comment)
Updated comment to make it clearer 👉 9d19ee9

llvm/lib/Transforms/Yk/ModuleClone.cpp

vext01 · 2024-10-11T13:21:47Z

Looking good! Just a handful of comments from me.

Pavel-Durov · 2024-10-11T13:24:12Z

Looking good! Just a handful of comments from me.

Thanks, will address 👍

vext01 · 2024-10-14T09:03:54Z

I think this looks good, bar some minor bits you can fix at squash-time.

@ptersilie would you mind glancing over this in case I've missed anything?

ptersilie · 2024-10-14T09:28:42Z

llvm/lib/Transforms/Yk/ModuleClone.cpp

+//     original names. This distinguishes them from the original functions.
+//
+// - **Clone Process:**
+//   - 1. The orignal model is cloned, creating two copies of the module.


Should this be s/orignal model/original module/ ?

updated 👉 86efe30

ptersilie · 2024-10-14T09:31:19Z

llvm/lib/Transforms/Yk/ModuleClone.cpp

+  }
+
+  // This pass duplicates functions within the module:
+  // the original and the cloned version.


This line seems redundant?

updated 👉 5f2716d

ptersilie · 2024-10-14T09:35:14Z

Apart from some typos and odd formatting in the comments (which hopefully yk_format_new_files.sh will sort out), this looks good to me.

Pavel-Durov · 2024-10-14T10:03:49Z

Apart from some typos and odd formatting in the comments (which hopefully yk_format_new_files.sh will sort out), this looks good to me.

What is odd with the formatting?
When running yk_format_new_files.sh I get this diff 👉 5388c52

vext01 · 2024-10-14T15:12:10Z

(Formatters often don't format comments)

vext01 · 2024-10-14T15:12:55Z

I'm going to let @ptersilie take it from here, since his comments are the only ones remaining.

Pavel-Durov · 2024-10-15T14:51:52Z

Comment format updated 👉 f43ee77

ptersilie · 2024-10-15T15:54:24Z

Let's squash this.

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.

Pavel-Durov · 2024-10-15T18:35:45Z

squashed 👉 64e9c76

Pavel-Durov mentioned this pull request Sep 14, 2024

LLVM module clone #201

Closed

Pavel-Durov closed this Sep 14, 2024

Pavel-Durov reopened this Sep 14, 2024

Pavel-Durov assigned vext01 Sep 14, 2024

Pavel-Durov assigned ltratt Sep 14, 2024

Pavel-Durov changed the title ~~Module clone Pass.~~ Module clone pass. Sep 15, 2024

Pavel-Durov marked this pull request as ready for review September 15, 2024 13:32

vext01 reviewed Sep 16, 2024

View reviewed changes

llvm/lib/Transforms/Yk/ModuleClone.cpp Outdated Show resolved Hide resolved

vext01 reviewed Sep 16, 2024

View reviewed changes

Pavel-Durov assigned ptersilie Sep 16, 2024

ptersilie reviewed Sep 27, 2024

View reviewed changes

Pavel-Durov commented Sep 28, 2024

View reviewed changes

ptersilie reviewed Sep 30, 2024

View reviewed changes

Pavel-Durov mentioned this pull request Oct 10, 2024

Enable llvm module clone pass for SWT. ykjit/yk#1421

Merged

vext01 reviewed Oct 11, 2024

View reviewed changes

llvm/lib/CodeGen/TargetPassConfig.cpp Outdated Show resolved Hide resolved

vext01 reviewed Oct 11, 2024

View reviewed changes

llvm/lib/Transforms/Yk/BasicBlockTracer.cpp Show resolved Hide resolved

vext01 reviewed Oct 11, 2024

View reviewed changes

llvm/lib/Transforms/Yk/ControlPoint.cpp Outdated Show resolved Hide resolved

vext01 reviewed Oct 11, 2024

View reviewed changes

llvm/lib/Transforms/Yk/ControlPoint.cpp Show resolved Hide resolved

vext01 reviewed Oct 11, 2024

View reviewed changes

llvm/lib/Transforms/Yk/ModuleClone.cpp Outdated Show resolved Hide resolved

ptersilie reviewed Oct 14, 2024

View reviewed changes

Pavel-Durov force-pushed the clone-llvm-module-pass branch from f43ee77 to 64e9c76 Compare October 15, 2024 17:50

vext01 added this pull request to the merge queue Oct 15, 2024

Merged via the queue into ykjit:main with commit 3ff37cb Oct 15, 2024
2 checks passed

Module clone pass. #202

Module clone pass. #202

Conversation

Pavel-Durov commented Sep 14, 2024 • edited Loading

Key changes

Function Cloning

Module Linking

ltratt commented Sep 14, 2024

ltratt commented Sep 14, 2024

Pavel-Durov commented Sep 14, 2024

Pavel-Durov commented Sep 14, 2024

Pavel-Durov commented Sep 15, 2024 • edited Loading

ltratt commented Sep 15, 2024

vext01 commented Sep 16, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Pavel-Durov Sep 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vext01 commented Sep 16, 2024

Pavel-Durov commented Sep 21, 2024

Pavel-Durov commented Sep 26, 2024

vext01 commented Sep 27, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Pavel-Durov commented Sep 28, 2024

Choose a reason for hiding this comment

Pavel-Durov commented Sep 29, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Pavel-Durov commented Oct 8, 2024

Pavel-Durov commented Oct 10, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vext01 commented Oct 11, 2024

Pavel-Durov commented Oct 11, 2024

vext01 commented Oct 14, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ptersilie commented Oct 14, 2024

Pavel-Durov commented Oct 14, 2024

vext01 commented Oct 14, 2024

vext01 commented Oct 14, 2024

Pavel-Durov commented Oct 15, 2024

ptersilie commented Oct 15, 2024

Pavel-Durov commented Oct 15, 2024

Pavel-Durov commented Sep 14, 2024 •

edited

Loading

Pavel-Durov commented Sep 15, 2024 •

edited

Loading

Pavel-Durov Sep 16, 2024 •

edited

Loading