IOUB

IOOpt: Automatic Derivation of I/O Complexity Bounds for Affine Programs (Olivry et al., PLDI 2021) 리뷰

IOLB가 “어떤 스케줄로도 이보다 적게 옮길 수 없다”는 하한을 증명했다면, IOUB는 “이 스케줄로 이만큼만 옮기면 된다”는 상한을 자동으로 유도한다. 둘을 합친 도구가 IOOpt이고, 부수적으로 타일 크기와 루프 순서 추천까지 뱉는다. 하한이 타이트한지 확인하는 유일한 방법은 그 값에 도달하는 구현을 실제로 제시해 상한을 맞춰보는 것뿐이다.

이 논문이 제시한 것들은 다음과 같다.

1. 파라메트릭 타일 크기를 가진 다차원 타일 코드의 데이터 이동량을 기호적으로 과대근사하는 최초의 알고리즘
2. 그 식의 최소화를 최적화 문제(NLP)로 자동 정식화
3. 작은 차원(small dimension) 이 있을 때의 하한 유도를 개선 — IOLB의 약점이던 convolution을 잡는다
4. 이를 통합해 연산 복잡도 / 증명된 I/O 상·하한 / 데이터 이동을 최소화하는 타일 코드를 함께 내놓는 도구

대상 프로그램

imperfectly nested affine loop이되, 직사각형 타일로 모든 차원을 타일링할 수 있는(fully tilable) 프로그램을 가정한다.

이론적으로는 임의의 affine 프로그램을 폴리헤드럴 컴파일러로 전처리해 fully permutable band를 만들 수 있지만, 실제로는 비정규 반복 도메인에서 복잡한 기호식을 단순화·풀이해야 해서 만만치 않다. 대신 이 가정에 들어오는 중요한 클래스가 충분히 많다 (convolution, 모든 종류의 tensor contraction, 행렬곱을 비롯한 상당수 선형대수 커널 등…)

메모리 모델은 IOLB와 같은 red-white pebble game이다. 상한 쪽도 같은 형식화를 쓰기 때문에 두 결과를 그대로 비교할 수 있다.

타일링 스케줄

타일링된 코드는 바깥의 inter-tile 루프와 안쪽의 intra-tile 루프로 나뉜다. 핵심 관찰은 intra-tile 루프의 순서는 이 비용 모델에 영향을 주지 않는다는 것이다. 타일 하나는 어차피 통째로 캐시에 올라가므로 안에서 어떤 순서로 도는지는 I/O에 무관하다. 따라서 스케줄은 이렇게만 적으면 된다.

(P, T) — P = inter-tile 루프 순서(바깥→안), T = {T_d} 차원별 타일 크기

행렬곱의 타일 코드는 ((i, j, k), {T_i=Ti, T_j=Tj, T_k=1})로 표현된다.

비용 모델

레벨 j의 sub-domain SD_{d_j}는 자기보다 바깥 루프들의 인덱스를 고정했을 때 남는 반복 공간이다. 여기에 배열 A의 접근 함수를 씌우면 sub-domain footprint SDF가 되고, 연속한 두 sub-domain의 footprint 교집합이 재사용량 SDR이 된다.

SDF_{A,j} = f_A(SD_{d_j})
SDR_{A,j} = f_A(SD_{d_j}(...,i_j)) ∩ f_A(SD_{d_j}(...,i_j - T_{d_j}))

“이전 sub-domain에서의 재사용”은 루프의 첫 번째 반복에는 적용되지 않는다. 그래서 반복 공간을 i_j = 0인 front와 i_j ≠ 0인 back으로 쪼갠다. 각각에 대해 inverse density를 정의한다.

ID_{A,j}(back)  = (SDF - SDR) / |SD|
ID_{A,j}(front) = SDF / |SD|

연산 하나당 필요한 데이터 이동량, 즉 inverse operational intensity의 과대근사다. 도메인이 직사각형이 아니면 식이 복잡해지므로 front/back 각각에 대해 최댓값을 취한다. 대부분의 경우 sub-domain 크기와 교집합이 루프 인덱스에 의존하지 않아서 최댓값을 따로 구할 필요도 없다.

