PEP 836: JIT Go Brrr: The Path to a Supported JIT Compiler for CPython

Thanks for your comments, Thomas.

The timeline and method JIT proposal are providing what the SC asked for, and we put a lot of thought and time into forming a cohesive PEP around them because they were explicitly requested.

From the SC’s post (extracting the two relevant bullets):

Also in the same post:

We are providing a concrete proposal for an alternative implementation that a lot of people have suggested over the years, with clear, measurable success metrics and timelines for that alternative implementation, as well as setting expectations about how stable and unstable the different parts of the current JIT architecture are in light of that transition.

So I think it would be inappropriate for the PEP to not include these. Further, shifting to the alternative design directly affects the timeline/maintainability/etc., so a separate PEP seems like it makes the process more complicated and opaque, not less. Here, we have a way to present a complete story.

(And if the SC feels differently, it can totally say “We accept the PEP, but we want to stick with tracing because of X, Y, and Z.”)

I think this is a misunderstanding.

Firstly, I want to re-emphasize that the uop IR is what’s currently in main, we are not introducing a new IR with this PEP. The main reason why we like the existing uop IR so much and don’t want to change it is because it’s literally how the base interpreter itself is written. It’s not just “understandable”, it’s what’s actually used today to define every single bytecode instruction. Furthermore, lowering is trivial (and mostly generated), keeping the Python frame in a sane state is easy (because the IR operates on the exact same stack as the bytecode)… there are many nice things about sharing definitions between the two parts of the runtime, and a concatenative, stack-based IR is a very straightforward way to do that in our experience.

So I would argue that yes, this is already understood well by anyone who works in the interpreter loop. More importantly, is it maintainable by a small team due to deriving from the interpreter. If you’re interested, take a look at What’s new in Python 3.15 — Python 3.16.0a0 documentation; I count 14 contributors to the JIT optimizer, of which only four are core devs. I think this is pretty strong evidence that the current JIT IR can be understood and is accessible to other members outside the core team.

That’s a huge part of why the method JIT is part of the proposal. There are a lot more people familiar with whole-function compiler architectures than tracing ones. It turns out the uops/stack also have nice SSA properties if you’re careful, and so tons of relevant theory is directly applicable to the bytecode-fragment IR we already have.

7 Likes

(Speaking here as a PEP co-author, not in my SC capacity; obviously I’ll recuse from any SC decision on this PEP.)

Outlining goals like we have in the PEP without being able to share the plan and everything we’ve thought through alongside it would make me uncomfortable. To me, goals without plans behind them aren’t really goals at all.

I also think it’s better to err on the side of transparency about our plans — both because the SC (not just this sitting SC, but previous SCs as well) has been asking for a crisper articulation of our plans and how the pieces fit together, and because hearing from the community could help us improve the plan.

7 Likes

Okay, I guess I understand where you’re coming from now. I still think we need to set the criteria for accepting a JIT, and I assumed this PEP would be that proposal (partly because it talks about goals in the Proposal section, I guess! :stuck_out_tongue: ) but after reading your response I took another look at the PEP title and, sure, it’s not actually supposed to be about setting the criteria. To me that still feels like it should be a separate PEP, though, albeit not one the SC explicitly asked for.

I will say, when I read[1] the SC’s request for a PEP for the JIT, I didn’t think we were talking about a PEP for a future JIT, but for the current JIT. If the plans in this PEP for the future JIT were rejected, that still leaves the fate of the current JIT unclear.


  1. and, admittedly, when I helped write ↩︎

Hi @thomas,

still think we need to set the criteria for accepting a JIT, and I assumed this PEP would be that proposal

You’re right here. It is that proposal! We do say in the PEP if the JIT doesn’t meet its goals we should reconsider removing it, and that the goals are the criteria for marking the JIT supported.

I didn’t think we were talking about a PEP for a future JIT, but for the current JIT.

I think there might be a misunderstanding here. This PEP isn’t proposing a new or future JIT. PEP 836 says:

The rest of the JIT (intermediate representation, middle-end/optimizer, and Copy and Patch backend), remain almost completely unchanged from CPython 3.15. In other words, only what the JIT selects to compile is evolving from traces to methods, nothing else is changing from CPython 3.15.

For people who are unfamiliar with recent JIT developments, the design we described might seem like a radical change, but most of it is already in the current CPython main branch! Again, the only change is growing support for methods.

