더블 버퍼링 스텐실 연산

🔬 세밀한 동기화: mbarrier vs barrier()

이 퍼즐은 이전 퍼즐에서 사용한 기본 barrier() 함수보다 훨씬 강력한 제어를 제공하는 명시적 메모리 배리어 API를 소개합니다.

기본 barrier()의 한계:

• 일회성 사용: 상태 추적 없이 단일 동기화 지점만 제공
• 블록 전체 전용: 블록의 모든 스레드가 동시에 참여해야 함
• 재사용 불가: 매 barrier() 호출이 새로운 동기화 이벤트를 생성
• 세밀도 부족: 메모리 순서와 타이밍에 대한 제한적 제어
• 정적 조정: 스레드 참여 패턴의 변화에 적응 불가

고급 mbarrier API의 기능:

• 정밀한 제어: mbarrier_init()로 특정 스레드 수를 지정하여 재사용 가능한 배리어 객체를 설정
• 상태 추적: mbarrier_arrive()로 개별 스레드 완료를 알리고 도착 횟수를 유지
• 유연한 대기: mbarrier_test_wait()로 특정 완료 상태를 기다릴 수 있음
• 재사용 가능한 객체: 동일한 배리어를 여러 반복에 걸쳐 재초기화하고 재사용 가능
• 다중 배리어: 서로 다른 동기화 지점(초기화, 반복, 마무리)에 서로 다른 배리어 객체 사용
• 하드웨어 최적화: GPU 하드웨어 동기화 기본 요소에 직접 매핑하여 더 나은 성능
• 메모리 의미론: 메모리 가시성과 순서 보장에 대한 명시적 제어

반복 알고리즘에서 왜 중요한가: 더블 버퍼링 패턴에서는 버퍼 교체 단계 간의 정밀한 조정이 필요합니다. 기본 barrier()로는 다음에 필요한 세밀한 제어를 제공할 수 없습니다:

• 버퍼 역할 교대: buffer_A에 대한 모든 쓰기가 완료된 후에야 buffer_A에서 읽기 시작되도록 보장
• 반복 경계: 단일 커널 내에서 여러 동기화 지점 조율
• 상태 관리: 어떤 스레드가 어떤 처리 단계를 완료했는지 추적
• 성능 최적화: 재사용 가능한 배리어 객체를 통해 동기화 오버헤드 최소화

이 퍼즐은 반복법, 시뮬레이션 프레임워크, 고성능 이미지 처리 파이프라인 등 실제 GPU 컴퓨팅 애플리케이션에서 사용되는 동기화 패턴을 보여줍니다.

개요

더블 버퍼링 공유 메모리를 사용하여 반복 스텐실 연산을 수행하는 커널을 구현합니다. 반복 간 안전한 버퍼 교체를 보장하기 위해 명시적 메모리 배리어로 조정합니다. 스텐실 연산은 배열의 각 요소 값을 이웃 요소들의 고정된 패턴을 기반으로 계산하는 연산 패턴입니다.

참고: 버퍼 역할이 교대합니다: buffer_A와 buffer_B가 매 반복마다 읽기와 쓰기 연산을 교대하며, mbarrier 동기화가 버퍼 교체 전에 모든 스레드의 쓰기 완료를 보장합니다.

알고리즘 아키텍처: 이 퍼즐은 두 개의 공유 메모리 버퍼가 여러 반복에 걸쳐 읽기와 쓰기 대상의 역할을 교대하는 더블 버퍼링 패턴을 구현합니다. 데이터를 한 번만 처리하는 단순한 스텐실 연산과 달리, 이 접근 방식은 버퍼 전환 중 경쟁 상태를 방지하기 위한 세심한 메모리 배리어 조정과 함께 반복적 개선을 수행합니다.

파이프라인 개념: 알고리즘은 반복적 스텐실 개선을 통해 데이터를 처리합니다. 각 반복은 하나의 버퍼에서 읽고 다른 버퍼에 쓰며, 버퍼들은 매 반복마다 역할을 교대하여 데이터 손상 없이 연속 처리를 가능하게 하는 핑퐁 패턴을 만듭니다.

데이터 의존성과 동기화: 각 반복은 이전 반복의 완성된 결과에 의존합니다:

• 반복 N → 반복 N+1: 현재 반복이 다음 반복이 소비하는 개선된 데이터를 생성
• 버퍼 조정: 읽기와 쓰기 버퍼가 매 반복마다 역할을 교환
• 메모리 배리어가 경쟁 상태를 방지: 새로 기록된 버퍼에서 읽기를 시작하기 전에 모든 쓰기가 완료되도록 보장

구체적으로, 더블 버퍼링 스텐실은 세 가지 수학 연산으로 구성된 반복적 스무딩 알고리즘을 구현합니다:

