Tile - Memory-Efficient Tiled Processing

comptime TILE_SIZE = 32


def tiled_elementwise_add[
    LayoutT: TensorLayout,
    dtype: DType,
    simd_width: Int,
    rank: Int,
    size: Int,
    tile_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 process_tiles[
        simd_width: Int, alignment: Int = align_of[dtype]()
    ](indices: Coord) capturing -> None:
        var tile_id = Int(indices[0].value())

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

        comptime for i in range(tile_size):
            var a_vec = a_tile.aligned_load[width=simd_width](Index(i))
            var b_vec = b_tile.aligned_load[width=simd_width](Index(i))
            var ret = a_vec + b_vec
            output_tile.store[simd_width](Index(i), ret)

    var 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
• Small per-thread footprint: Each thread touches only SIMD_WIDTH elements
• Coalesced access: Consecutive threads cover one contiguous run of memory

Tiled approach:

• Deep parallelism: Fewer threads, each doing substantial work
• Localized memory access: Each thread works on contiguous data
• Per-thread locality: Spatial locality within a thread, traded against warp-level coalescing (see below)

2. Tile organization and indexing

var tile_id = Int(indices[0].value())
var output_tile = output.tile[tile_size](tile_id).to_layout_tensor()
var a_tile = a.tile[tile_size](tile_id).to_layout_tensor()
var b_tile = b.tile[tile_size](tile_id).to_layout_tensor()

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

comptime for i in range(tile_size):
    var a_vec = a_tile.aligned_load[width=simd_width](Index(i))
    var b_vec = b_tile.aligned_load[width=simd_width](Index(i))
    var ret = a_vec + b_vec
    output_tile.store[simd_width](Index(i), ret)

Why sequential processing?

• Cache optimization: Consecutive memory accesses maximize cache hit rates
• Compiler optimization: comptime for loops unroll completely at compile-time
• Predictable addresses: Each thread’s addresses advance by one element per step, which the compiler folds into the unrolled loop
• Reduced coordination: One thread owns a whole tile, so there is nothing to synchronize

Execution pattern within one tile (TILE_SIZE=32):

Thread processes tile sequentially, one element per iteration:
Step 0:  Load/store element [0]
Step 1:  Load/store element [1]
Step 2:  Load/store element [2]
...
Step 31: Load/store element [31]
Total: 32 scalar operations per thread (comptime for i in range(tile_size))

The width here is 1, not SIMD_WIDTH. The inner process_tiles declares its own simd_width parameter, which shadows the enclosing function’s, and elementwise[process_tiles, 1, target="gpu"](num_tiles, ctx) instantiates it with 1. So aligned_load[width=simd_width] loads a single element and the loop walks the tile one element at a time. The two vectorized kernels later in this puzzle avoid the shadowing by naming their inner parameter num_threads_per_tile, which is why they really do load SIMD_WIDTH elements at a time.

4. Memory access pattern analysis

Cache behavior comparison:

Elementwise pattern:

Thread 0: accesses positions [0:4]                        ← One SIMD chunk
Thread 1: accesses positions [4:8]                        ← Next SIMD chunk
...
Result: Each thread touches one small chunk, and the chunks spread across
        the 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
• TLB efficiency: Fewer translation lookaside buffer misses
• Coalescing caveat: GPUs have no hardware prefetcher, and coalescing is a per-warp property, not a per-thread one. Because each thread walks its own tile, on any single iteration the lanes of a warp sit tile_size elements apart - this pattern buys per-thread sequentiality at the cost of warp-level coalescing

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: Each thread’s own stream of addresses is contiguous, though the warp’s are not
• 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.