Also my personal opinion: I don’t think the PEP must be strictly confined to the “current” implementation of the JIT with no room for slight changes. There’s no implicit blessing or formal acceptance of the current implementation in the first place (which, I guess is why we are all here now :upside_down_face: ). The burden is on the current implementation to show, through a PEP, that it should be kept in main was my understanding of the situation.

8 Likes

All future optimizations upon resuming JIT development will be reviewed with free-threading compatibility and performance impact required before merge. Optimizations that rely solely on the GIL build and break on the free-threaded build will be rejected.

It’s not only an issue of the free-threading runtime mode, but it’s going to be more intensified by free-threading. If the JIT compiler is expected to move memory reads/writes around (or eliminate redundant reads, move a read out of loop, …), which is a common and important optimization performed by any decent compiler, then Python needs some kind of memory model. Moreover, the memory reordering done by modern CPUs may already be an issue - the memory barriers included in within free-threaded refcounting probably help there, but one of the aims of the JIT optimizations will be eliminating the refcounting.

“Memory model” may seem like a big thing, but many programming languages have gotten a long way with pragmatic informal model along the lines of “use some annotation or language construct to prevent moving reads/writes around, otherwise anything can happen”. The important thing IMO is to clearly state that users should not rely on any ordering of concurrent reads/writes of plain memory locations and should use some abstraction provided by stdlib (recognized and intrinsified by the JIT - this can be easily polyfilled in older Python versions) or a language construct. 99.9% of the code doesn’t need to care and gets the benefit of advanced JIT optimizations. The 0.1% (or hopefully even less) of the code that performs unprotected concurrent memory accesses is inherently already doing complex concurrent programming, and the author should know very well what they are doing - ideally this kind of code would be found only in low level concurrency libraries/frameworks.

My concern is that specifying Python memory accesses as sequentially consistent could create a long-term constraint similar to exposing implementation details through the C API. These kinds of decisions may seem reasonable at the time, but by the time their drawbacks become apparent, it may be already too late.

Related discussion: What is the Long-Term Vision for a Parallel Python Programming Model?

6 Likes

@steve-s

Hi Stepan! Thanks for your very thoughtful post.

My concern is that specifying Python memory accesses as sequentially consistent could create a long-term constraint similar to exposing implementation details through the C API. These kinds of decisions may seem reasonable at the time, but by the time their drawbacks become apparent, it may be already too late.

To be clear: This PEP is not specifying that. For the avoidance of doubt: we’re saying that if your optimization currently only works on the GIL build and not on the FT build, it will be rejected. Not that we require sequential consistency. What CPython defines as “work” is currently underspecified, and even in some edge cases, it’s possible for the specializing interpreter (and indeed the normal interpreter itself), to load stale values or stale reads in attribute loads and other operations. For now, because it’s underspecified, we choose to not perform those optimizations and defer them to a future PEP. This does not mean the JIT (or any other JIT for that matter) can never perform those optimizations.

I will open a PR to clarify that we are not specifying JITs must guarantee sequential consistency, but rather that we (for now) will not introduce optimizations that only work on the GIL build.

Edit: PR up PEP 836: Clarify FT support != sequential consistency by Fidget-Spinner · Pull Request #5025 · python/peps · GitHub

3 Likes

The JIT should obey the memory model, not define it.

Having said that, I too am concerned that the memory model (mostly sequential consistency) of free-threading is going to be an issue for performance.

6 Likes

I agree that this seems like two PEPs. For now, we should have a PEP solely for establishing what are the goals for acceptance of the PEP, and what it would take to acheive those goals.

Choosing implementation details, like the nature of region selection, seems out of place in this first PEP. Even if it were PEP-worthy, the highly technical discussion it needs is very different from a
discussion about broad requirements, ease of support and timelines that a broad “JIT acceptance” PEP merits.

I too am a bit concerned about putting specific number on the goals as that could create some perverse incentives to not improve interpreter performance or to rush in features that should wait.

I can’t help but point out that past predictions for performance progress have rarely panned out precisely as predicted

Indeed GitHub - markshannon/faster-cpython: How to make CPython faster. · GitHub :slight_smile:

Regarding startup time and memory overhead, we’re already in good shape there.
So that’s more just a case of ticking a box than making any changes.

Warmup time is not well defined, and assumes behavior that might not exist so I think we
should leave that out [1602.00602] Virtual Machine Warmup Blows Hot and Cold

Regarding the understandability and maintainability of the JIT:
I think that the best way to improve that, by far, is with better documentation.
There are many complex parts of CPython, the parser, the GC, and the bytecode compiler, to
name a few. All of those are made much more understandable by having clear internal docs.

