☸️ 클러스터 전체 집합 연산

개요

이전 섹션의 기본 클러스터 조정을 바탕으로, 이 도전에서는 클러스터 전체 집합 연산을 구현하는 방법을 배웁니다 - Puzzle 27에서 익힌 block.sum 패턴을 여러 스레드 블록에 걸쳐 확장합니다.

도전 과제: 4개의 조정된 블록에 걸쳐 1024개 요소를 처리하고, 각 블록의 개별 리덕션을 하나의 전역 결과로 합치는 클러스터 전체 리덕션을 구현합니다.

핵심 학습: 전체 클러스터 조정을 위한 cluster_sync()와 효율적인 최종 리덕션을 위한 elect_one_sync()를 배웁니다.

문제: 대규모 전역 합산

단일 블록은 (Puzzle 27에서 배웠듯이) 스레드 수와 Puzzle 8의 공유 메모리 용량에 의해 제한됩니다. 단일 블록 리덕션을 넘어서는 대규모 데이터셋의 전역 통계(평균, 분산, 합계)를 구하려면 클러스터 전체 집합 연산이 필요합니다.

과제: 다음과 같은 클러스터 전체 합산 리덕션을 구현하세요:

1. 각 블록이 로컬 리덕션을 수행합니다 (Puzzle 27의 block.sum()과 유사)
2. Puzzle 29의 동기화를 사용하여 블록들이 부분 결과를 합칩니다
3. 선출된 하나의 스레드가 워프 선출 패턴을 사용하여 최종 전역 합계를 계산합니다

문제 명세

알고리즘 흐름:

1단계 - 로컬 리덕션 (각 블록 내부): \[R_i = \sum_{j=0}^{TPB-1} input[i \times TPB + j] \quad \text{for block } i\]

2단계 - 전역 집계 (클러스터 전체): \[\text{Global Sum} = \sum_{i=0}^{\text{CLUSTER_SIZE}-1} R_i\]

조정 요구사항:

1. 로컬 리덕션: 각 블록이 트리 리덕션으로 부분 합을 계산합니다
2. 클러스터 동기화: cluster_sync()로 모든 부분 결과가 준비되었는지 보장합니다
3. 최종 집계: 선출된 하나의 스레드가 모든 부분 결과를 합칩니다

설정

• 문제 크기: SIZE = 1024 요소
• 블록 설정: TPB = 256 블록당 스레드 수 (256, 1)
• 그리드 설정: CLUSTER_SIZE = 4 클러스터당 블록 수 (4, 1)
• 데이터 타입: DType.float32
• 메모리 레이아웃: 입력 Layout.row_major(SIZE), 출력 Layout.row_major(1)
• 임시 저장소: 부분 결과를 위한 Layout.row_major(CLUSTER_SIZE)

예상 결과: 수열 0, 0.01, 0.02, ..., 10.23의 합 = 523,776

완성할 코드

fn cluster_collective_operations[
    in_layout: Layout, out_layout: Layout, tpb: Int
](
    output: LayoutTensor[dtype, out_layout, MutAnyOrigin],
    input: LayoutTensor[dtype, in_layout, ImmutAnyOrigin],
    temp_storage: LayoutTensor[
        dtype, Layout.row_major(CLUSTER_SIZE), MutAnyOrigin
    ],
    size: Int,
):
    """Cluster-wide collective operations using real cluster APIs."""
    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)
    local_i = Int(thread_idx.x)

    # FILL IN (roughly 24 lines)

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

클러스터 API 참조

gpu.primitives.cluster 모듈:

• cluster_sync(): 전체 클러스터 동기화 - arrive/wait 패턴보다 강력
• elect_one_sync(): 효율적인 조정을 위해 워프 내에서 단일 스레드를 선출
• block_rank_in_cluster(): 클러스터 내 고유한 블록 식별자를 반환

트리 리덕션 패턴

Puzzle 27의 전통적인 내적에서 배운 트리 리덕션 패턴을 떠올려 보세요:

Stride 128: [T0] += [T128], [T1] += [T129], [T2] += [T130], ...
Stride 64:  [T0] += [T64],  [T1] += [T65],  [T2] += [T66],  ...
Stride 32:  [T0] += [T32],  [T1] += [T33],  [T2] += [T34],  ...
Stride 16:  [T0] += [T16],  [T1] += [T17],  [T2] += [T18],  ...
...
Stride 1:   [T0] += [T1] → Final result at T0

이제 이 패턴을 클러스터 규모로 확장합니다 - 각 블록이 하나의 부분 결과를 생성한 뒤, 블록 간에 결합합니다.

코드 실행

pixi run p34 --reduction

uv run poe p34 --reduction

예상 출력:

Testing Cluster-Wide Reduction
SIZE: 1024 TPB: 256 CLUSTER_SIZE: 4
Expected sum: 523776.0
Cluster reduction result: 523776.0
Expected: 523776.0
Error: 0.0
✅ Passed: Cluster reduction accuracy test
✅ Cluster-wide collective operations tests passed!

