Parameterization: compile-time metaprogramming

An introduction to parameters and compile-time metaprogramming.

Many languages have facilities for metaprogramming: that is, for writing code that generates or modifies code. Python has facilities for dynamic metaprogramming: features like decorators, metaclasses, and many more. These features make Python very flexible and productive, but since they’re dynamic, they come with runtime overhead. Other languages have static or compile-time metaprogramming features, like C preprocessor macros and C++ templates. These can be limiting and hard to use.

To support Modular’s work in AI, Mojo aims to provide powerful, easy-to-use metaprogramming with zero runtime cost. This compile-time metaprogramming uses the same language as runtime programs, so you don’t have to learn a new language—just a few new features.

The main new feature is parameters. You can think of a parameter as a compile-time variable that becomes a runtime constant. This usage of “parameter” is probably different from what you’re used to from other languages, where “parameter” and “argument” are often used interchangeably. In Mojo, “parameter” and “parameter expression” refer to compile-time values, and “argument” and “expression” refer to runtime values.

In Mojo, you can add parameters to a struct or function. You can also define named parameter expressions—aliases—that you can use as runtime constants.

Parameterized functions

To define a parameterized function, add parameters in square brackets ahead of the argument list. Each parameter is formatted just like an argument: a parameter name, followed by a colon and a type (which is required). In the following example, the function has a single parameter, count of type Int.

fn repeat[count: Int](msg: String):
    @unroll
    for i in range(count):
        print(msg)

(The @unroll directive shown here unrolls loops at compile time. The directive only works if the loop limits are compile-time constants.)

Calling a parameterized function, you provide values for the parameters, just like function arguments:

repeat[3]("Hello")
Hello
Hello
Hello

The compiler resolves the parameter values during compilation, and creates a concrete version of the repeat[]() function for each unique parameter value. After resolving the parameter values and unrolling the loop, the repeat[3]() function would be roughly equivalent to this:

fn repeat_3(msg: String):
    print(msg)
    print(msg)
    print(msg)

Note: This doesn’t represent actual code generated by the compiler. By the time parameters are resolved, Mojo code has already been transformed to an intermediate representation in MLIR.

If the compiler can’t resolve all parameter values to constant values, compilation fails.

Parameterized structs

You can also add parameters to structs. You can use parameterized structs to build generic containers. For example, a generic array type might include code like this:

struct GenericArray[T: AnyRegType]:
    var data: Pointer[T]
    var size: Int

    fn __init__(inout self, *elements: T):
        self.size = len(elements)
        self.data = Pointer[T].alloc(self.size)
        for i in range(self.size):
            self.data[i] = elements[i]

    fn __del__(owned self):
        self.data.free()

    fn __getitem__(self, i: Int) raises -> T:
        if (i < self.size):
            return self.data[i]
        else:
            raise Error("Out of bounds")

This struct has a single parameter, T, which is a placeholder for the data type you want to store in the array, sometimes called a type parameter. T is typed as AnyRegType, which is a metatype representing any register-passable type. This means our GenericArray can hold fixed-size data types like integers and floating-point numbers that can be passed in a machine register, but not dynamically allocated data like strings or vectors.

Using a capital T as the parameter doesn’t have any special significance to the compiler, but it’s a convention in a number of languages to use a short name for a type parameter, frequently defaulting to T (for “type”).

As with parameterized functions, you need to pass in parameter values when you use a parameterized struct. In this case, when you create an instance of GenericArray, you need to specify the type you want to store, like Int, or Float64. (This is a little confusing, because the parameter value you’re passing in this case is a type. That’s OK: a Mojo type is a valid compile-time value.)

You’ll see that T is used throughout the struct where you’d usually see a type name. For example, as the formal type for the elements in the constructor, and the return type of the __getitem__() method.

At the moment, this code only works with register-passable types, which is why the type parameter T is limited to AnyRegType. There’s also an AnyType metatype, which includes all Mojo types.

Here’s an example of using GenericArray:

var array = GenericArray[Int](1, 2, 3, 4)
for i in range(array.size):
    print_no_newline(array[i], " ")
1  2  3  4  

A parameterized struct can use the Self type to represent a concrete instance of the struct (that is, with all its parameters specified). For example, you could add a static factory method to GenericArray with the following signature:

