Make sure that unoptimised functions calls ONLY unoptimised functions. #216

Pavel-Durov · 2024-11-24T15:06:45Z

Ensure that calls in original functions are referencing only original functions.

I encountered this issue when I was building yklua with latest yk YKB_TRACER=swt.
The main issue here is with serialization in YkIRWriter, where we don't serialize cloned functions but we still might end up with original functions calling cloned functions.

Lua build error:

Function not found in function index map: __yk_clone_luaO_tostring
LLVM ERROR: Function not found in function index map

source

ltratt · 2024-11-24T15:42:12Z

llvm/lib/Transforms/Yk/ModuleClone.cpp

+   * - Function `f` calls function `g`.
+   * - Function `g` is cloned as `__yk_clone_g`.
+   * - Function `f` is not cloned because its address is taken.
+   * - As a result, function `f` calls `__yk_clone_g` instead of `g`.


IIUC the idea here is "we need to be able to trace f, but its address has taken, so we will always call f with tracing calls inserted"?

More like:
We need to be able to trace f, but its address has taken.
We need to make sure that any functions called from f are also non-cloned functions (functions that we can trace).

Updated tests 👉 0f3e89c

As a result, function f calls __yk_clone_g instead of g.

Is that what's already happening or what this function does? I feel like there's another bullet point missing.

Also, please remind me which functions have the yk_record_block calls, the original or the cloned ones?

Is that what's already happening or what this function does? I feel like there's another bullet point missing.

The comment example scenario is about what's currently happening (before this function in this PR)
I can update it to make it clearer though.

As a result, function f calls __yk_clone_g instead of g.

If f isn't cloned, why is it calling __yk_clone_g. We must have altered f somewhere. So it appears, we first alter f to call _yk_clone_g instead of g and then in this function we revert this change. Is that correct? If so, why don't we only update function calls to their cloned version if the caller function is also cloned?

I think what's happening here, is that we rename all functions in the cloned module before linking.
But then some of them are not renamed cause their address is taken, while sill referencing functions that were renamed.

let's say we have:

func a { } // function which address is taken func b { a(); }

After calling renameFunctions on a cloned module we have:

func __yk_clone_a { } // let b a function which address is taken func b { __yk_clone_a(); }

function a was renamed but function b wasn't (because address is taken).

And then when we link both modules,b from cloned module overrides the original b definition.

Example of actual module IR generated (before this PR):

; Function Attrs: noinline nounwind optnone uwtable define internal i32 @luaB_assert(ptr noundef %0) #0 { %2 = getelementptr inbounds %struct.lua_State, ptr %0, i64 0, i32 8 %3 = load ptr, ptr %2, align 8, !tbaa !27 %4 = load ptr, ptr %3, align 8, !tbaa !22 %5 = getelementptr inbounds %union.StackValue, ptr %4, i64 1 %6 = getelementptr inbounds %struct.lua_State, ptr %0, i64 0, i32 6 %7 = load ptr, ptr %6, align 8, !tbaa !22 %8 = icmp ult ptr %5, %7 br i1 %8, label %13, label %9 9: ; preds = %1 %10 = getelementptr inbounds %struct.lua_State, ptr %0, i64 0, i32 7 %11 = load ptr, ptr %10, align 8, !tbaa !10 %12 = getelementptr inbounds %struct.global_State, ptr %11, i64 0, i32 8 br label %13 13: ; preds = %9, %1 %14 = phi ptr [ %12, %9 ], [ %5, %1 ] %15 = getelementptr inbounds %struct.TValue, ptr %14, i64 0, i32 1 %16 = load i8, ptr %15, align 8, !tbaa !23 %17 = icmp eq i8 %16, 1 %18 = and i8 %16, 15 %19 = icmp eq i8 %18, 0 %20 = or i1 %17, %19 br i1 %20, label %21, label %24, !prof !5 21: ; preds = %13 tail call void @__yk_clone_luaL_checkany(ptr noundef %0, i32 noundef 1) #77 tail call void @__yk_clone_lua_rotate(ptr noundef %0, i32 noundef 1, i32 noundef -1) #77 tail call void @__yk_clone_lua_settop(ptr noundef %0, i32 noundef -2) #77 %22 = tail call ptr @__yk_clone_lua_pushstring(ptr noundef %0, ptr noundef nonnull @.str.26.676) #77 tail call void @__yk_clone_lua_settop(ptr noundef %0, i32 noundef 1) #77 %23 = tail call i32 @luaB_error(ptr noundef %0) unreachable

Or maybe the question is how these call sites in the function that is not cloned got renamed in the first place?

Looks like F.setName is actually updating the callsites as well as the function name.
I can see:
Before F.setName is called for luaL_checkany function:

21: ; preds = %13 tail call void @luaL_checkany(ptr noundef %0, i32 noundef 1) #77 ... unreachable

And just after F.setName is called for luaL_checkany function:

21: ; preds = %13 tail call void @__yk_clone_luaL_checkany(ptr noundef %0, i32 noundef 1) #77 .... unreachable