성공 기준:

• 완벽한 정확도: 결과가 예상 합계(523,776)와 정확히 일치합니다
• 클러스터 조정: 4개 블록 모두가 부분 합에 기여합니다
• 효율적인 최종 리덕션: 선출된 단일 스레드가 최종 결과를 계산합니다

솔루션

fn cluster_collective_operations[
    in_layout: Layout, out_layout: Layout, tpb: Int
](
    output: LayoutTensor[dtype, out_layout, MutAnyOrigin],
    input: LayoutTensor[dtype, in_layout, ImmutAnyOrigin],
    temp_storage: LayoutTensor[
        dtype, Layout.row_major(CLUSTER_SIZE), MutAnyOrigin
    ],
    size: Int,
):
    """Cluster-wide collective operations using real cluster APIs."""
    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)
    local_i = Int(thread_idx.x)
    my_block_rank = Int(block_rank_in_cluster())
    block_id = Int(block_idx.x)

    # Each thread accumulates its data
    var my_value: Float32 = 0.0
    if global_i < size:
        my_value = input[global_i][0]

    # Block-level reduction using shared memory
    shared_mem = LayoutTensor[
        dtype,
        Layout.row_major(tpb),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    shared_mem[local_i] = my_value
    barrier()

    # Tree reduction within block
    var stride = tpb // 2
    while stride > 0:
        if local_i < stride and local_i + stride < tpb:
            shared_mem[local_i] += shared_mem[local_i + stride]
        barrier()
        stride = stride // 2

    # FIX: Store block result using block_idx for reliable indexing
    if local_i == 0:
        temp_storage[block_id] = shared_mem[0]

    # Use cluster_sync() for full cluster synchronization
    cluster_sync()

    # Final cluster reduction (elect one thread to do the final work)
    if elect_one_sync() and my_block_rank == 0:
        var total: Float32 = 0.0
        for i in range(CLUSTER_SIZE):
            total += temp_storage[i][0]
        output[0] = total

var my_value: Float32 = 0.0
if global_i < size:
    my_value = input[global_i][0]  # Load with bounds checking
shared_mem[local_i] = my_value     # Store in shared memory
barrier()                          # Ensure all threads complete loading

var stride = tpb // 2  # Start with half the threads (128)
while stride > 0:
    if local_i < stride and local_i + stride < tpb:
        shared_mem[local_i] += shared_mem[local_i + stride]
    barrier()          # Synchronize after each reduction step
    stride = stride // 2

Step 1: stride=128  [T0]+=T128, [T1]+=T129, ..., [T127]+=T255
Step 2: stride=64   [T0]+=T64,  [T1]+=T65,  ..., [T63]+=T127
Step 3: stride=32   [T0]+=T32,  [T1]+=T33,  ..., [T31]+=T63
Step 4: stride=16   [T0]+=T16,  [T1]+=T17,  ..., [T15]+=T31
Step 5: stride=8    [T0]+=T8,   [T1]+=T9,   ..., [T7]+=T15
Step 6: stride=4    [T0]+=T4,   [T1]+=T5,   [T2]+=T6,  [T3]+=T7
Step 7: stride=2    [T0]+=T2,   [T1]+=T3
Step 8: stride=1    [T0]+=T1    → Final result at shared_mem[0]

if elect_one_sync() and my_block_rank == 0:
    var total: Float32 = 0.0
    for i in range(CLUSTER_SIZE):
        total += temp_storage[i][0]  # Sum: temp[0] + temp[1] + temp[2] + temp[3]
    output[0] = total

집합 연산 패턴 이해하기

이 퍼즐은 분산 컴퓨팅에서 사용되는 고전적인 2단계 리덕션 패턴을 보여줍니다:

1. 로컬 집계: 각 처리 단위(블록)가 자신의 데이터 영역을 리덕션합니다
2. 전역 조정: 처리 단위들이 동기화하고 결과를 교환합니다
3. 최종 리덕션: 선출된 하나의 단위가 모든 부분 결과를 합칩니다

단일 블록 방식과의 비교:

• 기존 block.sum(): 최대 256개 스레드 내에서만 동작합니다
• 클러스터 집합 연산: 여러 블록에 걸쳐 1000개 이상의 스레드로 확장됩니다
• 동일한 정확도: 둘 다 동일한 수학적 결과를 생성합니다
• 다른 규모: 클러스터 방식이 더 큰 데이터셋을 처리합니다

성능 이점:

• 더 큰 데이터셋: 단일 블록 용량을 초과하는 배열을 처리합니다
• 더 나은 활용률: 더 많은 GPU 연산 유닛을 동시에 사용합니다
• 확장 가능한 패턴: 복잡한 다단계 알고리즘의 기반이 됩니다

다음 단계: 최종 도전을 할 준비가 되셨나요? 고급 클러스터 알고리즘 으로 이동하여 워프 프로그래밍+블록 조정+클러스터 동기화를 결합한 계층적 패턴을 배워보세요. 성능 최적화 기법을 기반으로 합니다!