멀티 블록 조정 기초

개요

첫 번째 클러스터 프로그래밍 도전에 오신 것을 환영합니다! 이 섹션에서는 SM90+ 클러스터 API를 사용한 블록 간 조정의 기본 구성 요소를 소개합니다.

도전 과제: 4개의 스레드 블록이 조정하여 서로 다른 데이터 범위를 처리하고 결과를 공유 출력 배열에 저장하는 멀티 블록 히스토그램 알고리즘을 구현합니다.

핵심 학습: cluster_arrive() → 처리 → cluster_wait()라는 필수적인 클러스터 동기화 패턴을 배웁니다. Puzzle 29의 barrier()에서 배운 동기화 개념을 확장합니다.

문제: 멀티 블록 히스토그램 구간 분류

Puzzle 27과 같은 기존의 단일 블록 알고리즘은 하나의 블록이 가진 스레드 용량(예: 256개 스레드) 내에 들어오는 데이터만 처리할 수 있습니다. Puzzle 8의 공유 메모리 용량을 초과하는 더 큰 데이터셋의 경우, 여러 블록이 협력해야 합니다.

과제: 4개 블록 각각이 서로 다른 데이터 범위를 처리하고, 고유한 블록 순위로 값을 스케일링하며, Puzzle 29의 동기화 패턴을 사용하여 다른 블록들과 조정함으로써 모든 블록의 처리가 완료된 후에야 최종 결과를 읽을 수 있도록 하는 히스토그램을 구현하세요.

문제 명세

멀티 블록 데이터 분배:

• Block 0: 요소 0-255를 처리, 1배 스케일링
• Block 1: 요소 256-511을 처리, 2배 스케일링
• Block 2: 요소 512-767을 처리, 3배 스케일링
• Block 3: 요소 768-1023을 처리, 4배 스케일링

조정 요구사항:

1. 각 블록은 cluster_arrive()를 사용하여 완료를 알려야 합니다
2. 모든 블록은 cluster_wait()를 사용하여 다른 블록을 기다려야 합니다
3. 최종 출력은 각 블록의 처리된 합계를 4개 요소 배열로 보여줍니다

설정

• 문제 크기: SIZE = 1024 요소 (1D 배열)
• 블록 설정: TPB = 256 블록당 스레드 수 (256, 1)
• 그리드 설정: CLUSTER_SIZE = 4 클러스터당 블록 수 (4, 1)
• 데이터 타입: DType.float32
• 메모리 레이아웃: 입력 Layout.row_major(SIZE), 출력 Layout.row_major(CLUSTER_SIZE)

스레드 블록 분배:

• Block 0: 스레드 0-255 → 요소 0-255
• Block 1: 스레드 0-255 → 요소 256-511
• Block 2: 스레드 0-255 → 요소 512-767
• Block 3: 스레드 0-255 → 요소 768-1023

완성할 코드

