Mojo🔥 Enabling AI to be used by anyone, anywhere

Mojo is a new programming language that builds on top of Python’s usability that provides powerful features like systems programming and metaprogramming. Mojo is being developed by Modular an AI startup armed with some tech fire-power like Chris Lattner (also known for co-founding LLVM, Clang compiler, MLIR compiler infrastructure, and the Swift programming language).

To quickly summarize what Modular is doing right now, I will make a reference to the famous saying “During a gold rush, sell shovels”, because this is more or less what Modular is trying to build.

This “shovel” is not just any ordinary “shovel”…it is built on top of next-generation compiler technologies that allow significant performance gains (allegedly 35,000x faster than Python).

Python superset + C-level speed = 🔥

Now if you ever tried to program some low-level AI hardware, chances are that most likely some proficiency in C, C++, or CUDA was required. But not anymore, now you can write Python or scale all the way down to the hardware target such as CPUs, GPUs, TPUs, and other custom hardware accelerators with ease.

Mojo is a language designed for programming on AI Hardware (e.g. GPU running CUDA), which currently is accessible via a Jupyter interactive Notebook.

Mojo Notebook

The learning curve is not that steep since Mojo is a superset of Python, therefore you don’t need to learn another language to take advantage of it. Python is not designed nor does it excel for systems programming, but Mojo is, and it introduces struct type which is pretty similar to Python’s class type, but quite different in the sense that is static and bound at compile time.

Also, Mojo introduces var and let to encourage explicit declarations. If in Python is no difference between variable assignment and declaration, we can do:

x = 3

even if x x was not already declared, and also due to the dynamic typing nature of Python we can even do:

x = 3
type(x) # returns <class 'int'>

x = "string"
type(x) # returns <class 'str'>

In Mojo declaration must have either a type or an initializer, where var and let explicitly introduce a new scoped runtime value:

# mutable with type
var a: Int

# mutable with intializer
var b = 4

a = 2
a = 3 # it's allowed 

# immutable in assignment
let c = 434

c = 3 # will retrn error: expression must be mutable in assignment

Even tho you can still use dynamic types just like in Python, you can employ strong type-checking with struct which also provides safe abstractions on top of low-level data. Also, Mojo provides a fn declaration that is like a “strict mode” for def when creating functions.

“A key trick in Mojo is that you can opt-in at any time to a faster ‘mode’ as a developer, by using `fn` instead of `def` to create your function. In this mode, you have to declare exactly what the type of every variable is, and as a result, Mojo can create optimized machine code to implement your function.”

WIP

Small benchmark

What benchmark is more suited than matrix multiplication?

Matrices have the property of representing the data efficiently and are critical to ML systems.

matrix multiplication

The performance of an ML model is tightly coupled with the notion of FLOPS. Metric that is used to assess the performance of GPUs, and CPUs, and indicate their capability to compute numerical calculations (i.e. matrix multiplication).

Using Python for a 128x128 square matrices multiplication we need ~2s and 0.002 GFLOP/s.

Python

Using Mojo we run the same computation, in 0.14s and 0.02 GFLOP/s without syntactically modifying the Python code.

Mojo

The only modifications needed were importing utilities and defining the matrix:

from Benchmark import Benchmark
from DType import DType
from Intrinsics import strided_load
from List import VariadicList
from Math import div_ceil, min
from Memory import memset_zero
from Object import object, Attr
from Pointer import DTypePointer
from Random import rand, random_f64
from TargetInfo import dtype_sizeof, dtype_simd_width
....
def matrix_init(rows: Int, cols: Int) -> object:
    value = object([])
    return object(
        Attr("value", value), Attr("__getitem__", matrix_getitem), Attr("__setitem__", matrix_setitem), 
        Attr("rows", rows), Attr("cols", cols), Attr("append", matrix_append),
    )

In all honesty, this is merely scratching the surface and we cannot really know what happens behind the scenes but nonetheless, what Modular is trying to achieve is undeniably remarkable.

Last but not least, the amazing thing is that Mojo gained quite some attention without having a compiler released to the public 🤯.

mojo