완성 버전

# Kernel 1: Compute local prefix sums and store block sums in out
fn prefix_sum_local_phase[
    out_layout: Layout, in_layout: Layout
](
    output: LayoutTensor[dtype, out_layout, MutAnyOrigin],
    a: LayoutTensor[dtype, in_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()

    # Load data into shared memory
    # Example with SIZE_2=15, TPB=8, BLOCKS=2:
    # Block 0 shared mem: [0,1,2,3,4,5,6,7]
    # Block 1 shared mem: [8,9,10,11,12,13,14,uninitialized]
    # Note: The last position remains uninitialized since global_i >= size,
    # but this is safe because that thread doesn't participate in computation
    if global_i < size:
        shared[local_i] = a[global_i]

    barrier()

    # Compute local prefix sum using parallel reduction
    # This uses a tree-based algorithm with log(TPB) iterations
    # Iteration 1 (offset=1):
    #   Block 0: [0,0+1,2+1,3+2,4+3,5+4,6+5,7+6] = [0,1,3,5,7,9,11,13]
    # Iteration 2 (offset=2):
    #   Block 0: [0,1,3+0,5+1,7+3,9+5,11+7,13+9] = [0,1,3,6,10,14,18,22]
    # Iteration 3 (offset=4):
    #   Block 0: [0,1,3,6,10+0,14+1,18+3,22+6] = [0,1,3,6,10,15,21,28]
    #   Block 1 follows same pattern to get [8,17,27,38,50,63,77,???]
    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 < TPB:
            current_val = shared[local_i - offset]  # read

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

        barrier()
        offset *= 2

    # Write local results to output
    # Block 0 writes: [0,1,3,6,10,15,21,28]
    # Block 1 writes: [8,17,27,38,50,63,77,???]
    if global_i < size:
        output[global_i] = shared[local_i]

    # Store block sums in auxiliary space
    # Block 0: Thread 7 stores shared[7] == 28 at position size+0 (position 15)
    # Block 1: Thread 7 stores shared[7] == ??? at position size+1 (position 16).  This sum is not needed for the final output.
    # This gives us: [0,1,3,6,10,15,21,28, 8,17,27,38,50,63,77, 28,???]
    #                                                           ↑  ↑
    #                                                     Block sums here
    if local_i == TPB - 1:
        output[size + block_idx.x] = shared[local_i]


# Kernel 2: Add block sums to their respective blocks
fn prefix_sum_block_sum_phase[
    layout: Layout
](output: LayoutTensor[dtype, layout, MutAnyOrigin], size: UInt):
    global_i = block_dim.x * block_idx.x + thread_idx.x

    # Second pass: add previous block's sum to each element
    # Block 0: No change needed - already correct
    # Block 1: Add Block 0's sum (28) to each element
    #   Before: [8,17,27,38,50,63,77]
    #   After: [36,45,55,66,78,91,105]
    # Final result combines both blocks:
    # [0,1,3,6,10,15,21,28, 36,45,55,66,78,91,105]
    if block_idx.x > 0 and global_i < size:
        prev_block_sum = output[size + block_idx.x - 1]
        output[global_i] += prev_block_sum

이 솔루션은 여러 스레드 블록에 걸치는 배열을 처리하기 위해 2개의 커널을 사용하는 멀티 블록 누적 합을 구현합니다. 각 부분을 자세히 살펴보겠습니다:

블록 간 통신의 과제

GPU 프로그래밍의 근본적인 제약은 barrier()를 사용한 스레드 동기화가 블록 내부에서만 가능하다는 점입니다. 데이터가 여러 블록에 걸쳐 있을 때 다음과 같은 과제에 직면합니다: 블록이 부분 결과를 다른 블록에 어떻게 전달할 수 있을까?

메모리 레이아웃 시각화

테스트 케이스 SIZE_2 = 15, TPB = 8의 경우:

Input array:  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]

Block 0 처리: [0, 1, 2, 3, 4, 5, 6, 7]
Block 1 처리: [8, 9, 10, 11, 12, 13, 14] (유효 원소 7개)

블록 합계를 위한 공간을 포함하도록 출력 버퍼를 확장합니다:

확장 버퍼: [데이터 값 (15개)] + [블록 합계 (2개)]
           [0...14] + [block0_sum, block1_sum]

이 확장 버퍼의 크기: EXTENDED_SIZE = SIZE_2 + num_blocks = 15 + 2 = 17

1단계 커널: 로컬 누적 합

로컬 단계에서의 경쟁 상태 방지

로컬 단계는 기본 버전과 동일한 명시적 동기화 패턴을 사용하여 읽기-쓰기 충돌을 방지합니다:

• 읽기 단계: 모든 스레드가 먼저 필요한 값을 로컬 변수 current_val에 읽어둠
• 동기화: barrier()로 모든 읽기가 완료된 후에야 쓰기가 시작되도록 보장
• 쓰기 단계: 모든 스레드가 계산된 값을 안전하게 공유 메모리에 기록

이를 통해 병렬 리덕션 중 여러 스레드가 동시에 같은 공유 메모리 위치에 접근할 때 발생할 수 있는 경쟁 상태를 방지합니다.

Block 0 단계별 실행

1. 공유 메모리에 값 로드:

   shared = [0, 1, 2, 3, 4, 5, 6, 7]
   

2. 병렬 리덕션 반복 (\(\log_2(TPB) = 3\)회 반복):

   반복 1 (offset=1):

   읽기 단계: 각 활성 스레드가 필요한 값을 읽음:

   T₁ reads shared[0] = 0    T₅ reads shared[4] = 4
   T₂ reads shared[1] = 1    T₆ reads shared[5] = 5
   T₃ reads shared[2] = 2    T₇ reads shared[6] = 6
   T₄ reads shared[3] = 3
   

   동기화: barrier()로 모든 읽기 완료를 보장

   쓰기 단계: 각 스레드가 읽은 값을 더함:

   shared[0] = 0              (변경 없음)
   shared[1] = 1 + 0 = 1
   shared[2] = 2 + 1 = 3
   shared[3] = 3 + 2 = 5
   shared[4] = 4 + 3 = 7
   shared[5] = 5 + 4 = 9
   shared[6] = 6 + 5 = 11
   shared[7] = 7 + 6 = 13
   

   배리어 후: shared = [0, 1, 3, 5, 7, 9, 11, 13]

   반복 2 (offset=2):

   읽기 단계: 각 활성 스레드가 필요한 값을 읽음:

   T₂ reads shared[0] = 0    T₅ reads shared[3] = 5
   T₃ reads shared[1] = 1    T₆ reads shared[4] = 7
   T₄ reads shared[2] = 3    T₇ reads shared[5] = 9
   

   동기화: barrier()로 모든 읽기 완료를 보장

   쓰기 단계: 각 스레드가 읽은 값을 더함:

   shared[0] = 0              (변경 없음)
   shared[1] = 1              (변경 없음)
   shared[2] = 3 + 0 = 3      (변경 없음)
   shared[3] = 5 + 1 = 6
   shared[4] = 7 + 3 = 10
   shared[5] = 9 + 5 = 14
   shared[6] = 11 + 7 = 18
   shared[7] = 13 + 9 = 22
   

   배리어 후: shared = [0, 1, 3, 6, 10, 14, 18, 22]

   반복 3 (offset=4):

   읽기 단계: 각 활성 스레드가 필요한 값을 읽음:

   T₄ reads shared[0] = 0    T₆ reads shared[2] = 3
   T₅ reads shared[1] = 1    T₇ reads shared[3] = 6
   

   동기화: barrier()로 모든 읽기 완료를 보장

   쓰기 단계: 각 스레드가 읽은 값을 더함:

   shared[0] = 0              (변경 없음)
   shared[1] = 1              (변경 없음)
   shared[2] = 3              (변경 없음)
   shared[3] = 6              (변경 없음)
   shared[4] = 10 + 0 = 10    (변경 없음)
   shared[5] = 14 + 1 = 15
   shared[6] = 18 + 3 = 21
   shared[7] = 22 + 6 = 28
   

   배리어 후: shared = [0, 1, 3, 6, 10, 15, 21, 28]

3. 로컬 결과를 전역 메모리에 기록:

   output[0...7] = [0, 1, 3, 6, 10, 15, 21, 28]
   

4. 블록 합계를 보조 공간에 저장 (마지막 스레드만):

   output[15] = 28  // 위치: size + block_idx.x = 15 + 0
   

Block 1 단계별 실행

1. 공유 메모리에 값 로드:

   shared = [8, 9, 10, 11, 12, 13, 14, 미초기화]
   

   참고: 스레드 7은 global_i = 15 >= SIZE_2이므로 아무것도 로드하지 않아 shared[7]이 미초기화 상태로 남습니다. 스레드 7은 최종 출력에 참여하지 않으므로 안전합니다.

2. 병렬 리덕션 반복 (\(\log_2(TPB) = 3\)회 반복):

   실제로 연산에 참여하는 것은 처음 7개 스레드뿐입니다. 세 번의 반복을 거치면:

   shared = [8, 17, 27, 38, 50, 63, 77, 미초기화]
   

3. 로컬 결과를 전역 메모리에 기록:

   output[8...14] = [8, 17, 27, 38, 50, 63, 77]  // 유효 출력 7개만
   

4. 블록 합계를 보조 공간에 저장 (블록의 마지막 스레드만):

   output[16] = shared[7]  // 스레드 7 (TPB-1)이 shared[7]의 값을 저장
   

   참고: 스레드 7은 유효한 입력을 로드하지 않았지만, 블록 내 누적 합 연산에는 그대로 참여합니다. shared[7]은 병렬 리덕션을 거치며 갱신되지만, 미초기화 상태에서 시작했기 때문에 최종 값을 예측할 수 없습니다. 다만 Block 1이 마지막 블록이므로 이 블록 합계는 2단계에서 사용되지 않아 정확성에는 영향이 없습니다.

1단계 이후 출력 버퍼의 내용:

[0, 1, 3, 6, 10, 15, 21, 28, 8, 17, 27, 38, 50, 63, 77, 28, ???]
                                                        ^   ^
                                                블록 합계가 여기에 저장됨

참고: 마지막 블록 합계 (???) 는 미초기화 메모리에 기반하므로 예측할 수 없지만, 최종 결과에는 영향을 주지 않습니다.

호스트-디바이스 동기화: 실제로 필요한 시점

두 커널 단계는 명시적 동기화 없이 순차적으로 실행됩니다:

# 1단계: 로컬 누적 합
ctx.enqueue_function[prefix_sum_local_phase[...], prefix_sum_local_phase[...]](...)

# 2단계: 블록 합계 더하기 (자동으로 1단계 완료를 대기)
ctx.enqueue_function[prefix_sum_block_sum_phase[...], prefix_sum_block_sum_phase[...]](...)

핵심 통찰: Mojo의 DeviceContext는 단일 실행 스트림(NVIDIA GPU에서는 CUDA 스트림, AMD ROCm GPU에서는 HIP 스트림)을 사용하므로, 큐에 넣은 커널이 정확히 넣은 순서대로 실행됨을 보장합니다. 커널 간에 명시적 동기화가 필요 없습니다.

ctx.synchronize()가 필요한 시점:

# 두 커널 완료 후, 호스트에서 결과를 읽기 전
ctx.synchronize()  # 호스트가 GPU 완료를 대기

with out.map_to_host() as out_host:  # 이제 GPU 결과를 안전하게 읽을 수 있음
    print("out:", out_host)

ctx.synchronize() 호출의 역할:

• 호스트-디바이스 동기화: 결과에 접근하기 전에 호스트가 모든 GPU 작업의 완료를 대기하도록 보장
• 메모리 안전성: 연산이 끝나기 전에 GPU 메모리를 읽는 것을 방지

실행 모델: 블록 내부의 스레드를 동기화하는 barrier()와 달리, 커널 실행 순서는 Mojo의 단일 스트림 실행 모델에서 보장되며, ctx.synchronize()는 호스트-디바이스 간 조율을 담당합니다.

2단계 커널: 블록 합계 더하기

1. Block 0: 변경 불필요 (이미 올바른 상태).

2. Block 1: 각 스레드가 Block 0의 합계를 자기 원소에 더함:

   prev_block_sum = output[size + block_idx.x - 1] = output[15] = 28
   output[global_i] += prev_block_sum
   

   Block 1의 값이 변환됩니다:

   Before: [8, 17, 27, 38, 50, 63, 77]
   After:  [36, 45, 55, 66, 78, 91, 105]
   

성능 및 최적화 고려 사항

주요 구현 상세

로컬 단계 동기화 패턴: 블록 내 각 반복은 엄격한 읽기 → 동기화 → 쓰기 패턴을 따릅니다:

1. var current_val: out.element_type = 0 - 로컬 변수 초기화
2. current_val = shared[local_i - offset] - 읽기 단계 (조건 충족 시)
3. barrier() - 경쟁 상태 방지를 위한 명시적 동기화
4. shared[local_i] += current_val - 쓰기 단계 (조건 충족 시)
5. barrier() - 다음 반복 전 동기화

블록 간 동기화: 이 알고리즘은 두 단계의 동기화를 사용합니다:

• 블록 내부: 로컬 누적 합 연산 중 barrier()로 각 블록 내 스레드를 동기화
• 블록 간: DeviceContext가 큐에 넣은 커널을 순차 실행하여 1단계가 2단계 전에 완료되도록 보장. 결과를 읽기 전에 호스트-디바이스 동기화가 필요하면 ctx.synchronize()를 사용합니다.

경쟁 상태 방지: 로컬 단계에서 읽기와 쓰기를 명시적으로 분리하여, 병렬 리덕션 중 여러 스레드가 같은 공유 메모리 위치에 동시에 접근할 때 생길 수 있는 경쟁 상태를 방지합니다.

1. 작업 효율성: 이 구현의 작업 복잡도는 \(O(n \log n)\)이며, 순차 알고리즘은 \(O(n)\)입니다. 병렬 알고리즘에서 전형적인 공간-시간 트레이드오프입니다.

2. 메모리 오버헤드: 블록 합계를 위한 추가 공간은 아주 적습니다 (블록당 원소 하나).

이 2개 커널 접근 방식은 블록 간 통신이 필요한 GPU 알고리즘의 기본 패턴입니다. 기수 정렬, 히스토그램 계산, 리덕션 연산 등 다른 병렬 알고리즘에도 동일한 전략을 적용할 수 있습니다.