warp.prefix_sum() 하드웨어 최적화 병렬 스캔

워프 레벨 병렬 스캔 연산에서는 prefix_sum()을 사용하여 복잡한 공유 메모리 알고리즘을 하드웨어 최적화 기본 요소로 대체할 수 있습니다. 이 강력한 연산을 통해 수십 줄의 공유 메모리 및 동기화 코드가 필요했을 효율적인 누적 계산, 병렬 파티셔닝, 고급 조정 알고리즘을 구현할 수 있습니다.

핵심 통찰: prefix_sum() 연산은 하드웨어 가속 병렬 스캔을 활용하여 워프 레인에 걸쳐 \(O(\log n)\) 복잡도로 누적 연산을 수행하며, 복잡한 다단계 알고리즘을 단일 함수 호출로 대체합니다.

병렬 스캔이란? 병렬 스캔 (누적 합)은 데이터 요소에 걸쳐 누적 연산을 수행하는 기본적인 병렬 기본 요소입니다. 덧셈의 경우 [a, b, c, d]를 [a, a+b, a+b+c, a+b+c+d]로 변환합니다. 이 연산은 스트림 컴팩션, quicksort 파티셔닝, 병렬 정렬 같은 병렬 알고리즘에 필수적입니다.

핵심 개념

이 퍼즐에서 배울 내용:

• prefix_sum()을 활용한 하드웨어 최적화 병렬 스캔
• 포함(inclusive) vs 비포함(exclusive) 누적 합 패턴
• 데이터 재배치를 위한 워프 레벨 스트림 컴팩션
• 여러 워프 기본 요소를 결합한 고급 병렬 파티셔닝
• 복잡한 공유 메모리를 대체하는 단일 워프 알고리즘 최적화

이를 통해 다단계 공유 메모리 알고리즘이 우아한 단일 함수 호출로 변환되어, 명시적 동기화 없이 효율적인 병렬 스캔 연산이 가능합니다.

1. 워프 포함 누적 합

구성

• 벡터 크기: SIZE = WARP_SIZE (GPU에 따라 32 또는 64)
• 그리드 구성: (1, 1) 그리드당 블록 수
• 블록 구성: (WARP_SIZE, 1) 블록당 스레드 수
• 데이터 타입: DType.float32
• 레이아웃: Layout.row_major(SIZE) (1D row-major)

prefix_sum의 이점

기존 누적 합은 복잡한 다단계 공유 메모리 알고리즘이 필요합니다. Puzzle 14: 누적 합에서는 명시적 공유 메모리 관리로 이를 힘들게 구현했습니다:

fn prefix_sum_simple[
    layout: Layout
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    a: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    size: UInt,
):
    global_i = block_dim.x * block_idx.x + thread_idx.x
    local_i = thread_idx.x
    shared = LayoutTensor[
        dtype,
        Layout.row_major(TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    if global_i < size:
        shared[local_i] = a[global_i]

    barrier()

    offset = UInt(1)
    for i in range(Int(log2(Scalar[dtype](TPB)))):
        var current_val: output.element_type = 0
        if local_i >= offset and local_i < size:
            current_val = shared[local_i - offset]  # read

        barrier()
        if local_i >= offset and local_i < size:
            shared[local_i] += current_val

        barrier()
        offset *= 2

    if global_i < size:
        output[global_i] = shared[local_i]

기존 방식의 문제점:

• 메모리 오버헤드: 공유 메모리 할당이 필요
• 다중 배리어: 복잡한 다단계 동기화
• 복잡한 인덱싱: 수동 스트라이드 계산과 경계 검사
• 낮은 확장성: 각 단계 사이에 배리어가 필요한 \(O(\log n)\) 단계

prefix_sum()을 사용하면 병렬 스캔이 간단해집니다:

# 하드웨어 최적화 방식 - 단일 함수 호출!
current_val = input[global_i]
scan_result = prefix_sum[exclusive=False](current_val)
output[global_i] = scan_result

prefix_sum의 장점:

• 메모리 오버헤드 제로: 하드웨어 가속 연산
• 동기화 불필요: 단일 아토믹 연산
• 하드웨어 최적화: 전용 스캔 유닛 활용
• 완벽한 확장성: 모든 WARP_SIZE (32, 64 등)에서 동작

완성할 코드

하드웨어 최적화 prefix_sum() 기본 요소를 사용하여 포함 누적 합을 구현합니다.

수학적 연산: 각 레인이 자신의 위치까지 모든 요소의 합을 포함하는 누적 합을 계산합니다: \[\Large \text{output}[i] = \sum_{j=0}^{i} \text{input}[j]\]

입력 데이터 [1, 2, 3, 4, 5, ...]를 누적 합 [1, 3, 6, 10, 15, ...]으로 변환하며, 각 위치에 이전 모든 요소와 자기 자신의 합이 담깁니다.

fn warp_inclusive_prefix_sum[
    layout: Layout, size: Int
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
):
    """
    Inclusive prefix sum using warp primitive:
    Each thread gets sum of all elements up to and including its position.
    Compare this to Puzzle 12's complex shared memory + barrier approach.

    Puzzle 12 approach:
    - Shared memory allocation
    - Multiple barrier synchronizations
    - Log(n) iterations with manual tree reduction
    - Complex multi-phase algorithm

    Warp prefix_sum approach:
    - Single function call!
    - Hardware-optimized parallel scan
    - Automatic synchronization
    - O(log n) complexity, but implemented in hardware.

    NOTE: This implementation only works correctly within a single warp (WARP_SIZE threads).
    For multi-warp scenarios, additional coordination would be needed.
    """
    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)

    # FILL ME IN (roughly 4 lines)

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

워프 포함 누적 합 테스트:

pixi run p26 --prefix-sum

pixi run -e amd p26 --prefix-sum

pixi run -e apple p26 --prefix-sum

uv run poe p26 --prefix-sum

풀었을 때의 예상 출력:

WARP_SIZE:  32
SIZE:  32
output: [1.0, 3.0, 6.0, 10.0, 15.0, 21.0, 28.0, 36.0, 45.0, 55.0, 66.0, 78.0, 91.0, 105.0, 120.0, 136.0, 153.0, 171.0, 190.0, 210.0, 231.0, 253.0, 276.0, 300.0, 325.0, 351.0, 378.0, 406.0, 435.0, 465.0, 496.0, 528.0]
expected: [1.0, 3.0, 6.0, 10.0, 15.0, 21.0, 28.0, 36.0, 45.0, 55.0, 66.0, 78.0, 91.0, 105.0, 120.0, 136.0, 153.0, 171.0, 190.0, 210.0, 231.0, 253.0, 276.0, 300.0, 325.0, 351.0, 378.0, 406.0, 435.0, 465.0, 496.0, 528.0]
✅ Warp inclusive prefix sum test passed!

솔루션

fn warp_inclusive_prefix_sum[
    layout: Layout, size: Int
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
):
    """
    Inclusive prefix sum using warp primitive: Each thread gets sum of all elements up to and including its position.
    Compare this to Puzzle 12's complex shared memory + barrier approach.

    Puzzle 12 approach:
    - Shared memory allocation
    - Multiple barrier synchronizations
    - Log(n) iterations with manual tree reduction
    - Complex multi-phase algorithm

    Warp prefix_sum approach:
    - Single function call!
    - Hardware-optimized parallel scan
    - Automatic synchronization
    - O(log n) complexity, but implemented in hardware.

    NOTE: This implementation only works correctly within a single warp (WARP_SIZE threads).
    For multi-warp scenarios, additional coordination would be needed.
    """
    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)

    if global_i < size:
        current_val = input[global_i]

        # This one call replaces ~30 lines of complex shared memory logic from Puzzle 12!
        # But it only works within the current warp (WARP_SIZE threads)
        scan_result = prefix_sum[exclusive=False](
            rebind[Scalar[dtype]](current_val)
        )

        output[global_i] = scan_result

if global_i < size:
    current_val = input[global_i]

    # 이 한 줄이 Puzzle 14의 복잡한 공유 메모리 로직 ~30줄을 대체합니다!
    # 단, 현재 워프 (WARP_SIZE 스레드) 내에서만 동작합니다
    scan_result = prefix_sum[exclusive=False](
        rebind[Scalar[dtype]](current_val)
    )

    output[global_i] = scan_result