Pablo’s docs on the GC are an excellent example, making a complex subject much more approachable.

6 Likes

Hi there! :waving_hand: @thomas mentioned me so I tried to put together some thoughts.

First off, it looks like a ton of work went into this PEP and the PoC. Personally, I’m super excited to have the opportunity to discuss JITs in depth on a forum like this. So many thanks for all the work and for creating the opportunity to discuss it.

I agree with @thomas and @markshannon that this seems like two PEPs. I think that may be more a product of fulfilling what was being asked by the SC in the original announcement than any kind of oversight on your part. But now that it’s here, it seems like we would really want acceptance/rollback criteria to be defined in one PEP, especially if other JIT PEPs are on their way. From my own perspective, I would definitely want to see a few additional things either fully fleshed out or at least mentioned:

  • Most of the PEP talks about performance, which is great, but it seems to exclusively use geomean across platforms. I think there might be some missing discussion around the floor of performance. While it’s great to see a huge geomean, I worry more about the possibility of regressing whole classes of workloads.

  • As it stands the only real mention I see of memory is “may initially use more memory”. In my experience, method JITs in particular require more book-keeping (“Recording extra type profiling information”, “Path splitting”), holding invariants, pointers to ICs, previous compiled versions and code lifetime management (“Respecializing instructions”), etc. I would hope to see budgets for overall binary size, generated code size, and metadata memory size loosely discussed (even if the discussion is just “there will be a setting for this”).

  • I agree with @markshannon that compile latency/warmup is largely solved by the fact that you’ve chosen copy-and-patch, but I still think it should be explicitly called out outside the rejected ideas section of the PEP.

  • As mentioned by others, I would hope to see something like a real workload suite, not just pyperformance. At Meta we are struggling with this as well, as getting representative workloads is a bit of a sisyphean task. Nevertheless, I think it would be helpful to cover this even if the answer is “a small todo django app that we made”.

  • You already mentioned CET/BTI as per the Fedora request, so that’s good.

I realize this is a lot to ask, and maybe shouldn’t even have to be completed by the authors of this PEP specifically. But if there’s going to be acceptance criteria for a JIT within Python, I would hope to see most of those things covered.

Beyond the acceptance criteria, there are a couple of things I worry about with the direction this PEP is going in. I’ll go into detail below.

In my opinion, the PEP very much understates the sizable difference between a tracing JIT and a method-at-a-time JIT. It mentions “nothing else is changing from CPython 3.15” and that the middle-end “just needs to support merging type information at control-flow merge points”. I think this is glossing over a couple of important details. (Not to mention that the current incomplete PoC is already a 10K LoC change.)

A tracing JIT naturally captures the hot path, allows multiple entry points into a function by entering at back edges, and easily maintains branch history. That all must be replaced by value and control-flow merges, deopts, guard placement, and likely quite a few improvements to ICs. That’s not even to mention OSR (or some other explicit multi-entry/resume-entry design), which I would expect to become necessary for avoiding regressions on long-running loops because you’re replacing a tracing frontend and you will eventually need to support generators and coroutines. We’ve hit this problem with pyperformance and CinderX because a lot of the benchmarks are long-running loops.

The PEP mentions “Most optimizations can be implemented as local rewrites”. It’s totally possible I’m misreading this, or have a different definition of “rewrite”. But if I understand the intention correctly, I think this is contradictory to the other promised wins. Things like devirtualizing calls, lock removal, handling generators/coroutines are all outside of what I would say is “local”, and require having much more available context like control-flow and lifetime reasoning. Again, totally possible this is just a linguistic thing.

“The JIT can recover all of free-threading’s single-threaded performance losses and maybe even more.” I think this presumes the JIT will be able to eliminate many more locks and refcount operations than is actually possible. In CinderX we have a method-at-a-time JIT, and we still have a ton of these kinds of operations. It’s totally possible you can come up with some more clever optimizations than we have and eliminate some more, but from our own experience I can tell you I would be surprised if the JIT can eliminate most of them.

If I understand the PEP correctly, most of the promised performance gains are planned to be gained by the frontend change. As the PEP states multiple times in multiple ways, nothing else is changing. I don’t see how both of these things can be true. The backend is still going to be the hard limiting factor regardless. As I understand it, values across uops are limited to a small top-of-stack _PyStackRef cache rather than a general register-allocated/unboxed representation. The original copy-and-patch authors admit in their own paper “we conclude that copy-and-patch replaces the role of baseline compilers” and “we believe copy-and-patch can also be used as a fast profiling tier”. My concern is that this PEP lays out many person-years of work to change the frontend of a compiler architecture that explicitly states in its design that it is designed to be a baseline compiler, and then rejects a higher-tier JIT.

