Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Mojo’s SIMD[type, width]: width as a type parameter

What does SIMD look like when the TYPE SYSTEM, not a library, owns it? In Mojo, scalars are literally width-1 vectors, and vector width is a compile-time parameter you abstract over — the ergonomic ceiling that std::simd and wide approximate from below. Before the docs, this chapter builds the idea step by step — the ergonomics ladder, the width-1 unification, the parametric vectorize, and what none of it solves. Read it for the language-design angle; it’s the contrast that explains why our Rust experiments hand-write what Mojo’s vectorize generates.

The problem in one sentence

The same dot-product kernel must be written once for NEON’s 4 f32 lanes, again for AVX-512’s 16, and again as a scalar remainder loop — polars literally ships its AVX-512 filter twice for two element types — because in today’s languages vector width is a property of the code you wrote, not a parameter the compiler fills in.

The concepts, step by step

Step 1 — the ladder of SIMD ergonomics

SIMD (single instruction, multiple data — one instruction operating on a vector of W values, its lanes) can be programmed at three levels of abstraction, each trading control for portability:

 raw intrinsics      vfmaq_f32(acc, a, b)         per-ISA names, unsafe,
 (core::arch)                                     exact instruction control
        │
 portable library    acc = a.mul_add(b, acc)      one vocabulary, library
 (std::simd, wide)   Simd<f32, 4> / f32x4         picks instructions; width
        │                                         is a const generic bolted on
        │
 language type       SIMD[DType.float32, 4]       scalars ARE SIMD[T,1];
 (Mojo)              fn foo[w: Int](x: SIMD[T,w]) width is a first-class
                                                  compile-time parameter

Intrinsics (compiler-provided functions mapping 1:1 to machine instructions, like NEON’s vfmaq_f32) name the exact instruction — maximal control, zero portability. Portable libraries give one vocabulary and let the library pick instructions. Mojo moves the whole idea into the language.

Step 2 — scalars are width-1 vectors

Mojo’s unification move: Float32 is literally an alias for SIMD[DType.float32, 1]. There is no separate scalar world — every scalar function is already the width-1 instance of a width-generic function, so vectorizing an algorithm means changing a parameter, not rewriting the body. Two practical wins: no duplicate scalar/vector implementations to keep in sync, and the width-1 instance is a free test oracle — run the same body at w=1 and w=4 and diff the outputs (steal this for dot.rs’s tests).

Step 3 — parametric width in action: simdwidthof + vectorize

With width as a compile-time parameter, the machine’s natural width becomes a queryable constant and the boilerplate becomes generated code:

fn kernel[w: Int](v: SIMD[DType.float32, w]) -> SIMD[DType.float32, w]
vectorize[kernel, simdwidthof[DType.float32]()](n)

simdwidthof is a compile-time query of the TARGET (4 on NEON, 16 on AVX-512); vectorize instantiates the kernel at that width plus a scalar remainder at width 1 — the same body, monomorphized twice. The polars contrast makes the value concrete: polars hardcodes STRIPE=16 and writes the AVX-512 filter twice (u8 and u32) because Rust’s const generics can’t cleanly abstract “the best width for this type on this target.” Question: what stops wide/std::simd from doing this today? (The width is a const generic, but there’s no portable “native width” query, and no autogenerated remainder loop — you write both by hand in dot.rs.)

Step 4 — the matmul arc: four separable layers

The famous progression (numbers from Modular’s blog, M-series-class hardware, GFLOPS order-of-magnitude):

 Python baseline                ~0.002   ×1
 naive Mojo (same loops)        ~5       ×2000    compiled, typed
 + vectorize inner loop         ~25      SIMD lanes
 + parallelize outer loop       ~100     cores (topic 14, not 17)
 + tile + unroll                ~200+    cache blocking (topic 13)

Note the ORDER: types/compilation first, lanes second, cores third, cache blocking last — and each step is a decorator/parameter, not a rewrite. The lesson isn’t “Mojo is fast”; it’s that the four layers (scalar semantics → lanes → threads → tiles) are SEPARABLE when width is parametric. Our experiments walk the same rungs by hand.

Step 5 — what parametric width does NOT solve

The type system dissolves boilerplate, not microarchitecture:

  • compress/gather-shaped problems (simdjson’s LUT trick) still need per-ISA thought — a parametric width doesn’t conjure vpcompress on NEON; you still write the LUT or take the branchless store.
  • ports × latency (the accumulator count — reading-simsimd.md) is still yours: the matmul blog manually unrolls 4 accumulator vectors, exactly like SimSIMD’s 4-state API. No compiler infers “12 chains” for you — reassociation of floats stays illegal in every language.

The dividing line: anything expressible as “same op, any width” the language absorbs; anything about which instructions and how many independent chains stays engineering.

Step 6 — the translation table for our stack

Every Mojo construct has a hand-written stable-Rust equivalent — this table is the map from the ideal to what M17 actually ships:

Mojostable Rust (our experiments)
SIMD[DType.float32, 4]wide::f32x4
simdwidthof[T]()hardcode 4 on NEON (128-bit / 32)
vectorize[kernel, w](n)hand-written chunks_exact(w) + remainder
@unroll / accumulator params4 named accumulator variables
width-1 fallbackyour scalar fn, kept for dispatch

How to read the docs (with the concepts in hand)

  1. The SIMD stdlib page — skim the type’s surface with Step 2 in mind; the interesting part is what’s a parameter (dtype, width) versus what’s a method.
  2. The matmul blog arc — read in full against Step 4’s table; at each rung, name which topic of this course (11, 13, 14, 17) the speedup belongs to.
  3. Then go back to dot.rs and identify each hand-written line that vectorize would have generated — that’s Step 6 made concrete.

Questions for notes.md

  1. Float32 = SIMD[f32,1]: what does making scalars width-1 vectors buy for TESTING kernels (hint: run the same body at w=1 as the oracle for w=4 — steal this for dot.rs’s tests)?
  2. vectorize generates the remainder loop; where in our filter.rs is the equivalent, and why is the remainder the classic source of SIMD bugs (simdjson pads its input to 64 instead — compare)?
  3. Rust’s std::simd Simd<f32, N> has the parametric type but sits on nightly for years. What’s actually hard: the type, the portable ops (compress!), or stabilizing the ISA mapping?
  4. The matmul arc gains more from tiling than from SIMD. Reconcile with this topic’s “last 10× on a single core” framing — when is topic 13 the bigger lever than topic 17?
  5. For M17: our engine will hardcode NEON width 4. Write the one sentence justifying that (deployment target) and the one-line escape hatch if SVE servers arrive.

References

Papers & docs

  • Modular — Mojo stdlib docs for SIMD (docs.modular.com) — the parametric type itself; no clone needed
  • Modular — the “Matrix Multiplication in Mojo” blog arc (matmul.mojo) — the ×2000-to-tiled progression walked in §4