gh-145742: Manually emit _LOAD_FAST_BORROW to reduce stencil bloat

[corona10]

• Manually emit _LOAD_FAST_BORROW at JIT compile time, encoding the operand offset directly into the instruction instead of loading it from the GOT at runtime.
• This shrinks the generic case (oparg ≥ 8) from 28 bytes to 8 bytes and eliminates 27 stencil functions.
• I've compared machine code through godbolt:
  • x86-64 / aarch64: https://godbolt.org/z/3M9zKeosj
  • i686: https://godbolt.org/z/cdjdzev5Y

• Issue: Some (mostly) easy ways to reduce the size of JIT generated code #145742

[corona10]

For i686: https://godbolt.org/z/cdjdzev5Y

[diegorusso]

Some initial feedback on this:

• we should not to pollute jit.c with uops implementation. They should live in separate compile units and have the same signature of the other ones (e.g.: void emit__UOP_NAME(unsigned char *code, unsigned char *data, _PyExecutorObject *executor, const _PyUOpInstruction *instruction, jit_state *state)
• ifdefs can select the right architecture of the custom implementation
• in bytecodes.c we should have a way to tell the JIT machinery not to generated any code for a specific uops but the uops implemetation should be accounted in the table in the jit-stencils-*.h (static const StencilGroup stencil_groups[MAX_UOP_REGS_ID + 1])
• The linker later on will pick up our own version of the uops implementation.

[corona10]

Thanks, @diegorusso. I’ll keep working on this based on your feedback.

[markshannon]

A couple of other things:

• This PR asserts that the immediate value fits into the space given, but this will fail for larger opargs.
• I don't if this matters, but the x86 code is inferior to that generated by the stencils for oparg 0-5. For example, for LOAD_FAST_BORROW_1_r01 in the stencil generated code uses a 1 byte offset instead of the 4 byte offset this PR generates. For oparg > 5, the code is the same.

I think you need to split _LOAD_FAST_BORROW into two variants for the JIT. One for all normal opargs, that can use manual code generation, and a generated fallback for huge opargs.

[corona10]

@markshannon I would like to get feedback with current imolementation before starting to refactor as Diego suggested. :)

[markshannon]

The generated code looks nice and efficient.

You'll need to split the two parts into different uops and insert them in the fix up pass in uop_optimize, as the JIT can only handle fixed length stencils.

E.g you might need LOAD_FAST_BORROW_0_to_5, LOAD_FAST_BORROW_TO_500. Then replace LOAD_FAST_BORROW with them in uop_optimize. Mark them both as "manual" or something like that, so that the stencil generator knows to omit them.

[brandtbucher]

Are there benchmarks for this yet?

[brandtbucher]

In general this is a big increase in scope and complexity for the JIT backend, so I'd definitely want to make sure it's justified.

From a maintenance perspective, can I propose an alternate design that I figured would make sense if we ever wanted to do something like this? Basically, stick a directory full of .S files in the Tools/jit directory. So, we could add a file, Tools/jit/overrides/_LOAD_FAST_BORROW.S that contains syntax-highlighted, manually-written textual assembly. The JIT build scripts would just use that instead once we re-compile everything after the textual assembly pass, and the rest of the backend stays unchanged. We can also put a hash at the top of the file to make sure we check the file whenever the source of its corresponding uop changes.

That seems like a much more maintainable, scalable way of doing this, IMO.

[corona10]

Are there benchmarks for this yet?

With Ken Jin’s x86 benchmarks, it does not show any regression. That is not too surprising, though, since this change only reduces binary size.

From a maintenance perspective, can I propose an alternate design that I figured would make sense if we ever wanted to do something like this

I agree. If we want to do this in better way, we may need an assembly generator like ZJIT, or use DynASM. Alternatively, we could go with the approach you suggested.

[brandtbucher]

I don’t think we should do this if it isn’t faster.

[markshannon]

So, we could add a file, Tools/jit/overrides/_LOAD_FAST_BORROW.S that contains syntax-highlighted, manually-written textual assembly.

Sadly, that won't work. It isn't the assembly, but the patching that needs customization.

Take this example from the PR:

// AArch64: preserve_none CC: x21=frame, x24/x25/x26=cache0/1/2
// Small oparg (imm12 fits): ldr x8, [x21, #off] ; orr xDST, x8, #1  (8 bytes)

We need to compute the offset at JIT time, and somehow communicate to Clang that #off needs a custom relocation that it isn't aware of.

I don’t think we should do this if it isn’t faster.

This change will be too small to measure, but it is still an improvement. The assembly code is objectively better.
If the change is clean and makes it easy to add other similar improvements, then it will be worth it.

[brandtbucher]

Sadly, that won't work. It isn't the assembly, but the patching that needs customization.

Take this example from the PR:

// AArch64: preserve_none CC: x21=frame, x24/x25/x26=cache0/1/2
// Small oparg (imm12 fits): ldr x8, [x21, #off] ; orr xDST, x8, #1  (8 bytes)

We need to compute the offset at JIT time, and somehow communicate to Clang that #off needs a custom relocation that it isn't aware of.

This is a just 12-bit immediate for a load instruction, right? The existing JIT does this all the time. You can encode this by prefixing the symbol with :lo12:.

Here's what that looks like for your example. If we add "R_AARCH64_LDST64_ABS_LO12_NC": "patch_aarch64_12" to our relocation mapping in _stencils.py, the existing backend just works.

I don’t think we should do this if it isn’t faster.

This change will be too small to measure, but it is still an improvement. The assembly code is objectively better. If the change is clean and makes it easy to add other similar improvements, then it will be worth it.

I still disagree. The JIT backend's goal has never been to make the "objectively best" assembly, it's to make fast code with the right maintenance tradeoffs. It was (and still is) carefully designed to be as easy-to-reason-about and hard-to-break as possible, while still giving us quite performant code. This PR hacks some really subtle special cases onto that for no performance benefit. And I certainly don't want to make it "easy to add similar improvements", because I'm still not convinced that this is an improvement at all.

(I also have a hunch that the additional overhead from this patch will negate any potential wins in practice. But as you note, it's be impossible to tell, since it's all in the noise anyways. 🙃)