Mojo Vectorization

Vectorization rewrites a loop so that each iteration processes multiple data elements simultaneously using SIMD hardware instructions. Instead of computing one result per loop cycle, vectorized code computes 4, 8, or 16 results per cycle. This is one of Mojo's most powerful performance tools.

Scalar vs Vectorized Loop

Scalar loop (one element at a time):
  for i in range(8):
      result[i] = a[i] + b[i]

  Cycle 1: result[0] = a[0] + b[0]
  Cycle 2: result[1] = a[1] + b[1]
  ...
  Cycle 8: result[7] = a[7] + b[7]
  Total: 8 CPU additions

Vectorized loop (8 elements at once with SIMD width=8):
  result[0:8] = a[0:8] + b[0:8]

  Cycle 1: ALL 8 additions happen simultaneously
  Total: 1 CPU instruction
  Speedup: ~8×

Manual Vectorization with SIMD

from memory import UnsafePointer

fn add_arrays_scalar(a: UnsafePointer[Float32],
                     b: UnsafePointer[Float32],
                     result: UnsafePointer[Float32],
                     n: Int):
    for i in range(n):
        result[i] = a[i] + b[i]

fn add_arrays_vector(a: UnsafePointer[Float32],
                     b: UnsafePointer[Float32],
                     result: UnsafePointer[Float32],
                     n: Int):
    alias width = 8   # process 8 floats at once
    var i = 0
    while i + width <= n:
        var va = SIMD[DType.float32, width].load(a + i)
        var vb = SIMD[DType.float32, width].load(b + i)
        (va + vb).store(result + i)
        i += width
    # Handle remaining elements
    while i < n:
        result[i] = a[i] + b[i]
        i += 1

Vectorized loop iteration diagram (width=4):

  i=0:   a[0,1,2,3] + b[0,1,2,3] → result[0,1,2,3]   ← 1 instruction
  i=4:   a[4,5,6,7] + b[4,5,6,7] → result[4,5,6,7]   ← 1 instruction
  i=8:   scalar fallback for any remaining elements

Using the vectorize Higher-Order Function

Mojo's standard library provides a vectorize function that handles the SIMD width selection and loop structure automatically.

from algorithm import vectorize
from memory import UnsafePointer

fn scale_array(data: UnsafePointer[Float32], n: Int, factor: Float32):
    @parameter
    fn scale_simd[width: Int](i: Int):
        var chunk = SIMD[DType.float32, width].load(data + i)
        (chunk * factor).store(data + i)

    vectorize[scale_simd, 8](n)

fn main():
    let n = 16
    var arr = UnsafePointer[Float32].alloc(n)
    for i in range(n):
        arr.init_pointee_copy(Float32(i + 1))

    scale_array(arr, n, 2.0)

    for i in range(n):
        print(arr[i], end=" ")   # 2 4 6 8 10 12 14 16 ...

    for i in range(n):
        (arr + i).destroy_pointee()
    arr.free()

Choosing SIMD Width

Width | Bits used | Typical hardware support
------|-----------|-----------------------------
  4   | 128 bits  | SSE (Intel), NEON (ARM)
  8   | 256 bits  | AVX2 (modern Intel/AMD)
  16  | 512 bits  | AVX-512 (server CPUs)

Query the hardware width at compile time:
  from sys.info import simdwidthof
  alias W = simdwidthof[DType.float32]()   # auto-detect

Vectorized Dot Product

from algorithm import vectorize
from memory import UnsafePointer

fn dot_product(a: UnsafePointer[Float32],
               b: UnsafePointer[Float32],
               n: Int) -> Float32:
    var accumulator = SIMD[DType.float32, 8].splat(0.0)

    @parameter
    fn multiply_add[width: Int](i: Int):
        var va = SIMD[DType.float32, width].load(a + i)
        var vb = SIMD[DType.float32, width].load(b + i)
        accumulator += va * vb

    vectorize[multiply_add, 8](n)
    return accumulator.reduce_add()

Diagram: Vectorized dot product (width=4, n=8)

  Iteration 1 (i=0):
    a[0,1,2,3] × b[0,1,2,3] → partial[0,1,2,3]
    accumulator += partial

  Iteration 2 (i=4):
    a[4,5,6,7] × b[4,5,6,7] → partial[4,5,6,7]
    accumulator += partial

  Final: accumulator.reduce_add() → scalar sum

Vectorization and Memory Alignment

SIMD loads are fastest when data sits at an aligned memory address — a multiple of the SIMD vector size in bytes. Misaligned loads still work but run slower on some hardware.

from memory import UnsafePointer

# alloc() returns aligned memory by default in Mojo
var arr = UnsafePointer[Float32].alloc(64)   # aligned for SIMD

Real-World Impact

Operation         | Scalar time | Vectorized time | Speedup
------------------|-------------|-----------------|--------
Add 1M floats     | 4 ms        | 0.5 ms          | ~8×
Scale image pixels| 12 ms       | 1.5 ms          | ~8×
Dot product (1M)  | 6 ms        | 0.4 ms          | ~15×
(Higher speedups when memory fits in CPU cache)

Key Takeaways

Vectorization processes multiple data elements per CPU instruction using SIMD. Manual vectorization loads chunks into SIMD variables, performs operations, and stores results. The vectorize higher-order function automates SIMD width selection and loop structure. Match your SIMD width to the target hardware — 8 float32 values at a time covers most modern CPUs. Aligned memory maximizes SIMD throughput. Vectorization delivers 4–16× speedups on numerical workloads and is one of the primary reasons Mojo dramatically outperforms Python.