As a final point, I worry about the PEP’s focus on simplicity as a virtue of the JIT. In particular there are two quotes that stick out: “A single maintainer should be able to “fit” the entire system in their head” and “the JIT must remain accessible for contributors who do not work on it”. I’m uneasy with this framing. JITs and dynamic language runtimes are some of the more complex pieces of software out there. I promise I’m sympathetic to having too little energy and contributors, and it makes sense that you would want the things you mention. But I worry that eschewing complexity in favor of making a JIT compiler “fit” in someone’s head ultimately results in a JIT compiler not worth having.

Again, thank you very much for your work and thoughtfulness in putting together this PEP. I hope these comments are helpful in driving the conversation toward what some good acceptance criteria might look like and helping frame the discussion around the technical direction of this JIT in particular.

13 Likes

Hi! I’m just a random passerby whose interested in this project, and has been lurking in this thread for a while.

My question is: when can work restart on the JIT? Is it when this PEP is accepted? If this PEP gets split into multiple PEPs, will people have to wait for all of the PEPs to be accepted before work can continue, or do we just need the first PEP that explains the goals and performance targets before the PEP that explains the more technical side of thing.

Personally, I’m interested in either contributing or testing the JIT, so I’m really itching to see it get back on track, but I also understand that it’s really important to have an idea of where to go before continuing.

Also, another question, but, has there been an more indepth explanation of WHY copy and patch was chosen over other options? (Or, even, what copy and patch even IS) That might have been in the initial PEP (which I didn’t read), but I’d appreciate if an explanation of why copy and patch was chosen (over like LLVM ORC or Cranelift or whatever as a backend).

2 Likes

Hi @kddnewton, thanks for taking the time to share:

I agree with almost everything you said (regarding memory, performance floors, etc.)

I disagree however with the point that a tracing jit requires lots of changes to support methods for the current state of CPython. I agree that advanced versions of them diverge quite a bit. However, the current tracing JIT implementation in cpython is low-medium stage on the path to a full-blown tracing jit like say LuaJIT. We have not implemented the usual complex things like inter-trace optimizations, recompilation, unboxing, etc. This means the complexity to switch at this point is not high. It will be high if we had a more mature tracing runtime, but we do not. IMO, tracing and method runtimes only start diverging in implementation (though not conceptually) a lot when we implement the more mature optimizations. That’s why we put forth the PEP as it is for CPython right now.

Also I think some of the design difference points mentioned for a method JIT apply to a tracing jit in cpython as well?

  • Control-flow merges: You need this for loop peeling for LICM in tracing jits https://dl.acm.org/doi/10.1145/2384577.2384586
  • Guard placement: We already do this for the upstream tracing JIT.
  • Deopt: We also already do this for the upstream tracing JIT. This will also need to be done once we start doing serious optimizations regardless if it’s tracing or method?
  • Multi-entry resume/entry: We need this anyways to support generators. In tracing it’s an n-way branch keyed on instruction pointer for generator yields/resumes in a tracing JIT. The current tracing JIT cannot do that efficiently. I tried (again, I implemented the current trace recorder), and unless we implement self-modifying code or allow growing dispatch tables, the naiive implementation is just too slow. So there’s really not much difference here.
  • Branch history: This is done in the CPython interpreter, not the tracing JIT. So it works for both.
  • More advanced ICs: This I agree with. However, our IC infrastructure in CPython is unique in that we already have runtime specializers that can profile data at specialization time (Python/specialize.c). Instrumenting them to save extra info isn’t hard. So it’s not the same as in other runtimes.

In short, additional things you mentioned that a method JIT requires is already all conceptually required by a mature tracing one in CPython, though the actual implementation can differ quite a bit. I appreciate your expertise in CinderX and other JITs. However, the method frontend proposal sits in a unique spot not because of tracing vs method, but that the nuances in CPython make it easy to switch at this point of time (I can’t say it’s will be the same in the future!)