fn cluster_coordination_basics[
    in_layout: Layout, out_layout: Layout, tpb: Int
](
    output: LayoutTensor[dtype, out_layout, MutAnyOrigin],
    input: LayoutTensor[dtype, in_layout, ImmutAnyOrigin],
    size: Int,
):
    """Real cluster coordination using SM90+ cluster APIs."""
    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)
    local_i = thread_idx.x

    # Check what's happening with cluster ranks
    my_block_rank = Int(block_rank_in_cluster())
    block_id = Int(block_idx.x)

    shared_data = LayoutTensor[
        dtype,
        Layout.row_major(tpb),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    # FIX: Use block_idx.x for data distribution instead of cluster rank
    # Each block should process different portions of the data
    var data_scale = Float32(
        block_id + 1
    )  # Use block_idx instead of cluster rank

    # Phase 1: Each block processes its portion
    if global_i < size:
        shared_data[local_i] = input[global_i] * data_scale
    else:
        shared_data[local_i] = 0.0

    barrier()

    # Phase 2: Use cluster_arrive() for inter-block coordination
    # Signal this block has completed processing

    # FILL IN 1 line here

    # Block-level aggregation (only thread 0)
    if local_i == 0:
        # FILL IN 4 line here
        ...

    # Wait for all blocks in cluster to complete

    # FILL IN 1 line here

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

코드 실행

pixi run p34 --coordination

uv run poe p34 --coordination

예상 출력:

Testing Multi-Block Coordination
SIZE: 1024 TPB: 256 CLUSTER_SIZE: 4
Block coordination results:
  Block 0 : 127.5
  Block 1 : 255.0
  Block 2 : 382.5
  Block 3 : 510.0
✅ Multi-block coordination tests passed!

성공 기준:

• 4개 블록 모두 0이 아닌 결과를 생성합니다
• 결과가 스케일링 패턴을 보여줍니다: Block 1 > Block 0, Block 2 > Block 1 등
• 경쟁 상태나 조정 실패가 없어야 합니다

솔루션

fn cluster_coordination_basics[
    in_layout: Layout, out_layout: Layout, tpb: Int
](
    output: LayoutTensor[dtype, out_layout, MutAnyOrigin],
    input: LayoutTensor[dtype, in_layout, ImmutAnyOrigin],
    size: Int,
):
    """Real cluster coordination using SM90+ cluster APIs."""
    global_i = Int(block_dim.x * block_idx.x + thread_idx.x)
    local_i = thread_idx.x

    # Check what's happening with cluster ranks
    my_block_rank = Int(block_rank_in_cluster())
    block_id = Int(block_idx.x)

    shared_data = LayoutTensor[
        dtype,
        Layout.row_major(tpb),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    # FIX: Use block_idx.x for data distribution instead of cluster rank
    # Each block should process different portions of the data
    var data_scale = Float32(
        block_id + 1
    )  # Use block_idx instead of cluster rank

    # Phase 1: Each block processes its portion
    if global_i < size:
        shared_data[local_i] = input[global_i] * data_scale
    else:
        shared_data[local_i] = 0.0

    barrier()

    # Phase 2: Use cluster_arrive() for inter-block coordination
    cluster_arrive()  # Signal this block has completed processing

    # Block-level aggregation (only thread 0)
    if local_i == 0:
        var block_sum: Float32 = 0.0
        for i in range(tpb):
            block_sum += shared_data[i][0]
        # FIX: Store result at block_idx position (guaranteed unique per block)
        output[block_id] = block_sum

    # Wait for all blocks in cluster to complete
    cluster_wait()

global_i = block_dim.x * block_idx.x + thread_idx.x  # Global thread index
local_i = thread_idx.x                               # Local thread index within block
my_block_rank = Int(block_rank_in_cluster())         # Cluster rank (0-3)
block_id = Int(block_idx.x)                          # Block index for reliable addressing

if local_i == 0:
    var block_sum: Float32 = 0.0
    for i in range(tpb):
        block_sum += shared_data[i][0]  # Sum all elements in shared memory
    output[block_id] = block_sum        # Store result at unique block position

패턴 이해하기

클러스터 조정의 핵심 패턴은 단순하지만 강력한 구조를 따릅니다:

1. 1단계: 각 블록이 할당된 데이터 영역을 독립적으로 처리합니다
2. 신호: cluster_arrive()로 처리 완료를 알립니다
3. 2단계: 다른 블록의 결과에 의존하는 연산을 안전하게 수행할 수 있습니다
4. 동기화: cluster_wait()로 모든 블록이 완료된 후 다음으로 진행합니다

다음 단계: 더 고급 조정을 배울 준비가 되셨나요? 클러스터 전체 집합 연산 으로 이동하여 Puzzle 27의 block.sum() 패턴을 클러스터 규모로 확장하는 방법을 배워보세요. Puzzle 24의 워프 레벨 리덕션을 기반으로 합니다!