struct GenericArray[T: AnyRegType]:
    ...

    @staticmethod
    splat(count: Int, value: T) -> Self:
        # Create a new array with count instances of the given value

Here, Self is equivalent to writing GenericArray[T]. That is, you can call the splat() method like this:

GenericArray[Float64].splat(8, 0)

The method returns an instance of GenericArray[Float64].

Fully-bound, partially-bound, and unbound types

A parametric type with its parameters specified is said to be fully-bound. That is, all of its parameters are bound to values. As mentioned before, you can only instantiate a fully-bound type (sometimes called a concrete type).

However, parametric types can be unbound or partially bound in some contexts. For example, given the following type:

struct MyType[s: String, i: Int]:
    pass

It can appear in code in the following forms:

  • Fully bound, with all of its parameters specified. For example, MyType["Hello", 3].
  • Partially bound, with some but not all of its parameters specified. For example, MyType["Hola"].
  • Unbound, with no parameters specified. For example, MyType.

Partially-bound and unbound parametric types can be used in some contexts where the missing (unbound) parameters can be supplied later—such as in function arguments.

Case study: the SIMD type

For a real-world example of a parameterized type, let’s look at the SIMD type from Mojo’s standard library.

Single instruction, multiple data (SIMD) is a parallel processing technology built into many modern CPUs, GPUs, and custom accelerators. SIMD allows you to perform a single operation on multiple pieces of data at once. For example, if you want to take the square root of each element in an array, you can use SIMD to parallelize the work.

Processors implement SIMD using low-level vector registers in hardware that hold multiple instances of a scalar data type. In order to use the SIMD instructions on these processors, the data must be shaped into the proper SIMD width (data type) and length (vector size). Processors may support 512-bit or longer SIMD vectors, and support many data types from 8-bit integers to 64-bit floating point numbers, so it’s not practical to define all of the possible SIMD variations.

Mojo’s SIMD type (defined as a struct) exposes the common SIMD operations through its methods, and makes the SIMD data type and size values parametric. This allows you to directly map your data to the SIMD vectors on any hardware.

Here’s a cut-down (non-functional) version of Mojo’s SIMD type definition:

struct SIMD[type: DType, size: Int]:
    var value: … # Some low-level MLIR stuff here

    # Create a new SIMD from a number of scalars
    fn __init__(inout self, *elems: SIMD[type, 1]):  ...

    # Fill a SIMD with a duplicated scalar value.
    @staticmethod
    fn splat(x: SIMD[type, 1]) -> SIMD[type, size]: ...

    # Cast the elements of the SIMD to a different elt type.
    fn cast[target: DType](self) -> SIMD[target, size]: ...

    # Many standard operators are supported.
    fn __add__(self, rhs: Self) -> Self: ...

So you can create and use a SIMD vector like this:

var vector = SIMD[DType.int16, 4](1, 2, 3, 4)
vector = vector * vector
for i in range(4):
    print_no_newline(vector[i], " ")
1  4  9  16  

As you can see, a simple arithmetic operator like * applied to a pair of SIMD vector operates on the corresponding elements in each vector.

Defining each SIMD variant with parameters is great for code reuse because the SIMD type can express all the different vector variants statically, instead of requiring the language to pre-define every variant.

Because SIMD is a parameterized type, the self argument in its functions carries those parameters—the full type name is SIMD[type, size]. Although it’s valid to write this out (as shown in the return type of splat()), this can be verbose, so we recommend using the Self type (from PEP673) like the __add__ example does.

Overloading on parameters

Functions and methods can be overloaded on their parameter signatures. The overload resolution logic filters for candidates according to the following rules, in order of precedence:

  1. Candidates with the minimal number of implicit conversions (in both arguments and parameters).
  2. Candidates without variadic arguments.
  3. Candidates without variadic parameters.
  4. Candidates with the shortest parameter signature.
  5. Non-@staticmethod candidates (over @staticmethod ones, if available).

If there is more than one candidate after applying these rules, the overload resolution fails. For example:

@register_passable("trivial")
struct MyInt:
    """A type that is implicitly convertible to `Int`."""
    var value: Int

    @always_inline("nodebug")
    fn __init__(_a: Int) -> Self:
        return Self {value: _a}

