Analyze SIMD Usage Opportunities

Identify where SIMD (Single Instruction Multiple Data) can improve performance.

When to Use

• Performance-critical tensor operations
• Element-wise operations on large arrays
• Vectorization of loops processing multiple elements
• Optimizing matrix/vector operations
• Finding performance bottlenecks in ML code

Quick Reference

# Find loops processing arrays/tensors
grep -n "for.*in.*range\|@unroll\|@vectorize" *.mojo

# Find element-wise operations
grep -n "\.load\|\.store\|\.broadcast" *.mojo

# Check for SIMD parameters
grep -n "simd_width\|nelems\|\[.*:\]" *.mojo

# Identify candidates
grep -n "for i in range.*:" -A 10 *.mojo | grep -E "array\[i\]|tensor\[i\]"

SIMD Optimization Opportunities

Vectorizable Patterns:

• ✅ Element-wise addition: a[i] + b[i] for all i
• ✅ Scalar multiplication: a[i] * scalar for all i
• ✅ Unary operations: sin(a[i]), exp(a[i]) for all i
• ✅ Reduction operations: sum, max, min over array
• ❌ Dependent iterations: a[i] = a[i-1] + value (sequential)
• ❌ Conditional branches: if a[i] > threshold: (hard to vectorize)
• ❌ Function calls: unpredictable latency (avoid in tight loops)

SIMD Width Selection:

• @parameter fn[simd_width: Int] - Generic SIMD width
• simd_width=4 - Typically good for float32
• simd_width=8 - Optimal for many operations
• simd_width=16+ - For int32 or specialized ops
• Match hardware capabilities (AVX2=4-8, AVX512=8-16)

Vectorization Patterns:

• ✅ @vectorize decorator for simple loops
• ✅ @unroll for small loops (2-4 iterations)
• ✅ Manual SIMD with .load[] and .store[]
• ✅ Tensor operations with SIMD dimensions

Analysis Workflow

1. Profile code: Identify bottlenecks using time/memory metrics
2. Find loops: Locate loops processing large amounts of data
3. Check vectorizability: Verify no loop-carried dependencies
4. Estimate speedup: SIMD could provide 4-16x improvement
5. Implement SIMD: Use @vectorize, @unroll, or manual SIMD
6. Measure performance: Verify improvement with benchmarks
7. Document changes: Note what was optimized and why

Output Format

Report SIMD analysis with:

1. Hotspots - Functions/loops using most CPU time
2. Vectorization Potential - Operations that could use SIMD
3. Estimated Speedup - Expected performance improvement
4. Implementation Priority - High/medium/low impact
5. Technical Approach - How to implement SIMD
6. Risks - Potential issues with vectorization
7. Recommendations - Which optimizations to pursue first

Optimization Examples

Example 1: Element-wise Addition

# Before: scalar loop
fn add_scalar(a: Tensor, b: Tensor) -> Tensor:
    var result = Tensor(a.shape)
    for i in range(a.num_elements()):
        result._data[i] = a._data[i] + b._data[i]
    return result

# After: vectorized
@vectorize
fn add_simd[simd_width: Int](i: Int):
    result._data.store[simd_width](i,
        a._data.load[simd_width](i) + b._data.load[simd_width](i))

def add_vectorized(a: Tensor, b: Tensor) -> Tensor:
    var result = Tensor(a.shape)
    # 4x-8x speedup typical
    return result

Example 2: Reduction (Sum)

# Before: scalar loop
fn sum_scalar(tensor: Tensor) -> Float32:
    var total: Float32 = 0
    for i in range(tensor.num_elements()):
        total += tensor._data[i]
    return total

# After: SIMD reduction
fn sum_simd[simd_width: Int](tensor: Tensor) -> Float32:
    # Process simd_width elements at a time
    # Then reduce results - can be much faster
    return total

Error Handling

SIMD Decision Tree

• Does loop process large arrays? → YES → Check vectorizability
• Loop-carried dependencies? → YES → Can't vectorize, optimize differently
• Simple operations on many elements? → YES → Use @vectorize or @unroll
• Critical path (hot loop)? → YES → Worth optimizing
• Implement → Measure → Iterate

References

• See mojo-simd-optimize for implementation guidance
• See CLAUDE.md for SIMD code patterns
• See performance section in module documentation