다단계 파이프라인 조정

개요

조율된 3단계 파이프라인을 통해 이미지를 처리하는 커널을 구현합니다. 서로 다른 스레드 그룹이 특화된 처리 단계를 담당하고, 명시적 배리어로 동기화됩니다.

참고: 스레드 역할이 특화되어 있습니다: Stage 1 (스레드 0-127)은 데이터를 로드하고 전처리하며, Stage 2 (스레드 128-255)는 블러 연산을 적용하고, Stage 3 (전체 스레드)은 최종 스무딩을 수행합니다.

알고리즘 아키텍처: 이 퍼즐은 하나의 GPU 블록 안에서 서로 다른 스레드 그룹이 완전히 다른 알고리즘을 실행하는 생산자-소비자 파이프라인을 구현합니다. 모든 스레드가 서로 다른 데이터에 대해 동일한 알고리즘을 실행하는 전통적인 GPU 프로그래밍과 달리, 이 접근 방식은 스레드를 기능별로 특화하여 분할합니다.

파이프라인 개념: 알고리즘은 세 개의 구분된 단계를 통해 데이터를 처리하며, 각 단계에는 서로 다른 알고리즘을 실행하는 특화된 스레드 그룹이 있습니다. 각 단계는 다음 단계가 소비하는 데이터를 생성하여, 배리어로 신중하게 동기화해야 하는 명시적 생산자-소비자 관계를 만듭니다.

데이터 의존성과 동기화: 각 단계는 다음 단계가 소비하는 데이터를 생성합니다:

• Stage 1 → Stage 2: 첫 번째 단계가 블러 처리를 위한 전처리 데이터를 생성
• Stage 2 → Stage 3: 두 번째 단계가 최종 스무딩을 위한 블러 결과를 생성
• 배리어가 경쟁 상태를 방지: 의존하는 단계가 시작되기 전에 해당 단계가 완전히 완료되도록 보장

구체적으로, 다단계 파이프라인은 세 가지 수학 연산으로 구성된 조율된 이미지 처리 알고리즘을 구현합니다:

Stage 1 - 전처리 강화:

\[P[i] = I[i] \times 1.1\]

여기서 \(P[i]\)는 전처리된 데이터이고 \(I[i]\)는 입력 데이터입니다.

Stage 2 - 수평 블러 필터:

\[B[i] = \frac{1}{N_i} \sum_{k=-2}^{2} P[i+k] \quad \text{where} i+k \in [0, 255]\]

여기서 \(B[i]\)는 블러 결과이고, \(N_i\)는 타일 경계 내의 유효한 이웃 수입니다.

Stage 3 - 연쇄적 이웃 스무딩:

\[F[i] = \begin{cases} (B[i] + B[i+1]) \times 0.6 & \text{if } i = 0 \\ ((B[i] + B[i-1]) \times 0.6 + B[i+1]) \times 0.6 & \text{if } 0 < i < 255 \\ (B[i] + B[i-1]) \times 0.6 & \text{if } i = 255 \end{cases}\]

여기서 \(F[i]\)는 연쇄적 스무딩이 적용된 최종 출력입니다.

스레드 특화:

• 스레드 0-127: \(i \in \{0, 1, 2, \ldots, 255\}\)에 대해 \(P[i]\) 계산 (스레드당 2개 요소)
• 스레드 128-255: \(i \in \{0, 1, 2, \ldots, 255\}\)에 대해 \(B[i]\) 계산 (스레드당 2개 요소)
• 전체 256개 스레드: \(i \in \{0, 1, 2, \ldots, 255\}\)에 대해 \(F[i]\) 계산 (스레드당 1개 요소)

동기화 지점:

\[\text{barrier}_1 \Rightarrow P[i] \text{ complete} \Rightarrow \text{barrier}_2 \Rightarrow B[i] \text{ complete} \Rightarrow \text{barrier}_3 \Rightarrow F[i] \text{ complete}\]

핵심 개념

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

• 하나의 GPU 블록 안에서 스레드 역할 특화 구현
• 처리 단계 간 생산자-소비자 관계 조율
• 서로 다른 알고리즘 간의 동기화를 위한 배리어 사용 (동일한 알고리즘 내부뿐 아니라)

