Puzzle 8: Shared Memory

Overview

Implement a kernel that adds 10 to each position of a 1D TileTensor a and stores it in 1D TileTensor output.

Shared memory is fast, on-chip storage that is visible to all threads within the same block. Unlike global memory (which all blocks can access but is slow), shared memory lives on-chip, alongside the cores running the block, so an access costs on the order of tens of cycles rather than the hundreds a global-memory access can take. Each block gets its own private shared memory region—threads in one block cannot see the shared memory of another block. Because threads can read and write to the same shared memory locations, coordination via barrier() is required to prevent one thread from reading a value before another thread has finished writing it.

Note: You have fewer threads per block than the size of a.

Key concepts

In this puzzle, you’ll learn about:

• Using TileTensor’s shared memory features with address_space
• Thread synchronization with shared memory
• Block-local data management with TileTensor

The key insight is how TileTensor simplifies shared memory management while maintaining the performance benefits of block-local storage.

Configuration

• Array size: SIZE = 8 elements
• Threads per block: TPB = 4
• Number of blocks: 2
• Shared memory: TPB elements per block

Warning: Each block can only have a constant amount of shared memory that threads in that block can read and write to. This needs to be a compile-time constant (a comptime binding), not a runtime variable. After writing to shared memory you need to call barrier to ensure that threads do not cross.

Educational Note: In this specific puzzle, the barrier() isn’t strictly necessary since each thread only accesses its own shared memory location. However, it’s included to teach proper shared memory synchronization patterns for more complex scenarios where threads need to coordinate access to shared data.

Code to complete

comptime TPB = 4
comptime SIZE = 8
comptime BLOCKS_PER_GRID = (2, 1)
comptime THREADS_PER_BLOCK = (TPB, 1)
comptime dtype = DType.float32
comptime layout = row_major[SIZE]()
comptime LayoutType = type_of(layout)


def add_10_shared_tile_tensor(
    output: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
    size_dev: Int32,
):
    var size = Int(size_dev)
    # Allocate shared memory. Unlike most kernel variables, which are private to
    # each thread, shared memory (`AddressSpace.SHARED`) is scoped per-block.
    # So, each thread in a block can use the same "scratchpad" to store data.
    var shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB]())

    var global_i = block_dim.x * block_idx.x + thread_idx.x
    var local_i = thread_idx.x

    if global_i < size:
        shared[local_i] = a[global_i]

    # Note: barrier is not strictly needed here since each thread only accesses
    # its own shared memory location. However, it's included to teach proper
    # shared memory synchronization patterns for more complex scenarios where
    # threads need to coordinate access to shared data.
    barrier()

    # FILL ME IN (roughly 2 lines)

View full file: problems/p08/p08.mojo

Running the code

To test your solution, run the following command in your terminal:

pixi run p08

pixi run -e amd p08

pixi run -e apple p08

uv run poe p08

Your output will look like this if the puzzle isn’t solved yet:

out: HostBuffer([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([11.0, 11.0, 11.0, 11.0, 11.0, 11.0, 11.0, 11.0])

Solution

def add_10_shared_tile_tensor(
    output: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
    size_dev: Int32,
):
    var size = Int(size_dev)
    # Allocate shared memory. Unlike most kernel variables, which are private to
    # each thread, shared memory (`AddressSpace.SHARED`) is scoped per-block.
    # So, each thread in a block can use the same "scratchpad" to store data.
    var shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB]())

    var global_i = block_dim.x * block_idx.x + thread_idx.x
    var local_i = thread_idx.x

    if global_i < size:
        shared[local_i] = a[global_i]

    # Note: barrier is not strictly needed here since each thread only accesses
    # its own shared memory location. However, it's included to teach proper
    # shared memory synchronization patterns for more complex scenarios where
    # threads need to coordinate access to shared data.
    barrier()

    if global_i < size:
        output[global_i] = shared[local_i] + 10