Intro
Numba has become an incredibly popular package for JIT compiling numerically oriented functions in Python. We have reached over 10 million downloads per month from our large user base and active community. However, the popularity also brings the burden of maintenance. The complexity of the Numba code base has made introducing new maintainers or even just making new features a growing challenge. This problem has become apparent in the last few years, maintainers and contributors have felt that progress has slowed down. There are many factors contributing to its high complexity, but to highlight a few:
- Numba is a monolithic Python JIT compiler that contains re-implementations of core Python functions and a significant subset of the NumPy array library; it also targets multiple hardware architectures — CPU and GPUs.
- Numba was originally designed to be a straightforward numerical/array-oriented compiler but the community demands have over time transformed Numba into an extensible compiler project.
- Early design decisions introduced assumptions that are no longer useful, but these assumptions are scattered across the code base and are difficult to undo without breaking users in production.
- Ecosystem changes—CPython bytecode changes, NumPy API changes, LLVM updates etc.
There are many more factors but we will not go into the details in this document. The key note here is that Numba needs to refactor and redesign parts of itself in terms of both development process and code to sustain and grow into an extensible compiler toolkit.
Modularizing Numba
We need smaller, self-contained components that can be composed and rearranged to form both new compilers and replace portions of the existing Numba compiler. These new components add a separation of concerns in both code and developer teams, avoiding the need for maintainers and contributors to have a good understanding of every part of the compiler. To start with, we are proposing work on the following:
- A Regionalized Value-State Dependence Graphs (RVSDG) based bytecode analysis frontend.
- A JIT-friendly shared library format.
- A new intermediate representation (IR).
There are many more components needed as alluded to in the introduction. However, this document only lists the above components as these ideas and their immediate usefulness have become apparent.
A RVSDG bytecode frontend
A large portion of the effort in maintaining Numba is dealing with the ever changing Python bytecode. In the early days of Numba, the bytecode was quite stable. However, the faster-cpython project by the CPython developers has led to the bytecode repeatedly changing in a significant way and we envisage future related challenges in coming years. We cannot continue our current implementation strategy for analysis that relies heavily on CPython-version specific bytecode behavior. We need to move to a bytecode analyzer that is more resilient to bytecode changes. As part of making changes for Python 3.11, we have begun the work on a new bytecode analyzer based on concepts from RVSDG, which have useful properties that can simplify many of the analyses we do to untangle the Python bytecode.
A JIT-friendly shared library format
The user community has repeatedly asked for code caching and ahead-of-time (AOT) compilation. Code caching and AOT are two related concepts. Both require storing executable code on disk. A lot of the challenges for Numba to implement these features came from its JIT-oriented nature. However, we now realize it is much easier to support a JIT compiler when working from the base of an already existing AOT compiler. The cornerstone of a new AOT offering is through producing a shared library that embeds low-level compiler IR (such as LLVM IR) for JIT compilation and link-time optimization. This led to the idea of PIXIE (Portable Instructions eXchanged In Executable).
PIXIE is a Python package for creating specially-formatted shared libraries of functions that give the user the choice about portability and performance through the whole program lifetime. PIXIE tries to combine the best of AOT and JIT by putting the choice about execution-latency, compilation-performance and runtime-performance in the hands of the user. A PIXIE library typically contains multiple architectural variants of given functions, and also LLVM bitcode of the same. A PIXIE library often comes as a Python C-Extension and has a Python “overlay” at the module root that makes it trivial to introspect and call the contents of the shared library.
A new intermediate representation (IR)
The existing IR has numerous problem:
- It was not designed for efficient transformation.
- Accidental mutation of the IR is easy and often leads to malformed IR.
- The IR lacks a formal specification or grammar.
- It is very low level and reflects the CPython bytecode.
Extending the benefits from the RVSDG frontend to the rest of the compiler is another reason for a new IR. These benefits can simplify many compiler passes and reduce compilation overhead. (See Value Dependence Graphs: Representations Without Taxation, and RVSDG: An Intermediate Representation for Optimizing Compilers).
Near-term plan
The proposed new components are designed to complement and replace parts of the existing Numba compilation pipeline:
Potential milestones:
- To address the need from Python 3.11+ faster-cpython bytecode changes, we need to replace the “frontend” portion of the compiler pipeline with the new RVSDG bytecode analyzer. The code lifting passes (with-lifting and loop-lifting) will be rewritten to use the new information from the RVSDG forms. The new passes should become simpler due to the “regionalized” nature of RVSDG.
- Experimental PIXIE support can be added to further investigate the effectiveness of AOT compilation with JIT LTO. Targeting PIXIE libraries as output may need an experimental compiler pipeline to create them.
- Experimental pipelines can be added to investigate the usefulness of the new IR.
NOTE: This is a high level document. More detailed information and proposals will be provided in a separate document.