Mojo Metaprogramming

Metaprogramming is writing code that generates or transforms other code at compile time. Instead of repeating similar code for every type or size, you write a template once and the compiler produces specialized versions automatically. Mojo's metaprogramming system runs at zero runtime cost — the generated code is as fast as hand-written code for each specific case.

The Stamp Factory Analogy

  Without metaprogramming:
    Write add_int() by hand
    Write add_float() by hand
    Write add_double() by hand
    ... repeat for every type

  With metaprogramming:
    Write add[dtype]() once
    → Compiler stamps out add[Int]    automatically
    → Compiler stamps out add[Float32] automatically
    → Compiler stamps out add[Float64] automatically

  You write the template once. The compiler writes the copies.

Parameterized Functions

fn maximum[T: Comparable](a: T, b: T) -> T:
    if a > b:
        return a
    return b

fn main():
    print(maximum[Int](10, 20))           # 20
    print(maximum[Float64](3.14, 2.71))   # 3.14
    print(maximum[String]("apple", "mango"))   # mango
Compiler generates:
  maximum[Int]:     if a > b (Int comparison) ...
  maximum[Float64]: if a > b (Float comparison) ...
  maximum[String]:  if a > b (String comparison) ...

Three separate compiled functions, each optimized for its type.

Parameterized Structs

struct Stack[T: AnyType]:
    var items: List[T]
    var _size: Int

    fn __init__(inout self):
        self.items = List[T]()
        self._size = 0

    fn push(inout self, item: T):
        self.items.append(item)
        self._size += 1

    fn pop(inout self) -> T:
        self._size -= 1
        return self.items.pop()

    fn size(self) -> Int:
        return self._size

fn main():
    var int_stack = Stack[Int]()
    int_stack.push(1)
    int_stack.push(2)
    int_stack.push(3)
    print(int_stack.pop())    # 3
    print(int_stack.size())   # 2

    var str_stack = Stack[String]()
    str_stack.push("hello")
    str_stack.push("world")
    print(str_stack.pop())    # world

Compile-Time Type Introspection

Mojo lets you query properties of types at compile time and branch on them using @parameter if-statements.

fn describe_type[T: AnyType]():
    @parameter
    if T == Int:
        print("Integer type")
    elif T == Float64:
        print("64-bit floating point")
    elif T == String:
        print("String type")
    else:
        print("Unknown type")

fn main():
    describe_type[Int]()      # Integer type
    describe_type[Float64]()  # 64-bit floating point
    describe_type[String]()   # String type

Generating Code with alias

alias VECTOR_WIDTH = simdwidthof[DType.float32]()

fn process_chunk[width: Int](data: UnsafePointer[Float32], offset: Int):
    var v = SIMD[DType.float32, width].load(data + offset)
    (v * 2.0).store(data + offset)

fn process_all(data: UnsafePointer[Float32], n: Int):
    var i = 0
    while i + VECTOR_WIDTH <= n:
        process_chunk[VECTOR_WIDTH](data, i)
        i += VECTOR_WIDTH
    while i < n:
        data[i] *= 2.0
        i += 1
alias VECTOR_WIDTH = 8  (for AVX2 CPU)

The compiler generates:
  process_chunk[8] → SIMD load 8 floats, multiply, store
  Specialized for width=8 with no runtime branching.

Recursive Metaprogramming

Compile-time recursion generates sequences of code at compile time — useful for loop unrolling and building tables of values.

fn unrolled_sum[N: Int](data: UnsafePointer[Float32]) -> Float32:
    @parameter
    if N == 0:
        return 0.0
    else:
        return data[N - 1] + unrolled_sum[N - 1](data)

fn main():
    var arr = UnsafePointer[Float32].alloc(4)
    arr.init_pointee_copy(1.0)
    (arr + 1).init_pointee_copy(2.0)
    (arr + 2).init_pointee_copy(3.0)
    (arr + 3).init_pointee_copy(4.0)

    print(unrolled_sum[4](arr))   # 10.0

    for i in range(4):
        (arr + i).destroy_pointee()
    arr.free()
unrolled_sum[4] expands to:
  data[3] + data[2] + data[1] + data[0] + 0.0
  All at compile time — the loop disappears, replaced by fixed additions.

When to Use Metaprogramming

Good uses:
  ✓ Writing one function that works for all numeric types
  ✓ Fixed-size data structures (buffer, matrix, ring queue)
  ✓ Algorithm selection based on type (e.g., use SIMD for floats)
  ✓ Compile-time loop unrolling for small known sizes
  ✓ Zero-cost debugging flags

Avoid when:
  ✗ Logic is simple enough without it
  ✗ The compile-time parameter is rarely varied
  ✗ It makes code unreadable without benefit

Key Takeaways

Metaprogramming generates specialized code at compile time with zero runtime overhead. Parameterize functions and structs with [T: Trait] syntax to write once and apply to many types. Use @parameter if-statements to branch at compile time and eliminate dead code. The compiler generates separate, optimized implementations for each unique combination of compile-time parameters. Recursive metaprogramming unrolls loops and builds compile-time tables. Use metaprogramming when it eliminates repetition or enables optimizations that runtime code cannot achieve.

Leave a Comment

Your email address will not be published. Required fields are marked *