Puzzle 28: 비동기 메모리 연산과 복사 중첩

GPU 메모리 병목 현상: 실제 GPU 알고리즘 대부분은 좌절스러운 벽에 부딪힙니다 - 연산 능력이 아니라 메모리 대역폭에 의해 제한된다는 것입니다. 비싼 GPU 코어가 느린 DRAM에서 데이터가 도착하기를 기다리며 놀고 있는 것이죠.

GPU 프로그래밍에서 흔히 볼 수 있는 상황을 살펴보겠습니다:

# 성능의 적 - 순차적 메모리 연산
load_input_tile()     # ← DRAM 대기 500 사이클
load_kernel_data()    # ← 또 100 사이클 대기
barrier()             # ← 모든 스레드가 유휴 대기
compute()             # ← 드디어 실제 연산 50 사이클
# 총: 650 사이클, 연산 활용률 겨우 7.7%!

이렇게 할 수 있다면 어떨까요?

# 성능 개선 - 중첩 연산
launch_async_load()   # ← 백그라운드에서 500 사이클 전송 시작
load_small_data()     # ← 대기 중 유용한 작업 100 사이클
wait_and_compute()    # ← 나머지 ~400 사이클만 대기 후 연산
# 총: ~550 사이클, 45% 향상!

이것이 비동기 메모리 연산의 위력입니다 - 느린 알고리즘과 GPU의 잠재력을 최대한 발휘하는 알고리즘의 차이를 만들어 냅니다.

왜 중요한가

이 퍼즐에서는 Puzzle 13의 메모리 바운드 1D 합성곱을 연산 뒤에 메모리 지연 시간을 숨기는 고성능 구현으로 변환합니다. 단순한 학술적 연습이 아닙니다 - 이 패턴들은 다음 분야의 핵심입니다:

• 딥러닝: 가중치와 활성화값의 효율적 로딩
• 과학 연산: 스텐실 연산에서 데이터 전송 중첩
• 이미지 처리: 메모리 계층 구조를 통한 대규모 데이터셋 스트리밍
• 모든 메모리 바운드 알고리즘: 대기 시간을 생산적인 작업으로 전환

사전 준비

시작하기 전에 다음 내용을 확실히 이해하고 있어야 합니다:

필수 GPU 프로그래밍 개념:

• 공유 메모리 프로그래밍 (Puzzle 8, Puzzle 16) - matmul 패턴을 확장합니다
• 메모리 병합(coalescing) (Puzzle 21) - 최적의 비동기 전송에 필수
• 타일 기반 처리 (Puzzle 23) - 이 최적화의 기반

하드웨어 이해:

• GPU 메모리 계층 구조 (DRAM → 공유 메모리 → 레지스터)
• 스레드 블록 구성과 동기화
• 메모리 지연 시간 vs. 대역폭에 대한 기본 이해

API 숙지: Mojo GPU Memory Operations

⚠️ 하드웨어 호환성 참고: 이 퍼즐은 최신 GPU 아키텍처가 필요할 수 있는 비동기 복사 연산(copy_dram_to_sram_async, async_copy_wait_all)을 사용합니다. .async 수정자나 지원되지 않는 연산 관련 컴파일 오류가 발생하면 해당 GPU가 이 기능을 지원하지 않는 것일 수 있습니다. 그래도 메모리 최적화 패턴을 이해하는 데 개념은 여전히 유용합니다.

GPU 컴퓨팅 능력 확인:

nvidia-smi --query-gpu=name,compute_cap --format=csv,noheader,nounits

• SM_70 이상 (예: V100, T4, A10G, RTX 20+ 시리즈): 기본 비동기 복사 지원
• SM_80 이상 (예: A100, RTX 30+ 시리즈): 전체 비동기 복사 기능
• SM_90 이상 (예: H100, RTX 40+ 시리즈): 고급 TMA 연산 지원

학습 내용

이 퍼즐을 마치면 다음을 직접 경험하게 됩니다:

핵심 기법

• 비동기 복사 기본 요소: 백그라운드 DRAM→SRAM 전송 시작
• 지연 시간 은폐(latency hiding): 비용이 큰 메모리 연산을 유용한 연산과 중첩
• 스레드 레이아웃 최적화: 메모리 접근 패턴을 하드웨어에 맞추기
• 파이프라인 프로그래밍: 메모리 활용을 극대화하도록 알고리즘 구조화

주요 API

Puzzle 16의 관용적 matmul에서 소개한 비동기 복사 연산을 기반으로, 이제 메모리 최적화 잠재력에 집중합니다:

• copy_dram_to_sram_async(): 전용 복사 엔진을 사용하여 백그라운드 DRAM→SRAM 전송 시작
• async_copy_wait_all(): 공유 메모리 접근 전 전송 완료 동기화

