Thanks for organizing our thoughts on this. Okay if I assign you, since you expressed interest in working on it?

Implement trampoline with LDR of a PC relative literal (instead of movk). It saves

Interesting! Mind elaborating on this a bit more? I get that it saves memory, but I'm curious if it's expected to be faster too.

Generate trampoline at the end of the trace instead of at the end of every micro op and write a function to generate the trampoline.

I'd break this up into a couple of phases:

• A PR to move the trampolines from the "code" section of a micro-op to the "data" section, so it's out-of-line.
• A PR to emit all of the trampolines at the end of every trace, so that each opcode doesn't need its own copy of the trampolines it uses (this could waste some memory initially, but is a nice intermediate step).
• Once we have a slab allocator from JIT: improve memory allocation #119730, a PR use one set of trampolines per-slab rather than per-trace.

Also worth mentioning: we'll want to move to short jumps with trampolines on all platforms, not just AArch64 (AArch64 just sort of forces our hand right now since it only lets us use short jumps). So this work should also benefit other platforms too, which is nice.

Interesting! Mind elaborating on this a bit more? I get that it saves memory, but I'm curious if it's expected to be faster too.

I've updated the original comment saying that it saves 8 bytes. About the speed, I think we need to measure it somehow but I would think it would be the same. The other saving is that we will do only one relocation instead of four.

The code will be something like that:

ldr x8, [PC+8]
br x8
&_Py_Dealloc

So this work should also benefit other platforms too, which is nice.

Of course :)

The plan to implement the task "Emit all of the trampolines at the end of every trace, so that each opcode doesn't need its own copy of the trampolines it uses. Also write a function to generate the trampoline."

Instead of generating trampolines in the data section of the stencil, generate a call to a function to generate trampolines (like patch_aarch64_trampoline).
The signature of the function is something like:

void patch_trampoline(unsigned char *location, uint64_t value);

where:

• location is the pointer to the code to patch
• value is the address of the trampoline.

In jit.c implement such function. It will do mainly 2 things:

• if the trampoline for that symbol is not already emitted, emit (write in memory) the trampoline and store it in a table. If a trampoline for that symbol is already emitted (it is present the table), skip this step.
• patch the uop with the correct address of the trampoline

The function will contain the code of the trampoline

 f"{base + 4 * 0:x}: 58000048      ldr     x8, 8",
 f"{base + 4 * 1:x}: d61f0100      br      x8",
 f"{base + 4 * 2:x}: 00000000",
 f"{base + 4 * 2:016x}:  R_AARCH64_ABS64    {hole.symbol}",
 f"{base + 4 * 3:x}: 00000000",

Now that the trampolines are not in the stencils anymore, they will be generated and patched at runtime.
We need though space in memory to store them but we don't know how many we need to emit.

As first implementation we could assume the worst case scenario and allocate enough memory to store all the possible trampolines. Currently we emit 101 different trampolines across all the stencils, so we could allocate memory to store 150 trampolines. This allows some room for increasing the number of trampolines emitted by the JIT without incurring in
memory overflow.
The total size allocate by _PyJIT_Compile would be something like (POC):

#define TRAMPOLINE_SIZE (4 * 4 * 150)

...
...

size_t total_size = code_size + data_size + TRAMPOLINES_SIZE + padding;

In order to improve the allocation above, we could even emit exactly the correct number of trampolines. We can detect what's needed when building jit_stencils.h, and give the stencils metadata about what trampolines they
need (sort of like how we put their size now).
Then when we're doing that first loop in _PyJIT_Compile to get the size of the JIT code, we can also collect the trampolines that are needed at the same time.

For completeness a comment from Brandt:

"Btw, I don’t think the trampoline emit func needs to take a string. We can probably just generate an enum (or #define) with integer values for each symbol, and use those. That way we don’t need to parse strings, and we could use bit-flags later when we want to only emit the needed trampolines.
Also, we can count the unique symbols when generating the templates, and use that to define the max extra space needed (instead of just hardcoding 150).
In general I think we can really rely on the stencil generator here to make things easier.
And more precise."

The PR to generate trampolines at runtime has been created: #123872

The approach to allocate memory is the optimal one as we count how many trampolines we need for every trace and we allocate the right amount of memory for these trampolines.