개요

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

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

핵심 개념

이 퍼즐에서 배울 내용:

• 1D 벡터를 각각 다른 차원 방향으로 브로드캐스트하기
• 2D 스레드 인덱스로 브로드캐스트 연산 수행하기
• 브로드캐스트 패턴에서 경계 조건 처리하기

핵심은 두 1D 벡터의 원소들을 브로드캐스트로 2D 출력 행렬에 매핑하는 방법을 이해하고, 스레드 경계를 올바르게 처리하는 것입니다.

• 브로드캐스트: a의 각 원소가 b의 각 원소와 결합
• 스레드 매핑: \(2 \times 2\) 출력에 \((3 \times 3)\) 스레드 그리드 사용
• 벡터 접근: a와 b는 서로 다른 접근 패턴 사용
• 경계 검사: 행렬 범위를 벗어나는 스레드를 가드로 처리

완성할 코드

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


fn broadcast_add(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    b: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: UInt,
):
    row = thread_idx.y
    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])

솔루션

fn broadcast_add(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    b: 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[col] + b[row]