🔰 Elementwise - Basic GPU Functional Operations

fn elementwise_add[
    layout: Layout, dtype: DType, simd_width: Int, rank: Int, size: Int
](
    output: LayoutTensor[mut=True, dtype, layout, MutableAnyOrigin],
    a: LayoutTensor[mut=False, dtype, layout, MutableAnyOrigin],
    b: LayoutTensor[mut=False, dtype, layout, MutableAnyOrigin],
    ctx: DeviceContext,
) raises:
    @parameter
    @always_inline
    fn add[
        simd_width: Int, rank: Int
    ](indices: IndexList[rank]) capturing -> None:
        idx = indices[0]
        # Note: This is thread-local SIMD - each thread processes its own vector of data
        # we'll later better see this hierarchy in Mojo:
        # SIMD within threads, warp across threads, block across warps
        a_simd = a.load[simd_width](idx, 0)
        b_simd = b.load[simd_width](idx, 0)
        ret = a_simd + b_simd
        # print(
        #     "idx:", idx, ", a_simd:", a_simd, ", b_simd:", b_simd, " sum:", ret
        # )
        output.store[simd_width](idx, 0, ret)

    elementwise[add, SIMD_WIDTH, target="gpu"](a.size(), ctx)

The elementwise functional pattern in Mojo demonstrates several fundamental concepts for modern GPU programming:

1. Functional abstraction philosophy

The elementwise function represents a paradigm shift from traditional GPU programming:

Traditional CUDA/HIP approach:

// Manual thread management
idx = thread_idx.x + block_idx.x * block_dim.x
if idx < size:
    out[idx] = a[idx] + b[idx];  // Scalar operation

Mojo functional approach:

// Automatic management + SIMD vectorization
elementwise[add_function, simd_width, target="gpu"](size, ctx)

What elementwise abstracts away:

• Thread grid configuration: No need to calculate block/grid dimensions
• Bounds checking: Automatic handling of array boundaries
• Memory coalescing: Optimal memory access patterns built-in
• SIMD orchestration: Vectorization handled transparently
• GPU target selection: Works across different GPU architectures

2. Deep dive: nested function architecture

@parameter
@always_inline
fn add[simd_width: Int, rank: Int](indices: IndexList[rank]) capturing -> None:

Parameter Analysis:

• @parameter: This decorator enables compile-time specialization. The function is generated separately for each unique simd_width and rank, allowing aggressive optimization.
• @always_inline: Critical for GPU performance - eliminates function call overhead by embedding the code directly into the kernel.
• capturing: Enables lexical scoping - the inner function can access variables from the outer scope without explicit parameter passing.
• IndexList[rank]: Provides dimension-agnostic indexing - the same pattern works for 1D vectors, 2D matrices, 3D tensors, etc.

3. SIMD execution model deep dive

idx = indices[0]                          // Linear index: 0, 4, 8, 12... (GPU-dependent spacing)
a_simd = a.load[simd_width](idx, 0)      // Load: [a[0:4], a[4:8], a[8:12]...] (4 elements per load)
b_simd = b.load[simd_width](idx, 0)      // Load: [b[0:4], b[4:8], b[8:12]...] (4 elements per load)
ret = a_simd + b_simd                    // SIMD: 4 additions in parallel (GPU-dependent)
out.store[simd_width](idx, 0, ret)       // Store: 4 results simultaneously (GPU-dependent)

Execution Hierarchy Visualization:

GPU Architecture:
├── Grid (entire problem)
│   ├── Block 1 (multiple warps)
│   │   ├── Warp 1 (32 threads) --> We'll learn about Warp in the next Part VI
│   │   │   ├── Thread 1 → SIMD[4 elements]  ← Our focus (GPU-dependent width)
│   │   │   ├── Thread 2 → SIMD[4 elements]
│   │   │   └── ...
│   │   └── Warp 2 (32 threads)
│   └── Block 2 (multiple warps)

For a 1024-element vector with SIMD_WIDTH=4 (example GPU):

• Total SIMD operations needed: 1024 ÷ 4 = 256
• GPU launches: 256 threads (1024 ÷ 4)
• Each thread processes: Exactly 4 consecutive elements
• Memory bandwidth: SIMD_WIDTH× improvement over scalar operations

Note: SIMD width varies by GPU architecture (e.g., 4 for some GPUs, 8 for RTX 4090, 16 for A100).

4. Memory access pattern analysis

a.load[simd_width](idx, 0)  // Coalesced memory access

Memory Coalescing Benefits:

• Sequential access: Threads access consecutive memory locations
• Cache optimization: Maximizes L1/L2 cache hit rates
• Bandwidth utilization: Achieves near-theoretical memory bandwidth
• Hardware efficiency: GPU memory controllers optimized for this pattern

Example for SIMD_WIDTH=4 (GPU-dependent):

Thread 0: loads a[0:4]   → Memory bank 0-3
Thread 1: loads a[4:8]   → Memory bank 4-7
Thread 2: loads a[8:12]  → Memory bank 8-11
...
Result: Optimal memory controller utilization

5. Performance characteristics & optimization

Computational Intensity Analysis (for SIMD_WIDTH=4):

• Arithmetic operations: 1 SIMD addition per 4 elements
• Memory operations: 2 SIMD loads + 1 SIMD store per 4 elements
• Arithmetic intensity: 1 add ÷ 3 memory ops = 0.33 (memory-bound)

Why This Is Memory-Bound:

Memory bandwidth >>> Compute capability for simple operations

Optimization Implications:

• Focus on memory access patterns rather than arithmetic optimization
• SIMD vectorization provides the primary performance benefit
• Memory coalescing is critical for performance
• Cache locality matters more than computational complexity

6. Scaling and adaptability

Automatic Hardware Adaptation:

alias SIMD_WIDTH = simdwidthof[dtype, target = _get_gpu_target()]()

• GPU-specific optimization: SIMD width adapts to hardware (e.g., 4 for some cards, 8 for RTX 4090, 16 for A100)
• Data type awareness: Different SIMD widths for float32 vs float16
• Compile-time optimization: Zero runtime overhead for hardware detection

Scalability Properties:

• Thread count: Automatically scales with problem size
• Memory usage: Linear scaling with input size
• Performance: Near-linear speedup until memory bandwidth saturation

7. Advanced insights: why this pattern matters

Foundation for Complex Operations: This elementwise pattern is the building block for:

• Reduction operations: Sum, max, min across large arrays
• Broadcast operations: Scalar-to-vector operations
• Complex transformations: Activation functions, normalization
• Multi-dimensional operations: Matrix operations, convolutions

Compared to Traditional Approaches:

// Traditional: Error-prone, verbose, hardware-specific
__global__ void add_kernel(float* out, float* a, float* b, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];  // No vectorization
    }
}

// Mojo: Safe, concise, automatically vectorized
elementwise[add, SIMD_WIDTH, target="gpu"](size, ctx)

Benefits of Functional Approach:

• Safety: Automatic bounds checking prevents buffer overruns
• Portability: Same code works across GPU vendors/generations
• Performance: Compiler optimizations often exceed hand-tuned code
• Maintainability: Clean abstractions reduce debugging complexity
• Composability: Easy to combine with other functional operations

This pattern represents the future of GPU programming - high-level abstractions that don’t sacrifice performance, making GPU computing accessible while maintaining optimal efficiency.