핵심 통찰은 서로 다른 스레드 그룹이 완전히 다른 알고리즘을 실행하면서 전략적 배리어 배치를 통해 조율되는 다단계 파이프라인을 어떻게 설계하는지 이해하는 것입니다.

왜 중요한가: 대부분의 GPU 튜토리얼은 단일 알고리즘 내에서의 배리어 사용법 - 리덕션이나 공유 메모리 연산 중 스레드를 동기화하는 것 - 을 가르칩니다. 하지만 실제 GPU 알고리즘에서는 신중하게 조율해야 하는 여러 개의 구분된 처리 단계를 포함하는 아키텍처적 복잡성이 필요한 경우가 많습니다. 이 퍼즐은 단일체적 알고리즘을 특화되고 조율된 처리 파이프라인으로 변환하는 방법을 보여줍니다.

이전 퍼즐과 현재의 배리어 사용 비교:

• 이전 퍼즐 (P8, P12, P15): 모든 스레드가 동일한 알고리즘을 실행하고, 배리어는 알고리즘 단계 내에서 동기화
• 이 퍼즐: 서로 다른 스레드 그룹이 서로 다른 알고리즘을 실행하고, 배리어는 서로 다른 알고리즘 간의 조율

스레드 특화 아키텍처: 스레드가 데이터 인덱스만 다른 데이터 병렬 처리와 달리, 이 퍼즐은 파이프라인에서의 역할에 따라 스레드가 근본적으로 다른 코드 경로를 실행하는 알고리즘 병렬 처리를 구현합니다.

구성

시스템 매개변수:

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

스레드 특화 아키텍처:

• Stage 1 스레드: STAGE1_THREADS = 128 (스레드 0-127, 블록의 전반부)

  • 역할: 전역 메모리에서 입력 데이터를 로드하고 전처리 적용
  • 작업 분배: 효율적인 부하 균형을 위해 스레드당 2개 요소 처리
  • 출력: input_shared[256]에 전처리된 데이터 채우기

• Stage 2 스레드: STAGE2_THREADS = 128 (스레드 128-255, 블록의 후반부)

  • 역할: 전처리된 데이터에 수평 블러 필터 적용
  • 작업 분배: 스레드당 2개의 블러 연산 처리
  • 출력: blur_shared[256]에 블러 결과 채우기

• Stage 3 스레드: 전체 256개 스레드 협력

  • 역할: 최종 스무딩 및 전역 메모리로 출력
  • 작업 분배: 일대일 매핑 (스레드 i가 요소 i를 처리)
  • 출력: 글로벌 output 배열에 최종 결과 기록

완성할 코드

comptime TPB = 256  # Threads per block for pipeline stages
comptime SIZE = 1024  # Image size (1D for simplicity)
comptime BLOCKS_PER_GRID = (4, 1)
comptime THREADS_PER_BLOCK = (TPB, 1)
comptime dtype = DType.float32
comptime layout = row_major[SIZE]()
comptime LayoutType = type_of(layout)

# Multi-stage processing configuration
comptime STAGE1_THREADS = TPB // 2
comptime STAGE2_THREADS = TPB // 2
comptime BLUR_RADIUS = 2


def multi_stage_image_blur_pipeline(
    output: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    input: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
    size: Int,
):
    """Multi-stage image blur pipeline with barrier coordination.

    Stage 1 (threads 0-127): Load input data and apply 1.1x preprocessing
    Stage 2 (threads 128-255): Apply 5-point blur with BLUR_RADIUS=2
    Stage 3 (all threads): Final neighbor smoothing and output
    """

    # Shared memory buffers for pipeline stages
    var input_shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB]())
    var blur_shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB]())

    var global_i = block_dim.x * block_idx.x + thread_idx.x
    var local_i = thread_idx.x

    # Stage 1: Load and preprocess (threads 0-127)

    # FILL ME IN (roughly 10 lines)

    barrier()  # Wait for Stage 1 completion

    # Stage 2: Apply blur (threads 128-255)

    # FILL ME IN (roughly 25 lines)

    barrier()  # Wait for Stage 2 completion

    # Stage 3: Final smoothing (all threads)

    # FILL ME IN (roughly 7 lines)

    barrier()  # Ensure all writes complete

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

