JIT: map uops with code generated by the JIT

Feature or enhancement

Proposal:

This is a follow up feature of https://github.com/python/cpython/issues/117958. That feature exposes just the bare JIT Code of the executor object.

To improve the debug experience of the JIT implementation, a map between the UOp and the generated code should be implemented.
@brandtbucher suggested the following in his comment

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.

Originally posted by @brandtbucher in https://github.com/python/cpython/issues/117959#issuecomment-2088028660

Has this already been discussed elsewhere?

I have already discussed this feature proposal on Discourse

Links to previous discussion of this feature:

https://discuss.python.org/t/jit-mapping-bytecode-instructions-and-assembly/50809

https://github.com/python/cpython/issues/117958

首先阅读相关 issue 117958 和 117959、Discourse 讨论以及 jit_stencils.h。跟踪 executor 如何公开其 JIT 代码，以及 stencil_groups 如何提供代码和数据大小。完成的标志是公开一个基地址，以及按 uop 划分的代码和数据字节字符串列表，并且这些列表可以映射回 executor 的 uop。