fn foo[x: MyInt, a: Int]():
    print("foo[x: MyInt, a: Int]()")

fn foo[x: MyInt, y: MyInt]():
    print("foo[x: MyInt, y: MyInt]()")

fn bar[a: Int](b: Int):
    print("bar[a: Int](b: Int)")

fn bar[a: Int](*b: Int):
    print("bar[a: Int](*b: Int)")

fn bar[*a: Int](b: Int):
    print("bar[*a: Int](b: Int)")

fn parameter_overloads[a: Int, b: Int, x: MyInt]():
    # `foo[x: MyInt, a: Int]()` is called because it requires no implicit
    # conversions, whereas `foo[x: MyInt, y: MyInt]()` requires one.
    foo[x, a]()

    # `bar[a: Int](b: Int)` is called because it does not have variadic
    # arguments or parameters.
    bar[a](b)

    # `bar[*a: Int](b: Int)` is called because it has variadic parameters.
    bar[a, a, a](b)

parameter_overloads[1, 2, MyInt(3)]()

struct MyStruct:
    fn __init__(inout self):
        pass

    fn foo(inout self):
        print("calling instance menthod")

    @staticmethod
    fn foo():
        print("calling static method")

fn test_static_overload():
    var a = MyStruct()
    # `foo(inout self)` takes precedence over a static method.
    a.foo()
foo[x: MyInt, a: Int]()
bar[a: Int](b: Int)
bar[*a: Int](b: Int)

Using parameterized types and functions

You can instantiate parametric types and functions by passing values to the parameters in square brackets. For example, for the SIMD type above, type specifies the data type and size specifies the length of the SIMD vector (it must be a power of 2):

# Make a vector of 4 floats.
let small_vec = SIMD[DType.float32, 4](1.0, 2.0, 3.0, 4.0)

# Make a big vector containing 1.0 in float16 format.
let big_vec = SIMD[DType.float16, 32].splat(1.0)

# Do some math and convert the elements to float32.
let bigger_vec = (big_vec+big_vec).cast[DType.float32]()

# You can write types out explicitly if you want of course.
let bigger_vec2 : SIMD[DType.float32, 32] = bigger_vec

print('small_vec type:', small_vec.element_type, 'length:', len(small_vec))
print('bigger_vec2 type:', bigger_vec2.element_type, 'length:', len(bigger_vec2))
small_vec type: float32 length: 4
bigger_vec2 type: float32 length: 32

Note that the cast() method also needs a parameter to specify the type you want from the cast (the method definition above expects a target parametric value). Thus, just as the SIMD struct is a generic type definition, the cast() method is a generic method definition that gets instantiated at compile-time instead of runtime, based on the parameter value.

The code above shows the use of concrete types (that is, it instantiates SIMD using known type values), but the major power of parameters comes from the ability to define parametric algorithms and types (code that uses the parameter values). For example, here’s how to define a parametric algorithm with SIMD that is type- and width-agnostic:

from math import sqrt

fn rsqrt[dt: DType, width: Int](x: SIMD[dt, width]) -> SIMD[dt, width]:
    return 1 / sqrt(x)

print(rsqrt[DType.float16, 4](42))
[0.154296875, 0.154296875, 0.154296875, 0.154296875]

Notice that the x argument is actually a SIMD type based on the function parameters. The runtime program can use the value of the parameters, because the parameters are resolved at compile-time before they are needed by the runtime program (but compile-time parameter expressions cannot use runtime values).

The Mojo compiler is also smart about type inference with parameters. Note that the above function is able to call the parametric sqrt[]() function without specifying the parameters—the compiler infers its parameters based on the parametric x value passed into it, as if you wrote sqrt[dt, width](x) explicitly. Also note that rsqrt() chose to name its second parameter width even though the SIMD type names it size, and there is no problem.

Optional parameters and keyword parameters

Just as you can specify optional arguments in function signatures, you can also define an optional parameter by giving it a default value. You can also pass parameters by keyword. For a function or struct with multiple optional parameters, using keywords allows you to pass only the parameters you want to specify, regardless of their position in the function signature.

For example, here’s a function with two parameters, each with a default value:

fn speak[a: Int = 3, msg: StringLiteral = "woof"]():
    print(msg, a)

