Puzzle 2: Zip

개요

벡터 a와 벡터 b의 각 위치를 더해 output에 저장하는 커널을 구현해 보세요.

참고: 각 위치마다 스레드 1개가 배정됩니다.

핵심 개념

이 퍼즐에서 배우는 내용:

• 여러 입력 배열의 병렬 처리
• 여러 입력에 대한 요소별 연산
• 배열 간 스레드-데이터 매핑
• 여러 배열의 메모리 접근 패턴

각 스레드 \(i\)에 대해: \[\Large output[i] = a[i] + b[i]\]

메모리 접근 패턴

Thread 0:  a[0] + b[0] → output[0]
Thread 1:  a[1] + b[1] → output[1]
Thread 2:  a[2] + b[2] → output[2]
...

💡 참고: 이제 커널에서 세 개의 배열(a, b, output)을 다루고 있습니다. 연산이 복잡해질수록 여러 배열에 대한 접근을 관리하기가 점점 어려워집니다.

완성할 코드

comptime SIZE = 4
comptime BLOCKS_PER_GRID = 1
comptime THREADS_PER_BLOCK = SIZE
comptime dtype = DType.float32


fn add(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    b: UnsafePointer[Scalar[dtype], MutAnyOrigin],
):
    i = thread_idx.x
    # FILL ME IN (roughly 1 line)

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

코드 실행

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

pixi run p02

pixi run -e amd p02

pixi run -e apple p02

uv run poe p02

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

out: HostBuffer([0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([0.0, 2.0, 4.0, 6.0])

솔루션

fn add(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    b: UnsafePointer[Scalar[dtype], MutAnyOrigin],
):
    i = thread_idx.x
    output[i] = a[i] + b[i]

앞으로 다룰 내용

직접 인덱싱은 간단한 요소별 연산에서 잘 작동하지만, 다음 상황을 생각해 보세요:

• 배열의 레이아웃이 서로 다르다면?
• 한 배열을 다른 배열에 브로드캐스트해야 한다면?
• 여러 배열에서 병합(coalesced) 접근을 어떻게 보장할 수 있을까?

이러한 질문들은 Puzzle 4의 LayoutTensor 알아보기에서 다룹니다.