입력: [1, 2, 3, 4, 5, 6, 7, 8, ...]

사이클 1: 모든 레인이 동시에 값을 로드
  Lane 0: current_val = 1
  Lane 1: current_val = 2
  Lane 2: current_val = 3
  Lane 3: current_val = 4
  ...
  Lane 31: current_val = 32

사이클 2: prefix_sum[exclusive=False] 실행 (하드웨어 가속)
  Lane 0: scan_result = 1 (요소 0~0의 합)
  Lane 1: scan_result = 3 (요소 0~1의 합: 1+2)
  Lane 2: scan_result = 6 (요소 0~2의 합: 1+2+3)
  Lane 3: scan_result = 10 (요소 0~3의 합: 1+2+3+4)
  ...
  Lane 31: scan_result = 528 (요소 0~31의 합)

사이클 3: 결과 저장
  Lane 0: output[0] = 1
  Lane 1: output[1] = 3
  Lane 2: output[2] = 6
  Lane 3: output[3] = 10
  ...

2. 워프 파티션

구성

• 벡터 크기: SIZE = WARP_SIZE (GPU에 따라 32 또는 64)
• 그리드 구성: (1, 1) 그리드당 블록 수
• 블록 구성: (WARP_SIZE, 1) 블록당 스레드 수

완성할 코드

shuffle_xor과 prefix_sum 기본 요소를 모두 사용하여 단일 워프 병렬 파티셔닝을 구현합니다.

수학적 연산: 피벗 값을 기준으로 요소를 분할하여, < pivot인 요소는 왼쪽에, >= pivot인 요소는 오른쪽에 배치합니다: \[\Large \text{output} = [\text{elements} < \text{pivot}] \,|\, [\text{elements} \geq \text{pivot}]\]

고급 알고리즘: 이 알고리즘은 두 가지 정교한 워프 기본 요소를 결합합니다:

1. shuffle_xor(): 왼쪽 요소 개수를 세기 위한 워프 레벨 버터플라이 리덕션
2. prefix_sum(): 각 파티션 내 위치 계산을 위한 비포함 스캔

이는 단일 워프 내에서 여러 워프 기본 요소를 결합하여 복잡한 병렬 알고리즘을 구현하는 강력함을 보여줍니다.

fn warp_partition[
    layout: Layout, size: Int
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    pivot: Float32,
):
    """
    Single-warp parallel partitioning using BOTH shuffle_xor AND prefix_sum.
    This implements a warp-level quicksort partition step that places elements < pivot
    on the left and elements >= pivot on the right.

    ALGORITHM COMPLEXITY - combines two advanced warp primitives:
    1. shuffle_xor(): Butterfly pattern for warp-level reductions
    2. prefix_sum(): Warp-level exclusive scan for position calculation.

    This demonstrates the power of warp primitives for sophisticated parallel algorithms
    within a single warp (works for any WARP_SIZE: 32, 64, etc.).

    Example with pivot=5:
    Input:  [3, 7, 1, 8, 2, 9, 4, 6]
    Result: [3, 1, 2, 4, 7, 8, 9, 6] (< pivot | >= pivot).
    """
    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)

    if global_i < size:
        current_val = input[global_i]

        # FILL ME IN (roughly 13 lines)

워프 파티션 테스트:

uv run poe p26 --partition

pixi run p26 --partition

풀었을 때의 예상 출력:

WARP_SIZE:  32
SIZE:  32
output: HostBuffer([3.0, 1.0, 2.0, 4.0, 0.0, 3.0, 1.0, 4.0, 3.0, 1.0, 2.0, 4.0, 0.0, 3.0, 1.0, 4.0, 7.0, 8.0, 9.0, 6.0, 10.0, 11.0, 12.0, 13.0, 7.0, 8.0, 9.0, 6.0, 10.0, 11.0, 12.0, 13.0])
expected: HostBuffer([3.0, 1.0, 2.0, 4.0, 0.0, 3.0, 1.0, 4.0, 3.0, 1.0, 2.0, 4.0, 0.0, 3.0, 1.0, 4.0, 7.0, 8.0, 9.0, 6.0, 10.0, 11.0, 12.0, 13.0, 7.0, 8.0, 9.0, 6.0, 10.0, 11.0, 12.0, 13.0])
pivot: 5.0
✅ Warp partition test passed!