fn use_defaults() raises:
    speak()             # prints 'woof 3'
    speak[5]()          # prints 'woof 5'
    speak[7, "meow"]()  # prints 'meow 7'
    speak[msg="baaa"]() # prints 'baaa 3'

Recall that when a parametric function is called, Mojo can infer the parameter values. That is, it can use the parameter values attached to an argument value (see the sqrt[]() example above). If the parametric function also has a default value defined, then the inferred parameter type takes precedence.

For example, in the following code, we update the parametric speak[]() function to take an argument with a parametric type. Although the function has a default parameter value for a, Mojo instead uses the inferred a parameter value from the bar argument (as written, the default a value can never be used, but this is just for demonstration purposes):

@value
struct Bar[v: Int]:
    pass

fn foo[a: Int = 3, msg: StringLiteral = "woof"](bar: Bar[a]):
    print(msg, a)

fn use_inferred():
    foo(Bar[9]())  # prints 'woof 9'

As mentioned above, you can also use optional parameters and keyword parameters in a struct:

struct KwParamStruct[greeting: String = "Hello", name: String = "🔥mojo🔥"]:
    fn __init__(inout self):
        print(greeting, name)

fn use_kw_params():
    let a = KwParamStruct[]()                 # prints 'Hello 🔥mojo🔥'
    let b = KwParamStruct[name="World"]()     # prints 'Hello World'
    let c = KwParamStruct[greeting="Hola"]()  # prints 'Hola 🔥mojo🔥'

Note: Mojo currently includes only partial support for keyword parameters, so some features such as keyword-only parameters and variadic keyword parameters (for example, **kwparams) are not supported yet.

Parameter expressions are just Mojo code

A parameter expression is any code expression (such as a+b) that occurs where a parameter is expected. Parameter expressions support operators and function calls, just like runtime code, and all parameter types use the same type system as the runtime program (such as Int and DType).

Because parameter expressions use the same grammar and types as runtime Mojo code, you can use many “dependent type” features. For example, you might want to define a helper function to concatenate two SIMD vectors:

fn concat[ty: DType, len1: Int, len2: Int](
        lhs: SIMD[ty, len1], rhs: SIMD[ty, len2]) -> SIMD[ty, len1+len2]:

    var result = SIMD[ty, len1 + len2]()
    for i in range(len1):
        result[i] = SIMD[ty, 1](lhs[i])
    for j in range(len2):
        result[len1 + j] = SIMD[ty, 1](rhs[j])
    return result

let a = SIMD[DType.float32, 2](1, 2)
let x = concat[DType.float32, 2, 2](a, a)

print('result type:', x.element_type, 'length:', len(x))
result type: float32 length: 4

Note that the resulting length is the sum of the input vector lengths, and this is expressed with a simple + operation.

Powerful compile-time programming

While simple expressions are useful, sometimes you want to write imperative compile-time logic with control flow. You can even do compile-time recursion. For instance, here is an example “tree reduction” algorithm that sums all elements of a vector recursively into a scalar:

fn slice[ty: DType, new_size: Int, size: Int](
        x: SIMD[ty, size], offset: Int) -> SIMD[ty, new_size]:
    var result = SIMD[ty, new_size]()
    for i in range(new_size):
        result[i] = SIMD[ty, 1](x[i + offset])
    return result

fn reduce_add[ty: DType, size: Int](x: SIMD[ty, size]) -> Int:
    @parameter
    if size == 1:
        return x[0].to_int()
    elif size == 2:
        return x[0].to_int() + x[1].to_int()

    # Extract the top/bottom halves, add them, sum the elements.
    alias half_size = size // 2
    let lhs = slice[ty, half_size, size](x, 0)
    let rhs = slice[ty, half_size, size](x, half_size)
    return reduce_add[ty, half_size](lhs + rhs)

let x = SIMD[DType.index, 4](1, 2, 3, 4)
print(x)
print("Elements sum:", reduce_add[DType.index, 4](x))
[1, 2, 3, 4]
Elements sum: 10

This makes use of the @parameter decorator to create a parametric if condition, which is an if statement that runs at compile-time. It requires that its condition be a valid parameter expression, and ensures that only the live branch of the if statement is compiled into the program.

Mojo types are just parameter expressions

While we’ve shown how you can use parameter expressions within types, type annotations can themselves be arbitrary expressions (just like in Python). Types in Mojo have a special metatype type, allowing type-parametric algorithms and functions to be defined.

