elementwise - 기본 GPU 함수형 연산

fn elementwise_add[
    layout: Layout, dtype: DType, simd_width: Int, rank: Int, size: Int
](
    output: LayoutTensor[mut=True, dtype, layout, MutAnyOrigin],
    a: LayoutTensor[mut=False, dtype, layout, MutAnyOrigin],
    b: LayoutTensor[mut=False, dtype, layout, MutAnyOrigin],
    ctx: DeviceContext,
) raises:
    @parameter
    @always_inline
    fn add[
        simd_width: Int, rank: Int, alignment: Int = align_of[dtype]()
    ](indices: IndexList[rank]) capturing -> None:
        idx = indices[0]
        # Note: This is thread-local SIMD - each thread processes its own vector of data
        # we'll later better see this hierarchy in Mojo:
        # SIMD within threads, warp across threads, block across warps
        a_simd = a.aligned_load[width=simd_width](Index(idx))
        b_simd = b.aligned_load[width=simd_width](Index(idx))
        ret = a_simd + b_simd
        # print(
        #     "idx:", idx, ", a_simd:", a_simd, ", b_simd:", b_simd, " sum:", ret
        # )
        output.store[simd_width](Index(idx), ret)

    elementwise[add, SIMD_WIDTH, target="gpu"](a.size(), ctx)

Mojo의 elementwise 함수형 패턴은 현대 GPU 프로그래밍을 위한 몇 가지 기본 개념을 소개합니다:

1. 함수형 추상화 철학

elementwise 함수는 기존 GPU 프로그래밍에서의 패러다임 전환을 나타냅니다:

전통적인 CUDA/HIP 방식:

# 수동 스레드 관리
idx = thread_idx.x + block_idx.x * block_dim.x
if idx < size:
    output[idx] = a[idx] + b[idx];  // 스칼라 연산

Mojo 함수형 방식:

# 자동 관리 + SIMD 벡터화
elementwise[add_function, simd_width, target="gpu"](size, ctx)

elementwise가 추상화하는 것들:

• 스레드 그리드 구성: 블록/그리드 차원을 계산할 필요 없음
• 경계 검사: 배열 경계를 자동으로 처리
• 메모리 병합: 최적의 메모리 접근 패턴이 내장
• SIMD 오케스트레이션: 벡터화를 투명하게 처리
• GPU 타겟 선택: 다양한 GPU 아키텍처에서 동작

2. 심층 분석: 중첩 함수 아키텍처

@parameter
@always_inline
fn add[simd_width: Int, rank: Int](indices: IndexList[rank]) capturing -> None:

매개변수 분석:

• @parameter: 이 데코레이터는 컴파일 타임 특수화를 제공합니다. 각 고유한 simd_width와 rank에 대해 함수가 별도로 생성되어 적극적인 최적화가 가능합니다.
• @always_inline: GPU 성능에 매우 중요합니다 - 코드를 커널에 직접 내장하여 함수 호출 오버헤드를 제거합니다.
• capturing: 렉시컬 스코핑을 활성화합니다 - 내부 함수가 명시적 매개변수 전달 없이 외부 스코프의 변수에 접근할 수 있습니다.
• IndexList[rank]: 차원 무관 인덱싱을 제공합니다 - 동일한 패턴이 1D 벡터, 2D 행렬, 3D 텐서 등에서 작동합니다.

3. SIMD 실행 모델 심층 분석

idx = indices[0]                                    # 선형 인덱스: 0, 4, 8, 12... (GPU 의존적 간격)
a_simd = a.aligned_load[simd_width](Index(idx))     # 로드: [a[0:4], a[4:8], a[8:12]...] (로드당 4개 요소)
b_simd = b.aligned_load[simd_width](Index(idx))     # 로드: [b[0:4], b[4:8], b[8:12]...] (로드당 4개 요소)
ret = a_simd + b_simd                               # SIMD: 4개 덧셈을 병렬 수행 (GPU 의존적)
output.store[simd_width](Index(global_start), ret)  # 저장: 4개 결과를 동시 저장 (GPU 의존적)

실행 계층 구조 시각화:

GPU 아키텍처:
├── Grid (전체 문제)
│   ├── Block 1 (여러 Warp)
│   │   ├── Warp 1 (32개 스레드) --> Warp는 다음 Part VI에서 학습
│   │   │   ├── Thread 1 → SIMD[4개 요소]  ← 현재 초점 (GPU 의존적 폭)
│   │   │   ├── Thread 2 → SIMD[4개 요소]
│   │   │   └── ...
│   │   └── Warp 2 (32개 스레드)
│   └── Block 2 (여러 Warp)

