Elementwise - Basic GPU Functional Operations

def elementwise_add[
    LayoutT: TensorLayout, dtype: DType, simd_width: Int, rank: Int, size: Int
](
    output: TileTensor[mut=True, dtype, LayoutT, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutT, MutAnyOrigin],
    ctx: DeviceContext,
) raises:
    @parameter
    @always_inline
    def add[
        simd_width: Int, alignment: Int = align_of[dtype]()
    ](indices: Coord) capturing -> None:
        var idx = Int(indices[0].value())
        # Convert inside GPU kernel to avoid host-captured LayoutTensor issues
        var a_lt = a.to_layout_tensor()
        var b_lt = b.to_layout_tensor()
        var out_lt = output.to_layout_tensor()
        # 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
        var a_simd = a_lt.aligned_load[width=simd_width](Index(idx))
        var b_simd = b_lt.aligned_load[width=simd_width](Index(idx))
        var ret = a_simd + b_simd
        out_lt.store[simd_width](Index(idx), ret)

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

The elementwise functional pattern in Mojo introduces 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
var idx = thread_idx.x + block_idx.x * block_dim.x
if idx < size:
    output[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
• Tail handling: Elements left over when SIMD_WIDTH doesn’t divide the shape are re-invoked one at a time, with width=1
• Memory coalescing: Threads walk the array with a grid stride, so neighboring threads touch neighboring chunks
• SIMD orchestration: Vectorization handled transparently
• GPU target selection: Works across different GPU architectures

2. Deep dive: nested function architecture

@parameter
@always_inline
def add[
    simd_width: Int, alignment: Int = align_of[dtype]()
](indices: Coord) capturing -> None:

Parameter Analysis:

• @parameter: This decorator provides compile-time specialization. The function is generated separately for each unique simd_width, 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.
• Coord: Carries the per-dimension indices for the SIMD chunk being processed; indices[0] is the linear start position for 1D operations.

3. SIMD execution model deep dive

var idx = Int(indices[0].value())                            # Linear index: 0, 4, 8, 12... (GPU-dependent spacing)
var a_lt = a.to_layout_tensor()                              # LayoutTensor views for vectorized access
var b_lt = b.to_layout_tensor()
var out_lt = output.to_layout_tensor()
var a_simd = a_lt.aligned_load[width=simd_width](Index(idx))  # Load: [a[0:4], a[4:8], a[8:12]...] (4 elements per load)
var b_simd = b_lt.aligned_load[width=simd_width](Index(idx))  # Load: [b[0:4], b[4:8], b[8:12]...] (4 elements per load)
var ret = a_simd + b_simd                                    # SIMD: 4 additions in parallel (GPU-dependent)
out_lt.store[simd_width](Index(idx), 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 Part VII
│   │   │   ├── 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
• Body invocations: 256, one per SIMD chunk
• Each invocation processes: Exactly 4 consecutive elements
• Threads: However many elementwise decides it needs to saturate the GPU - the 256 chunks are handed out across that grid, so a thread may run the body once or several times

Note: SIMD width varies with the target’s vector register width and the data type. Current NVIDIA and AMD GPU targets report a 128-bit vector width, which gives SIMD_WIDTH = 4 for float32.

4. Memory access pattern analysis

a_lt.aligned_load[width=simd_width](Index(idx))  # 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]   → bytes 0-15
Thread 1: loads a[4:8]   → bytes 16-31
Thread 2: loads a[8:12]  → bytes 32-47
...
Result: A 32-lane warp covers one contiguous 512-byte span, so the memory
        controller fetches whole cache lines instead of scattered words

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:

comptime SIMD_WIDTH = simd_width_of[dtype, target=get_gpu_target()]()

• GPU-specific optimization: SIMD width adapts to hardware - it is the target’s SIMD register width divided by the size of dtype, so a 128-bit vector register gives 4 for float32 and 8 for float16
• 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* output, float* a, float* b, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        output[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: The tail is handled for you, so a shape that isn’t a multiple of SIMD_WIDTH doesn’t need a hand-written remainder branch
• Portability: Same code works across GPU vendors/generations
• Performance: Grid sizing and the SIMD width come from the target, so the same source adapts instead of being retuned by hand
• 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.