배열이 하나일 때 비용은 이렇게 나온다. 먼저 타일 하나의 footprint가 캐시에 들어가야 한다 (SDF_{A,1} ≤ S). 그다음 outermost reuse dimension l — P에서 SDF_{A,l} ≤ S를 만족하는 가장 바깥 차원 — 을 찾으면

IO_A = ID_{A,l}(front) × |I_front| + ID_{A,l}(back) × |I_back|

가 유효한 상한이다. 각 sub-domain은 시작 전에 필요한 데이터를 한 번만 가져오면 실행할 수 있고, back의 sub-domain은 직전 것이 올려둔 데이터를 재사용할 수 있기 때문이다.

배열이 여러 개면 캐시를 배열별 영역으로 자른다. S₁ + … + S_s = S로 나누고 각 배열에 대해 같은 조건을 적용한다.

다단계 캐시로의 확장은 간단하다. 캐시 레벨마다 타일링 band를 하나씩 두고 같은 논리를 독립적으로 적용하면 된다.

루프 순서

타일 크기는 기호식이 있으니 NLP 솔버에 넘기면 되는데, 루프 순열의 탐색 공간은 차원 수와 캐시 깊이에 대해 지수적으로 늘어난다. 그래서 가지치기가 필요하다.

두 가지 관찰을 쓴다. 첫째, I/O 비용이 같은 순열들이 많다 (conv-1D에서 바깥 두 루프 c1과 f1을 바꿔도 비용이 같다). 둘째, 어떤 순열은 명백히 나쁘다 (가장 안쪽 차원이 연속한 타일 간에 아무 재사용도 만들지 못하는 순열은 탐색할 가치가 없다)

재사용 차원을 이렇게 정의한다. 차원 d를 가장 안쪽에 놓았을 때 레벨 2의 footprint가 레벨 1에 비해 크게 늘지 않으면(SDF_{A,2} − SDF_{A,1} ≪ SDF_{A,1}) d는 배열 A에 대한 재사용 차원이다.

conv-1D의 Out[f][x]는 명확하다. c를 안쪽에 놓으면 SDF가 그대로지만(재사용 차원), f를 놓으면 (Nf − Tf) × Tx만큼 늘어난다(재사용 차원 아님). 그런데 Image[x+w][c]처럼 첨자가 복잡하면 애매해진다 — w를 안쪽에 놓는 판단이 Tw − 1과 Tx의 크기 비교로 귀결되는데, 이건 어떤 차원이 작고 어떤 차원이 큰지 알려주는 오라클(사용자) 없이는 결정할 수 없다.

알고리즘 1은 이 오라클을 가정하고 안쪽부터 바깥으로 순열을 짜 나간다. 남은 차원 집합과 각 차원이 재사용을 만드는 배열 집합을 추적하면서, 재사용 배열 집합이 다른 차원의 진부분집합인 차원은 가지친다. conv-1D(4차원, 24가지 순열)에서 실제로 3개까지 줄어든다. c를 가장 안쪽에 놓으면 Out에만 재사용이 생기지만 w를 놓으면 Out과 Image 둘 다에 생기므로, c 쪽 가지는 잘린다.

TileOpt와 기호식 정리

IOUB가 순열별로 (I/O 비용식, 타일 크기 제약)을 만들면 TileOpt가 IPOPT(NLP 솔버)로 각 순열의 최적 타일 크기를 찾고, 전체 최솟값을 고른다. 구현은 ISL과 Sympy 위에 올렸다. 부산물로 해당 타일링을 구현한 기본 코드도 나온다.

문제는 이 결과가 타일 크기에 의존하는 식이라는 것이다. 하한과 비교하려면 프로그램 파라미터와 캐시 크기 S만으로 표현해야 한다. 일반적으로는 풀기 어렵지만, 벤치마크 범위에서는 가정 두 개로 해결된다 — 타일을 정사각으로 두고, 타일이 캐시를 정확히 꽉 채운다고 본다. 최선이 아닐 수는 있어도 여전히 유효한 상한이다.

행렬곱으로 보면, IOUB가 내놓는 식과 제약이

