This guest post is by Brian Grenier, a C++ developer working in cardiac image processing and an active contributor to the Mojo standard library. He maintains two open-source Mojo libraries: EmberJSON, a pure Mojo JSON parser, and Kelvin, a type-safe dimensional analysis toolkit. You can often find him sharing insights and answering questions in the Modular forum and Discord server (@eggsquad).
Mojo's parameter system is one of its most powerful and expressive features. In this post, I will take a closer look at a focused subset of Mojo's metaprogramming capabilities. For a breadth-first introduction to these features, please refer to the official Modular documentation.
Many of this article's examples include snippets of LLVM-IR to help illustrate what the Mojo compiler is doing under the hood. If you are not familiar with LLVM-IR, you can find a helpful language reference here.
Parameters for compile-time calculations
Metaprogramming is a technique widely used in many programming languages where we write code that modifies itself or other code. Python does this dynamically at runtime, whereas compiled languages like Rust, C++, and Mojo provide utilities for directing the compiler to generate and modify code at compile time. If you are already familiar with Rust generics or C++ templates, the Mojo parameter system will likely feel familiar.
A common use case for metaprogramming is creating generic functions and types; Python can do this implicitly via duck typing. The Python interpreter is generally unconcerned with the type of a variable until it comes time to perform some operation on it, so as long as the type you provide to a function satisfies those definitions, everything will move along just fine.
Without using generics, statically typed languages like Mojo need to rely on creating multiple definitions of that function in order to handle different input types:
We can solve the problem of duplicating functions like this by using generics, which in Mojo are implemented as part of the parameter system. In Mojo, the set of expected parameters for a struct or function is defined using a square bracket-enclosed list. All parameters must be types or expressions that are known at compile time. Mojo uses traits to inform the compiler about the kinds of types that a type parameter can bind to. Using traits, we can create an Addable trait like so:
With this in place, we can write a new doubler function that operates on both of these types by telling the compiler, "This function operates on all types that have the Addable trait and as a result, implement the __add__ method."
By diving into the LLVM-IR output from the Mojo compiler, we can see that the compiler has done the work for us of generating a different function definition for each type:
Function parameters vs. arguments
Before we dive deeper, how is a “parameter” different from an “argument”? Mojo allows functions to define two different kinds of inputs: one category that must be known at compile time when a function is instantiated, and one that is allowed to be unknown until the function is called at runtime. We need words to reason about this so we say that “parameters” are compile time inputs to a function - you might say that functions can be “parameterized” - and that arguments are the runtime inputs to the function.
For example, this function takes width as a parameter which must be known at compile time. It is used in the type of arg1 and arg2 (which are known at runtime), and supports math like multiply:
Parameter values are always immutable, and can be further used in other compile-time expressions. Note the use of width in a comparison expression against 4 in the constrained function call. constrained allows us to make compile-time assertions (like static_assert in other languages). In this case, we arbitrarily assert that the value of width is always greater than four.
Parameter values can also "hop the fence" into the runtime domain–observe how the constant 30 has been dropped directly into the multiplication expression in the IR output:
The alias keyword
Mojo provides a very powerful alias keyword that allows us to create new names for types, as well as define compile-time constant values. In the latter case, alias functions as the compile-time equivalent of var for runtime variables. Since values defined using alias are known at compile time, they can also be used in a parameter list:
Looking at the IR output, we can see the value of MyConstants has been dropped directly into the expression:
Thanks to Mojo's compile-time interpreter, you can use alias to precompute non-trivial expressions as well, which enables offloading runtime compute to the compiler:
And just like that, this function call has been folded into a store instruction:
When using alias , Mojo guarantees the expression will be completely compile-time folded. Any instructions that are not compile-time computable (such as special hardware intrinsics) will result in compilation failure.
You can even create an alias to a parameterized function, referring back to our multiplier function–we can create a function that always multiplies the input by two:
The compiler is also sometimes capable of entirely eliding the runtime materialization of a compile-time variable. For instance, if you create an alias to an InlineArray and index it using a compile-time-known value, the compiler will simply replace that array access with the appropriate value:
Compile-time branching
Through the power of parameter values, we can even make branching decisions at compile time by using @parameter if statements:
We can see in the IR that the if statement has been folded into a constant value branch. Mojo guarantees unreachable code is not compiled into the final binary. So this is effectively equivalent to an #ifdef expression in C/C++, but much more powerful, including the added benefit of both branches being type checked. Meaning that developers writing code that targets multiple devices, such as GPUs and CPUs, have type safety through the entire codebase. Regardless of which platform they are currently compiling for:
More recently, the is_compile_time function was added to the Mojo standard library, which allows invoking different behavior at compile time and runtime:
This is generally useful to provide fallback implementations of functions that would otherwise invoke platform-specific instructions which cannot be handled by the MLIR interpreter, or other expressions that don't have meaning at compile time, such as this example from the Mojo standard library:
In this case, the interpreter would throw an error upon trying to evaluate llvm_intrinsic["llvm.assume", NoneType](val), so we ignore it completely in that case.
Loop unrolling
Loop unrolling is an optimization technique used by compilers to save the overhead of incrementing loop variables and handling branch instructions. In performance-critical applications, this is even done manually by the programmer. In Mojo, we can ensure that the compiler always unrolls a particular loop using the @parameter decorator:
In the IR, we can see the loop has been inlined:
Dependent types
As the name suggests, "dependent types" are types whose definition depends on a value. For a concrete example, let’s take a look at a struct backed by Int:
Adder takes a value parameter I ; therefore, the particular value of I is now a part of the type definition of Adder . Note in the definition of Adder.__add___() that we use this property to enforce the compile-time evaluation of the addition by encoding it into a type system transformation. The @always_inline('builtin') decorator is used to ensure that the method body can be fully compile-time folded so expressions like Adder[1]() + Adder[3]() are in this case exactly equal to Adder[4](). Refer to the feature proposal for a more in depth explanation.
Encoding information in this way is also useful for enforcing that a function argument is parameterized with a particular value:
Since the addition operation has been moved into the type system, the resulting value I is simply materialized into a runtime store instruction of the value 12:
Representing complex information in the type system
To show a more complex example of what we can represent in a parameter, we're going to dive into a library I wrote called Kelvin. It is a type-safe dimensional analysis library in which all the safety checks are done at compile time to ensure minimal performance overhead over raw numbers. Let's explore how we can represent "10 meters per second" in code.
First, we need a definition for a physical "dimension" – here we define a struct Dimension that has 3 parameters:
Z: The value of the dimension, eg. the value ofZfor the length dimension of area would be2R: The scale of the dimension. This is used to differentiate between different scales of the same unit, eg.MetersvsKilometerssuffix: The string suffix for the value, e.g.km,K,mol, etc
We also have a distinct Angle struct since angle isn't a true dimension and behaves differently:
The Mojo parameter system allows us to write functions where the type of the return value is dependent on the function inputs. This allows us to build the following example where we define methods for Dimension that subtract the Z value of an input dimension, and return a new dimension type parameterized:
This behavior empowers us to do some cool stuff later on.
As another example, there are 7 dimensions in the SI unit system, so we can compose 7 dimension units plus an angle unit to create the Dimensions of a value:
Now we can define the result of dividing dimensions using the __sub__ method we just created. Units that live in the denominator are defined as negative Z values, e.g. m / s is represented by m^1 s^-1:
With that plumbing in place, we can define a Quantity struct that represents our actual value:
Now we can define particular units easily using alias:
Since we've defined a Dimensions / Dimensions operation, we can now define derivative units such as velocity in a single line of code:
Using the power of compile-time programming, we can now define operations on Quantity that respect the rules of the SI unit system, which are all verified at compile time. Starting with a simple example (since addition/subtraction only makes sense between the quantity of matching dimensions), we can use the special Self type to ensure that self and other are the exact same type:
Things get a bit more complex for multiplication and division since those operations are allowed on non-matching units, and the return type is dependent on those input types - we'll take division as our example here:
Here at compile-time, we have a static assertion that the dimensions of the two types have matching scale, e.g. meters per second, and miles per hour have matching dimensions, but are incompatible due to differences in unit scale. In the output argument, we compute the derivative dimensions using D / OD; therefore, the output of the method is the exact type we would expect from this operation:
Inspecting the IR, we can see all this parameter complexity has been resolved, and all that remains is a simple division operation:
Using this power responsibly
Understanding exactly how these abstractions work is critical to using them effectively; for instance, using a @parameter for loop will always result in the entire loop being unrolled, potentially resulting in binary bloat and negative performance implications:
Wrapping up
Mojo offers powerful tools for executing code seamlessly at compile time, allowing programmers to build powerful abstractions into the type system. This completely removes the associated complexity from runtime, and progress is continuously made (such as parameteric aliases).
I hope you found this article helpful in understanding how to use these features, and feel free to reach out on Discord (@eggsquad) or the Modular forum (@bgreni) if you have any questions! Head over to my GitHub (@bgreni) to follow my progress on Kelvin and other Mojo libraries I work on.
