Tinygrad Unroll vs Upcast

26 Mar, 2025

1. Unroll (루프 펼치기)

목적: 커널 내의 루프(loop) 를 최적화합니다. 루프의 반복 실행 오버헤드(조건 검사, 인덱스 증가, 분기 등)를 줄여 성능을 향상시키는 것이 주 목적입니다.
작동 방식: 루프의 반복 코드를 명시적으로 여러 번 풀어서 작성합니다. 예를 들어 4번 반복하는 루프를 2번 unroll하면, 루프는 2번만 돌지만 루프 본문 코드는 한 번에 2개의 원래 반복에 해당하는 연산을 수행하도록 변경됩니다. (float2 사용 예시) 완전히 unroll하면 루프 자체가 사라지고 모든 반복이 순차적인 코드로 대체됩니다.
적용 대상: 주로 리덕션 축(reduction axis) 과 같이 커널 스레드 내부에서 반복적으로 수행되는 계산에 적용됩니다.
효과:
- 루프 제어 오버헤드 감소.
- 명시적인 코드로 인해 컴파일러가 추가적인 최적화(예: 명령어 스케줄링, 벡터화)를 수행할 가능성 증가.
- 코드 크기가 증가할 수 있음.
예시: Opt(OptOps.UNROLL, 0, 2)는 마지막 축(axis 0)에 대한 루프를 2번씩 펼칩니다. 원래 4번 반복했다면, 이제 루프는 2번만 돌고 각 반복마다 float2를 사용하여 2개의 값을 처리합니다.

2. Upcast (업캐스팅)

목적: 각 커널 스레드가 더 많은 데이터 요소를 한 번에 처리하도록 하여 전체 커널 실행의 효율성을 높입니다. 메모리 대역폭 활용률을 높이고 커널 실행(launch) 오버헤드를 줄이는 것이 주 목적입니다.
작동 방식: 특정 축(주로 병렬 처리되는 축)에 대해 각 스레드가 처리하는 요소의 수를 늘립니다. 예를 들어, 16개의 요소를 16개의 스레드가 각각 1개씩 처리하는 대신, 2개의 스레드가 각각 8개씩 처리하도록 변경합니다. 이를 위해 더 큰 데이터 타입(예: float4)을 사용하여 여러 요소를 동시에 로드/저장/계산하는 경우가 많습니다.
적용 대상: 주로 병렬 처리 축(parallel axis) 에 적용됩니다. 리덕션 축에는 적용할 수 없습니다.
효과:
- 실행되는 총 커널 스레드 수 감소 -> 커널 실행 오버헤드 감소.
- 벡터 데이터 타입(float4 등)을 활용하여 메모리 접근 및 계산 효율성 증대 (메모리 대역폭 활용률 향상).
- 각 스레드가 더 많은 작업을 수행하게 됨.
예시: Opt(OptOps.UPCAST, 0, 8)는 16개 요소를 처리하는 작업을, 각 스레드가 8개씩 처리하도록 변경합니다. 따라서 총 스레드 수는 2개(16 / 8)로 줄어듭니다. 각 스레드는 float4를 두 번 사용하여 총 8개의 float 요소를 처리합니다.

요약:

특징	Unroll (루프 펼치기)	Upcast (업캐스팅)
주 목적	커널 내 루프 최적화, 반복 횟수 풀어내어 루프 오버헤드 감소	스레드당 작업량 증가, 메모리 효율 증대, 여러 요소를 동시에 처리하여 작업 밀도를 높임
작동 방식	루프 본문 코드를 명시적으로 반복, 반복문 제거 및 인라인 코드 증가	스레드당 처리 요소 수를 늘림, 커널이 다중 요소를 처리하도록 데이터 변화
적용 대상/축	주로 리덕션 축(Reduction Axis)의 루프, 주로 관계된 축	주로 병렬 처리 축(Parallel Axis), 연속적 데이터 축 (예: elementwise 축)
스레드 수	변경 없음	감소
스레드 작업	루프 구조 변경, 작업량 자체는 동일	처리하는 요소 수 증가
주요 효과	루프 오버헤드 감소, 추가 최적화 기회, 파이프라인 최적화	커널 실행 오버헤드 감소, 메모리 대역폭 활용, 메모리 접근 최적화, 벡터화로 처리량 증가
사용 예시	`Opt(OptOps.UNROLL, axis, factor)`	`Opt(OptOps.UPCAST, axis, amount)`
적용 조건	루프가 특정 값으로 반복 가능해야 함 (예: `amt=2`)	데이터가 연속적(contiguous)이고 reduce 축이 아니어야 함