tile - 메모리 효율적인 타일링 처리

comptime TILE_SIZE = 32


fn tiled_elementwise_add[
    layout: Layout,
    dtype: DType,
    simd_width: 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:
    @parameter
    @always_inline
    fn process_tiles[
        simd_width: Int, rank: Int, alignment: Int = align_of[dtype]()
    ](indices: IndexList[rank]) capturing -> None:
        tile_id = indices[0]

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

        @parameter
        for i in range(tile_size):
            a_vec = a_tile.load[simd_width](Index(i))
            b_vec = b_tile.load[simd_width](Index(i))
            ret = a_vec + b_vec
            output_tile.store[simd_width](Index(i), ret)

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

타일링 처리 패턴은 GPU 프로그래밍을 위한 고급 메모리 최적화 기법을 보여줍니다:

1. 타일링 철학과 메모리 계층 구조

타일링은 병렬 처리에 대한 사고 방식의 근본적인 전환을 나타냅니다:

요소별 방식:

• 넓은 병렬성: 많은 스레드가 각각 최소한의 작업 수행
• 전역 메모리 부하: 스레드들이 전체 배열에 분산
• 캐시 미스: 스레드 경계를 넘나드는 낮은 공간 지역성

타일링 방식:

• 깊은 병렬성: 더 적은 스레드가 각각 상당한 작업 수행
• 지역화된 메모리 접근: 각 스레드가 연속적인 데이터에서 작업
• 캐시 최적화: 우수한 공간 및 시간 지역성

2. 타일 구성과 인덱싱

tile_id = indices[0]
out_tile = output.tile[tile_size](tile_id)
a_tile = a.tile[tile_size](tile_id)
b_tile = b.tile[tile_size](tile_id)

타일 매핑 시각화 (TILE_SIZE=32):

원본 배열: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, ..., 1023]

Tile 0 (thread 0): [0, 1, 2, ..., 31]      ← 요소 0-31
Tile 1 (thread 1): [32, 33, 34, ..., 63]   ← 요소 32-63
Tile 2 (thread 2): [64, 65, 66, ..., 95]   ← 요소 64-95
...
Tile 31 (thread 31): [992, 993, ..., 1023] ← 요소 992-1023

핵심 인사이트:

• tile[size](id)는 원본 텐서에 대한 뷰를 생성합니다
• 뷰는 제로 카피로 동작합니다 - 데이터를 복사하지 않고 포인터 연산만 수행
• 타일 경계는 항상 tile_size 단위로 정렬됩니다

3. 순차 처리 심층 분석

@parameter
for i in range(tile_size):
    a_vec = a_tile.load[simd_width](i, 0)
    b_vec = b_tile.load[simd_width](i, 0)
    ret = a_vec + b_vec
    out_tile.store[simd_width](i, 0, ret)

왜 순차 처리인가?

• 캐시 최적화: 연속적인 메모리 접근이 캐시 히트율을 극대화
• 컴파일러 최적화: @parameter 루프가 컴파일 타임에 완전히 전개됨
• 메모리 대역폭: 순차 접근이 메모리 컨트롤러 설계에 부합
• 조정 비용 감소: SIMD 그룹 간 동기화가 불필요

하나의 타일 내 실행 패턴 (TILE_SIZE=32, SIMD_WIDTH=4):

스레드가 타일을 순차 처리:
Step 0: 요소 [0:4]를 SIMD로 처리
Step 1: 요소 [4:8]를 SIMD로 처리
Step 2: 요소 [8:12]를 SIMD로 처리
...
Step 7: 요소 [28:32]를 SIMD로 처리
합계: 스레드당 8회 SIMD 연산 (32 ÷ 4 = 8)

4. 메모리 접근 패턴 분석

캐시 동작 비교:

요소별 패턴:

Thread 0: 글로벌 위치 [0, 4, 8, 12, ...] 접근    ← Stride = SIMD_WIDTH
Thread 1: 글로벌 위치 [4, 8, 12, 16, ...] 접근   ← Stride = SIMD_WIDTH
...
결과: 메모리 접근이 전체 배열에 분산

Tiled 패턴:

