개요

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

참고: 블록당 스레드 수가 a의 크기보다 작습니다.

핵심 개념

이 퍼즐에서 배울 내용:

• 스레드 블록 내에서 공유 메모리 사용하기
• 배리어(barrier)로 스레드 동기화하기
• 블록 로컬 데이터 저장소 관리하기

핵심은 공유 메모리가 블록 내 모든 스레드가 접근할 수 있는 빠른 로컬 저장소라는 점, 그리고 이를 사용할 때 스레드 간 조율이 필요하다는 점을 이해하는 것입니다.

구성

• 배열 크기: SIZE = 8 원소
• 블록당 스레드 수: TPB = 4
• 블록 수: 2
• 공유 메모리: 블록당 TPB개 원소

참고:

• 공유 메모리: 블록 내 스레드들이 함께 사용하는 빠른 저장소
• 스레드 동기화: barrier()를 사용한 조율
• 메모리 범위: 공유 메모리는 블록 내에서만 보임
• 접근 패턴: 로컬 인덱스 vs 전역 인덱스

주의: 각 블록이 가질 수 있는 공유 메모리 크기는 상수 로 정해져야 합니다. 이 값은 변수가 아닌 리터럴 Python 상수여야 합니다. 공유 메모리에 쓴 후에는 barrier를 호출하여 스레드들이 서로 앞서가지 않도록 해야 합니다.

학습 참고: 이 퍼즐에서는 각 스레드가 자신의 공유 메모리 위치에만 접근하므로 barrier()가 엄밀히 필요하지 않습니다. 하지만 더 복잡한 상황에서 필요한 올바른 동기화 패턴을 익히기 위해 포함되어 있습니다.

완성할 코드

comptime TPB = 4
comptime SIZE = 8
comptime BLOCKS_PER_GRID = (2, 1)
comptime THREADS_PER_BLOCK = (TPB, 1)
comptime dtype = DType.float32


fn add_10_shared(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: UInt,
):
    shared = stack_allocation[
        TPB,
        Scalar[dtype],
        address_space = AddressSpace.SHARED,
    ]()
    global_i = block_dim.x * block_idx.x + thread_idx.x
    local_i = thread_idx.x
    # Load local data into shared memory
    if global_i < size:
        shared[local_i] = a[global_i]

    # wait for all threads to complete
    # works within a thread block
    barrier()

    # FILL ME IN (roughly 2 lines)

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

코드 실행

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

pixi run p08

pixi run -e amd p08

pixi run -e apple p08

uv run poe p08

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

out: HostBuffer([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([11.0, 11.0, 11.0, 11.0, 11.0, 11.0, 11.0, 11.0])

솔루션

fn add_10_shared(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: UInt,
):
    shared = stack_allocation[
        TPB,
        Scalar[dtype],
        address_space = AddressSpace.SHARED,
    ]()
    global_i = block_dim.x * block_idx.x + thread_idx.x
    local_i = thread_idx.x
    # Load local data into shared memory
    if global_i < size:
        shared[local_i] = a[global_i]

    # Wait for all threads to complete (works within a thread block).
    # Note: barrier is not strictly needed here since each thread only accesses
    # its own shared memory location. However, it's included to teach proper
    # shared memory synchronization patterns for more complex scenarios where
    # threads need to coordinate access to shared data.
    barrier()

    # process using shared memory
    if global_i < size:
        output[global_i] = shared[local_i] + 10