반복 패턴 - 버퍼 교대:

\[\text{Iteration } i: \begin{cases} \text{Read from buffer_A, Write to buffer_B} & \text{if } i \bmod 2 = 0 \\ \text{Read from buffer_B, Write to buffer_A} & \text{if } i \bmod 2 = 1 \end{cases}\]

스텐실 연산 - 3점 평균:

\[S^{(i+1)}[j] = \frac{1}{N_j} \sum_{k=-1}^{1} S^{(i)}[j+k] \quad \text{where } j+k \in [0, 255]\]

여기서 \(S^{(i)}[j]\)는 반복 \(i\) 이후 위치 \(j\)에서의 스텐실 값이고, \(N_j\)는 유효한 이웃 수입니다.

메모리 배리어 조정:

\[\text{mbarrier_arrive}() \Rightarrow \text{mbarrier_test_wait}() \Rightarrow \text{buffer swap} \Rightarrow \text{next iteration}\]

최종 출력 선택:

\[\text{Output}[j] = \begin{cases} \text{buffer_A}[j] & \text{if STENCIL_ITERATIONS } \bmod 2 = 0 \\ \text{buffer_B}[j] & \text{if STENCIL_ITERATIONS } \bmod 2 = 1 \end{cases}\]

핵심 개념

이 퍼즐에서는 다음을 배웁니다:

• 반복 알고리즘을 위한 더블 버퍼링 패턴 구현
• mbarrier API를 사용한 명시적 메모리 배리어 조정
• 반복에 걸쳐 교대하는 읽기/쓰기 버퍼 역할 관리

핵심 통찰은 읽기와 쓰기 연산 사이의 경쟁 상태가 적절히 동기화되지 않으면 데이터를 손상시킬 수 있는 반복 알고리즘에서 버퍼 교체를 안전하게 조율하는 방법을 이해하는 것입니다.

왜 중요한가: 대부분의 GPU 튜토리얼은 단순한 단일 패스 알고리즘을 보여주지만, 실제 애플리케이션에서는 데이터에 대한 다중 패스를 수행하는 반복적 개선이 필요한 경우가 많습니다. 더블 버퍼링은 각 반복이 이전 반복의 완성된 결과에 의존하는 반복법, 이미지 처리 필터, 시뮬레이션 업데이트 같은 알고리즘에 필수적입니다.

이전 퍼즐과 현재의 동기화 비교:

• 이전 퍼즐 (P8, P12, P15): 단일 패스 알고리즘을 위한 단순 barrier() 호출
• 이 퍼즐: 버퍼 교체 타이밍에 대한 정밀한 제어를 위한 명시적 mbarrier API

메모리 배리어 특화: 기본적인 스레드 동기화와 달리, 이 퍼즐은 메모리 연산이 언제 완료되는지에 대한 세밀한 제어를 제공하는 명시적 메모리 배리어를 사용하며, 이는 복잡한 메모리 접근 패턴에 필수적입니다.

구성

시스템 매개변수:

• 이미지 크기: SIZE = 1024 요소 (간소화를 위해 1D)
• 블록당 스레드 수: TPB = 256 스레드, (256, 1) 블록 차원으로 구성
• 그리드 구성: 전체 이미지를 타일 단위로 처리하기 위한 (4, 1) 블록 (총 4개 블록)
• 데이터 타입: 모든 연산에 DType.float32

반복 매개변수:

• 스텐실 반복 횟수: STENCIL_ITERATIONS = 3 개선 패스
• 버퍼 수: BUFFER_COUNT = 2 (더블 버퍼링)
• 스텐실 커널: 반지름 1의 3점 평균

버퍼 아키텍처:

• buffer_A: 주 공유 메모리 버퍼 ([256] 요소)
• buffer_B: 보조 공유 메모리 버퍼 ([256] 요소)
• 역할 교대: 매 반복마다 버퍼가 읽기 소스와 쓰기 대상 사이를 교체

처리 요구사항:

초기화 단계:

• 버퍼 설정: buffer_A를 입력 데이터로, buffer_B를 0으로 초기화
• 배리어 초기화: 동기화 지점을 위한 mbarrier 객체 설정
• 스레드 조정: 모든 스레드가 초기화에 참여

반복 처리:

• 짝수 반복 (0, 2, 4…): buffer_A에서 읽고 buffer_B에 쓰기
• 홀수 반복 (1, 3, 5…): buffer_B에서 읽고 buffer_A에 쓰기
• 스텐실 연산: 3점 평균 \((\text{left} + \text{center} + \text{right}) / 3\)
• 경계 처리: 버퍼 가장자리의 요소에 대해 적응적 평균 사용

메모리 배리어 조정:

• mbarrier_arrive(): 각 스레드가 쓰기 단계 완료를 알림
• mbarrier_test_wait(): 모든 스레드가 쓰기를 완료할 때까지 대기
• 버퍼 교체 안전성: 다른 스레드가 아직 쓰고 있는 동안 버퍼에서 읽는 것을 방지
• 배리어 재초기화: 반복 간에 배리어 상태를 재설정

출력 단계:

• 최종 버퍼 선택: 반복 횟수의 홀짝에 따라 활성 버퍼 선택
• 전역 메모리 쓰기: 최종 결과를 출력 배열에 복사
• 완료 배리어: 블록 종료 전 모든 쓰기 완료 보장

완성할 코드

# Double-buffered stencil configuration
comptime STENCIL_ITERATIONS = 3
comptime BUFFER_COUNT = 2


fn double_buffered_stencil_computation[
    layout: Layout
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    size: Int,
):
    """Double-buffered stencil computation with memory barrier coordination.

    Iteratively applies 3-point stencil using alternating buffers.
    Uses mbarrier APIs for precise buffer swap coordination.
    """

    # Double-buffering: Two shared memory buffers
    buffer_A = LayoutTensor[
        dtype,
        Layout.row_major(TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    buffer_B = LayoutTensor[
        dtype,
        Layout.row_major(TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    # Memory barriers for coordinating buffer swaps
    init_barrier = LayoutTensor[
        DType.uint64,
        Layout.row_major(1),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    iter_barrier = LayoutTensor[
        DType.uint64,
        Layout.row_major(1),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    final_barrier = LayoutTensor[
        DType.uint64,
        Layout.row_major(1),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)
    local_i = Int(thread_idx.x)

    # Initialize barriers (only thread 0)
    if local_i == 0:
        mbarrier_init(init_barrier.ptr, TPB)
        mbarrier_init(iter_barrier.ptr, TPB)
        mbarrier_init(final_barrier.ptr, TPB)

    # Initialize buffer_A with input data

    # FILL ME IN (roughly 4 lines)

    # Wait for buffer_A initialization
    _ = mbarrier_arrive(init_barrier.ptr)
    _ = mbarrier_test_wait(init_barrier.ptr, TPB)

    # Iterative stencil processing with double-buffering
    @parameter
    for iteration in range(STENCIL_ITERATIONS):

        @parameter
        if iteration % 2 == 0:
            # Even iteration: Read from A, Write to B

            # FILL ME IN (roughly 12 lines)
            ...

        else:
            # Odd iteration: Read from B, Write to A

            # FILL ME IN (roughly 12 lines)
            ...

        # Memory barrier: wait for all writes before buffer swap
        _ = mbarrier_arrive(iter_barrier.ptr)
        _ = mbarrier_test_wait(iter_barrier.ptr, TPB)

        # Reinitialize barrier for next iteration
        if local_i == 0:
            mbarrier_init(iter_barrier.ptr, TPB)

    # Write final results from active buffer
    if local_i < TPB and global_i < size:

        @parameter
        if STENCIL_ITERATIONS % 2 == 0:
            # Even iterations end in buffer_A
            output[global_i] = buffer_A[local_i]
        else:
            # Odd iterations end in buffer_B
            output[global_i] = buffer_B[local_i]

    # Final barrier
    _ = mbarrier_arrive(final_barrier.ptr)
    _ = mbarrier_test_wait(final_barrier.ptr, TPB)

전체 파일 보기: problems/p29/p29.mojo

코드 실행

솔루션을 테스트하려면 터미널에서 다음 명령을 실행합니다:

pixi run p29 --double-buffer

pixi run -e amd p29 --double-buffer

uv run poe p29 --double-buffer

퍼즐을 성공적으로 완료하면 다음과 유사한 출력이 표시됩니다:

Puzzle 29: GPU Synchronization Primitives
==================================================
TPB: 256
SIZE: 1024
STENCIL_ITERATIONS: 3
BUFFER_COUNT: 2

Testing Puzzle 29B: Double-Buffered Stencil Computation
============================================================
Double-buffered stencil completed
Input sample: 1.0 1.0 1.0
GPU output sample: 1.0 1.0 1.0
✅ Double-buffered stencil test PASSED!

솔루션

fn double_buffered_stencil_computation[
    layout: Layout
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    size: Int,
):
    """Double-buffered stencil computation with memory barrier coordination.

    Iteratively applies 3-point stencil using alternating buffers.
    Uses mbarrier APIs for precise buffer swap coordination.
    """

    # Double-buffering: Two shared memory buffers
    buffer_A = LayoutTensor[
        dtype,
        Layout.row_major(TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    buffer_B = LayoutTensor[
        dtype,
        Layout.row_major(TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    # Memory barriers for coordinating buffer swaps
    init_barrier = LayoutTensor[
        DType.uint64,
        Layout.row_major(1),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    iter_barrier = LayoutTensor[
        DType.uint64,
        Layout.row_major(1),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    final_barrier = LayoutTensor[
        DType.uint64,
        Layout.row_major(1),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)
    local_i = Int(thread_idx.x)

    # Initialize barriers (only thread 0)
    if local_i == 0:
        mbarrier_init(init_barrier.ptr, TPB)
        mbarrier_init(iter_barrier.ptr, TPB)
        mbarrier_init(final_barrier.ptr, TPB)

    # Initialize buffer_A with input data
    if local_i < TPB and global_i < size:
        buffer_A[local_i] = input[global_i]
    else:
        buffer_A[local_i] = 0.0

    # Wait for buffer_A initialization
    _ = mbarrier_arrive(init_barrier.ptr)
    _ = mbarrier_test_wait(init_barrier.ptr, TPB)

    # Iterative stencil processing with double-buffering
    @parameter
    for iteration in range(STENCIL_ITERATIONS):

        @parameter
        if iteration % 2 == 0:
            # Even iteration: Read from A, Write to B
            if local_i < TPB:
                var stencil_sum: Scalar[dtype] = 0.0
                var stencil_count: Int = 0

                # 3-point stencil: [i-1, i, i+1]
                for offset in range(-1, 2):
                    sample_idx = local_i + offset
                    if sample_idx >= 0 and sample_idx < TPB:
                        stencil_sum += rebind[Scalar[dtype]](
                            buffer_A[sample_idx]
                        )
                        stencil_count += 1

                if stencil_count > 0:
                    buffer_B[local_i] = stencil_sum / stencil_count
                else:
                    buffer_B[local_i] = buffer_A[local_i]

        else:
            # Odd iteration: Read from B, Write to A
            if local_i < TPB:
                var stencil_sum: Scalar[dtype] = 0.0
                var stencil_count: Int = 0

                # 3-point stencil: [i-1, i, i+1]
                for offset in range(-1, 2):
                    sample_idx = local_i + offset
                    if sample_idx >= 0 and sample_idx < TPB:
                        stencil_sum += rebind[Scalar[dtype]](
                            buffer_B[sample_idx]
                        )
                        stencil_count += 1

                if stencil_count > 0:
                    buffer_A[local_i] = stencil_sum / stencil_count
                else:
                    buffer_A[local_i] = buffer_B[local_i]

        # Memory barrier: wait for all writes before buffer swap
        _ = mbarrier_arrive(iter_barrier.ptr)
        _ = mbarrier_test_wait(iter_barrier.ptr, TPB)

        # Reinitialize barrier for next iteration
        if local_i == 0:
            mbarrier_init(iter_barrier.ptr, TPB)

    # Write final results from active buffer
    if local_i < TPB and global_i < size:

        @parameter
        if STENCIL_ITERATIONS % 2 == 0:
            # Even iterations end in buffer_A
            output[global_i] = buffer_A[local_i]
        else:
            # Odd iterations end in buffer_B
            output[global_i] = buffer_B[local_i]

    # Final barrier
    _ = mbarrier_arrive(final_barrier.ptr)
    _ = mbarrier_test_wait(final_barrier.ptr, TPB)

# 왼쪽, 중심, 오른쪽 이웃과의 평균
stencil_sum = buffer[i-1] + buffer[i] + buffer[i+1]
result[i] = stencil_sum / 3.0

# 유효한 이웃만 평균에 포함
stencil_count = 0
for neighbor in valid_neighbors:
    stencil_sum += buffer[neighbor]
    stencil_count += 1
result[i] = stencil_sum / stencil_count

# 스레드 A가 buffer_B[10]에 쓰기
buffer_B[10] = stencil_result_A

# 스레드 B가 스텐실 연산을 위해 buffer_B[10]을 즉시 읽기
# 경쟁 상태: 스레드 B가 스레드 A의 쓰기가 완료되기 전에 이전 값을 읽을 수 있음
stencil_input = buffer_B[10]  // 미정의 동작!

# 모든 스레드가 결과를 쓰기
buffer_B[local_i] = stencil_result

# 쓰기 완료 알림
mbarrier_arrive(barrier)

# 모든 스레드의 쓰기 완료까지 대기
mbarrier_test_wait(barrier, TPB)

# 이제 읽기 안전 - 모든 쓰기 완료 보장
stencil_input = buffer_B[neighbor_index]  // 항상 올바른 값을 읽음

@parameter
if STENCIL_ITERATIONS % 2 == 0:
    # 짝수 총 반복 횟수는 buffer_A에서 종료
    output[global_i] = buffer_A[local_i]
else:
    # 홀수 총 반복 횟수는 buffer_B에서 종료
    output[global_i] = buffer_B[local_i]