For example, we can create a simplified Array that supports arbitrary types of elements (via the AnyRegType parameter):

struct Array[T: AnyRegType]:
    var data: Pointer[T]
    var size: Int

    fn __init__(inout self, size: Int, value: T):
        self.size = size
        self.data = Pointer[T].alloc(self.size)
        for i in range(self.size):
            self.data[i] = value

    fn __getitem__(self, i: Int) -> T:
        return self.data[i]

    fn __del__(owned self):
        self.data.free()

var v = Array[Float32](4, 3.14)
print(v[0], v[1], v[2], v[3])
3.1400001049041748 3.1400001049041748 3.1400001049041748 3.1400001049041748

Notice that the T parameter is being used as the formal type for the value arguments and the return type of the __getitem__() function. Parameters allow the Array type to provide different APIs based on the different use-cases.

There are many other cases that benefit from more advanced use of parameters. For example, you can execute a closure N times in parallel, feeding in a value from the context, like this:

fn parallelize[func: fn (Int) -> None](num_work_items: Int):
    # Not actually parallel: see the 'algorithm' module for real implementation.
    for i in range(num_work_items):
        func(i)

Another example where this is important is with variadic generics, where an algorithm or data structure may need to be defined over a list of heterogeneous types such as for a tuple. Right now, this is not fully supported in Mojo and requires writing some MLIR by hand. In the future, this will be possible in pure Mojo.

alias: named parameter expressions

It is very common to want to name compile-time values. Whereas var defines a runtime value, and let defines a runtime constant, we need a way to define a compile-time temporary value. For this, Mojo uses an alias declaration.

For example, the DType struct implements a simple enum using aliases for the enumerators like this (the actual DType implementation details vary a bit):

struct DType:
    var value : UI8
    alias invalid = DType(0)
    alias bool = DType(1)
    alias int8 = DType(2)
    alias uint8 = DType(3)
    alias int16 = DType(4)
    alias int16 = DType(5)
    ...
    alias float32 = DType(15)

This allows clients to use DType.float32 as a parameter expression (which also works as a runtime value) naturally. Note that this is invoking the runtime constructor for DType at compile-time.

Types are another common use for aliases. Because types are compile-time expressions, it is handy to be able to do things like this:

alias Float16 = SIMD[DType.float16, 1]
alias UInt8 = SIMD[DType.uint8, 1]

var x : Float16   # Float16 works like a "typedef"

Like var and let, aliases obey scope, and you can use local aliases within functions as you’d expect.

By the way, the special types None, AnyType, and AnyRegType are defined as type aliases.

Automatic parameterization of functions

Mojo supports “automatic” parameterization of functions. If a function argument type is parametric but the function signature doesn’t specify parameters, the unbound parameters are automatically added as input parameters on the function. This is easier to understand with an example:

fn print_params(vec: SIMD):
    print(vec.type)
    print(vec.size)

fn main():
    let v = SIMD[DType.float64, 4](1.0, 2.0, 3.0, 4.0)
    print_params(v)

In the above example, the print_params function is automatically parameterized. The vec argument takes an argument of type SIMD. This is an unbound parameterized type—that is, it doesn’t specify any parameter values for the type. Mojo treats the unbound parameters on vec as implicit parameters on the function. This is roughly equivalent to the following code, which includes explicit input parameters:

fn print_params[t: DType, s: Int](vec: SIMD[t, s]):
    print(vec.type)
    print(vec.size)

When you call print_params() you must pass it a concrete instance of the SIMD type—that is, one with all of its parameters specified, like SIMD[DType.float64, 4]. The Mojo compiler infers the unbound parameter values from the input argument.

With a manually parameterized function, you can access the input parameters by name (for example, t and s in the previous example). For an automatically parameterized function, you can access the parameters as attributes on the input argument (for example, vec.type).

This ability to access a type’s input parameters is not specific to automatically parameterized functions, you can use it anywhere. You can access the input parameters of a parameterized type as attributes on the type itself:

fn on_type():
    print(SIMD[DType.float32, 2].size) # prints 2

Or as attributes on an instance of the type:

fn on_instance():
    let x = SIMD[DType.int32, 2](4, 8)
    print(x.type) # prints int32

