Puzzle 7: 2D 블록

개요

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

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

핵심 개념

이 퍼즐에서 배울 내용:

• 여러 블록과 함께 TileTensor 사용하기
• 2D 블록 구성으로 큰 행렬 처리하기
• 블록 인덱싱과 TileTensor 접근 결합하기

핵심은 TileTensor가 2D 인덱싱을 단순화해 주지만, 큰 행렬에서는 여전히 블록 간 조율이 필요하다는 점입니다.

🔑 2D 스레드 인덱싱 방식

Puzzle 4: 2D Map의 블록 기반 인덱싱을 2D로 확장합니다:

전역 위치 계산:
row = block_dim.y * block_idx.y + thread_idx.y
col = block_dim.x * block_idx.x + thread_idx.x

예를 들어, 4×4 그리드에서 2×2 블록을 사용하면:

Block (0,0):   Block (1,0):
[0,0  0,1]     [0,2  0,3]
[1,0  1,1]     [1,2  1,3]

Block (0,1):   Block (1,1):
[2,0  2,1]     [2,2  2,3]
[3,0  3,1]     [3,2  3,3]

각 위치는 해당 스레드의 전역 인덱스 (row, col)를 나타냅니다. 블록 차원과 인덱스가 함께 작동하여 다음을 보장합니다:

• 2D 공간 전체를 빈틈없이 처리
• 블록 간 겹침 없음
• 효율적인 메모리 접근 패턴

구성

• 행렬 크기: \(5 \times 5\) 원소
• 레이아웃 처리: TileTensor가 행 우선 구성 관리
• 블록 조율: 여러 블록으로 전체 행렬 커버
• 2D 인덱싱: 경계 검사와 함께 자연스러운 \((i,j)\) 접근
• 총 스레드 수: \(25\)개 원소에 대해 \(36\)개
• 스레드 매핑: 각 스레드가 행렬 원소 하나씩 처리

완성할 코드

comptime SIZE = 5
comptime BLOCKS_PER_GRID = (2, 2)
comptime THREADS_PER_BLOCK = (3, 3)
comptime dtype = DType.float32
comptime out_layout = row_major[SIZE, SIZE]()
comptime a_layout = row_major[SIZE, SIZE]()
comptime OutLayout = type_of(out_layout)
comptime ALayout = type_of(a_layout)


def add_10_blocks_2d(
    output: TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, ALayout, ImmutAnyOrigin],
    size: Int,
):
    var row = block_dim.y * block_idx.y + thread_idx.y
    var col = block_dim.x * block_idx.x + thread_idx.x
    # FILL ME IN (roughly 2 lines)

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

코드 실행

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

pixi run p07

pixi run -e amd p07

pixi run -e apple p07

uv run poe p07

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

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

솔루션

def add_10_blocks_2d(
    output: TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, ALayout, ImmutAnyOrigin],
    size: Int,
):
    var row = block_dim.y * block_idx.y + thread_idx.y
    var col = block_dim.x * block_idx.x + thread_idx.x
    if row < size and col < size:
        output[row, col] = a[row, col] + 10.0