⚡ Tile - Memory-Efficient Tiled Processing

alias TILE_SIZE = 32


fn tiled_elementwise_add[
    layout: Layout,
    dtype: DType,
    simd_width: Int,
    rank: Int,
    size: Int,
    tile_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 process_tiles[
        simd_width: Int, rank: Int
    ](indices: IndexList[rank]) capturing -> None:
        tile_id = indices[0]

        output_tile = output.tile[tile_size](tile_id)
        a_tile = a.tile[tile_size](tile_id)
        b_tile = b.tile[tile_size](tile_id)

        @parameter
        for i in range(tile_size):
            a_vec = a_tile.load[simd_width](i, 0)
            b_vec = b_tile.load[simd_width](i, 0)
            ret = a_vec + b_vec
            output_tile.store[simd_width](i, 0, ret)

    num_tiles = (size + tile_size - 1) // tile_size
    elementwise[process_tiles, 1, target="gpu"](num_tiles, ctx)

The tiled processing pattern demonstrates advanced memory optimization techniques for GPU programming:

1. Tiling philosophy and memory hierarchy

Tiling represents a fundamental shift in how we think about parallel processing:

Elementwise approach:

• Wide parallelism: Many threads, each doing minimal work
• Global memory pressure: Threads scattered across entire array
• Cache misses: Poor spatial locality across thread boundaries

Tiled approach:

• Deep parallelism: Fewer threads, each doing substantial work
• Localized memory access: Each thread works on contiguous data
• Cache optimization: Excellent spatial and temporal locality

2. Tile organization and indexing

tile_id = indices[0]
out_tile = out.tile[tile_size](tile_id)
a_tile = a.tile[tile_size](tile_id)
b_tile = b.tile[tile_size](tile_id)

Tile mapping visualization (TILE_SIZE=32):

Original array: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, ..., 1023]

Tile 0 (thread 0): [0, 1, 2, ..., 31]      ← Elements 0-31
Tile 1 (thread 1): [32, 33, 34, ..., 63]   ← Elements 32-63
Tile 2 (thread 2): [64, 65, 66, ..., 95]   ← Elements 64-95
...
Tile 31 (thread 31): [992, 993, ..., 1023] ← Elements 992-1023

Key insights:

• Each tile[size](id) creates a view into the original tensor
• Views are zero-copy - no data movement, just pointer arithmetic
• Tile boundaries are always aligned to tile_size boundaries

3. Sequential processing deep dive

@parameter
for i in range(tile_size):
    a_vec = a_tile.load[simd_width](i, 0)
    b_vec = b_tile.load[simd_width](i, 0)
    ret = a_vec + b_vec
    out_tile.store[simd_width](i, 0, ret)

Why sequential processing?

• Cache optimization: Consecutive memory accesses maximize cache hit rates
• Compiler optimization: @parameter loops unroll completely at compile-time
• Memory bandwidth: Sequential access aligns with memory controller design
• Reduced coordination: No need to synchronize between SIMD groups

Execution pattern within one tile (TILE_SIZE=32, SIMD_WIDTH=4):

Thread processes tile sequentially:
Step 0: Process elements [0:4] with SIMD
Step 1: Process elements [4:8] with SIMD
Step 2: Process elements [8:12] with SIMD
...
Step 7: Process elements [28:32] with SIMD
Total: 8 SIMD operations per thread (32 ÷ 4 = 8)

4. Memory access pattern analysis

Cache behavior comparison:

Elementwise pattern:

Thread 0: accesses global positions [0, 4, 8, 12, ...]    ← Stride = SIMD_WIDTH
Thread 1: accesses global positions [4, 8, 12, 16, ...]   ← Stride = SIMD_WIDTH
...
Result: Memory accesses spread across entire array

Tiled pattern:

Thread 0: accesses positions [0:32] sequentially         ← Contiguous 32-element block
Thread 1: accesses positions [32:64] sequentially       ← Next contiguous 32-element block
...
Result: Perfect spatial locality within each thread

Cache efficiency implications:

• L1 cache: Small tiles often fit better in L1 cache, reducing cache misses
• Memory bandwidth: Sequential access maximizes effective bandwidth
• TLB efficiency: Fewer translation lookbook buffer misses
• Prefetching: Hardware prefetchers work optimally with sequential patterns

5. Thread configuration strategy

elementwise[process_tiles, 1, target="gpu"](num_tiles, ctx)

Why 1 instead of SIMD_WIDTH?

• Thread count: Launch exactly num_tiles threads, not num_tiles × SIMD_WIDTH
• Work distribution: Each thread handles one complete tile
• Load balancing: More work per thread, fewer threads total
• Memory locality: Each thread’s work is spatially localized

Performance trade-offs:

• Fewer logical threads: May not fully utilize all GPU cores at low occupancy
• More work per thread: Better cache utilization and reduced coordination overhead
• Sequential access: Optimal memory bandwidth utilization within each thread
• Reduced overhead: Less thread launch and coordination overhead

Important note: “Fewer threads” refers to the logical programming model. The GPU scheduler can still achieve high hardware utilization by running multiple warps and efficiently switching between them during memory stalls.

6. Performance characteristics

When tiling helps:

• Memory-bound operations: When memory bandwidth is the bottleneck
• Cache-sensitive workloads: Operations that benefit from data reuse
• Complex operations: When compute per element is higher
• Limited parallelism: When you have fewer threads than GPU cores

When tiling hurts:

• Highly parallel workloads: When you need maximum thread utilization
• Simple operations: When memory access dominates over computation
• Irregular access patterns: When tiling doesn’t improve locality

For our simple addition example (TILE_SIZE=32):

• Thread count: 32 threads instead of 256 (8× fewer)
• Work per thread: 32 elements instead of 4 (8× more)
• Memory pattern: Sequential vs strided access
• Cache utilization: Much better spatial locality

7. Advanced tiling considerations

Tile size selection:

• Too small: Poor cache utilization, more overhead
• Too large: May not fit in cache, reduced parallelism
• Sweet spot: Usually 16-64 elements for L1 cache optimization
• Our choice: 32 elements balances cache usage with parallelism

Hardware considerations:

• Cache size: Tiles should fit in L1 cache when possible
• Memory bandwidth: Consider memory controller width
• Core count: Ensure enough tiles to utilize all cores
• SIMD width: Tile size should be multiple of SIMD width

Comparison summary:

Elementwise: High parallelism, scattered memory access
Tiled:       Moderate parallelism, localized memory access

The choice between elementwise and tiled patterns depends on your specific workload characteristics, data access patterns, and target hardware capabilities.