솔루션

fn warp_partition[
    layout: Layout, size: Int
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    pivot: Float32,
):
    """
    Single-warp parallel partitioning using BOTH shuffle_xor AND prefix_sum.
    This implements a warp-level quicksort partition step that places elements < pivot
    on the left and elements >= pivot on the right.

    ALGORITHM COMPLEXITY - combines two advanced warp primitives:
    1. shuffle_xor(): Butterfly pattern for warp-level reductions
    2. prefix_sum(): Warp-level exclusive scan for position calculation.

    This demonstrates the power of warp primitives for sophisticated parallel algorithms
    within a single warp (works for any WARP_SIZE: 32, 64, etc.).

    Example with pivot=5:
    Input:  [3, 7, 1, 8, 2, 9, 4, 6]
    Result: [3, 1, 2, 4, 7, 8, 9, 6] (< pivot | >= pivot).
    """
    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)

    if global_i < size:
        current_val = input[global_i]

        # Phase 1: Create warp-level predicates
        predicate_left = Float32(1.0) if current_val < pivot else Float32(0.0)
        predicate_right = Float32(1.0) if current_val >= pivot else Float32(0.0)

        # Phase 2: Warp-level prefix sum to get positions within warp
        warp_left_pos = prefix_sum[exclusive=True](predicate_left)
        warp_right_pos = prefix_sum[exclusive=True](predicate_right)

        # Phase 3: Get total left count using shuffle_xor reduction
        warp_left_total = predicate_left

        # Butterfly reduction to get total across the warp: dynamic for any WARP_SIZE
        offset = WARP_SIZE // 2
        while offset > 0:
            warp_left_total += shuffle_xor(warp_left_total, offset)
            offset //= 2

        # Phase 4: Write to output positions
        if current_val < pivot:
            # Left partition: use warp-level position
            output[Int(warp_left_pos)] = current_val
        else:
            # Right partition: offset by total left count + right position
            output[Int(warp_left_total + warp_right_pos)] = current_val

if global_i < size:
    current_val = input[global_i]

    # 1단계: 워프 레벨 프레디케이트 생성
    predicate_left = Float32(1.0) if current_val < pivot else Float32(0.0)
    predicate_right = Float32(1.0) if current_val >= pivot else Float32(0.0)

    # 2단계: 워프 레벨 누적 합으로 워프 내 위치 계산
    warp_left_pos = prefix_sum[exclusive=True](predicate_left)
    warp_right_pos = prefix_sum[exclusive=True](predicate_right)

    # 3단계: shuffle_xor 버터플라이 리덕션으로 왼쪽 총 개수 구하기
    warp_left_total = predicate_left

    # 워프 전체의 합산을 위한 버터플라이 리덕션: 모든 WARP_SIZE에 동적 대응
    offset = WARP_SIZE // 2
    while offset > 0:
        warp_left_total += shuffle_xor(warp_left_total, offset)
        offset //= 2

    # 4단계: 출력 위치에 기록
    if current_val < pivot:
        # 왼쪽 파티션: 워프 레벨 위치 사용
        output[Int(warp_left_pos)] = current_val
    else:
        # 오른쪽 파티션: 왼쪽 총 개수 + 오른쪽 위치로 offset
        output[Int(warp_left_total + warp_right_pos)] = current_val

초기 상태:
  Lane 0: current_val=3 (< 5)  Lane 1: current_val=7 (>= 5)
  Lane 2: current_val=1 (< 5)  Lane 3: current_val=8 (>= 5)
  Lane 4: current_val=2 (< 5)  Lane 5: current_val=9 (>= 5)
  Lane 6: current_val=4 (< 5)  Lane 7: current_val=6 (>= 5)

1단계: 프레디케이트 생성
  Lane 0: predicate_left=1.0, predicate_right=0.0
  Lane 1: predicate_left=0.0, predicate_right=1.0
  Lane 2: predicate_left=1.0, predicate_right=0.0
  Lane 3: predicate_left=0.0, predicate_right=1.0
  Lane 4: predicate_left=1.0, predicate_right=0.0
  Lane 5: predicate_left=0.0, predicate_right=1.0
  Lane 6: predicate_left=1.0, predicate_right=0.0
  Lane 7: predicate_left=0.0, predicate_right=1.0

