There have been a lot of articles written following Modular’s announcement of
Mojo, a Python “successor” language that
will bring high performance computing (especially for AI applications) to
Python users. Here at Anaconda, we’ve been working on the problem of HPC and
Python for over a decade now. As part of that work, we created
Numba in 2012, a Python just-in-time compiler
specifically for numerical computing, and have been improving it ever since.
We’ve learned a lot during that time about what it means to compile and
accelerate Python, and thought it would be useful to share some of that
perspective here. Along the way, you’ll learn about some of the design
tradeoffs Numba has had to navigate, and maybe get some context for the
challenges that Mojo will face as it takes a different (but potentially
complementary) approach to Python compilation.
When is Python useful for HPC?
The difficulty in deciding how Python should be best used for high performance
computing (HPC) hinges on the difficulty in defining the purpose of HPC. HPC
is often discussed in terms of maximizing execution speed of large, complex
calculations. For these purposes, Python is not very useful, except maybe as
a job configuration or control language delegating nearly all computation to
highly optimized components. (This is the traditional “glue language” role
that Python is very good at.)
However, the economies of scale in software can confuse our understanding of
priorities. Some kinds of software, for example linear algebra algorithms,
are so ubiquitous across disciplines and so frequently used that they merit
spending a huge amount of developer effort on performance. There are almost
no situations where you should implement matrix multiplication yourself.
Instead, you should use MKL, or OpenBLAS, or cuBLAS, or some other library
that has had a tremendous amount of effort from hardware experts invested into
it. The time spent developing and debugging these sorts of libraries is tiny
compared to the amount of time the library will be used by the entire software
community. Additionally, it is often possible to find resources to fund these
fundamental libraries from hardware vendors and non-profit funding sources
because these applications are pivotal. This is the “HPC mass production”
scenario.
But many software projects are not in this “mass production” category. They
are research projects where a small team of developers needs to iterate
quickly to discover a solution to a particular problem. Or it is an
established project where maintenance costs and barriers to contribution need
to be minimized because the team is small, or turnover (as in many academic
research contexts) is high. The incredible diversity of use cases in the
Python numerical computing ecosystem is its strength, but also means that
development processes optimized for “mass production” are not necessarily the
best choice. Speed matters, but total time matters more, including developer
onboarding time, feature development time, debugging time, maintenance time,
packaging time, as well as execution time.
This is the vast world of “high-enough performance computing,” and is full of
tradeoffs. Numba seeks to find a balance that reduces total person time in
the numerical computing space, of which execution time is one component. Mojo
is trying to do a similar thing, but with a different set of tradeoffs.
(Literally anyone who isn’t hand-coding assembly language for critical code
paths is trading execution speed for something.) Hopefully I can help
illustrate some of these tradeoffs in the following sections.
Python isn’t a language, it is a community
When people talk about the characteristics (like “speed” or “ease of use”) of
“Python,” they are usually mashing together a bunch of distinct parts that are
important to think about separately when designing a strategy to improve
Python in some way. These parts include:
- The Python language specification
- The Python interpreter (usually CPython for most users)
- The Python standard library
- The huge collection of third party packages
- The network of users and developers who help create, maintain, and teach others about all of the above
When a new solution to “the Python performance problem” is announced, I find
it useful to understand how it impacts each of these items separately. For
example, Numba makes the following choices:
-
Python language: No changes. Numba uses standard Python language
features, like decorators and context managers, to intercept execution and
annotate code for compilation.
-
Python interpreter: No changes, although Numba is CPython-specific. Numba
is specifically designed to work inside CPython as a regular module to
minimize barriers to use and maximize application compatibility.
-
Python standard library: No changes, but Numba can only accelerate a small
number of standard Python types and built-in functions. Applications are
free to use the entire standard library, but only outside of Numba-optimized
functions (with some exceptions).
-
Third party packages: Numba does not prevent usage of third party packages
in an application because Numba is confined only to specific functions the
developer chooses to compile. However, Numba cannot compile code that calls
most third party libraries either. Numba includes specific optimization for
usage of NumPy arrays and functions, and support for
other libraries can be added using Numba’s extension mechanism (for example
Awkward Array).
-
Users and developers: Numba tries to be easy to pick up for anyone
familiar with Python and NumPy, although more advanced capabilities require
learning a few new concepts. Numba also strives to be easy to integrate
into other Python packages, avoiding a compile step during packaging, in
some cases allowing packages to be distributed as pure Python source code
which is compiled for the user’s specific hardware at runtime.
In general, Numba is trying to strike a balance of minimizing the barrier to adopting Numba in a numerical application by existing simply as a module that is compatible with the standard CPython interpreter. The application author does need to actively opt-in to Numba for each function that needs to be compiled, but we find this minimizes surprises and leads to more successful usage in the long run.
Mojo takes a different approach:
-
Python language: Mojo adds significant new syntax to the Python language, aspiring to be a “successor language” for Python that offers more low-level types and performance control. In particular, Mojo defines a new kind of function (using the “fn” keyword) that will be more strictly typed and less dynamic so that faster code can be generated. Regular functions (defined with the standard “def” keyword) have more dynamic and implicit behavior, which will make them feel more familiar to Python developers, but they are not “Python functions” in a strict sense.
-
Python interpreter: Mojo has its own runtime interpreter (called “mojo”) which must be used to call Mojo code, and a compiler for standalone binaries will be available as well. Mojo also links to the CPython runtime, which it uses to allow it to import Python modules and call methods on Python objects.
-
Python standard library: Mojo has its own standard library, which currently is very low-level (at least in the public documentation). Mojo modules can access the Python standard library, via the Python interop functionality.
-
Third party packages: As the Mojo compiler is not generally available, no one has made any Mojo packages. It is unclear how they will be distributed, or what the packaging infrastructure and tooling will look like. Mojo can access 3rd party Python packages via the interop interface, same as the Python standard library.
-
Users and developers: Mojo’s target audience currently seems to be developers of what I would call “application code.” Mojo is not yet targeting anyone creating reusable modules for redistribution, and not Python package developers either.
The best way I’ve found to think about Mojo is that it is trying to be a language that looks like Python, feels somewhat familiar to Python developers, and is designed to interoperate with Python out of the box. It is not a drop-in replacement for Python, however, unless you are writing a new application from scratch that makes limited use of other Python modules.
Mojo fn/def vs. Numba’s object/nopython modes
One of the more interesting decisions in Mojo is to create two kinds of functions:
-
def functions: These functions are designed to feel like standard Python
functions. The dynamic behavior we expect from Python should continue to
work here, although Mojo def functions can go beyond Python via creation of
immutable variables, define type signatures using Mojo types, and so on.
-
fn functions: These functions have more strict rules about the contents of
the function, such as required type declarations for arguments and local
variables, and explicit declaration of exceptions.
It is clear from the restrictions that fn functions are designed to be
statically typed by a compiler and can be compiled ahead of time with the
performance one would expect from a statically typed language.
Numba takes a somewhat similar approach, and has two different compilation
modes, which are called “object mode” and “nopython” mode. These two compiler
modes have the following behavior:
-
object mode: Does not attempt type inference on the function, treating all
values as opaque Python objects. Object mode simulates the Python
interpreter at compile time, emitting calls to the CPython C API to perform
all operations on the Python objects in the function. This results in a
negligible performance improvement to the code. Due to limitations in the
bytecode analysis front-end of Numba, there are some Python language
constructs not supported. Execution is still subject to the Python Global
Interpreter Lock (GIL).
-
nopython mode: This performs a complete type-inference pass on the
function, and, for the subset of Python that is supported, eliminates all
direct manipulation of Python objects. Scalar values are unboxed into
native machine types, and operations on NumPy arrays directly modify the
data buffer of those arrays. Most reference counting is completely
eliminated, although a few instances may remain depending on whether new
NumPy arrays are allocated. Numba, however, uses its own thread-safe
reference counting internally, so the GIL can be released and multiple
threads can execute Numba functions simultaneously.
In practice, we’ve found object mode is no longer very useful. In early
Numba, nopython mode was very limited, so object mode was paired with a
technique called “loop lifting.” The Numba compiler would find loops within a
function that could be compiled to nopython mode and would “lift” them out of
the function body, and the rest of the function would be compiled with object
mode. As nopython mode has gained more capabilities (including memory
allocation and access to typed Python-like containers), we find the use cases
for object mode compilation of functions have diminished. In fact, we’ve even
inverted the situation, and now have support for object mode blocks within a
nopython mode function, which is useful for things like supporting Python
callbacks.
Since Numba exists within the CPython interpreter, the best practice for
functions that cannot be compiled with nopython mode is to not compile them at
all! Thanks to the Faster CPython effort, these sorts of generic functions
are getting faster with every release anyway, and Numba is most effective on
the numerical core of the application, where the compiler can strip away all
of the Python overhead. For this reason, Siu Kwan Lam (one of Numba’s
creators) often calls Numba a “performance linter” because the limitations of
nopython mode guide Numba users into using Python in a way that can be most
effectively sped up by a compiler, usually between 2x and 200x. If we can’t
achieve that, there isn’t much point to compilation in the first place.
At first glance, Numba’s object mode seems like Mojo’s def functions, but the
difference is that Mojo’s def functions feel like Python functions, but are
not Python functions. Typing is optional, but new Mojo declarations (like var
and let) are allowed, and the types are still Mojo types, not Python types.
Access to Python modules needs to happen via special wrappers created by the
Mojo Python interface, which can import Python modules. Mojo uses CPython as
a library to handle these calls, but the body of the def function itself is
still executed by the Mojo runtime. In principle, this could make Mojo behave
a lot like Cython, but we still don’t know much about how convenient it will
be to use Python from Mojo in practice, and we know nothing about whether
Python will be able to call Mojo functions easily.
Array Types and Operations
The most important part of any numerical computing system in 2023 is the
multidimensional array. Multidimensional arrays are extremely useful and
memory efficient containers for numerical data, and concepts like
broadcasting, universal functions, and advanced indexing have dramatically
simplified the application of numerical computing to many situations,
including machine learning. Moreover, arrays are very compiler-friendly data
structures, which has created opportunities for many compilers and compiler
technologies. In the Python space, NumPy has become the standard array, and
without NumPy, there would be no Numba.
The most surprising part of the Mojo announcement for me was the complete lack
of a built-in multidimensional array type. Given Mojo’s target use case of
machine learning workloads, I can only assume information about the
multidimensional array was held back for a future update, because it
absolutely has to exist for Mojo to be useful. The absence was doubly
surprising because MLIR (which the Mojo
documentation mentions as part of its core compiler design) includes a tensor
type, which is very similar
to the NumPy array, and also defines a generic linear algebra
operation (linalg.generic) in
one of the MLIR dialects that is extremely powerful. I would describe the
linalg.generic operation as a superset of the generalized
ufunc
that both NumPy and Numba support (via the
@guvectorize
decorator), which can represent a surprisingly wide range of array operations
in a very natural and compact way. Not seeing linalg.generic mentioned in
the Mojo docs was very disappointing, and I hope we learn more about it in a
future update.
With no array type to operate on, Mojo is currently lacking in all the array
functions that one would expect from a numerical computing language. These
can surely be added, and even implemented directly in the Mojo language, but
they are not there now. Hopefully Modular will learn from the recently
created Python Array API standard
and consider using that as the basis for their user-facing array API.
Even beyond the array type and functions, it will be extremely important for
Mojo to have a clear mechanism to pass array-like data between it and other
languages, especially Python. NumPy’s C
API is just as
important as its Python API, as it enables passing data in and out of the
Python runtime to 3rd party libraries written in other languages. Mojo will
have a long road to broader adoption, so anything that eases interoperability
with existing numerical computing libraries will be essential.
Extensibility and MLIR
Whether or not the Mojo compiler needs to be extensible is unclear. Many
languages are self-contained, in that you are expected to “extend” the
system by writing even more code in that language. That seems to be how Mojo
is being implemented now, as most of the syntax and standard library of Mojo
is focused on very low-level types and concepts, like SIMD operations and
object lifetime management. With those ingredients, many of the essential
higher level features needed in Mojo can be implemented in Mojo.
Python (and Numba) are not like this at all. Python’s superpower (and curse
at times) is that it makes it very easy to extend Python with other languages,
typically C or C++, but recently Rust is also growing in popularity. Every
popular numerical computing package in Python either includes some non-Python
code in it, or depends on one that does. For this reason, the C API of Python
is very important for enabling extension of the interpreter to new types and
modules. In a similar way, Numba’s limitations mean that there are some
low-level operations that cannot be generated by Numba when it compiles a
function with its @jit decorator. For these cases, Numba has a extension API
which allows a variety of things, including:
- Generation of machine code (technically LLVM IR) to implement specific low level operations.
- Custom overrides of existing Python functions to define how they should be compiled for specific data types.
- Definition of custom “boxing” and “unboxing” methods for passing new data types through Numba-compiled code.
While Mojo does not necessarily need these things due to its design, I believe
the Python interoperability feature in Mojo will need an extensible mechanism
to define how to translate data structures across the Python / Mojo boundary.
We know from Numba experience that anything like a NumPy array is very easy to
move back and forth, and we believe that Apache Arrow data structures offer a
similar possibility for more heterogeneous data (dataframes and beyond). Numba
has even created a few typed data containers specifically to facilitate
movement of statically-typed
lists
and statically-typed
dictionaries
across the Numba / Python boundary. Manipulating Python data through a wrapper
mechanism (as is currently described in the Mojo docs) is a fine general
solution, but faster mechanisms will be needed if Mojo is going to coexist
with non-trivial Python code.
Another area of extensibility that could be important to Mojo and is
definitely important to the future of Numba is interoperability with other
MLIR-based tools. MLIR is designed to be a high level compiler intermediate
representation
(IR) with some unique features. MLIR is extensible by defining new dialects,
which allows the same IR to be used on a wide range of compiler use cases,
and, in principle, permits different compiler tools that support MLIR to work
together. The Numba team is very excited about the potential for MLIR and
plans to incorporate it into the next generation of Numba. However, we can’t
tell how Mojo plans to expose its MLIR-based internals. Would something like
Numba be able to target the Mojo compiler pipeline directly, or possibly the
reverse: could future Numba consume the functions produced by the Mojo
frontend and parser?
Packaging and Composability
The final area we are curious about Mojo is the packaging approach. One feature of Mojo that was emphasized was the ability of the Mojo compiler to produce fully self-contained executables that could be deployed wherever they were needed. Application distribution is a very important use case that has not gotten as much attention in the Python ecosystem as package distribution. Python’s plethora of packaging tools focus on allowing users to install Python components along with their complex web of dependencies. The Python Package Index (and adjacent package repositories, like conda-forge and the Anaconda Distribution) is truly one of Python’s superpowers, and if Mojo wants to jumpstart a broad community of developers, it will also need a packaging story. There are so many unanswered questions at this point:
- Will hybrid Mojo-Python projects be possible, with both Mojo and Python
dependencies?
- Will Mojo packages only distribute source code, with the intent that all
Mojo applications must be compiled by the application builder along with all
of their dependencies?
- Or, will Mojo allow precompiled modules to be distributed, in which case,how
will inter-module optimization be enabled?
Aside: Although this may be a long shot, I think it would be great if Mojo
would consider conda
as a package manager, as it already has excellent Python support, but is also
cross-platform and multi-language packaging system. Mojo doesn't need to
build yet another package manager, and conda would be an excellent choice for
the complex hybrid Mojo/Python projects that Modular hopes will exist in the
future. Give us a call, Chris, if you want to chat. 😀
The packaging story for Numba today is much simpler, as it is just a Python
package that does just-in-time compilation. Projects that use Numba
compilation do so by taking Numba as a dependency, and then can choose to
distribute their own package as pure source code. The translation from Python
to machine code happens at runtime on the end user’s system. This
dramatically simplifies the distribution of those packages downstream of
Numba, who do not need to build wheels for each target platform unless they
have other compiled code aside from Numba-decorated functions.
However, the Numba team definitely appreciates that JIT compilation is not
suitable for all situations, either because of the increased startup time, or
limitations of the target platform, which might forbid JIT compilation
entirely. Because of this, in our planning for the next generation of
Numba,
we’ve decided to first focus on an ahead-of-time compiler (like Mojo) and then
will treat just-in-time compilation as a special case. Ahead-of-time
compilation is much trickier because you have to make choices at build time
for which there are downstream consequences to consumers of your package. You
have to worry about ABI stability, supporting multiple variants of input
types, and different machine instruction set variants (like AVX, AVX2,
AVX-512, etc). At the same time, we do not want ahead-of-time compilation to
preclude JIT optimization across modules that were distributed separately. We
have been working on techniques to handle all of these situations in a new
library called
PIXIE,
which allows additional metadata to be inserted into library files to enable
dynamic dispatch, CPU feature selection, and optional future recompilation of
functions with a JIT.
Conclusions and a Wishlist for Mojo
I want to conclude by reiterating that my goal with this article is not to
dampen enthusiasm for Mojo, but to help the reader be more curious about the
details when they read about new Python and Python-adjacent compiler projects.
Impressive looking speedup factors are not where the most difficult challenges
are.
I think it is an interesting idea to try to blend a Python-like syntax with
the capabilities of MLIR to target a wide range of potential hardware. Mojo’s
Python interoperability features could expand the capabilities of the Python
ecosystem, which would be a great thing. But, there is a tremendous amount of
work left to do on Mojo and a lot of design details have not been shared. We
simply don’t know enough to decide how Mojo will impact Python, or how best to
interact with it.
Given that, I want to conclude by reiterating the wishlist of things I hope we see in Mojo in the coming months:
-
A native multidimensional array type that exposes the same power as MLIR’s tensor type and linalg dialect. That will enable the same kind of array programming we are familiar with from NumPy.
-
More explanation and infrastructure for managing data motion across the Mojo / Python boundary. NumPy arrays should be usable by Mojo native code (hopefully that tensor type we’re asking for in the previous item) without data copies. Similarly, being able to use Apache Arrow data structures across the two languages would be very powerful. In general, we need something extensible so that new data translators can be added.
-
Permit an inversion of control, so that Mojo modules can be imported from the Python interpreter. Mojo could be an interesting alternative to Cython if that was possible, and if NumPy arrays could be passed to Mojo functions with minimal overhead.
-
Expose the MLIR internals of Mojo via some kind of interface that lets us retarget the compiler. A true Python-to-Mojo compiler pipeline could be possible then, and that would allow a lot more experimentation. How should we use Mojo with new or custom MLIR dialects?
-
More details about Mojo packaging in general. Will there be a package index for Mojo-specific modules by third parties? How should mixed Mojo and Python environments be handled? (Seriously, please take a look at conda. 🙂)
Of course, we need to see an open source Mojo compiler and runtime before many
of these things will be possible. Hopefully we’ll get more details on that in
the future as well.