hook for JIT to support `_PyFrame_GetFrameObject`

[carljm]

A JIT will likely not want to maintain the full state of a _PyInterpreterFrame (most importantly, stack/locals; maybe also prev_instr) during execution of JITted code. In general, it should be the JIT's responsibility to ensure the interpreter frame is in the correct state before it returns control to the interpreter.

However, the functions PyEval_GetFrame and PyThreadState_GetFrame, which return a full PyFrameObject for the current interpreter frame, are public C API. PyThreadState_GetFrame is used internally by the coroutine_origin_tracking_depth feature and by PyErr_WriteUnraisable.

sys._getframe, sys._getframemodulename,sys._current_frames, and construction of tracebacks also reify one or more PyFrameObject from interpreter frames.

All of the above ultimately use _PyFrame_GetFrameObject, which takes a _PyInterpreterFrame and returns a PyFrameObject.

A JIT that does not constantly maintain the full state of the executing _PyInterpreterFrame will need some way to be notified when a full PyFrameObject is needed for one of its frames, so that it can populate the frame object appropriately based on its own internal execution state.

I'd ideally like to find a solution here that will work both for Cinder JIT and for future faster-cpython JIT, even given their differing designs.

So, questions:

1. Is the feature request sensible? Is there an entirely different way that a JIT should handle this?
2. What shape should the API take? Perhaps the simplest possibility is a single interpreter-wide function pointer that can be set, whose default value is _PyFrame_GetFrameObject.

[carljm]

For context, for a long time there was a similar hook (added for the sake of Psyco way back in 2002!), but it was removed in Python 3.9 (python/cpython#84127). But the implementation of frames has changed a lot since then.

That hook allowed overriding PyThreadState_GetFrame entirely. That's more capable (because it allows the hook setter to track their frame stack entirely separately from the interpreter; I don't think Cinder needs that capability). But with today's implementation of frames, that form of hook wouldn't be sufficient, because e.g. sys._getframe doesn't use PyThreadState_GetFrame. If it did, it would have to walk the chain of PyFrameObject to the one requested, materializing every frame along the way, which is a lot less efficient than what it does today, which is walk the stack of _PyInterpreterFrame and only create a single PyFrameObject for the requested frame.

[markshannon]

The whole point of the _PyInterpreterFrame is that it is cheap to create. So, in theory, Cinder should be able to use it and not worry about frame objects at all, as the VM would take care of them.

How could we make _PyInterpreterFrame even cheaper, so that you can do that?

You'll need a link, so previous has to remain.
The f_globals and f_builtins fields could be dropped, but that might hurt the performance of LOAD_GLOBAL.
Alternatively we could restrict their use to the interpreter, so that Cinder could ignore them.
We do need a way to access the globals and builtins, though, so f_funcobj seems useful.
f_locals is the "slow" locals, so just needs to be NULLed.
prev_instr is needed to find the line number. I guess we could allow it to be NULL, but setting it to NULL isn't much cheaper than setting to the start of the function.

There's a word of offsets and flags. They are introspected, so can be ignored.

That leaves the larger part of the frame: the locals and stack.
The stack isn't visible outside of the interpreter, so there is no need to set the stack.
That leaves the local variables. You'll need to store the values somewhere across calls (and there aren't many callee save registers left after the VM state), so why not put them in the frame?
They can be kept in register between calls.

[carljm]

Yep, I'd already come to the same conclusion that a) we should just use _PyInterpreterFrame in 3.12, and b) populating most of it shouldn't be a big issue for us. The locals are indeed the sticking point.

That leaves the local variables. You'll need to store the values somewhere across calls (and there aren't many callee save registers left after the VM state), so why not put them in the frame? They can be kept in register between calls.

This is an interesting suggestion. We may need to do some experimentation to see how much worse it would be than our current approach (we spill to C stack only values stored in caller-saved registers.) Off the top of my head there's a few components that could potentially make it worse: a) locality and code size effects of spilling to Python stack instead of C stack, b) needing to always spill all locals, not just those stored in caller-saved registers, c) needing to consult runtime metadata that maps registers to local indices, which we currently only have to consult when we deopt.

[carljm]

