개요

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

참고: 스레드 수가 행렬의 위치 수보다 많습니다.

핵심 개념

이 퍼즐에서 배울 내용:

• 2D 스레드 인덱스 다루기 (thread_idx.x, thread_idx.y)
• 2D 좌표를 1D 메모리 인덱스로 변환하기
• 2차원에서 경계 검사 처리하기

핵심은 2D 스레드 좌표 \((i,j)\)를 크기 \(n \times n\)인 행 우선 행렬의 원소로 매핑하는 방법을 이해하는 것입니다. 동시에 스레드 인덱스가 범위를 벗어나지 않는지도 확인해야 합니다.

• 2D 인덱싱: 각 스레드가 고유한 \((i,j)\) 위치를 가짐
• 메모리 레이아웃: 행 우선 순서로 2D를 1D 메모리에 매핑
• 가드 조건: 두 차원 모두 경계 검사 필요
• 스레드 범위: 스레드 \((3 \times 3)\)가 행렬 원소 \((2 \times 2)\)보다 많음

완성할 코드

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


fn add_10_2d(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: UInt,
):
    row = thread_idx.y
    col = thread_idx.x
    # FILL ME IN (roughly 2 lines)

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

코드 실행

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

pixi run p04

pixi run -e amd p04

pixi run -e apple p04

uv run poe p04

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

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

솔루션

fn add_10_2d(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: UInt,
):
    row = thread_idx.y
    col = thread_idx.x
    if row < size and col < size:
        output[row * size + col] = a[row * size + col] + 10.0