전역 메모리를 사용한 기본 버전

개요

정방 행렬 \(A\)와 \(B\)를 곱하여 결과를 \(\text{output}\)에 저장하는 커널을 구현하세요. 각 스레드가 출력 행렬의 원소 하나를 계산하는 가장 기본적인 구현입니다.

핵심 개념

이 퍼즐에서 다루는 내용:

• 행렬 연산을 위한 2D 스레드 구성
• 전역 메모리 접근 패턴
• 행 우선(row-major) 레이아웃에서의 행렬 인덱싱
• 스레드와 출력 원소 간 매핑

핵심은 2D 스레드 인덱스를 행렬 원소에 매핑하고, 내적을 병렬로 계산하는 방법을 이해하는 것입니다.

구성

• 행렬 크기: \(\text{SIZE} \times \text{SIZE} = 2 \times 2\)
• 블록당 스레드 수: \(\text{TPB} \times \text{TPB} = 3 \times 3\)
• 그리드 차원: \(1 \times 1\)

레이아웃 구성:

• 입력 A: Layout.row_major(SIZE, SIZE)
• 입력 B: Layout.row_major(SIZE, SIZE)
• 출력: Layout.row_major(SIZE, SIZE)

완성할 코드

from gpu import thread_idx, block_idx, block_dim, barrier
from gpu.memory import AddressSpace
from layout import Layout, LayoutTensor


comptime TPB = 3
comptime SIZE = 2
comptime BLOCKS_PER_GRID = (1, 1)
comptime THREADS_PER_BLOCK = (TPB, TPB)
comptime dtype = DType.float32
comptime layout = Layout.row_major(SIZE, SIZE)


fn naive_matmul[
    layout: Layout, size: UInt
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    a: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    b: LayoutTensor[dtype, layout, ImmutAnyOrigin],
):
    row = block_dim.y * block_idx.y + thread_idx.y
    col = block_dim.x * block_idx.x + thread_idx.x
    # FILL ME IN (roughly 6 lines)

전체 파일 보기: problems/p16/p16.mojo

코드 실행

솔루션을 테스트하려면 터미널에서 다음 명령어를 실행하세요:

pixi run p16 --naive

pixi run -e amd p16 --naive

pixi run -e apple p16 --naive

uv run poe p16 --naive

퍼즐을 아직 풀지 않았다면 출력은 다음과 같습니다:

out: HostBuffer([0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([4.0, 6.0, 12.0, 22.0])

솔루션

fn naive_matmul[
    layout: Layout, size: UInt
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    a: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    b: LayoutTensor[dtype, layout, ImmutAnyOrigin],
):
    row = block_dim.y * block_idx.y + thread_idx.y
    col = block_dim.x * block_idx.x + thread_idx.x

    if row < size and col < size:
        var acc: output.element_type = 0

        @parameter
        for k in range(size):
            acc += a[row, k] * b[k, col]

        output[row, col] = acc

Matrix A:          Matrix B:                   Output C:
[a[0,0] a[0,1]]    [b[0,0] b[0,1]]             [c[0,0] c[0,1]]
[a[1,0] a[1,1]]    [b[1,0] b[1,1]]             [c[1,0] c[1,1]]