Thread 0: 위치 [0:32]를 순차 접근               ← 연속적인 32개 요소 블록
Thread 1: 위치 [32:64]를 순차 접근             ← 다음 연속적인 32개 요소 블록
...
결과: 각 스레드 내에서 완벽한 공간 지역성

캐시 효율 시사점:

• L1 캐시: 작은 타일이 L1 캐시에 더 잘 맞아 캐시 미스 감소
• 메모리 대역폭: 순차 접근이 유효 대역폭을 극대화
• TLB 효율: TLB 미스 감소 (역주: TLB(Translation Lookaside Buffer)는 가상 주소를 물리 주소로 변환하는 캐시로, 미스가 줄면 메모리 접근이 빨라집니다)
• 프리페칭: 하드웨어 프리페처가 순차 패턴에서 최적으로 동작

5. 스레드 구성 전략

elementwise[process_tiles, 1, target="gpu"](num_tiles, ctx)

왜 SIMD_WIDTH 대신 1인가?

• 스레드 수: num_tiles × SIMD_WIDTH가 아닌 정확히 num_tiles개의 스레드만 실행
• 작업 분배: 각 스레드가 하나의 완전한 타일을 처리
• 로드 밸런싱: 스레드당 더 많은 작업, 전체적으로 더 적은 스레드
• 메모리 지역성: 각 스레드의 작업이 공간적으로 지역화

성능 트레이드오프:

• 더 적은 논리적 스레드: 낮은 점유율에서 모든 GPU 코어를 활용하지 못할 수 있음
• 스레드당 더 많은 작업: 더 나은 캐시 활용과 조정 오버헤드 감소
• 순차 접근: 각 스레드 내에서 최적의 메모리 대역폭 활용
• 오버헤드 감소: 스레드 실행 및 조정 오버헤드 감소

중요 참고: “더 적은 스레드“는 논리적 프로그래밍 모델을 의미합니다. GPU 스케줄러는 여러 워프를 실행하고 메모리 지연 시 효율적으로 전환하여 높은 하드웨어 활용률을 달성할 수 있습니다.

6. 성능 특성

타일링이 도움이 되는 경우:

• 메모리 바운드 연산: 메모리 대역폭이 병목인 경우
• 캐시 민감 워크로드: 데이터 재사용의 이점이 있는 연산
• 복잡한 연산: 요소당 연산량이 많은 경우
• 제한된 병렬성: GPU 코어보다 스레드가 적은 경우

타일링이 불리한 경우:

• 고도로 병렬적인 워크로드: 최대 스레드 활용이 필요한 경우
• 단순한 연산: 메모리 접근이 연산보다 지배적인 경우
• 불규칙적 접근 패턴: 타일링이 지역성을 개선하지 못하는 경우

단순 덧셈 예시 (TILE_SIZE=32):

• 스레드 수: 256개 대신 32개 (8배 적음)
• 스레드당 작업량: 4개 대신 32개 요소 (8배 많음)
• 메모리 패턴: 순차 vs 스트라이드 접근
• 캐시 활용: 훨씬 나은 공간 지역성

7. 고급 타일링 고려 사항

타일 크기 선택:

• 너무 작으면: 캐시 활용이 떨어지고, 오버헤드가 증가
• 너무 크면: 캐시에 맞지 않을 수 있고, 병렬성이 감소
• 최적 지점: L1 캐시 최적화를 위해 보통 16-64개 요소
• 현재 선택: 32개 요소로 캐시 활용과 병렬성의 균형 달성

하드웨어 고려 사항:

• 캐시 크기: 가능하면 타일이 L1 캐시에 맞아야 함
• 메모리 대역폭: 메모리 컨트롤러 폭을 고려
• 코어 수: 모든 코어를 활용하기에 충분한 타일 확보
• SIMD 폭: 타일 크기는 SIMD 폭의 배수여야 함

비교 요약:

Elementwise: 높은 병렬성, 분산된 메모리 접근
Tiled:       적당한 병렬성, 지역화된 메모리 접근

요소별 패턴과 타일링 패턴 간의 선택은 특정 워크로드 특성, 데이터 접근 패턴, 대상 하드웨어 능력에 따라 달라집니다.