Puzzle 5: 브로드캐스트

개요

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

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

핵심 개념

이 퍼즐에서 배울 내용:

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

핵심은 TileTensor가 서로 다른 텐서 크기 \((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 = row_major[SIZE, SIZE]()
comptime a_layout = row_major[1, SIZE]()
comptime b_layout = row_major[SIZE, 1]()
comptime OutLayout = type_of(out_layout)
comptime ALayout = type_of(a_layout)
comptime BLayout = type_of(b_layout)


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

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

코드 실행

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

pixi run p05

pixi run -e amd p05

pixi run -e apple p05

uv run poe p05

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

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

솔루션

def broadcast_add(
    output: TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, ALayout, ImmutAnyOrigin],
    b: TileTensor[mut=False, dtype, BLayout, ImmutAnyOrigin],
    size: Int,
):
    var row = thread_idx.y
    var col = thread_idx.x
    if row < size and col < size:
        output[row, col] = a[0, col] + b[row, 0]