Puzzle 6: 블록

개요

벡터 a의 각 위치에 10을 더해 output에 저장하는 커널을 구현해 보세요.

스레드 블록(thread block) (또는 줄여서 블록)은 하나의 GPU 멀티프로세서에서 함께 실행되는 스레드 묶음입니다. 같은 블록 안의 모든 스레드는 공유 메모리를 함께 사용하고 서로 동기화할 수 있습니다. 데이터가 한 블록이 처리할 수 있는 범위보다 크면 GPU는 여러 블록을 스케줄링하고, 각 블록은 자기 몫의 데이터를 독립적으로 처리합니다. 스레드의 전역 위치는 블록 내 위치(thread_idx.x)와 소속 블록(block_idx.x)을 합쳐 계산합니다: global_i = block_dim.x * block_idx.x + thread_idx.x.

참고: 블록당 스레드 수가 a의 크기보다 작습니다.

핵심 개념

이 퍼즐에서 다루는 내용:

스레드 블록 크기보다 큰 데이터 처리
여러 블록의 스레드 조율
전역 스레드 위치 계산

여기서 핵심은 여러 스레드 블록이 협력하여 단일 블록 용량보다 큰 데이터를 처리하면서도, 요소와 스레드 간 올바른 매핑을 유지하는 원리를 이해하는 것입니다.

완성할 코드

comptime SIZE = 9
comptime BLOCKS_PER_GRID = (3, 1)
comptime THREADS_PER_BLOCK = (4, 1)
comptime dtype = DType.float32


def add_10_blocks(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: Int,
):
    var i = block_dim.x * block_idx.x + thread_idx.x
    # FILL ME IN (roughly 2 lines)

전체 코드 보기: problems/p06/p06.mojo

참고: 이 퍼즐의 TileTensor 버전은 거의 동일하므로 독자에게 맡깁니다.

팁

전역 인덱스 계산: i = block_dim.x * block_idx.x + thread_idx.x
가드 추가: if i < size
가드 내부: output[i] = a[i] + 10.0

코드 실행

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

pixi run p06

pixi run -e amd p06

pixi run -e apple p06

uv run poe p06

퍼즐을 아직 풀지 않았다면 출력이 다음과 같이 나타납니다:

out: HostBuffer([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0])

솔루션

def add_10_blocks(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: Int,
):
    var i = block_dim.x * block_idx.x + thread_idx.x
    if i < size:
        output[i] = a[i] + 10.0

이 솔루션은 블록 기반 GPU 처리의 핵심 개념을 다룹니다:

전역 스레드 인덱싱
- 블록 인덱스와 스레드 인덱스를 결합: block_dim.x * block_idx.x + thread_idx.x
- 각 스레드를 고유한 전역 위치에 매핑
- 블록당 3개 스레드 예시:
```
Block 0: [0 1 2]
Block 1: [3 4 5]
Block 2: [6 7 8]
```
블록 조율
- 각 블록은 연속된 데이터 청크를 처리
- 블록 크기(3) < 데이터 크기(9)이므로 여러 블록 필요
- 블록 간 자동 작업 분배:
```
Data:    [0 1 2 3 4 5 6 7 8]
Block 0: [0 1 2]
Block 1:       [3 4 5]
Block 2:             [6 7 8]
```
경계 검사
- 가드 조건 i < size로 경계 케이스 처리
- 데이터 크기가 블록 크기로 나누어 떨어지지 않을 때 범위를 벗어난 접근 방지
- 데이터 끝부분의 불완전한 블록 처리에 필수
메모리 접근 패턴
- 병합(coalesced) 메모리 접근: 블록 내 스레드들이 연속된 메모리에 접근
- 각 스레드가 하나의 요소 처리: output[i] = a[i] + 10.0
- 블록 수준 병렬성으로 메모리 대역폭을 효율적으로 활용

이 패턴은 단일 스레드 블록 크기를 초과하는 대규모 데이터셋 처리의 기초가 됩니다.