SIMD_WIDTH=4인 1024개 요소 벡터의 경우 (GPU 예시):

• 필요한 총 SIMD 연산 수: 1024 ÷ 4 = 256
• GPU 실행: 256개 스레드 (1024 ÷ 4)
• 각 스레드가 처리하는 양: 정확히 4개의 연속 요소
• 메모리 대역폭: 스칼라 연산 대비 SIMD_WIDTH배 향상

참고: SIMD 폭은 GPU 아키텍처에 따라 다릅니다 (예: 일부 GPU는 4, RTX 4090은 8, A100은 16).

4. 메모리 접근 패턴 분석

a.aligned_load[simd_width](Index(idx))  // 병합 메모리 접근

메모리 병합의 이점:

• 순차적 접근: 스레드들이 연속적인 메모리 위치에 접근
• 캐시 최적화: L1/L2 캐시 히트율 극대화
• 대역폭 활용: 이론적 메모리 대역폭에 근접하는 성능 달성
• 하드웨어 효율: GPU 메모리 컨트롤러가 이 패턴에 최적화되어 있음

SIMD_WIDTH=4 (GPU 의존적) 예시:

Thread 0: a[0:4] 로드   → 메모리 뱅크 0-3
Thread 1: a[4:8] 로드   → 메모리 뱅크 4-7
Thread 2: a[8:12] 로드  → 메모리 뱅크 8-11
...
결과: 최적의 메모리 컨트롤러 활용

5. 성능 특성 및 최적화

산술 강도 분석 (SIMD_WIDTH=4 기준):

• 산술 연산: 4개 요소당 1회 SIMD 덧셈
• 메모리 연산: 4개 요소당 2회 SIMD 로드 + 1회 SIMD 저장
• 산술 강도: 1 덧셈 ÷ 3 메모리 연산 = 0.33 (메모리 바운드)

이것이 메모리 바운드인 이유:

단순 연산에서는 메모리 대역폭 >>> 연산 능력

최적화 시사점:

• 산술 최적화보다 메모리 접근 패턴에 집중해야 함
• SIMD 벡터화가 주요 성능 이점을 제공
• 메모리 병합이 성능에 매우 중요
• 연산 복잡도보다 캐시 지역성이 더 중요

6. 확장성과 적응성

자동 하드웨어 적응:

comptime SIMD_WIDTH = simd_width_of[dtype, target = _get_gpu_target()]()

• GPU별 최적화: SIMD 폭이 하드웨어에 맞게 조정됨 (예: 일부 카드는 4, RTX 4090은 8, A100은 16)
• 데이터 타입 인식: float32와 float16에 대해 서로 다른 SIMD 폭 적용
• 컴파일 타임 최적화: 하드웨어 감지에 대한 런타임 오버헤드 없음

확장성 특성:

• 스레드 수: 문제 크기에 따라 자동 확장
• 메모리 사용량: 입력 크기에 비례하여 선형 확장
• 성능: 메모리 대역폭 포화 시점까지 거의 선형적인 속도 향상

7. 고급 인사이트: 이 패턴이 중요한 이유

복잡한 연산의 기초: 이 elementwise 패턴은 다음 연산들의 기반이 됩니다:

• 리덕션 연산: 대규모 배열에서의 합계, 최댓값, 최솟값
• 브로드캐스트 연산: 스칼라-벡터 연산
• 복잡한 변환: 활성화 함수, 정규화
• 다차원 연산: 행렬 연산, 합성곱

전통적인 방식과의 비교:

// 전통적: 오류 발생 가능, 장황함, 하드웨어 종속적
__global__ void add_kernel(float* output, float* a, float* b, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        output[idx] = a[idx] + b[idx];  // 벡터화 없음
    }
}

// Mojo: 안전, 간결, 자동 벡터화
elementwise[add, SIMD_WIDTH, target="gpu"](size, ctx)

함수형 접근법의 이점:

• 안전성: 자동 경계 검사로 버퍼 오버플로우 방지
• 이식성: 동일한 코드가 다양한 GPU 벤더/세대에서 동작
• 성능: 컴파일러 최적화가 수동 튜닝 코드를 종종 능가
• 유지보수성: 깔끔한 추상화로 디버깅 복잡도 감소
• 조합성: 다른 함수형 연산과 쉽게 결합 가능

이 패턴은 GPU 프로그래밍의 미래를 나타냅니다 - 성능을 희생하지 않는 고수준 추상화로, 최적의 효율성을 유지하면서 GPU 컴퓨팅을 더 쉽게 접근할 수 있게 합니다.