Unlocking SIMD: A Hands-On Guide with Matrix Ops in Mojo
Single Instruction, Multiple Data (SIMD) represents one of the most powerful optimization techniques available to modern programmers. With Mojo's zero-cost abstraction and direct hardware access, we can harness SIMD's full potential while maintaining code readability. Today, we'll explore SIMD optimization through a practical example: implementing efficient matrix operations.
What is SIMD?
SIMD allows your CPU to perform the same operation on multiple data elements simultaneously. Instead of processing one floating-point number at a time, modern processors can handle 4, 8, or even 16 values in a single instruction cycle. This parallelism can deliver dramatic performance improvements, often 4x to 16x speedups for suitable workloads.
The key insight is that many computational problems are embarrassingly parallel at the data level. Matrix operations, image processing, and numerical computations all exhibit this characteristic, making them ideal candidates for SIMD optimization.
SIMD in Action: Simple Examples
Let's start with a concrete example to see SIMD in action. Imagine you want to add two arrays of numbers:
Array A: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]
Array B: [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5]
Result: [1.5, 3.5, 5.5, 7.5, 9.5, 11.5, 13.5, 15.5]
The Scalar Approach (Traditional)
Without SIMD, your CPU processes one pair at a time:
fn add_arrays_scalar(a: UnsafePointer[Float32], b: UnsafePointer[Float32],
result: UnsafePointer[Float32], size: Int):
for i in range(size):
result[i] = a[i] + b[i] # One addition per cycle
CPU Instructions:
- Load a[0], load b[0] → add → store result[0]
- Load a[1], load b[1] → add → store result[1]
- Load a[2], load b[2] → add → store result[2]
- ... (8 separate cycles for 8 elements)
The SIMD Approach (Vectorized)
With SIMD, your CPU can process multiple elements simultaneously:
fn add_arrays_simd(a: UnsafePointer[Float32], b: UnsafePointer[Float32],
result: UnsafePointer[Float32], size: Int):
alias simd_width = 8 # Process 8 float32s at once
@parameter
fn vectorized_add[width: Int](i: Int):
var vec_a = a.load[width=width](i) # Load 8 elements from A
var vec_b = b.load[width=width](i) # Load 8 elements from B
var vec_result = vec_a + vec_b # Add all 8 pairs at once!
result.store[width=width](i, vec_result) # Store 8 results
vectorize[vectorized_add, simd_width](size)
CPU Instructions (with AVX2):
- Load 8 floats from A, load 8 floats from B → add all 8 pairs → store 8 results Result: 8x fewer cycles!
Visual Representation
Here's what happens inside the CPU:
Scalar Processing (8 cycles):
Cycle 1: [1.0] + [0.5] = [1.5]
Cycle 2: [2.0] + [1.5] = [3.5]
Cycle 3: [3.0] + [2.5] = [5.5]
... (continues for each element)
SIMD Processing (1 cycle):
Cycle 1: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] +
[0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5] =
[1.5, 3.5, 5.5, 7.5, 9.5, 11.5, 13.5, 15.5]
A More Practical Example: Scaling Values
Let's see another common pattern, multiplying an array by a constant:
fn scale_array_naive(data: UnsafePointer[Float32], size: Int, factor: Float32):
for i in range(size):
data[i] = data[i] * factor
fn scale_array_simd(data: UnsafePointer[Float32], size: Int, factor: Float32):
alias simd_width = simdwidthof[DType.float32]()
@parameter
fn vectorized_scale[width: Int](i: Int):
var vec_data = data.load[width=width](i)
var vec_factor = SIMD[DType.float32, width](factor) # Broadcast scalar to vector
var scaled = vec_data * vec_factor
data.store[width=width](i, scaled)
vectorize[vectorized_scale, simd_width](size)
In this example, the scalar factor get broadcast to all lanes of the SIMD vector, then all multiplication happen simultaneously.
When SIMD Works Best
SIMD shines when you have:
- Regular data patterns (arrays, matrices)
- Same operation on multiple data (the "Single Instruction" part)
- Sufficient data parallelism (at least 4-16 elements to process)
- Memory-bound or compute-bound workloads (not branch-heavy code) Now that we understand the fundamental, let's see how these principles apply to matrix operations.
Building Our Foundation: The Matrix Structure
Let's start with a basic matrix implementation in Mojo:
struct Matrix[dtype: DType]():
var rows: Int
var cols: Int
var data: UnsafePointer[SIMD[dtype, 1]]
fn __init__(out self, rows: Int, cols: Int):
self.rows = rows
self.cols = cols
self.data = UnsafePointer[SIMD[dtype, 1]].alloc(rows * cols)
fn __del__(owned self):
self.data.free()
fn __getitem__(self, row: Int, col: Int) -> SIMD[dtype, 1]:
return self.data[row * self.cols + col]
fn __setitem__(self, row: Int, col: Int, value: SIMD[dtype, 1]):
self.data[row * self.cols + col] = value
This matrix uses row-major storage, elements are stored row by row in contiguous memory. This layout choice will become crucial when we optimize for cache performance later.
The Naive Approach: Row Mean Calculation
Let's implement a straightforward row mean calculation:
fn row_mean_naive[dtype: DType](
data: UnsafePointer[SIMD[dtype, 1]],
rows: Int,
cols: Int
) -> UnsafePointer[SIMD[dtype, 1]]:
var result = UnsafePointer[SIMD[dtype, 1]].alloc(rows)
for row in range(rows):
var sum = SIMD[dtype, 1](0)
for col in range(cols):
sum += data[row * cols + col]
result[row] = sum / cols
return result
This implementation processes one element at a time, exactly what SIMD is designed to accelerate. Each addition operation works on a single SIMD[dtype, 1] value, leaving the processor's vector units largely unused.
SIMD Optimization: Vectorization and Parallelization
Now let's unlock SIMD's power with Mojo's vectorize and parallelize functions:
fn row_mean_optim[dtype: DType](
data: UnsafePointer[SIMD[dtype, 1]],
rows: Int,
cols: Int,
) -> UnsafePointer[SIMD[dtype, 1]]:
var result = UnsafePointer[SIMD[dtype, 1]].alloc(rows)
alias simd_width = simdwidthof[dtype]()
@parameter
fn row_mean_parallelized(row: Int):
var sum = SIMD[dtype, 1](0)
@parameter
fn row_mean_vectorized[simd_width: Int](col: Int):
var data_ptr = data + row * cols + col
var simd_vals = data_ptr.load[width=simd_width]()
sum += simd_vals.reduce_add()
vectorize[row_mean_vectorized, simd_width](size=cols)
result[row] = sum / cols
parallelize[row_mean_parallelized](rows, rows)
return result
Breaking Down the Optimization
1. SIMD Width Detection:simdwidthof[dtype]() automatically determines the optimal vector width for your CPU and data type. For float32 on AVX2, this returns 8, meaning we can process 8 floats simultaneously.
2. Vectorization: The inner vectorize function processes multiple columns at once:
var simd_vals = data_ptr.load[width=simd_width]()
sum += simd_vals.reduce_add()
Instead of 8 separate load-and-add operations, we perform one vectorized load and one vectorized reduction.
3. Parallelization: The outer parallelization function distributes rows across CPU cores, enabling thread-level parallelism alongside data-level parallelism.
The Column Mean Challenge: Memory Layout Matters
Column operation present a different challenge. Let's look at the naive approach:
fn col_mean_naive[dtype: DType](
data: UnsafePointer[SIMD[dtype, 1]],
rows: Int,
cols: Int,
) -> UnsafePointer[SIMD[dtype, 1]]:
var result = UnsafePointer[SIMD[dtype, 1]].alloc(cols)
for col in range(cols):
var sum = SIMD[dtype, 1](0.0)
for row in range(rows):
sum += data[row * cols + col] # Strided access!
result[col] = sum / rows
return result
The problem here is memory access pattern. While row operations access contiguous memory (cache-friendly), column operations access elements separated by cols positions in memory. This strided access pattern can cause cache misses and prevent effective SIMD vectorization.
Advanced Optimization: Matrix Multiplication
fn col_mean_optim[dtype: DType](
data: UnsafePointer[SIMD[dtype, 1]],
rows: Int,
cols: Int,
) -> UnsafePointer[SIMD[dtype, 1]]:
var result = UnsafePointer[SIMD[dtype, 1]].alloc(cols)
# Create transposed data for better cache performance
var transposed_data = UnsafePointer[SIMD[dtype, 1]].alloc(rows * cols)
@parameter
fn transposed_parallelized(row: Int):
for col in range(cols):
transposed_data[col * rows + row] = data[row * cols + col]
parallelize[transposed_parallelized](rows, rows)
alias simd_width = simdwidthof[dtype]()
@parameter
fn col_mean_parallelize(col: Int):
var sum = SIMD[dtype, 1](0.0)
@parameter
fn col_mean_vectorized[simd_width: Int](row: Int):
var data_ptr = transposed_data + col * rows + row
var simd_vals = data_ptr.load[width=simd_width]()
sum += simd_vals.reduce_add()
vectorize[col_mean_vectorized, simd_width](size=rows)
result[col] = sum / rows
parallelize[col_mean_parallelize](cols, cols)
transposed_data.free()
return result
By transposing the matrix first, we transform the stripped column access into contiguous memory access, enabling efficient SIMD vectorization. This trades some additional memory and transpose cost for dramatically improved vectorization efficiency.
Benchmarking: Measuring Real Performance
Proper benchmarking is crucial for optimization work. Let's build a complete benchmarking harness to measure our SIMD optimization accurately.
The Complete Benchmarking Code
Here's the full main() function that demonstrates how to properly benchmark SIMD optimizations:
fn print_sample[dtype: DType](data: UnsafePointer[SIMD[dtype, 1]], size: Int, name: String):
print(name + ":")
if size > 10:
print("First 10 elements:")
for i in range(10):
print(" ", data[i])
else:
for i in range(size):
print(" ", data[i])
print()
fn main() raises:
# Set a seed for reproducible results
seed()
alias matrix_size = 1024 # 1024x1024 matrix = ~4MB of float32 data
var mat = Matrix[DType.float32](matrix_size, matrix_size)
print("Filling matrix with random values...")
mat.rand_init()
print("Matrix filled. Running benchmarks...")
# First, verify our implementations produce correct results
var result_naive = row_mean_naive[DType.float32](mat.data, mat.rows, mat.cols)
print_sample(result_naive, 10, "Row Mean Naive")
var result_optim = row_mean_optim[DType.float32](mat.data, mat.rows, mat.cols)
print_sample(result_optim, 10, "Row Mean Optimized")
# Define benchmark functions with anti-optimization measures
@parameter
fn bench_row_mean_naive():
var result = row_mean_naive[DType.float32](mat.data, mat.rows, mat.cols)
# Force the compiler to keep the computation
var sum = SIMD[DType.float32, 1](0)
for i in range(mat.rows):
sum += result[i] # "Use" the result
if sum > 1e10: # Impossible condition
print("Benchmark result:", sum) # Never executes, but compiler can't prove it
result.free()
@parameter
fn bench_row_mean_optimized():
var result = row_mean_optim[DType.float32](mat.data, mat.rows, mat.cols)
# Same anti-optimization pattern
var sum = SIMD[DType.float32, 1](0)
for i in range(mat.rows):
sum += result[i]
if sum > 1e10:
print("Row benchmark result:", sum)
result.free()
@parameter
fn bench_col_mean_naive():
var result = col_mean_naive[DType.float32](mat.data, mat.rows, mat.cols)
# Same protection needed for column operations
var sum = SIMD[DType.float32, 1](0)
for i in range(mat.cols):
sum += result[i]
if sum > 1e10:
print("Column benchmark result:", sum)
result.free()
@parameter
fn bench_col_mean_optimized():
var result = col_mean_optim[DType.float32](mat.data, mat.rows, mat.cols)
var sum = SIMD[DType.float32, 1](0)
for i in range(mat.cols):
sum += result[i]
if sum > 1e10:
print("Column optimized result:", sum)
result.free()
# Run the actual benchmarks
print("Running benchmarks...")
var report_row_mean_naive = benchmark.run[bench_row_mean_naive]()
var report_row_mean_optimized = benchmark.run[bench_row_mean_optimized]()
var report_col_mean_naive = benchmark.run[bench_col_mean_naive]()
var report_col_mean_optimized = benchmark.run[bench_col_mean_optimized]()
# Print detailed results
print("\n" + "="*50)
print("BENCHMARK RESULTS")
print("="*50)
print("\nRow Mean Operations:")
print("-" * 30)
print("Naive Implementation:")
report_row_mean_naive.print()
print("\nOptimized Implementation:")
report_row_mean_optimized.print()
print("Row Mean Speedup: ",
report_row_mean_naive.mean() / report_row_mean_optimized.mean(), "x")
print("\nColumn Mean Operations:")
print("-" * 30)
print("Naive Implementation:")
report_col_mean_naive.print()
print("Optimized Implementation:")
report_col_mean_optimized.print()
print("Column Mean Speedup: ",
report_col_mean_naive.mean() / report_col_mean_optimized.mean(), "x")
# Clean up sample results
result_naive.free()
result_optim.free()
Key Benchmarking Principles
1. Anti-Optimization Protection: The most critical aspect is preventing the compiler from optimizing away your computation:
var sum = SIMD[DType.float32, 1](0)
for i in range(mat.rows):
sum += result[i] # Force usage of results
if sum > 1e10: # Impossible condition
print("Benchmark result:", sum) # Never executes
This pattern ensures the compiler must actually compute and store results, but the impossible condition means no printing overhead affects timing.
2. Reproducible Results: Always seed your random number generator:
seed() # Ensures consistent data across benchmark runs
3. Verification Before Benchmarking: Print sample outputs to verify correctness:
var result_naive = row_mean_naive[DType.float32](mat.data, mat.rows, mat.cols)
print_sample(result_naive, 10, "Row Mean Naive")
4. Multiple Measurements: Mojo's benchmark.run() automatically handles multiple iterations and statistical analysis.
Sample Output
When you run this benchmark, you might see output like:
Filling matrix with random values...
Matrix filled. Running benchmarks...
Row Mean Naive:
First 10 elements:
0.489234
0.523891
0.456782
...
Running benchmarks...
==================================================
BENCHMARK RESULTS
==================================================
Row Mean Operations:
------------------------------
Naive Implementation:
--------------------------------------------------------------------------------
Benchmark Report (s)
--------------------------------------------------------------------------------
Mean: 0.0006902331518987342
Total: 2.399250436
Iters: 3476
Warmup Total: 0.000865924
Fastest Mean: 0.0006902331518987342
Slowest Mean: 0.0006902331518987342
Optimized Implementation:
--------------------------------------------------------------------------------
Benchmark Report (s)
--------------------------------------------------------------------------------
Mean: 0.00012104455985985986
Total: 2.418470306
Iters: 19980
Warmup Total: 0.000254114
Fastest Mean: 0.00012104455985985986
Slowest Mean: 0.00012104455985985986
Row Mean Speedup: 5.7023062638903905 x
Column Mean Operations:
------------------------------
Naive Implementation:
--------------------------------------------------------------------------------
Benchmark Report (s)
--------------------------------------------------------------------------------
Mean: 0.0032827859138166897
Total: 2.399716503
Iters: 731
Warmup Total: 0.003365046
Fastest Mean: 0.0032827859138166897
Slowest Mean: 0.0032827859138166897
Optimized Implementation:
--------------------------------------------------------------------------------
Benchmark Report (s)
--------------------------------------------------------------------------------
Mean: 0.002123317511
Total: 2.123317511
Iters: 1000
Warmup Total: 0.003421487
Fastest Mean: 0.002123317511
Slowest Mean: 0.002123317511
Column Mean Speedup: 1.5460645413650949 x
Understanding the Results
- Row operations typically show higher speedups (5-8x) because they access contiguous memory
- Column operations show moderate speedups (2-5x) due to transpose overhead
- Standard deviation should be low for reliable measurements
- Minimum times often reflect best-case performance without system interference
The benchmarking code demonstrates that proper measurement is as important as the optimization itself. Without accurate timing, you can't validate your performance gains or identify bottlenecks for further improvement.
Performance Expectations
For a 1024x1024 matrix of float32 values, you can expect:
- Row operations: 4-8x speedup from SIMD + parallelization
- Column operation: 2-4x speedup despite transpose overhead
- Memory bandwidth: Often becomes the limiting factor at large matrix sizes The exact speedup depends on your CPU's SIMD capabilities, core count, and memory subsystem performance.
Key Takeaways
- SIMD shines with regular, parallel computations like matrix operations, image processing, and numerical algorithms.
- Memory layout matters immensely, contiguous access patterns enable efficient vectorization, while strided patterns often require algorithm changes.
- Mojo's
vectorizeandparallelizefunction provide elegant abstractions for SIMD and thread parallelism without sacrificing performance. - Sometimes algorithmic changes (like matrix transposition) can unlock better SIMD utilization despite additional overhead.
- Always benchmark real workloads, theoretical speedups don't always translate to practice due to memory bandwidth, cache effects, and other system limitations.
Going Further
This example demonstrates fundamental SIMD principles, but there's much more to explore:
- Memory prefetching for even better cache utilization
- Loop tiling for large matrices that don't fit in cache
- Mixed precision computations for memory bandwidth optimization
- SIMD-aware algorithms designed from the ground up for vectorization
SIMD optimization is both an art and a science. With Mojo's powerful abstractions and direct hardware access, you have the tools to achieve remarkable performance while maintaining code clarity. The key is understanding your data patterns, measuring real performance, and letting the hardware guide your optimization decisions.
Start with these matrix operations, measure the speedups on your system, and use these patterns as building blocks for more complex optimizations. The performance gains are worth the effort!