2단계: 위치 계산을 위한 비포함 누적 합
  warp_left_pos:  [0, 0, 1, 1, 2, 2, 3, 3]
  warp_right_pos: [0, 0, 0, 1, 1, 2, 2, 3]

3단계: 왼쪽 총 개수를 위한 버터플라이 리덕션
  초기값: [1, 0, 1, 0, 1, 0, 1, 0]
  리덕션 후: 모든 레인이 warp_left_total = 4를 가짐

4단계: 출력 위치에 기록
  Lane 0: current_val=3 < pivot → output[0] = 3
  Lane 1: current_val=7 >= pivot → output[4+0] = output[4] = 7
  Lane 2: current_val=1 < pivot → output[1] = 1
  Lane 3: current_val=8 >= pivot → output[4+1] = output[5] = 8
  Lane 4: current_val=2 < pivot → output[2] = 2
  Lane 5: current_val=9 >= pivot → output[4+2] = output[6] = 9
  Lane 6: current_val=4 < pivot → output[3] = 4
  Lane 7: current_val=6 >= pivot → output[4+3] = output[7] = 6

최종 결과: [3, 1, 2, 4, 7, 8, 9, 6] (< pivot | >= pivot)

요약

prefix_sum() 기본 요소는 복잡한 다단계 알고리즘을 단일 함수 호출로 대체하는 하드웨어 가속 병렬 스캔 연산을 가능하게 합니다. 두 가지 문제를 통해 다음을 배웠습니다:

핵심 누적 합 패턴

1. 포함 누적 합 (prefix_sum[exclusive=False]):

   • 하드웨어 가속 누적 연산
   • 공유 메모리 코드 ~30줄을 단일 함수 호출로 대체
   • 전용 하드웨어 최적화를 동반한 \(O(\log n)\) 복잡도

2. 고급 다중 기본 요소 조정 (prefix_sum + shuffle_xor 결합):

   • 단일 워프 내 정교한 병렬 알고리즘
   • 위치 계산을 위한 비포함 스캔 + 총합을 위한 버터플라이 리덕션
   • 최적의 병렬 효율성을 가진 복잡한 파티셔닝 연산

핵심 알고리즘 통찰

하드웨어 가속의 이점:

• prefix_sum()이 현대 GPU의 전용 스캔 유닛을 활용
• 기존 방식 대비 공유 메모리 오버헤드 제로
• 명시적 배리어 없는 자동 동기화

다중 기본 요소 조정:

# 1단계: 파티션 소속을 위한 프레디케이트 생성
predicate = 1.0 if condition else 0.0

# 2단계: 로컬 위치를 위한 prefix_sum 사용
local_pos = prefix_sum[exclusive=True](predicate)

# 3단계: 전역 총합을 위한 shuffle_xor 사용
global_total = butterfly_reduce(predicate)

# 4단계: 최종 위치 결정을 위한 결합
final_pos = local_pos + partition_offset

성능 이점:

• 하드웨어 최적화: 소프트웨어 구현 대비 전용 스캔 유닛
• 메모리 효율성: 공유 메모리 할당 대비 레지스터 전용 연산
• 확장 가능한 복잡도: 하드웨어 가속을 동반한 \(O(\log n)\)
• 단일 워프 최적화: WARP_SIZE 한도 내 알고리즘에 최적

실용적 활용

이 누적 합 패턴들의 기반이 되는 분야:

• 병렬 스캔 연산: 누적 합, 누적 곱, min/max 스캔
• 스트림 컴팩션: 병렬 필터링과 데이터 재배치
• Quicksort 파티셔닝: 병렬 정렬 알고리즘의 핵심 빌딩 블록
• 병렬 알고리즘: 부하 분산, 작업 분배, 데이터 재구조화

prefix_sum()과 shuffle_xor()의 결합은 현대 GPU 워프 기본 요소가 최소한의 코드 복잡도와 최적의 성능 특성으로 정교한 병렬 알고리즘을 어떻게 구현할 수 있는지를 보여줍니다.