LayoutTensor 버전

개요

1D LayoutTensor a와 b를 브로드캐스트로 더해 2D LayoutTensor output에 저장하는 커널을 구현해 보세요.

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

핵심 개념

이 퍼즐에서 배울 내용:

• 브로드캐스트 연산에 LayoutTensor 사용하기
• 서로 다른 텐서 크기 다루기
• LayoutTensor로 2D 인덱싱 처리하기

핵심은 LayoutTensor가 서로 다른 텐서 크기 \((1, n)\)와 \((n, 1)\)을 \((n,n)\)으로 자연스럽게 브로드캐스트할 수 있다는 점입니다. 그러면서도 경계 검사는 여전히 필요합니다.

• 텐서 크기: 입력 벡터의 크기는 \((1, n)\)과 \((n, 1)\)
• 브로드캐스트: 두 차원을 결합해 \((n,n)\) 출력 생성
• 가드 조건: 출력 크기에 대한 경계 검사는 여전히 필요
• 스레드 범위: 텐서 원소 \((2 \times 2)\)보다 스레드 \((3 \times 3)\)가 많음

완성할 코드

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


fn broadcast_add[
    out_layout: Layout,
    a_layout: Layout,
    b_layout: Layout,
](
    output: LayoutTensor[dtype, out_layout, MutAnyOrigin],
    a: LayoutTensor[dtype, a_layout, ImmutAnyOrigin],
    b: LayoutTensor[dtype, b_layout, ImmutAnyOrigin],
    size: UInt,
):
    row = thread_idx.y
    col = thread_idx.x
    # FILL ME IN (roughly 2 lines)

전체 코드 보기: problems/p05/p05_layout_tensor.mojo

코드 실행

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

pixi run p05_layout_tensor

pixi run -e amd p05_layout_tensor

pixi run -e apple p05_layout_tensor

uv run poe p05_layout_tensor

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

out: HostBuffer([0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([1.0, 2.0, 11.0, 12.0])

솔루션

fn broadcast_add[
    out_layout: Layout,
    a_layout: Layout,
    b_layout: Layout,
](
    output: LayoutTensor[dtype, out_layout, MutAnyOrigin],
    a: LayoutTensor[dtype, a_layout, ImmutAnyOrigin],
    b: LayoutTensor[dtype, b_layout, ImmutAnyOrigin],
    size: UInt,
):
    row = thread_idx.y
    col = thread_idx.x
    if row < size and col < size:
        output[row, col] = a[0, col] + b[row, 0]