Puzzle 16과 다른 점은? Puzzle 16에서는 matmul의 깔끔한 타일 로딩을 위해 비동기 복사를 사용했다면, 이 퍼즐은 지연 시간 은폐에 집중합니다 - 비용이 큰 메모리 연산과 유용한 연산 작업을 중첩하도록 알고리즘을 구조화하는 것입니다.

성능 효과

이 기법들은 다음과 같은 방식으로 메모리 바운드 알고리즘의 성능을 크게 향상시킵니다:

• DRAM 지연 시간 숨기기: 유휴 대기를 생산적인 연산 시간으로 전환
• 대역폭 극대화: 최적의 메모리 접근 패턴으로 캐시 미스 방지
• 파이프라인 효율: 메모리 전송이 병렬로 일어나는 동안 연산 유닛을 바쁘게 유지

비동기 복사 연산이란? 비동기 복사 연산은 GPU 블록이 다른 작업을 계속하는 동안 백그라운드에서 실행되는 메모리 전송을 시작할 수 있게 해줍니다. 이를 통해 연산과 메모리 이동을 중첩할 수 있으며, 이는 메모리 바운드 알고리즘의 근본적인 최적화 기법입니다.

💡 성공 팁: 이것을 GPU 메모리를 위한 파이프라인 프로그래밍으로 생각하세요 - 단계를 중첩하고, 지연 시간을 숨기고, 처리량을 극대화합니다. 목표는 데이터가 백그라운드에서 이동하는 동안 비싼 연산 유닛을 바쁘게 유지하는 것입니다.

헤일로 영역 이해하기

비동기 복사 연산으로 들어가기 전에, 합성곱과 같은 스텐실 연산의 타일 기반 처리에 필수적인 헤일로 영역(ghost cell 또는 guard cell이라고도 함)을 이해하는 것이 중요합니다.

헤일로 영역이란?

헤일로 영역은 스텐실 연산에 필요한 이웃 데이터를 제공하기 위해 처리 타일의 경계를 넘어 확장되는 추가 요소입니다. 타일 가장자리 근처의 요소를 처리할 때, 스텐실 연산은 인접 타일의 데이터에 접근해야 합니다.

헤일로 영역이 필요한 이유

타일에서 5점 커널을 사용하는 1D 합성곱을 생각해 봅시다:

원본 데이터:      [... | a b c d e f g h i j k l m n o | ...]
처리 타일:              [c d e f g h i j k l m n o]
                            ^                 ^
                      왼쪽 타일에서        오른쪽 타일에서
                      이웃 필요           이웃 필요

헤일로 포함:       [a b | c d e f g h i j k l m n o | p q]
                 ^^^                               ^^^
                 왼쪽 헤일로                     오른쪽 헤일로

주요 특성:

• 헤일로 크기: 일반적으로 각 측면에 KERNEL_SIZE // 2개 요소
• 목적: 타일 경계에서 정확한 스텐실 연산 가능
• 내용: 이웃 타일의 데이터 복사본 또는 경계 조건
• 메모리 오버헤드: 큰 연산 이점을 위한 적은 추가 저장 공간

합성곱에서의 헤일로 영역

5점 합성곱 커널 \([k_0, k_1, k_2, k_3, k_4]\)의 경우:

• 중심 요소: \(k_2\)가 현재 처리 요소와 정렬
• 왼쪽 이웃: \(k_0, k_1\)은 왼쪽 2개 요소 필요
• 오른쪽 이웃: \(k_3, k_4\)은 오른쪽 2개 요소 필요
• 헤일로 크기: 각 측면에 HALO_SIZE = 5 // 2 = 2개 요소

헤일로 영역 없이:

• 타일 경계 요소에서 전체 합성곱을 수행할 수 없음
• 잘못된 출력이나 복잡한 경계 처리 로직이 필요
• 분산된 메모리 접근 패턴으로 성능 저하

헤일로 영역 사용 시:

• 모든 타일 요소가 로컬 데이터를 사용하여 전체 합성곱 수행 가능
• 예측 가능한 메모리 접근으로 간결하고 효율적인 연산
• 더 나은 캐시 활용과 메모리 병합

이 개념은 비동기 복사 연산을 구현할 때 특히 중요합니다. 헤일로 영역을 올바르게 로딩하고 동기화해야 여러 타일에 걸친 정확한 병렬 연산을 보장할 수 있습니다.

비동기 복사 중첩을 활용한 1D 합성곱

Puzzle 13 기반: 이 퍼즐은 Puzzle 13의 1D 합성곱을 다시 다루지만, 이번에는 비동기 복사 연산으로 메모리 지연 시간을 연산 뒤에 숨기는 최적화를 적용합니다. 단순한 동기식 메모리 접근 대신, 하드웨어 가속을 사용하여 비용이 큰 DRAM 전송과 유용한 작업을 중첩합니다.