IO = Ni·Nj·Nk·(1/Ti + 1/Tj + 1/Nk)
Ti + Tj + Ti·Tj ≤ S

이고, Ti = Tj = T로 두고 제약을 등식으로 보면 2T + T² = S → T = √(S+1) − 1. 이걸 대입하면

UB = Ni·Nj·( 2Nk / (√(S+1) − 1) + 1 )

이 된다. 최고차항이 2·Ni·Nj·Nk/√S로, IOLB가 준 하한 LB = 2Ni + Nk + 2Nj + 2·Ni·Nj·Nk/√S의 최고차항과 정확히 일치한다. 점근적으로 타이트함이 증명된 것이다.

tensor contraction으로 일반화할 때는 타일 크기의 크기 순서에 대한 정보를 조금 더 준다. TC의 차원들은 세 그룹으로 묶으면 행렬곱과 동형이 되므로(Fig. 5의 shared dimension), 각 그룹 안의 타일 크기 곱이 서로 같다는 조건을 걸면 된다.

IOUB만으로는 부족한 케이스가 convolution이다. 기존 IOLB가 주는 하한이 상한보다 점근 차수 자체가 몇 단계 아래였다. 원인이 두 가지라서 각각 고친다.

reduction 인식. convolution은 c, h, w 세 차원에 걸친 합 reduction이다. 그런데 코드는 루프로 쓰여 있으니 CDAG에는 reduce 차원을 따라 순차적 의존성 사슬이 생긴다. reduce 차원이 3개 이상이면 이 사슬이 경로 분석을 망가뜨려 추출되는 준동형 φ_j의 품질이 나빠진다.

합은 결합·교환법칙을 만족하므로 순서가 아무래도 상관없다. IOUB는 affine 의존성에 패턴 매칭을 걸어 (reduce 차원이 직사각 도메인이고 사전식 순서로 합해지는) 단순한 reduction을 검출하고, 순차 사슬을 두 벌의 broadcast 의존성으로 교체한다 — 최종 결과를 쓰는 계산에서 모든 reduction 노드로 가는 것, 그리고 초기화에서 모든 노드로 가는 것.

convolution에서 이 교체는 φ₁을 (b,c,f,x,y,h,0)에서 (b,0,f,x,y,0,0)으로 바꾼다. 커널이 커지고 s_j 제약이 좋아진다.

작은 차원(small dimension). 필터 크기 H, W처럼 캐시 크기 S보다 자릿수로 작은 차원들이 있다. 여기에 대응하는 준동형 φ_sd(작은 차원들로의 사영)를 Brascamp-Lieb에 하나 더 넣으면 |φ_sd(E)| ≤ N_sd라는 훨씬 빡빡한 제약이 붙는다.

|E| ≤ Π_j |φ_j(E)|^{s_j} · |φ_sd(E)|^{s_sd} ≤ K^σ · N_sd^{s_sd},  σ = Σ_j s_j (s_sd 제외)

s_sd가 기여하는 만큼 나머지 s_j의 제약이 느슨해져서 더 작은 s_j를 찾을 수 있다. 최적화는 σ를 먼저 최소화하고 그다음 s_sd를 최소화한다.

2D convolution에서 효과가 극적이다. 파라미터를 전부 N으로 두고 보면

단계	하한
기존 IOLB	O(N⁴)
+ reduction 인식	O(N⁷/S)
+ 작은 차원	O(√(HW)·N⁵/√S) — 점근적으로 타이트

Brascamp-Lieb 제약을 실제로 풀어보면, 작은 차원 없이는 s_j = 2/3, σ = 2 (즉 |E| ≤ K²)이고, φ_sd를 넣으면 s_j = 1/2, s_sd = 1/2, σ = 3/2 (즉 |E| ≤ K^{3/2}·(HW)^{1/2})이 된다.

Demmel et al.도 CNN에 대해 비슷한 트릭으로 타이트한 하한을 얻었지만 스케일링 상수를 주지 못한다. IOUB의 기여는 이 아이디어를 IOLB의 틀 안으로 가져와서, 점근 부등식이 아니라 상수가 붙은 정확한 식을 유지한 채 적용했다는 점이다.