코드 실행

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

pixi run p29 --multi-stage

pixi run -e amd p29 --multi-stage

uv run poe p29 --multi-stage

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

Puzzle 29: GPU Synchronization Primitives
==================================================
TPB: 256
SIZE: 1024
STAGE1_THREADS: 128
STAGE2_THREADS: 128
BLUR_RADIUS: 2

Testing Puzzle 29A: Multi-Stage Pipeline Coordination
============================================================
Multi-stage pipeline blur completed
Input sample: 0.0 1.01 2.02
Output sample: 1.6665002 2.3331003 3.3996604
✅ Multi-stage pipeline coordination test PASSED!

솔루션

def multi_stage_image_blur_pipeline(
    output: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    input: TileTensor[mut=False, dtype, LayoutType, MutAnyOrigin],
    size: Int,
):
    """Multi-stage image blur pipeline with barrier coordination.

    Stage 1 (threads 0-127): Load input data and apply 1.1x preprocessing
    Stage 2 (threads 128-255): Apply 5-point blur with BLUR_RADIUS=2
    Stage 3 (all threads): Final neighbor smoothing and output
    """

    # Shared memory buffers for pipeline stages
    var input_shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB]())
    var blur_shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB]())

    var global_i = block_dim.x * block_idx.x + thread_idx.x
    var local_i = thread_idx.x

    # Stage 1: Load and preprocess (threads 0-127)
    if local_i < STAGE1_THREADS:
        if global_i < size:
            input_shared[local_i] = input[global_i] * 1.1
            # Each thread loads 2 elements
            if local_i + STAGE1_THREADS < size:
                input_shared[local_i + STAGE1_THREADS] = (
                    input[global_i + STAGE1_THREADS] * 1.1
                )
        else:
            # Zero-padding for out-of-bounds
            input_shared[local_i] = 0.0
            if local_i + STAGE1_THREADS < TPB:
                input_shared[local_i + STAGE1_THREADS] = 0.0

    barrier()  # Wait for Stage 1 completion

    # Stage 2: Apply blur (threads 128-255)
    if local_i >= STAGE1_THREADS:
        var blur_idx = local_i - STAGE1_THREADS
        var blur_sum: Scalar[dtype] = 0.0
        blur_count = 0

        # 5-point blur kernel
        for offset in range(-BLUR_RADIUS, BLUR_RADIUS + 1):
            sample_idx = blur_idx + offset
            if sample_idx >= 0 and sample_idx < TPB:
                blur_sum += rebind[Scalar[dtype]](input_shared[sample_idx])
                blur_count += 1

        if blur_count > 0:
            blur_shared[blur_idx] = blur_sum / Scalar[dtype](blur_count)
        else:
            blur_shared[blur_idx] = 0.0

        # Process second element
        var second_idx = blur_idx + STAGE1_THREADS
        if second_idx < TPB:
            blur_sum = 0.0
            blur_count = 0
            for offset in range(-BLUR_RADIUS, BLUR_RADIUS + 1):
                sample_idx = second_idx + offset
                if sample_idx >= 0 and sample_idx < TPB:
                    blur_sum += rebind[Scalar[dtype]](input_shared[sample_idx])
                    blur_count += 1

            if blur_count > 0:
                blur_shared[second_idx] = blur_sum / Scalar[dtype](blur_count)
            else:
                blur_shared[second_idx] = 0.0

    barrier()  # Wait for Stage 2 completion

    # Stage 3: Final smoothing (all threads)
    if global_i < size:
        final_value = blur_shared[local_i]

        # Neighbor smoothing with 0.6 scaling
        if local_i > 0:
            final_value = (final_value + blur_shared[local_i - 1]) * 0.6
        if local_i < TPB - 1:
            final_value = (final_value + blur_shared[local_i + 1]) * 0.6

        output[global_i] = final_value

    barrier()  # Ensure all writes complete