Is that what's already happening or what this function does? I feel like there's another bullet point missing.

Updated comment for clarity 👉 3462947

Pavel-Durov · 2024-11-25T19:25:08Z

Added test in 👉 ykjit/yk#1477

ptersilie · 2024-11-26T11:20:06Z

llvm/lib/Transforms/Yk/ModuleClone.cpp

+   * - As a result, function `f` calls `__yk_clone_g` instead of `g`.
+   *
+   * **Reasoning:**
+   * In `YkIRWriter` we only serialise non-cloned functions. We want


AFAICS this isn't really related to the serialiser. The real reason we want unoptimised functions to only call other unoptimised functions is because otherwise block recording won't work. In other words, if f calls g_opt, then f will record the blocks, but g_opt won't, we get incorrect traces.

youre right.

ptersilie · 2024-11-26T11:32:21Z

Okay, I had to remind myself how this all works, but I think I get it now. Please let me know if I said anything wrong. The way we make the optimised module is:

We use LLVM's CloneModule API to create a copy of the current module.
We then iterate over all functions in the cloned module and rename them, unless they have their address taken.
Annoyingly, this also updates the callees of unoptimised functions, which means we end up with unoptimised functions calling optimised functions. Which is bad, because it messes with the block recording.
So we need to go back and reverse the callees inside unoptimised functions to only call unoptimised functions.
Finally, we link the two modules together to get one module.

Question:
Was there a reason why we use LLVM's CloneModule? From what I can see, we are only really interested in cloning some functions to make optimised versions of them. Wouldn't it be easier to do the following:

Iterate over the original module
Clone and rename (in one step) each function that doesn't have it's address taken, and add it to the same module.
Built up a VMap for the clones.
Once all functions have been cloned, iterate over all cloned functions and update the callees using the vmap.

This seems a bit more straight forward then what we are doing now, where things get updated that we don't want to and then we have to get back and fix it, due to the way CloneModule works. But there might be a valid reason we are using CloneModule that I have forgotten.

Separate suggestion (for a separate PR): I keep forgetting what the cloned and uncloned version of a function is. We should use unoptimised and optimised instead. I.e. the unoptimised function is the original (containing the yk_record_block calls). After cloning we get an optimised version without those calls. We should also rename the prefix to reflect that, maybe something like __yk_opt_foo.

Pavel-Durov · 2024-11-29T18:23:55Z

Okay, I had to remind myself how this all works, but I think I get it now. Please let me know if I said anything wrong. The way we make the optimised module is:
1. We use LLVM's `CloneModule` API to create a copy of the current module.

2. We then iterate over all functions in the cloned module and rename them, unless they have their address taken.

3. Annoyingly, this also updates the callees of unoptimised functions, which means we end up with unoptimised functions calling optimised functions. Which is bad, because it messes with the block recording.

4. So we need to go back and reverse the callees inside unoptimised functions to only call unoptimised functions.

5. Finally, we link the two modules together to get one module.
Question: Was there a reason why we use LLVM's CloneModule? From what I can see, we are only really interested in cloning some functions to make optimised versions of them. Wouldn't it be easier to do the following:
* Iterate over the original module

* Clone and rename (in one step) each function that doesn't have it's address taken, and add it to the same module.

* Built up a VMap for the clones.

* Once all functions have been cloned, iterate over all cloned functions and update the callees using the vmap.
This seems a bit more straight forward then what we are doing now, where things get updated that we don't want to and then we have to get back and fix it, due to the way CloneModule works. But there might be a valid reason we are using CloneModule that I have forgotten.

Separate suggestion (for a separate PR): I keep forgetting what the cloned and uncloned version of a function is. We should use unoptimised and optimised instead. I.e. the unoptimised function is the original (containing the yk_record_block calls). After cloning we get an optimised version without those calls. We should also rename the prefix to reflect that, maybe something like __yk_opt_foo.

Addressed here 👉 12fc9d7

ptersilie · 2024-12-02T15:25:19Z

llvm/lib/Transforms/Yk/ModuleClone.cpp

+      for (const Argument &OrigArg : F.args()) {
+        DestArgIt->setName(OrigArg.getName());
+        VMap[&OrigArg] = &*DestArgIt++;
+      }


Is this needed? From what I can see llvm::CloneFunction already copies the arguments. And the VMap isn't actually read anywhere here.

ok will remove.

removed 👉 4563f62f6b46eca4d91cd71192c4a32356690560

I don't see the commit.

ptersilie · 2024-12-02T15:31:26Z

Aha, this looks good. I actually wasn't sure if there was a reason we did things the other way (and was half expecting you to tell me why my approach won't work). But I'm glad this works as it's a bit easier to think about and we don't need to go back and fix things up as with the earlier approach.

Do we need to update the PR title now since our approach has changed, e.g. something like "Avoid updating unoptimised functions with optimised callees by processing optimised functions in isolation."?

