gh-119726: generate and patch AArch64 trampolines

[diegorusso]

AArch64 trampolines are now generated at runtime at the end of every trace.

In the PR there is a TODO which I'd like to have an opinion of.

• Issue: JIT: improve AArch64 code generation #119726

[brandtbucher]

This is great! I'd like to see this mechanism used across all platforms for symbols in the "global offset table" too (but that can be a separate project).

It's a clear improvement for the affected platforms:

• aarch64-unknown-linux-gnu: 0.8% faster, 0.6% less memory.
• aarch64-apple-darwin: 1.5% faster, 0.8% less memory.

Here's an initial review of the Python code. I'll dig into the C code tomorrow.

[brandtbucher]

This is a lot of braces, haha. Maybe move the braces to _get_trampoline_mask?

Suggested change

	return f"{{emit_{opname}, {len(self.code.body)}, {len(self.data.body)}, {{{self._get_trampoline_mask()}}}}}"
	return f"{{emit_{opname}, {len(self.code.body)}, {len(self.data.body)}, {self._get_trampoline_mask()}}}"

...and then turn ", ".join(trampoline_mask) to "{" + ", ".join(trampoline_mask) + "}" above. That makes it a bit clearer that it's returning valid C code.

[diegorusso]

Done.

[brandtbucher]

I had a bit of trouble following this, so I tried rewriting it. What do you think? This way we don't need to hardcode SYMBOL_MASK_SIZE, and can just delete it.

Suggested change

	if bitmask:
	trampoline_mask = [
	f"0x{(bitmask >> i*32) & ((1 << 32) - 1):x}"
	for i in range(0, SYMBOL_MASK_SIZE)
	]
	while bitmask:
	word = bitmask & ((1 << 32) - 1)
	trampoline_mask.append(f"{word:#04x}")
	bitmask >>= 32

[diegorusso]

Done.

[brandtbucher]

Is this accessed from outside of the class itself? If not, let's add an underscore:

Suggested change

	trampolines: set[int] = dataclasses.field(default_factory=set, init=False)
	_trampolines: set[int] = dataclasses.field(default_factory=set, init=False)

[diegorusso]

Done.

[brandtbucher]

Let's try to get rid of this global cache, since I would ideally like the classes to be able to used multiple times in the same process without sharing state like this. I'd need to think it through more, but maybe it should be moved to the _Target class?

[diegorusso]

let me see what I can place it...

[diegorusso]

I put it in the Target class. Have a look.

[brandtbucher]

Also not a huge fan of this magic constant, since it'll almost certainly need to be updated in the future (and may even be platform-dependent). This should be the same as math.ceil(len(known_symbols) / 32) once known_symbols is populated, right?

[diegorusso]

Yes, you are right. With time this changes (likely increases) and we need to keep updating this number. Also I agree, this could be platform dependent.

[diegorusso]

Removed.

[diegorusso]

It's a clear improvement for the affected platforms:

• aarch64-unknown-linux-gnu: 0.8% faster, 0.6% less memory.
• aarch64-apple-darwin: 1.5% faster, 0.8% less memory.

Well better than nothing :)