vectorize - SIMD 제어

개요

이 퍼즐에서는 수동 벡터화와 vectorize를 사용하여 GPU 커널 내에서 SIMD 연산을 정밀하게 제어하는 고급 벡터화 기법을 탐구합니다. 벡터화된 연산에 대해 두 가지 다른 접근법을 구현합니다:

1. 수동 벡터화: 명시적 인덱스 계산을 통한 직접적인 SIMD 제어
2. Mojo의 vectorize 함수: 자동 경계 검사를 포함한 고수준 벡터화

두 접근법 모두 타일링 개념을 기반으로 하지만, 제어, 안전성, 성능 최적화 간의 트레이드오프가 다릅니다.

핵심 통찰: 벡터화 전략은 성능 요구 사항과 복잡도 수준에 따라 달리 선택해야 합니다.

핵심 개념

이 퍼즐에서 배울 내용:

• 명시적 인덱스 관리를 통한 수동 SIMD 연산
• 안전하고 자동적인 벡터화를 위한 Mojo의 vectorize 함수
• 최적의 SIMD 정렬을 위한 청크 기반 메모리 구성
• 경계 조건을 위한 경계 검사 전략
• 수동 제어와 안전성 간의 성능 트레이드오프

이전과 동일한 수학적 연산: \[\Large \text{output}[i] = a[i] + b[i]\]

하지만 최대 성능을 위한 정교한 벡터화 전략을 사용합니다.

설정

• 벡터 크기: SIZE = 1024
• 타일 크기: TILE_SIZE = 32
• 데이터 타입: DType.float32
• SIMD 폭: GPU 의존적
• 레이아웃: Layout.row_major(SIZE) (1D 행 우선)

1. 수동 벡터화 방식

완성할 코드

