LayoutTensor 버전

개요

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

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

핵심 개념

이 퍼즐에서 배울 내용:

• 2D 배열 접근에 LayoutTensor 사용하기
• tensor[i, j]로 직접 2D 인덱싱하기
• LayoutTensor에서 경계 검사 처리하기

핵심은 LayoutTensor가 자연스러운 2D 인덱싱 인터페이스를 제공하여 내부 메모리 레이아웃을 추상화한다는 점입니다. 그러면서도 경계 검사는 여전히 필요합니다.

• 2D 접근: LayoutTensor로 자연스러운 \((i,j)\) 인덱싱
• 메모리 추상화: 수동 행 우선 계산 불필요
• 가드 조건: 두 차원 모두 경계 검사 필요
• 스레드 범위: 스레드 \((3 \times 3)\)가 텐서 원소 \((2 \times 2)\)보다 많음

완성할 코드

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


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

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

코드 실행

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

pixi run p04_layout_tensor

pixi run -e amd p04_layout_tensor

pixi run -e apple p04_layout_tensor

uv run poe p04_layout_tensor

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

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: LayoutTensor[dtype, layout, MutAnyOrigin],
    a: LayoutTensor[dtype, layout, MutAnyOrigin],
    size: UInt,
):
    row = thread_idx.y
    col = thread_idx.x
    if col < size and row < size:
        output[row, col] = a[row, col] + 10.0