Mojo Python Acceleration: SIMD Optimization and Parallel Processing for AI Workloads

Mojo bridges Python's simplicity with systems-level performance through native SIMD vectorization and parallel processing capabilities. This guide covers essential optimization patterns for AI workloads.

SIMD Vectorization Fundamentals

SIMD (Single Instruction, Multiple Data) enables parallel element-wise operations on contiguous memory blocks. Mojo's SIMD[DType, size] type maps directly to CPU vector registers.

from math import reduce_add

# Create SIMD vectors using splat or element-wise initialization
var a = SIMD[DType.float32, 8].splat(0.0)
a = SIMD[DType.float32, 8](1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0)
var b = SIMD[DType.float32, 8](10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0)

# Element-wise operations execute in parallel
var c = a * b + 5.0
var sum_result = reduce_add(c)

Handling Non-Power-of-2 Sizes

Real-world data rarely aligns to power-of-2 boundaries. Use the vectorize decorator for clean SIMD iteration:

from math import reduce_add
from algorithm import vectorize

fn process_with_tail(data: UnsafePointer[Float32], size: Int) -> Float32:
    let simd_width = 8
    var total = Float32(0.0)
    
    # Process full SIMD chunks using vectorize
    fn process_chunk[simd_width_param: Int](idx: Int) -> None:
        @unroll
        for j in range(simd_width_param):
            total += (data + idx + j).load() * (data + idx + j).load()
    
    vectorize[process_chunk, simd_width](size)
    
    # Handle remaining elements
    for i in range((size // simd_width) * simd_width, size):
        total += (data + i).load() * (data + i).load()
    
    return total

AI Workload Optimization Patterns

Element-wise Operations with Autotune

Use autotune to automatically select optimal SIMD width for your hardware:

from algorithm import vectorize, autotune
from math import reduce_add

fn vectorized_elementwise(
    a: UnsafePointer[Float32],
    b: UnsafePointer[Float32],
    result: UnsafePointer[Float32],
    size: Int
):
    fn compute_chunk[simd_width: Int](idx: Int) -> None:
        let va = SIMD[DType.float32, simd_width].load(a + idx)
        let vb = SIMD[DType.float32, simd_width].load(b + idx)
        let vresult = va * vb + (va * 0.1)
        vresult.store(result + idx)
    
    # Autotune selects best simd_width at compile time
    alias simd_width = autotune(1, 2, 4, 8, 16, 32)
    vectorize[compute_chunk, simd_width](size)

Tiled Matrix Multiplication

True matrix multiplication leverages SIMD through tiled computation for cache efficiency:

from math import reduce_add

fn matrix_multiply(
    a: UnsafePointer[Float32],
    b: UnsafePointer[Float32],
    result: UnsafePointer[Float32],
    m: Int, n: Int, k: Int
):
    let tile_size = 8
    let simd_width = 8
    
    for row in range(m):
        for col_tile in range(0, k, tile_size):
            var tile_sum = SIMD[DType.float32, tile_size].splat(0.0)
            
            # Compute dot product with tiled SIMD accumulation
            for idx in range(n):
                let a_val = (a + row * n + idx).load()
                let a_vec = SIMD[DType.float32, tile_size].splat(a_val)
                
                # Load tile from B matrix
                if col_tile + tile_size <= k:
                    let b_vec = SIMD[DType.float32, tile_size].load(b + idx * k + col_tile)
                    tile_sum = tile_sum + a_vec * b_vec
                else:
                    # Handle partial tile
                    for j in range(col_tile, min(col_tile + tile_size, k)):
                        tile_sum[j - col_tile] += a_val * (b + idx * k + j).load()
            
            # Store computed tile
            if col_tile + tile_size <= k:
                tile_sum.store(result + row * k + col_tile)
            else:
                for j in range(col_tile, min(col_tile + tile_size, k)):
                    (result + row * k + j).load() = tile_sum[j - col_tile]

Activation Functions

from math import exp

fn relu_vectorized(x: SIMD[DType.float32, 16]) -> SIMD[DType.float32, 16]:
    return x.max(SIMD[DType.float32, 16].splat(0.0))

fn sigmoid_vectorized(x: SIMD[DType.float32, 16]) -> SIMD[DType.float32, 16]:
    let neg_x = SIMD[DType.float32, 16].splat(0.0) - x
    let exp_neg_x = exp(neg_x)
    let ones = SIMD[DType.float32, 16].splat(1.0)
    return ones / (ones + exp_neg_x)

Parallel Processing

Mojo's parallelize distributes work across available CPU cores. The worker function must match the signature fn(Int) capturing -> None:

from algorithm import parallelize

fn parallel_batch_process(data: UnsafePointer[Float32], size: Int):
    let simd_width = 8
    
    fn worker(start: Int) -> None:
        let chunk_size = 64
        var end = min(start + chunk_size, size)
        for i in range(start, end, simd_width):
            if i + simd_width <= size:
                var vec = SIMD[DType.float32, simd_width].load(data + i)
                var activated = vec.max(SIMD[DType.float32, simd_width].splat(0.0))
                activated.store(data + i)
            else:
                # Handle tail elements
                for j in range(i, end):
                    if (data + j).load() < 0.0:
                        (data + j).store(0.0)
    
    parallelize[worker](size)

Memory Layout Optimization

AI workloads benefit from contiguous memory layouts that maximize cache efficiency:

struct Tensor[size: Int]:
    var data: UnsafePointer[Float32]
    
    fn __init__(out self):
        self.data = alloc[Float32](size)
    
    fn __del__(owned self):
        free(self.data)
    
    fn batch_process(inout self):
        let simd_width = 8
        
        fn worker(start: Int) -> None:
            let chunk_size = 64
            var end = min(start + chunk_size, size)
            for i in range(start, end, simd_width):
                if i + simd_width <= size:
                    var vec = SIMD[DType.float32, simd_width].load(self.data + i)
                    var activated = vec.max(SIMD[DType.float32, simd_width].splat(0.0))
                    activated.store(self.data + i)
        
        parallelize[worker](size)

Performance Benchmarks

SIMD operations demonstrate 10-50x speedup over scalar loops. Create test data outside the benchmark loop for fair comparison:

from benchmark import Benchmark
from math import reduce_add

fn benchmark_simd_vs_scalar():
    let iterations = 100000
    let simd_width = 8
    
    # Pre-create test data
    var data = SIMD[DType.float32, 8](1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0)
    
    fn scalar_version() -> Float32:
        var sum = Float32(0.0)
        for _ in range(iterations):
            @unroll
            for i in range(simd_width):
                sum += data[i] * data[i]
        return sum
    
    fn simd_version() -> Float32:
        var sum = Float32(0.0)
        for _ in range(iterations):
            var squared = data * data
            sum += reduce_add(squared)
        return sum
    
    let scalar_report = Benchmark.run(scalar_version)
    let simd_report = Benchmark.run(simd_version)
    print("Scalar time:", scalar_report.mean())
    print("SIMD time:", simd_report.mean())

Implementation Guidelines

1. Vector Width Selection: Use autotune for portable SIMD width selection, or match to CPU architecture (AVX-512: 16 float32, AVX2: 8 float32)
2. Memory Alignment: Ensure 32-byte alignment for optimal SIMD performance
3. Loop Unrolling: Apply @unroll to inner loops to reduce branch overhead
4. Data Type Consistency: Maintain uniform DType within SIMD operations
5. Tail Handling: Always process remainder elements when sizes are not power-of-2
6. Parallelization: Use parallelize[func](num_work_items) for multi-core scaling

Getting Started

1. Replace Python loops with SIMD operations using vectorize decorator
2. Profile hot paths using Mojo's Benchmark.run() function
3. Optimize memory layouts for cache-friendly access patterns
4. Use parallelize for multi-core scaling of independent work items
5. Apply autotune for portable performance across different CPU architectures

This approach delivers near-C performance while maintaining Python's development velocity for AI workloads.