Pavel-Durov · 2024-12-02T15:57:38Z

Aha, this looks good. I actually wasn't sure if there was a reason we did things the other way (and was half expecting you to tell me why my approach won't work). But I'm glad this works as it's a bit easier to think about and we don't need to go back and fix things up as with the earlier approach.

Do we need to update the PR title now since our approach has changed, e.g. something like "Avoid updating unoptimised functions with optimised callees by processing optimised functions in isolation."?

I think we went with cloning approach just because that's what we initially drawn on the whiteboard :)

Pavel-Durov · 2024-12-02T17:57:14Z

ready for review :)
@ptersilie

ptersilie · 2024-12-06T10:48:49Z

Please squash.

SWT. Switch from module linking to function cloning for simplicity and changed the function cloning prefix from `__yk_clone_` to `__yk_opt_` to better reflect the optimization intent in ModuleClone pass.

Pavel-Durov · 2024-12-06T18:00:14Z

Squashed 👉 453f961

Pavel-Durov assigned ltratt, ptersilie and vext01 Nov 24, 2024

ltratt reviewed Nov 24, 2024

View reviewed changes

This was referenced Nov 25, 2024

Add test that validates swt clone module function references. Pavel-Durov/yk-fork#17

Open

Add test that validates swt clone module function references. ykjit/yk#1477

Open

Pavel-Durov closed this Nov 25, 2024

Pavel-Durov reopened this Nov 25, 2024

ptersilie reviewed Nov 26, 2024

View reviewed changes

ptersilie reviewed Dec 2, 2024

View reviewed changes

Pavel-Durov changed the title ~~Update function call references in YkModuleClone.~~ Make sure that unoptiised functions calls only unoptiised functions. Dec 2, 2024

Pavel-Durov changed the title ~~Make sure that unoptiised functions calls only unoptiised functions.~~ Make sure that unoptimised functions calls ONLY unoptimised functions. Dec 2, 2024

Ensure that unoptimised functions call only unoptimised functions in

453f961

SWT. Switch from module linking to function cloning for simplicity and changed the function cloning prefix from `__yk_clone_` to `__yk_opt_` to better reflect the optimization intent in ModuleClone pass.

Pavel-Durov force-pushed the fix-function-reference-yk-module-clone-pass branch from 94a1f50 to 453f961 Compare December 6, 2024 17:59

ltratt added this pull request to the merge queue Dec 7, 2024

Merged via the queue into ykjit:main with commit b5ce1fa Dec 7, 2024
2 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make sure that unoptimised functions calls ONLY unoptimised functions. #216

Make sure that unoptimised functions calls ONLY unoptimised functions. #216

Pavel-Durov commented Nov 24, 2024 •

edited

Loading

ltratt Nov 24, 2024

Pavel-Durov Nov 24, 2024

Pavel-Durov Nov 24, 2024

ptersilie Nov 25, 2024

Pavel-Durov Nov 25, 2024

ptersilie Nov 25, 2024

Pavel-Durov Nov 25, 2024 •

edited

Loading

Pavel-Durov Nov 25, 2024 •

edited

Loading

Pavel-Durov Nov 25, 2024

Pavel-Durov Nov 25, 2024

Pavel-Durov commented Nov 25, 2024

ptersilie Nov 26, 2024 •

edited

Loading

Pavel-Durov Nov 29, 2024

ptersilie commented Nov 26, 2024 •

edited

Loading

Pavel-Durov commented Nov 29, 2024

ptersilie Dec 2, 2024

Pavel-Durov Dec 2, 2024

Pavel-Durov Dec 2, 2024

ptersilie Dec 2, 2024

ptersilie commented Dec 2, 2024

Pavel-Durov commented Dec 2, 2024

Pavel-Durov commented Dec 2, 2024

ptersilie commented Dec 6, 2024

Pavel-Durov commented Dec 6, 2024

Make sure that unoptimised functions calls ONLY unoptimised functions. #216

Make sure that unoptimised functions calls ONLY unoptimised functions. #216

Conversation

Pavel-Durov commented Nov 24, 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

Choose a reason for hiding this comment

Pavel-Durov Nov 25, 2024 • edited Loading

Choose a reason for hiding this comment

Pavel-Durov Nov 25, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Pavel-Durov commented Nov 25, 2024

ptersilie Nov 26, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ptersilie commented Nov 26, 2024 • edited Loading

Pavel-Durov commented Nov 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

ptersilie commented Dec 2, 2024

Pavel-Durov commented Dec 2, 2024

Pavel-Durov commented Dec 2, 2024

ptersilie commented Dec 6, 2024

Pavel-Durov commented Dec 6, 2024

Pavel-Durov commented Nov 24, 2024 •

edited

Loading

Pavel-Durov Nov 25, 2024 •

edited

Loading

Pavel-Durov Nov 25, 2024 •

edited

Loading

ptersilie Nov 26, 2024 •

edited

Loading

ptersilie commented Nov 26, 2024 •

edited

Loading