Mojo Parallelization

Parallelization splits work across multiple CPU cores so that different parts of a computation run at the same time. A modern CPU has 4, 8, 16, or more cores sitting idle during a single-threaded loop. Parallelization puts them all to work, cutting wall-clock time proportionally.

Single-Core vs Multi-Core

Single-core (sequential):
  Core 0: [chunk 0]──[chunk 1]──[chunk 2]──[chunk 3]
  Time:    ─────────────────────────────────────────→ 4 units

Multi-core (parallel, 4 cores):
  Core 0: [chunk 0]
  Core 1: [chunk 1]
  Core 2: [chunk 2]
  Core 3: [chunk 3]
  Time:    ─────────→ 1 unit   ← 4× faster

The parallelize Function

Mojo's parallelize function from the standard library distributes loop iterations across all available CPU cores automatically.

from algorithm import parallelize
from memory import UnsafePointer

fn main():
    let n = 1_000_000
    var data = UnsafePointer[Float64].alloc(n)
    for i in range(n):
        data.init_pointee_copy(Float64(i))

    # Square every element in parallel
    @parameter
    fn square_element(i: Int):
        data[i] = data[i] * data[i]

    parallelize[square_element](n)   # Mojo decides thread count

    print(data[0])   # 0.0
    print(data[4])   # 16.0  (4² = 16)

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

Parallelism with a Fixed Number of Workers

from algorithm import parallelize

fn main():
    let n = 16
    var results = UnsafePointer[Int].alloc(n)
    for i in range(n):
        results.init_pointee_copy(0)

    @parameter
    fn compute(i: Int):
        # Each worker computes its own chunk independently
        results[i] = i * i

    parallelize[compute](n, 4)   # use 4 threads

    for i in range(n):
        print(results[i], end=" ")
    print("")

    for i in range(n):
        (results + i).destroy_pointee()
    results.free()
Thread assignment (4 threads, 16 iterations):
  Thread 0: i = 0,1,2,3
  Thread 1: i = 4,5,6,7
  Thread 2: i = 8,9,10,11
  Thread 3: i = 12,13,14,15
  All run simultaneously on separate CPU cores.

Combining Parallelism and Vectorization

The greatest throughput comes from using both at once — parallelize across cores and vectorize within each core's chunk.

from algorithm import parallelize, vectorize
from memory import UnsafePointer

fn parallel_vector_scale(data: UnsafePointer[Float32], n: Int, factor: Float32):
    alias vector_width = 8

    @parameter
    fn process_chunk(chunk_start: Int):
        @parameter
        fn scale_simd[w: Int](offset: Int):
            var i = chunk_start + offset
            var v = SIMD[DType.float32, w].load(data + i)
            (v * factor).store(data + i)
        vectorize[scale_simd, vector_width](n // 4)

    parallelize[process_chunk](4)   # 4 CPU cores
Combined diagram:

  Core 0 handles indices 0–249,999:
    SIMD chunk [0:8], [8:16], [16:24] ...  (vectorized)

  Core 1 handles indices 250,000–499,999:
    SIMD chunk [250000:250008] ...          (vectorized)

  Core 2 handles indices 500,000–749,999:
    SIMD chunk ...                          (vectorized)

  Core 3 handles indices 750,000–999,999:
    SIMD chunk ...                          (vectorized)

  All cores run at the same time → 4 cores × 8 SIMD = 32× throughput

Avoiding Data Races

A data race occurs when two threads read and write the same memory location at the same time. The result is unpredictable. Design parallel workloads so each thread operates on its own exclusive portion of data.

SAFE — each index written by exactly one thread:
  @parameter
  fn compute(i: Int):
      results[i] = heavy_computation(i)   ← only thread i writes results[i]

UNSAFE — multiple threads write to the same location:
  @parameter
  fn accumulate(i: Int):
      shared_sum += data[i]   ← all threads write shared_sum → data race!

To safely accumulate across threads, compute a partial sum per thread, then combine partial sums after all threads finish.

When to Parallelize

Good candidates for parallelization:
  ✓ Large array transformations (scale, filter, normalize)
  ✓ Image processing (each pixel or row is independent)
  ✓ Monte Carlo simulations (independent random trials)
  ✓ Matrix multiplication (independent row × column products)

Poor candidates:
  ✗ Sequential algorithms where step N depends on step N-1
  ✗ Very small datasets (thread overhead exceeds the savings)
  ✗ Workloads bottlenecked by a single resource (disk, network)

Amdahl's Law

If 80% of your program can be parallelized and 20% must be sequential:

  Speedup with infinite cores = 1 / 0.20 = 5×

  No matter how many cores you add, you cannot exceed 5× because
  the sequential 20% sets the floor for total runtime.

  Lesson: Identify and minimize the sequential bottleneck first.

Key Takeaways

Parallelization distributes loop iterations across multiple CPU cores using parallelize. Each iteration must be independent — no shared mutable state between iterations. Combine parallelize and vectorize for maximum throughput on large numerical workloads. Avoid data races by ensuring each thread writes to its own memory region. Use partial results per thread and merge after completion for safe parallel reductions. Amdahl's Law caps the achievable speedup based on the fraction of work that remains sequential.

Leave a Comment

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