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

Bắt đầu bằng cách đọc các issue liên quan 117958 và 117959, cuộc thảo luận trên Discourse và jit_stencils.h. Theo dõi cách executor cung cấp mã JIT của nó và cách stencil_groups cung cấp kích thước mã và dữ liệu. Công việc được xem là hoàn tất khi cung cấp được một địa chỉ cơ sở cùng các danh sách chuỗi byte mã và dữ liệu theo từng uop, có thể ánh xạ ngược về các uop của executor.