gh-117958: Expose JIT code via access method in experimental UOpExecutor

[tonybaloney]

Adds the get_jit_code() access method to the UOp Executor along with the existing access methods.

This is only accessible via internal C APIs but would be helpful for testing and debugging.

• Issue: Expose jit_code field for UOp Executor #117958

[tonybaloney]

@brandtbucher copy of the original PR to your JIT branch.

[tonybaloney]

This could return an empty string instead of a error

[gvanrossum]

I think the exception makes more sense -- though maybe you should only check for jit_code == NULL.

[gvanrossum]

This LGTM. I have a nit for the news item and a suggestion for the actual code.

Curious what you're planning to do with this?

We might also worry about security implications -- while this doesn't allow writing the JIT code, it might give an attacker an easier way to analyze the JIT code and look for vulnerabilities. Though they can access this using ctypes as well, of course.

[gvanrossum]

I think the exception makes more sense -- though maybe you should only check for jit_code == NULL.

[tonybaloney]

Curious what you're planning to do with this?

I found a feature like this really useful for debugging in Pyjion and other JITs like ryuJIT. I know there are some other compiler folks who would use this too. There was a discuss thread where someone else was asking for this linked in the issue.

1. dumping the machine code into a file and disassembling it to do some analysis
2. looking at CFGs to understand the control flow and compare it with other JITs
3. (hopefully) emitting some debug symbols in future or at least markers for which offsets relates to the higher level instructions

Some security teams may also want the ability to export the JIT code for analysis, beyond what you can gather by looking at the C templates.

[brandtbucher]

I'm slightly leaning towards ditching the exceptions and returning None in the case where either _Py_JIT is not defined or jit_code == NULL, and an empty string if jit_size == 0.

I'm thinking about things like regression tests or runtime introspection (where it's more ergonomic to check for None instead of catching an exception in cases where the JIT was built but is disabled, etc.). Not a huge deal, but I think I'd personally rather see an empty string if the JIT code is empty or a None if no JIT code exists when debugging a buggy JIT. :)

Otherwise, this looks good. What do you think?

[tonybaloney]

@brandtbucher amended with your feedback.

• Removed the extra check (descriptors check seems to catch anyone trying to call it on another type)
• Now returns None instead of raising an exception if there is no JIT code.

[diegorusso]

Hello, thanks for the PR! It certainly does the job of capturing the machine code generated by the JIT but I was hoping to have a map between the uop byte code and the related machine code similarly to what I was envisaging here

[gvanrossum]

I'm happy to merge at this point, but did you consider putting this line inside #ifdef _Py_JIT instead? (And then the entire function definition as well.) That would make it possible to test whether this functionality exists without calling it, which is generally Pythonic API design.

[brandtbucher]

Hello, thanks for the PR! It certainly does the job of capturing the machine code generated by the JIT but I was hoping to have a map between the uop byte code and the related machine code similarly to what I was envisaging here

So, I've thought about this, and it should be possible with a couple of tweaks.

Basically, this current PR returns a byte string, which consists of the code for each instruction in sequence, followed by the auxiliary data for each instruction in sequence.

Meaning, for a trace of:

[A, B, C, D]

It returns:

b"".join([<A code>, <B code>, <C code>, <D code>, <A data>, <B data>, <C data>, <D data>, <padding>])

However, the executor knows the uops that make up its trace. If we #include "jit_stencils.h", we should be able to use stencil_groups[instruction->opcode].code.body_size and stencil_groups[instruction->opcode].data.body_size to compute these chunks.

Maybe @tonybaloney and @diegorusso can confirm, but it seems like the most useful info to return would be a 3-tuple of base address, a list of code byte strings (corresponding to uops) and a list of data byte strings (again, corresponding to uops).

So, for the above example, the return value would be:

(
    <base address>,
    [<A code>, <B code>, <C code>, <D code>],
    [<A data>, <B data>, <C data>, <D data>],
)

(I think base address is needed for some absolute addressing that we use in places.)

So each of the code or data lists can be zip'd with the executor to map them to individual uops. And if I want the raw string of data that this PR returns now, I can just take this tuple and do b"".join(result[1] + result[2]).

Would this meet everyone's needs, or am I overthinking it? Even though it's internal, I don't want to tweak this too much after the beta freeze on Monday, so I'm leaning towards providing more information rather than less.

[markshannon]

I'd be inclined to keep it simple (just returning a bytes object) for 3.13 as feature freeze is imminent.
We can always implement a richer API for 3.14.

Unless, someone really needs the fancier API and is able and willing to implement it in the next two or three days.

[diegorusso]

Hello, thanks for the follow up. I was going through a different route by adding a couple of fields in the executor struct

+    size_t *instruction_starts;
+    size_t trace_length;

and then work out where every instruction starts.

Anyway, because the feature freeze is imminent, I would vote for accepting this PR as it is, improve it in the next cycle and dedicate more thinking to the API. Better something good enough than nothing perfect :)
I will create a new issue with what @brandtbucher has suggested in his comment so we don't lose track of it.

Also it helps the fact that I'm off for a few days and I would miss anyway the feature freeze deadline.

[gvanrossum]

Okay, then I'll merge it as is.