Key concepts

In this puzzle, you’ll learn about:

• Basic GPU kernel structure

• Thread indexing with thread_idx.x

• Simple parallel operations

• Parallelism: Every thread runs the same kernel body concurrently

• Thread indexing: Access element at position i = thread_idx.x

• Memory access: Read from a[i] and write to output[i]

• Data independence: Each output depends only on its corresponding input

Code to complete

comptime SIZE = 4
comptime BLOCKS_PER_GRID = 1
comptime THREADS_PER_BLOCK = SIZE
comptime dtype = DType.float32


def add_10(
    output: Pointer[Scalar[dtype], MutAnyOrigin],
    a: Pointer[Scalar[dtype], MutAnyOrigin],
):
    var i = thread_idx.x
    # FILL ME IN (roughly 1 line)

View full file: problems/p01/p01.mojo

Note: The skeleton declares var i = thread_idx.x for you, but nothing uses i until you fill in the kernel. Until then the compiler warns assignment to 'i' was never used; assign to '_' instead?. That’s expected, and it clears as soon as your line reads from a[i] and writes to output[i]. Later puzzles scaffold their variables the same way.

Running the code

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

pixi run p01

pixi run -e amd p01

pixi run -e apple p01

uv run poe p01

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

out: HostBuffer([0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([10.0, 11.0, 12.0, 13.0])

Solution

def add_10(
    output: Pointer[Scalar[dtype], MutAnyOrigin],
    a: Pointer[Scalar[dtype], MutAnyOrigin],
):
    var i = thread_idx.x
    output[unsafe_offset=i] = a[unsafe_offset=i] + 10.0