구성

• 벡터 크기: VECTOR_SIZE = 16384 (여러 블록에 걸친 16K 요소)
• 타일 크기: CONV_TILE_SIZE = 256 (처리 타일 크기)
• 블록 구성: 블록당 (256, 1) 스레드
• 그리드 구성: 그리드당 (VECTOR_SIZE // CONV_TILE_SIZE, 1) 블록 (64개 블록)
• 커널 크기: KERNEL_SIZE = 5 (Puzzle 13과 동일한 간단한 1D 합성곱)
• 데이터 타입: DType.float32
• 레이아웃: Layout.row_major(VECTOR_SIZE) (1D row-major)

비동기 복사의 기회

Puzzle 16 기반: matmul에서 깔끔한 타일 로딩을 위해 copy_dram_to_sram_async를 사용하는 것을 이미 보셨습니다. 이제 고성능 메모리 바운드 알고리즘의 핵심인 지연 시간 은폐 기능에 집중합니다.

기존의 동기식 메모리 로딩은 전송 중 연산 유닛을 유휴 상태로 대기하게 합니다. 비동기 복사 연산은 전송과 유용한 작업의 중첩을 가능하게 합니다:

# 동기식 접근 - 비효율적:
for i in range(CONV_TILE_SIZE):
    input_shared[i] = input[base_idx + i]  # 각 로드가 DRAM을 기다림
for i in range(KERNEL_SIZE):
    kernel_shared[i] = kernel[i]           # DRAM 추가 대기
barrier()  # 연산 시작 전 모든 스레드 대기
# ↑ 총 시간 = input_transfer_time + kernel_transfer_time

# 비동기 복사 접근 - 효율적:
copy_dram_to_sram_async[thread_layout](input_shared, input_tile)  # 백그라운드 전송 시작
# 입력이 백그라운드에서 전송되는 동안, 커널을 동기식으로 로딩
for i in range(KERNEL_SIZE):
    kernel_shared[i] = kernel[i]  # 비동기 입력 전송과 중첩
async_copy_wait_all()  # 두 연산이 모두 완료될 때만 대기
# ↑ 총 시간 = MAX(input_transfer_time, kernel_transfer_time)

비동기 복사가 잘 동작하는 이유:

• 전용 복사 엔진: 최신 GPU는 레지스터를 우회하고 진정한 연산-메모리 중첩을 가능하게 하는 전용 하드웨어를 갖추고 있습니다 (Puzzle 16에서 설명)
• 지연 시간 은폐: GPU 스레드가 다른 연산을 실행하는 동안 메모리 전송이 이루어집니다
• 최적의 병합: 스레드 레이아웃이 효율적인 DRAM 접근 패턴을 보장합니다
• 리소스 활용: 연산 유닛이 유휴 대기 대신 계속 바쁘게 동작합니다

완성할 코드

Puzzle 16의 matmul 구현 패턴을 따라, 비동기 복사 연산으로 메모리 전송과 연산을 중첩하는 1D 합성곱을 구현하세요.

수학적 연산: 비동기 복사를 활용하여 대규모 벡터에 대한 1D 합성곱을 효율적으로 계산합니다: \[\text{output}[i] = \sum_{k=0}^{\text{KERNEL_SIZE}-1} \text{input}[i+k-\text{HALO_SIZE}] \times \text{kernel}[k]\]

비동기 복사 알고리즘:

1. 비동기 타일 로딩: 입력 데이터의 백그라운드 DRAM→SRAM 전송 시작
2. 중첩 연산: 입력 전송 중 작은 커널 데이터 로딩
3. 동기화: 전송 완료 대기 후 공유 메모리를 사용하여 연산

comptime VECTOR_SIZE = 16384
comptime CONV_TILE_SIZE = 256
comptime KERNEL_SIZE = 5
comptime HALO_SIZE = KERNEL_SIZE // 2  # Halo elements needed for boundary
comptime BUFFER_SIZE = CONV_TILE_SIZE + 2 * HALO_SIZE  # Include halo for boundary conditions
comptime BLOCKS_PER_GRID_ASYNC = (
    VECTOR_SIZE + CONV_TILE_SIZE - 1
) // CONV_TILE_SIZE
comptime THREADS_PER_BLOCK_ASYNC = 256
comptime dtype = DType.float32
comptime layout_async = Layout.row_major(VECTOR_SIZE)


fn async_copy_overlap_convolution[
    dtype: DType, layout: Layout
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    kernel: LayoutTensor[dtype, Layout.row_major(KERNEL_SIZE), ImmutAnyOrigin],
):
    """Demonstrates async copy operations building on p14 patterns.

    This shows how to use copy_dram_to_sram_async and async_copy_wait_all
    for efficient memory transfers, extending the patterns from p14 matmul.
    """

    # Shared memory buffers (like p14, but without .fill(0) to avoid race)
    input_shared = LayoutTensor[
        dtype,
        Layout.row_major(CONV_TILE_SIZE),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    kernel_shared = LayoutTensor[
        dtype,
        Layout.row_major(KERNEL_SIZE),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    # FILL IN HERE (roughly 19 lines)

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

비동기 복사 중첩 테스트:

pixi run p28

pixi run -e amd p28

pixi run -e apple p28

uv run poe p28

솔루션

fn async_copy_overlap_convolution[
    dtype: DType, layout: Layout
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
    kernel: LayoutTensor[dtype, Layout.row_major(KERNEL_SIZE), ImmutAnyOrigin],
):
    """Demonstrates async copy operations building on p14 patterns.

    This shows how to use copy_dram_to_sram_async and async_copy_wait_all
    for efficient memory transfers, extending the patterns from p14 matmul.
    """

    # Shared memory buffers (like p14, but without .fill(0) to avoid race)
    input_shared = LayoutTensor[
        dtype,
        Layout.row_major(CONV_TILE_SIZE),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    kernel_shared = LayoutTensor[
        dtype,
        Layout.row_major(KERNEL_SIZE),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    local_i = Int(thread_idx.x)

    # Phase 1: Launch async copy for input tile
    # Note: tile() does NOT perform bounds checking - ensure valid tile bounds
    input_tile = input.tile[CONV_TILE_SIZE](Int(block_idx.x))

    # Use async copy with thread layout matching p14 pattern
    comptime load_layout = Layout.row_major(THREADS_PER_BLOCK_ASYNC)
    copy_dram_to_sram_async[thread_layout=load_layout](input_shared, input_tile)

    # Phase 2: Load kernel synchronously (small data)
    if local_i < KERNEL_SIZE:
        kernel_shared[local_i] = kernel[local_i]

    # Phase 3: Wait for async copy to complete
    async_copy_wait_all()  # Always wait since we always do async copy
    barrier()  # Sync all threads

    # Phase 4: Compute convolution
    global_i = Int(block_idx.x) * CONV_TILE_SIZE + local_i
    if local_i < CONV_TILE_SIZE and global_i < output.shape[0]():
        var result: output.element_type = 0

        # Simple convolution avoiding boundary issues
        if local_i >= HALO_SIZE and local_i < CONV_TILE_SIZE - HALO_SIZE:
            # Full convolution for center elements
            for k in range(KERNEL_SIZE):
                input_idx = local_i + k - HALO_SIZE
                if input_idx >= 0 and input_idx < CONV_TILE_SIZE:
                    result += input_shared[input_idx] * kernel_shared[k]
        else:
            # For boundary elements, just copy input (no convolution)
            result = input_shared[local_i]

        output[global_i] = result

# Phase 1: Launch async copy for input tile
input_tile = input.tile[CONV_TILE_SIZE](block_idx.x)
comptime load_layout = Layout.row_major(THREADS_PER_BLOCK_ASYNC)
copy_dram_to_sram_async[thread_layout=load_layout](input_shared, input_tile)

# Phase 2: Load kernel synchronously (small data)
if local_i < KERNEL_SIZE:
    kernel_shared[local_i] = kernel[local_i]

# Phase 3: Wait for async copy to complete
async_copy_wait_all()  # Always wait since we always do async copy
barrier()  # Sync all threads

# Phase 4: Compute convolution
global_i = block_idx.x * CONV_TILE_SIZE + local_i
if local_i < CONV_TILE_SIZE and global_i < output.shape[0]():
    var result: output.element_type = 0

    if local_i >= HALO_SIZE and local_i < CONV_TILE_SIZE - HALO_SIZE:
        # Full convolution for center elements
        for k in range(KERNEL_SIZE):
            input_idx = local_i + k - HALO_SIZE
            if input_idx >= 0 and input_idx < CONV_TILE_SIZE:
                result += input_shared[input_idx] * kernel_shared[k]
    else:
        # For boundary elements, just copy input (no convolution)
        result = input_shared[local_i]

    output[global_i] = result

Total Time = Input_Transfer_Time + Kernel_Transfer_Time + Compute_Time
           = Large_DRAM_transfer + Small_DRAM_transfer + convolution
           = Major_latency + Minor_latency + computation_work

Total Time = MAX(Input_Transfer_Time, Kernel_Transfer_Time) + Compute_Time
           = MAX(Major_latency, Minor_latency) + computation_work
           = Major_latency + computation_work