IMO, PEPs are only required IMO if we change a fundamental or specified part of CPython. The current tracing JIT is neither a core part of CPython nor is it specified in the first place. I don’t understand why the burden is on a (genuinely, for our current state at least) minor change to produce two PEPs, when the current design has no community consensus nor formal approval to begin with. If anything, applying the same formal process in CPython means that the current upstream JIT must have two PEPs (one justifying its implementation, one criteria), same as any other implementation that is experimental. However, I don’t think the original SC announcement requested for that, and I don’t want to put the burden of producing two PEPs on the current upstream JIT as well.

6 Likes

As laid out in the PEP, I believe that the tracing-vs-method transition is a key part of “what it would take to achieve those goals”. It has a direct impact on both ease of support and timelines. In fact, the entire first phase of the PEP’s proposed roadmap is basically “ratchet up support while pruning the sharp edges of the implementation”. This is a big part of that.

I disagree that better documentation is the best way to improve the understandability and maintainability of the JIT, though it’s certainly a big part of it. I spent four years working full time on this JIT and giving lots of talks on how it works, and explaining it only got harder with time, not easier. If the system no longer fits in my head, that should at least be minor cause for alarm regarding the implementation itself. The super-bespoke nature of our frontend definitely isn’t doing us any favors in that department, especially because it offers so little in the way of abstraction. The only way to really understand how it works at a high level is to understand how it works at a low level.

Even more concerning than the complexity of the system itself is its behavior, which affects both maintainers and users. The more complex the behavior is, the harder it is to anticipate, debug, and fix the growing surface area of edge cases.

As a hopefully-illustrative example that I think anyone either familiar or unfamiliar with the JIT’s implementation can ponder, consider:

from graphlib import TopologicalSorter

list(TopologicalSorter(LARGE_GRAPH).static_order())

main’s tracing JIT requires over 40,000 nodes in LARGE_GRAPH before it reaches a “steady state” (where the only things it continues to compile are repeated failures), and during that time it compiles over 30 unique traces. Figuring out what those traces represent at a source level, and how they are rooted and connected together, is a “fun” way to burn a couple of hours (if you try to do it manually, at least).

A method JIT with similar warmup thresholds requires only ~8,000 nodes to reach a steady state, and compiles a grand total of four functions (plus three more if we want to implement on-stack-replacement in the future). I bet you could open the graphlib source code and quickly make a pretty accurate guess what gets compiled, and in roughly what order.

5 Likes

I believe that’s what the SC said.

Depends on what the split is for. The SC wants a PEP specifying the acceptance criteria for the JIT, so if there’s just one PEP that covers that need then that’s the PEP blocking work on main.

2 Likes

Hey @kj0, thanks for the detailed reply.

Totally see what you’re saying about rewriting the frontend before doing inter-trace opts, recompilation, unboxing, etc. Personally it still seems like a large change to me, but I will defer to your expertise and we can agree to disagree on that.

I agree that the different frontend definitely share a lot of the same concerns, but disagree that the implementations could be shared. To respond to your points directly:

  • Control-flow merges: thanks for the link to the paper, I had not heard of loop-peeling before. Reading through, I’m not sure this is the same kind of control-flow merge I am talking about though. It looks like from their optimization, it produces just one merge between the loop preamble and the body, and the types are already the same (since it relies on other opts like CSE running in order to be beneficial). In a method frontend every join in the CFG would be a type merge, and those are both bigger and harder to get right, which is why the splitting mentioned in the PEP is necessary at all. To be clear, I’m not saying this is or should be a blocker, all I’m saying is I don’t think the PEP should reduce this to “just needs to support merging type information at control-flow merge points”, as that undersells the work fairly heavily.
  • Guard placement/deopt: Yeah for sure, you have to have these for tracing. I’m saying more that these have to change to accomodate the new frontend.
  • Multi-entry: So there’s two things here: generator/resume and loop OSR. I will totally take your word for the generator/resume stuff. But I don’t see (or maybe I missed) the PEP discuss how it would replicate the current tracing JIT being about to enter a hot loop because it starts on a back-edge, whereas a method JIT starts at main entry. This is another one where I’m not saying this is or should be a blocker, but it seems like something the PEP should call out.
  • Branch history: Makes sense.
  • More advanced ICs: I see. My point here was that there would of course be an increased memory load because of this, but that gets wrapped up in the other points.

Beyond those points, there are a couple of things I’d still like your take on from my original post. Specifically the copy-and-patch backend and the part about free-threading.

As a final thing regarding process: I definitely agree with you, and that’s why I wrote that you probably shouldn’t even have to write the acceptance criteria PEP yourself. I think that would actually make this PEP more easily reviewed, because then it wouldn’t be trying to be both a design document and an answer to the SC’s questions.