After looking at this a bit, I think it's clear that spilling locals to the _PyInterpreterFrame around all calls is going to be too expensive to be feasible. It would be much better to find a solution that doesn't constrain an optimizer to always (on calls) maintaining locals (or ideally, anything else it doesn't need) in the _PyInterpreterFrame.

I also think a global hook maybe isn't a great idea; it makes it hard for multiple optimizers to coexist.

A better idea might be to reuse f_funcobj, since it's currently only needed internally to the interpreter, for fetching the function closure. So we could set that instead to a vector-callable object that knows how to fully update the _PyInterpreterFrame, and then we'd just need _PyFrame_GetFrameObject to do something like if (!PyFunction_Check(frame->f_funcobj)) { PyObject_Vectorcall(frame->f_funcobj, frame, 1, NULL); } before going ahead and making the full frame object. Or if the PyFunction_Check is too implicit, there's room for a few new bitflags on _PyInterpreterFrame without actually making it bigger.

[carljm]

After looking more into this, it seems like we could be OK with having f_locals on a JITted frame just be an empty dict. But there is one other problem: prev_instr. That ends up responsible for f_lineno on a frame, and we do need these to be correct. Updating prev_instr on the frame before every call (and every instruction that can trigger arbitrary code execution) would not be ideal for performance. So I think we still don't see a good way around having some kind of "frame reification" hook :/

[DinoV]

So I've been looking into this and it seems like the approach @carljm suggested (making f_funcobj a vectorcallable to initialize the frame) and it works pretty well. On 3.12 it's a 94 line change to check for the object and invoke it when we're potentially inspecting a less-than-fully-initialized frame. It's then the responsibility of the object to get the frame into a fully initialized state and restore the original function.

The end result is we only need to initialize f_funcobj, f_code, previous and prev_instr. We also need to store the function object somewhere (we don't want to store it in the object - each JITed frame can be associated with multiple function objects).

Something that's a little bit ugly is that to pass the _PyInterpreterFrame * via the vectorcall we need to convert it into a long and pass the address. Something else is to pass the test suite we need to disable trace malloc around this allocation otherwise we can end up an extra unexpected trace.

I'm also a little worried about using f_funcobj, it seems like one of the fields someone would be interested in (FWIW it appears PyTorch is already looking at this field and assuming it's a function). There may be other code out in the wild looking at it.

Another possibility would be doing this with frame_obj. It leaves us with the right place to store the function object and it's a somewhat logical place - a frequent reason to access this is to reify the frame. It is strongly typed to PyFrameObject right now but that's an easy change given it's all internal.

[DinoV]

Okay, getting back to this... Shortly after this was added @markshannon added python/cpython#100987 so that seems like the natural place for this to finally land. I'm thinking what could be done here it to add a new executable type like:

typedef int (*_PyFrame_Reifier)(struct _PyInterpreterFrame *, PyObject *state);

typedef struct {
    PyObject_HEAD
    PyCodeObject *je_code;
    _PyFrame_Reifier je_reifier;
} PyUnstable_PyJitExecutable;

PyAPI_DATA(PyTypeObject) PyUnstable_JITExecutable_Type;

#define PyUnstable_JITExecutable_Check(op) Py_IS_TYPE((op), &PyUnstable_JITExecutable_Type)

PyAPI_FUNC(PyObject *) PyUnstable_MakeJITExecutable(_PyFrame_Reifier reifier, PyCodeObject *code);

...

const PyTypeObject *const PyUnstable_ExecutableKinds[PyUnstable_EXECUTABLE_KINDS+1] = {
    [PyUnstable_EXECUTABLE_KIND_SKIP] = &_PyNone_Type,
    [PyUnstable_EXECUTABLE_KIND_PY_FUNCTION] = &PyCode_Type,
    [PyUnstable_EXECUTABLE_KIND_BUILTIN_FUNCTION] = &PyMethod_Type,
    [PyUnstable_EXECUTABLE_KIND_METHOD_DESCRIPTOR] = &PyMethodDescr_Type,
    [PyUnstable_EXECUTABLE_KIND_JIT] = &PyUnstable_JITExecutable_Type,
    [PyUnstable_EXECUTABLE_KINDS] = NULL,
};

The code object is there to support typical introspection. The callback is there to support reification of the frame - so when someone calls sys._getframe(), wants the lasti, etc... the runtime will call that and let the JIT update the current frame state.

Ideally we'd also add a new FRAME_OWNED_BY_... to indicate that this frame doesn't have locals (so stackpointer, locals, etc...) don't need to be populated as well.