You can even use this syntax in the function’s signature to define a function’s arguments and return type based on an argument’s parameters. For example, if you want your function to take two SIMD vectors with the same type and size, you can write code like this:

fn interleave(v1: SIMD, v2: SIMD[v1.type, v1.size]) -> SIMD[v1.type, v1.size*2]:
    var result = SIMD[v1.type, v1.size*2]()
    for i in range(v1.size):
        result[i*2] = SIMD[v1.type, 1](v1[i])
        result[i*2+1] = SIMD[v1.type, 1](v2[i])
    return result

let a = SIMD[DType.int16, 4](1, 2, 3, 4)
let b = SIMD[DType.int16, 4](0, 0, 0, 0)
let c = interleave(a, b)
for i in range(c.size):
    print_no_newline(c[i], " ")
1  0  2  0  3  0  4  0  

Partial automatic parameterization

Mojo also supports partial automatic parameterization: you can declare an argument that takes a partially-bound parameterized type (that is, a type with some but not all of the parameters specified).

For example, suppose we have a Fudge struct with three parameters:

@value
struct Fudge[sugar: Int, cream: Int, chocolate: Int = 7](Stringable):
    fn __str__(self) -> String:
        let values = StaticIntTuple[3](sugar, cream, chocolate)
        return str("Fudge") + String(values)

We can write a function that takes a Fudge argument with just one bound parameter (it’s partially bound):


fn eat(f: Fudge[5]):
    print("Ate " + str(f))

The eat() function takes a Fudge struct with the first parameter (sugar) bound to the value 5. The second parameter, cream, is unbound. (The third parameter, chocolate has a default value, which we’ll discuss shortly.)

The unbound cream parameter becomes an implicit input parameter on the eat function. In practice, this is roughly equivalent to writing:

fn eat[c: Int](f: Fudge[5, c]):
    print("Ate " + str(f))

In both cases, we can call the function by passing in an instance with the cream parameter bound:

eat(Fudge[5, 5]())
eat(Fudge[5, 8]())
Ate Fudge(5, 5, 7)
Ate Fudge(5, 8, 7)

If you try to pass in an argument with a sugar value other than 5, compilation fails, because it doesn’t match the argument type:

eat(Fudge[12, 5]()) 
# ERROR: invalid call to 'eat': argument #0 cannot be converted from 'Fudge[12, 5, 7]' to 'Fudge[5, 5, 7]'

We haven’t talked about the third parameter, chocolate. Currently, a parameter with a default value is considered bound to the default value if it’s not explicitly bound or unbound.

This means that the following code won’t compile with either of the previous versions of the eat() function, since it doesn’t use the default value for chocolate:

eat(Fudge[5, 5, 9]())
# ERROR: invalid call to 'eat': argument #0 cannot be converted from 'Fudge[5, 5, 9]' to 'Fudge[5, 5, 7]'

TODO: We believe this handling of default values is incorrect, and it will be changed in a future release.

You can also explicitly identify which parameters are unbound. This lets you unbind a parameter that has a default value, and gives you more freedom in specifying unbound parameters.

For example, you might want to let the user specify values for sugar and chocolate, and leave cream constant. To do this, replace each unbound parameter value with a single underscore (_):

fn devour(f: Fudge[_, 6, _]):
    print(str("Devoured ") + str(f))

Again, the unbound parameters (sugar and chocolate) are added as implicit input parameters on the function. This version is roughly equivalent to the following code, where these two values are explicitly bound to the input parameters, su and ch:

fn devour[su: Int, ch: Int](f: Fudge[su, 6, ch]):
        print(str("Devoured ") + str(f))

You can also specify parameters by keyword, or mix positional and keyword parameters, so the following function is roughly equivalent to the previous one: the first parameter, sugar is explicitly unbound with the underscore character. The chocolate parameter is unbound using the keyword syntax, chocolate=_. And cream is explicitly bound to the value 6:

fn devour(f: Fudge[_, chocolate=_, cream=6]):
    print(str("Devoured ") + str(f))

All three versions of the devour() function work with the following calls:

devour(Fudge[3, 6, 9]())
devour(Fudge[4, 6, 8]())
Devoured Fudge(3, 6, 9)
Devoured Fudge(4, 6, 8)