fn manual_vectorized_tiled_elementwise_add[
    layout: Layout,
    dtype: DType,
    simd_width: Int,
    num_threads_per_tile: Int,
    rank: Int,
    size: Int,
    tile_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:
    # Each tile contains tile_size groups of simd_width elements
    comptime chunk_size = tile_size * simd_width

    @parameter
    @always_inline
    fn process_manual_vectorized_tiles[
        num_threads_per_tile: Int, rank: Int, alignment: Int = align_of[dtype]()
    ](indices: IndexList[rank]) capturing -> None:
        tile_id = indices[0]
        print("tile_id:", tile_id)
        output_tile = output.tile[chunk_size](tile_id)
        a_tile = a.tile[chunk_size](tile_id)
        b_tile = b.tile[chunk_size](tile_id)

        # FILL IN (7 lines at most)

    # Number of tiles needed: each tile processes chunk_size elements
    num_tiles = (size + chunk_size - 1) // chunk_size
    elementwise[
        process_manual_vectorized_tiles, num_threads_per_tile, target="gpu"
    ](num_tiles, ctx)

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

comptime chunk_size = tile_size * simd_width  # 32 * 4 = 청크당 128개 요소

global_start = tile_id * chunk_size + i * simd_width

a_vec = a.load[simd_width](global_start, 0)     # 전역 텐서에서 로드
output.store[simd_width](global_start, 0, ret)  # 전역 텐서에 저장

수동 벡터화 실행

pixi run p23 --manual-vectorized

pixi run -e amd p23 --manual-vectorized

pixi run -e apple p23 --manual-vectorized

uv run poe p23 --manual-vectorized

퍼즐이 아직 풀리지 않은 경우 다음과 같이 출력됩니다:

SIZE: 1024
simd_width: 4
tile size: 32
tile_id: 0
tile_id: 1
tile_id: 2
tile_id: 3
tile_id: 4
tile_id: 5
tile_id: 6
tile_id: 7
out: HostBuffer([0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0])
expected: HostBuffer([1.0, 5.0, 9.0, ..., 4085.0, 4089.0, 4093.0])

수동 벡터화 풀이

fn manual_vectorized_tiled_elementwise_add[
    layout: Layout,
    dtype: DType,
    simd_width: Int,
    num_threads_per_tile: Int,
    rank: Int,
    size: Int,
    tile_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:
    # Each tile contains tile_size groups of simd_width elements
    comptime chunk_size = tile_size * simd_width

    @parameter
    @always_inline
    fn process_manual_vectorized_tiles[
        num_threads_per_tile: Int, rank: Int, alignment: Int = align_of[dtype]()
    ](indices: IndexList[rank]) capturing -> None:
        tile_id = indices[0]

        output_tile = output.tile[chunk_size](tile_id)
        a_tile = a.tile[chunk_size](tile_id)
        b_tile = b.tile[chunk_size](tile_id)

        @parameter
        for i in range(tile_size):
            global_start = tile_id * chunk_size + i * simd_width

            a_vec = a.aligned_load[simd_width](Index(global_start))
            b_vec = b.aligned_load[simd_width](Index(global_start))
            ret = a_vec + b_vec
            # print("tile:", tile_id, "simd_group:", i, "global_start:", global_start, "a_vec:", a_vec, "b_vec:", b_vec, "result:", ret)

            output.store[simd_width](Index(global_start), ret)

    # Number of tiles needed: each tile processes chunk_size elements
    num_tiles = (size + chunk_size - 1) // chunk_size
    elementwise[
        process_manual_vectorized_tiles, num_threads_per_tile, target="gpu"
    ](num_tiles, ctx)

comptime chunk_size = tile_size * simd_width  # 32 * 4 = 128

원본 배열: [0, 1, 2, 3, ..., 1023]

청크 0 (thread 0): [0:128]    ← 128개 요소 = 4개씩 32개 SIMD 그룹
청크 1 (thread 1): [128:256]  ← 다음 128개 요소
청크 2 (thread 2): [256:384]  ← 다음 128개 요소
...
청크 7 (thread 7): [896:1024] ← 마지막 128개 요소

@parameter
for i in range(tile_size):  # i = 0, 1, 2, ..., 31
    global_start = tile_id * chunk_size + i * simd_width
    # tile_id=0일 때: global_start = 0, 4, 8, 12, ..., 124
    # tile_id=1일 때: global_start = 128, 132, 136, 140, ..., 252

2. Mojo vectorize 방식

완성할 코드

fn vectorize_within_tiles_elementwise_add[
    layout: Layout,
    dtype: DType,
    simd_width: Int,
    num_threads_per_tile: Int,
    rank: Int,
    size: Int,
    tile_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:
    # Each tile contains tile_size elements (not SIMD groups)
    @parameter
    @always_inline
    fn process_tile_with_vectorize[
        num_threads_per_tile: Int, rank: Int, alignment: Int = align_of[dtype]()
    ](indices: IndexList[rank]) capturing -> None:
        tile_id = indices[0]
        tile_start = tile_id * tile_size
        tile_end = min(tile_start + tile_size, size)
        actual_tile_size = tile_end - tile_start
        print(
            "tile_id:",
            tile_id,
            "tile_start:",
            tile_start,
            "tile_end:",
            tile_end,
            "actual_tile_size:",
            actual_tile_size,
        )

        # FILL IN (9 lines at most)

    num_tiles = (size + tile_size - 1) // tile_size
    elementwise[
        process_tile_with_vectorize, num_threads_per_tile, target="gpu"
    ](num_tiles, ctx)

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

tile_start = tile_id * tile_size
tile_end = min(tile_start + tile_size, size)
actual_tile_size = tile_end - tile_start

fn vectorized_add[
  width: Int
](i: Int) unified {read tile_start, read a, read b, mut output}:
    global_idx = tile_start + i
    if global_idx + width <= size:  # 경계 검사
        # SIMD 연산 코드

vectorize[simd_width](actual_tile_size, vectorized_add)

Mojo vectorize 실행

uv run poe p23 --vectorized

pixi run p23 --vectorized

퍼즐이 아직 풀리지 않은 경우 다음과 같이 출력됩니다:

SIZE: 1024
simd_width: 4
tile size: 32
tile_id: 0 tile_start: 0 tile_end: 32 actual_tile_size: 32
tile_id: 1 tile_start: 32 tile_end: 64 actual_tile_size: 32
tile_id: 2 tile_start: 64 tile_end: 96 actual_tile_size: 32
tile_id: 3 tile_start: 96 tile_end: 128 actual_tile_size: 32
...
tile_id: 29 tile_start: 928 tile_end: 960 actual_tile_size: 32
tile_id: 30 tile_start: 960 tile_end: 992 actual_tile_size: 32
tile_id: 31 tile_start: 992 tile_end: 1024 actual_tile_size: 32
out: HostBuffer([0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0])
expected: HostBuffer([1.0, 5.0, 9.0, ..., 4085.0, 4089.0, 4093.0])

Mojo vectorize 풀이

fn vectorize_within_tiles_elementwise_add[
    layout: Layout,
    dtype: DType,
    simd_width: Int,
    num_threads_per_tile: Int,
    rank: Int,
    size: Int,
    tile_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:
    # Each tile contains tile_size elements (not SIMD groups)
    @parameter
    @always_inline
    fn process_tile_with_vectorize[
        num_threads_per_tile: Int, rank: Int, alignment: Int = align_of[dtype]()
    ](indices: IndexList[rank]) capturing -> None:
        tile_id = indices[0]
        tile_start = tile_id * tile_size
        tile_end = min(tile_start + tile_size, size)
        actual_tile_size = tile_end - tile_start

        fn vectorized_add[
            width: Int
        ](i: Int) unified {read tile_start, read a, read b, mut output}:
            global_idx = tile_start + i
            if global_idx + width <= size:
                a_vec = a.aligned_load[width](Index(global_idx))
                b_vec = b.aligned_load[width](Index(global_idx))
                result = a_vec + b_vec
                output.store[width](Index(global_idx), result)

        # Use vectorize within each tile
        vectorize[simd_width](actual_tile_size, vectorized_add)

    num_tiles = (size + tile_size - 1) // tile_size
    elementwise[
        process_tile_with_vectorize, num_threads_per_tile, target="gpu"
    ](num_tiles, ctx)

tile_start = tile_id * tile_size    # 0, 32, 64, 96, ...
tile_end = min(tile_start + tile_size, size)
actual_tile_size = tile_end - tile_start

fn vectorized_add[
  width: Int
](i: Int) unified {read tile_start, read a, read b, mut output}:
    global_idx = tile_start + i
    if global_idx + width <= size:
        # 자동 SIMD 최적화

Tile 0 처리:
  vectorize 호출 0: 요소 [0:4]를 SIMD_WIDTH=4로 처리
  vectorize 호출 1: 요소 [4:8]를 SIMD_WIDTH=4로 처리
  ...
  vectorize 호출 7: 요소 [28:32]를 SIMD_WIDTH=4로 처리
  합계: 8회 자동 SIMD 연산

성능 비교와 모범 사례

각 접근법의 선택 기준

수동 벡터화를 선택할 때:

• 최대 성능이 중요한 경우
• 예측 가능하고 정렬된 데이터 패턴이 있는 경우
• 메모리 접근에 대한 전문가 수준의 제어가 필요한 경우
• 수동으로 경계 안전성을 보장할 수 있는 경우
• 하드웨어별 최적화가 필요한 경우

Mojo vectorize를 선택할 때:

• 개발 속도와 안전성이 우선인 경우
• 불규칙하거나 동적인 데이터 크기를 다루는 경우
• 수동 경계 조건 관리 대신 자동 나머지 처리를 원하는 경우
• 경계 검사 복잡도가 오류를 유발할 수 있는 경우
• 수동 루프 관리보다 깔끔한 벡터화 패턴을 선호하는 경우

고급 최적화 인사이트

메모리 대역폭 활용:

수동:      8 스레드 × 32 SIMD 연산 = 총 256회 SIMD 연산
vectorize: 32 스레드 × 8 SIMD 연산 = 총 256회 SIMD 연산

둘 다 비슷한 총 처리량을 달성하지만, 병렬성 전략이 다릅니다.

캐시 동작:

• 수동: 대형 청크가 L1 캐시를 초과할 수 있지만, 완벽한 순차 접근
• vectorize: 작은 타일이 캐시에 더 잘 맞고, 자동 나머지 처리

하드웨어 매핑:

• 수동: 워프 활용과 SIMD 유닛 매핑에 대한 직접 제어
• vectorize: 자동 루프 및 나머지 관리를 통한 간소화된 벡터화

모범 사례 요약

수동 벡터화 모범 사례:

• 인덱스 계산을 항상 신중하게 검증
• 가능하면 chunk_size에 컴파일 타임 상수 사용
• 캐시 최적화를 위해 메모리 접근 패턴 프로파일링
• 최적의 SIMD 성능을 위한 정렬 요구 사항 고려

Mojo vectorize 모범 사례:

• 데이터와 하드웨어에 적합한 SIMD 폭 선택
• 미세 최적화보다 알고리즘의 명확성에 집중
• 깔끔한 벡터화 로직을 위해 중첩 파라미터 함수 사용
• 경계 조건에는 자동 경계 검사와 나머지 처리 신뢰

두 접근법 모두 GPU 성능 최적화 도구 모음에서 유효한 전략입니다. 수동 벡터화는 최대한의 제어를, Mojo의 vectorize는 안전성과 자동 나머지 처리를 제공합니다.

다음 단계

세 가지 기본 패턴을 모두 이해했다면:

• 🧠 GPU 스레딩 vs SIMD 개념: 실행 계층 구조 이해
• 📊 Mojo 벤치마킹: 성능 분석과 최적화

💡 핵심 요약: 벡터화 전략은 성능 요구 사항에 따라 달리 선택해야 합니다. 수동 벡터화는 최대한의 제어를, Mojo의 vectorize 함수는 안전성과 자동 나머지 처리를 제공합니다. 구체적인 성능 요구 사항과 개발 제약 조건에 따라 선택하세요.