IOLB

Automated Derivation of Parametric Data Movement Lower Bounds for Affine Programs (Olivry et al., 2019) 리뷰

affine 프로그램을 입력받아, 어떤 스케줄로 실행하더라도 반드시 발생하는 메모리 로드 횟수의 하한을 컴파일 타임에 자동으로 유도하는 도구 IOLB를 제안한다. 기존 연구들이 몇몇 알고리즘에 대해 사람 손으로 점근적 하한만 구했던 것을, 스케일링 상수까지 포함한 파라메트릭 식으로, 버튼 한 번에 뽑아낸다.

문제

지금 프로세서에서 연산 자체의 비용은 데이터를 메모리에서 ALU까지 옮기는 비용보다 훨씬 싸다. 지연시간은 파이프라이닝으로 가릴 수 있지만 처리량(throughput)은 못 가린다. 그래서 성능과 에너지의 근본적인 한계는 데이터 이동량이 결정한다.

이걸 재는 지표가 operational intensity (OI) — 연산 수 / 메모리 전송량(word). 이걸 기계의 machine balance (MB) — 피크 연산속도 / 피크 전송속도 — 와 비교하면, OI < MB인 프로그램은 메모리 대역폭에 묶인다. MB는 세대가 지날수록 계속 커지고 있으므로, 예전엔 compute-bound였던 코드가 점점 bandwidth-bound로 넘어간다.

문제는 OI가 알고리즘의 고정된 성질이 아니라는 것이다. N×N 행렬곱은 타일링하지 않으면 여섯 가지 루프 순서 중 무엇을 골라도 N² > S(캐시 용량)일 때 최소 N³ word를 옮겨야 하므로 OI가 2를 넘지 못한다. 반면 타일링하면 2N³/√S로 줄어 OI가 √S가 된다.

VTune이나 Intel SDE 같은 도구는 지금 구현이 달성한 OI만 알려준다. 그 값이 이 알고리즘이 낼 수 있는 최대치에 가까운 건지, 아니면 스케줄을 바꿔서 한참 더 올릴 수 있는 건지는 알려주지 못한다. IOLB는 이 질문에 답한다. 데이터 이동의 하한 Q_low를 구하면 #ops / Q_low가 달성 가능한 OI의 상한이 되기 때문이다.

IOLB는 일종의 증명 환경으로 볼 수 있다. 입력이 affine 코드, 출력이 그 코드의 모든 유효한 스케줄에 대한 I/O 하한이고, 증명 과정 자체를 출력에서 따라 읽을 수 있다.

CDAG와 red-white pebble game

Hong & Kung(1981)의 형식화를 따른다. 알고리즘을 CDAG(Computational DAG)로 추상화한다 — 정점은 연산 인스턴스나 입력값, 간선은 데이터 흐름 의존성이다. 여기엔 메모리 위치 개념이 없다. CDAG는 가능한 모든 유효한 스케줄을 한꺼번에 표현하며, 유일한 제약은 선행 정점이 먼저 실행되어야 한다는 것뿐이다.

데이터 이동은 red-white pebble game으로 모델링한다. 빨간 돌은 크기 S의 빠른 메모리에 있는 값, 흰 돌은 이미 계산된 값이다.

• (R1) 흰 돌이 있는 정점에 빨간 돌을 놓을 수 있다 (= 로드)
• (R2) 흰 돌이 없고 모든 직전 선행자에 빨간 돌이 있으면 빨간 돌을 놓을 수 있다 (= 계산). 흰 돌도 같이 놓인다
• (R3) 빨간 돌은 아무 때나 치울 수 있다 (= 방출)

게임의 비용은 (R1)의 적용 횟수, 즉 로드 횟수다. I/O 복잡도 Q(G)는 모든 정점에 흰 돌이 놓인 상태로 끝나는 게임의 최소 비용이다.

원본 red-blue 게임과 다른 점이 두 가지다.

• 재계산을 허용하지 않는다. 흰 돌은 한 번 놓이면 치워지지 않는다. 이 제약이 있어야 복잡한 CDAG를 부분 영역으로 쪼개서 하한을 합산할 수 있다.
• store를 세지 않고 load만 센다. 둘 다 세는 모델에 대해서도 여전히 유효한 하한이고, 대부분의 계산에서 load가 store를 압도하므로 타이트함에 큰 손해가 없다.

(S+T)-partition

Hong & Kung의 핵심 아이디어는 게임의 실행 시퀀스와 CDAG 정점 분할 사이의 대응이다.

실행을 정확히 T번의 로드를 포함하는 연속 구간들로 자른다. 한 구간이 시작될 때 빠른 메모리에는 최대 S개의 값이 있고 구간 안에서 T개를 더 로드하므로, 그 구간에서 계산된 정점 집합 P의 In-set — P 바깥에 있으면서 P 안에 후속자를 갖는 정점들 — 의 크기는 S+T 이하다. 이런 집합을 (S+T)-bounded라 한다.

여기서 (S+T)-bounded 집합의 크기 상한 U를 구할 수 있으면, 구간이 최소 |V\I|/U개 필요하므로 하한이 나온다.

Q(G) ≥ T · ⌈|V \ I| / U⌉

원래 Hong & Kung은 T = S로 고정해 2S-partition을 썼는데, Smith et al.(2019)이 T를 자유롭게 두는 일반화를 제안했고 IOLB는 이를 받아 T를 최적화한다. 이것만으로 많은 경우 하한이 더 타이트해진다.

입력 정점이 지정되지 않은 CDAG(분해할 때 생긴다)에는 Elango et al.의 변형을 쓴다. 일부 정점을 인위적으로 입력으로 태깅해서 하한을 구한 뒤, 추가된 입력 개수만큼 빼주면 된다.

Brascamp-Lieb 부등식으로 U 구하기

그럼 U — K-bounded 집합의 크기 상한 — 는 어떻게 구하나. 여기가 자동화의 핵심이다.

CDAG의 정점을 다차원 공간 Z^d의 점으로 매핑한다(보통 차원이 루프 인덱스다). 그러면 규칙적인 데이터 의존성은 저차원 공간으로의 사영(projection) 이 된다. “P가 K-bounded”라는 그래프 조건이 “ρ(P)의 사영들의 크기가 K 이하”라는 기하 조건으로 바뀐다.

주어진 사영들의 크기 제약으로부터 원래 집합의 크기를 제한하는 도구가 이산 Brascamp-Lieb 부등식이다. φ_j: Z^d → Z^{d_j}가 군 준동형이고 0 ≤ s_j ≤ 1 일 때, 모든 부분군 H에 대해

rank(H) ≤ Σ_j s_j · rank(φ_j(H))

가 성립하면

|E| ≤ Π_j |φ_j(E)|^{s_j}

이다. 축 방향 정규 사영에 s_j = 1/(d-1)인 특수 케이스가 우리가 아는 Loomis-Whitney 부등식이다.

Fig. 1 예제(for t: for i: A[i] = A[i] * C[t])로 보면, 2차원 격자 위 점 집합의 In-set은 두 축 사영을 포함하므로, T = S로 두면 각 사영이 2S 이하 → |E| ≤ 4S² = U → Q ≥ S·⌊MN/4S²⌋ ≈ MN/4S가 나온다.

실무적인 걸림돌이 두 가지 있다.

• 조건을 모든 부분군에 대해 확인해야 한다. 가능한 부등식의 개수 자체는 유한(계수가 정수이고 d로 유계)하고 admissible한 s_j들이 이루는 영역이 볼록 다면체라는 것도 알려져 있지만, 이를 계산하는 알고리즘은 조합적이고 복잡하다. IOLB는 대신 Valdimarsson의 결과를 쓴다 — Ker(φ_j)들이 생성하는 부분군 격자에 대해서만 확인해도 충분하다. 이 격자도 유한하다는 보장이 없어서 타임아웃을 두고, 사영을 하나씩 추가하며 격자를 갱신한다. 타임아웃이 나도 실패가 아니라 하한이 덜 타이트해질 뿐이다. Ker(φ_j)들이 선형독립이면 격자를 계산할 필요 없이 각 커널만 개별 확인하면 된다.
• s_j를 어떻게 고르나. |E| ≤ K^{Σs_j}이므로 Σ s_j를 최소화하면 된다. 제약이 선형부등식이라 선형계획으로 푼다.

DFG — CDAG의 압축 표현

CDAG는 한 번의 동적 실행 전체를 펼친 그래프라 매우 크다. IOLB는 폴리헤드럴 모델의 Data-flow graph(DFG) 로 압축해서 다룬다. 정점은 정적 statement 또는 입력 배열이고 파라메트릭 반복 도메인(ISL set)을 갖는다. 간선은 흐름 의존성이고 affine 관계(ISL map)를 갖는다. DFG 하나가 파라미터 값에 따라 여러 크기의 CDAG를 표현한다.

분석의 기본 단위는 DFG-path, 즉 DFG 위의 유향 경로다. 경로의 관계는 간선 관계들의 합성이며, 둘만 관심 대상이다.

• chain circuit: 한 정점에서 자기 자신으로 가는 순환이고 관계가 평행이동 S[x] → S[x + b]인 것. 재사용 방향을 뜻한다
• broadcast path: 첫 간선 외에 전부 단사 간선인 경로로, 역관계가 full-rank가 아닌 affine 함수인 것. 한 값이 여러 곳으로 퍼지는 패턴이다

broadcast path는 관계가 곧바로 사영이 되고, chain circuit은 평행이동 벡터 δ에 수직인 초평면으로의 정규 사영이 된다. 이렇게 그래프 위의 재사용 패턴이 기하 사영으로 번역된다.

간섭 그래프

두 경로 Q₁, Q₂에 대해 In-set에서 각각이 담당하는 부분이 서로소이면, 더 강한 부등식이 성립한다.

|φ_1(ρ(P))| + |φ_2(ρ(P))| ≤ K

곱이 아니라 합으로 묶이므로 훨씬 타이트하다. IOLB는 경로들을 정점으로 하고 독립적인 쌍을 간선으로 하는 간섭 그래프를 만들어, 모든 정점을 덮는 maximal clique들을 찾고 대응 부등식을 더해 일반형 Σ β_j |φ_j(E)| ≤ K를 얻는다. 여기에 라그랑주 승수로 최적화한 Lemma 5.2를 적용하면

|E| ≤ (K / Σs_j)^{Σs_j} · Π_j (s_j / β_j)^{s_j}

가 된다. 스케일링 상수가 여기서 나온다. 행렬곱 하한을 사람이 손으로 개선했던 기법(Dongarra, Lowery & Langou, Smith & van de Geijn)의 일반화다.

CDAG 분해

복잡한 프로그램은 부분 영역으로 쪼개고 각 영역의 하한을 더해야 한다. 그냥 서로소로 쪼개도 되지만, 논문은 더 일반적인 분해 보조정리를 제시한다.

부분 CDAG G|V_i에서 no-spill set은 (1) 나가는 간선이 없거나 (2) 들어오는 간선이 없고 나가는 간선이 최대 하나인 정점들이다. 이런 정점은 필요한 순간에 계산해서 바로 쓰고 버리면 되므로 로드로 세지 않아도 된다. 나머지가 may-spill set이다.

Lemma 4.2 V₁, …, V_k의 may-spill set들이 쌍마다 서로소이면 Q(G) ≥ Σ Q(G|V_i)

no-spill 정점은 여러 부분 영역에 겹쳐 들어가도 된다는 게 이 보조정리의 힘이다. IOLB는 이걸 두 방식으로 쓴다.

• bounded combination: 유한한 개수의 부분 CDAG를 조합한다. 알고리즘 1은 탐욕적으로, 복잡도가 가장 높은 것부터 고르고, 다음 것은 겹치는 부분을 빼고 복잡도를 다시 계산해서 더한다. “어느 쪽이 더 높은가”는 파라미터 값의 구체적 인스턴스를 넣어 비교한다 — 이건 휴리스틱 결정에만 쓰이고, 최종 식은 모든 파라미터 값에 대해 유효하다
• loop parametrization: 바깥 루프 인덱스를 파라미터 Ω로 고정해서 반복 공간을 슬라이스하고, 무한(파라메트릭) 개의 영역을 합산한다. ISL에서 도메인의 루프 인덱스를 파라미터로 바꾸는 것으로 구현된다. 결과가 Ω에 의존하면 다항식 합 공식으로 더한다

Floyd-Warshall처럼 하나의 statement 인스턴스를 부분 도메인들로 쪼개서 각각의 하한을 더해야 타이트해지는 경우도 자동으로 처리한다.

wavefront 기법

K-partition만으로는 못 잡는 케이스가 있다. 두 번째 무기는 Elango et al.의 wavefront다.

실행 중 어느 시점에서 이미 계산됐지만 아직 후속자가 남아 있는 정점들(live set)을 wavefront라 한다. 이 크기가 S를 넘으면 일부는 반드시 느린 메모리로 밀려났다가 다시 로드되어야 한다.

Q ≥ w_min - S

w_min은 일반적으로 계산 불가능하므로, 논문은 CDAG 구조에 조건을 붙인 따름정리를 쓴다. 서로소 집합 V₁, V₂가 있어 V₂의 모든 정점이 V₁의 모든 정점에서 도달 가능하고, V₁에서 시작해 V₂에서 끝나는 서로소 경로가 m개 있으면 w_min ≥ m이다.

이게 결정적인 경우가 있다. 타일링이 원천적으로 불가능한 커널(adi, durbin)은 기하학적 추론으로는 약한 하한밖에 못 얻는데, wavefront는 “달성 가능한 최선의 OI가 상수”임을 증명한다. 기하 추론으로 얻을 수 있는 어떤 하한보다도 최소 √S배 좋다.

전체 파이프라인

프론트엔드는 PET다. #pragma scop / #pragma endscop으로 감싼 SCoP을 받아 폴리헤드럴 표현을 뽑고, 거기서 DFG를 만든다. 배열 접근은 alias하지 않는다고 가정하고, 모든 스칼라 데이터는 원자적이고 같은 크기로 본다(CDAG에 가중치가 없다 — 개념적 한계가 아니라 구현 한계다).

메인 알고리즘은 루프 깊이 d × statement S의 이중 루프다.

1. 깊이 d의 바깥 인덱스들을 파라미터 Ω_d로 고정 (loop parametrization)
2. S로 끝나는 DFG-path들을 커널 차원 오름차순으로 모으면서, 부분군 격자를 갱신할 수 있는 것만 채택
3. 그 경로 조합으로 K-partition 하한을 계산, may-spill set을 CDAG 복사본에서 제거하고 반복 (한 statement에 대해 여러 개의 겹치지 않는 부분 CDAG가 나올 수 있다)
4. 같은 statement에 대해 wavefront 하한도 계산
5. 모인 하한들을 combine_subQ로 합산
6. compulsory miss를 반영해 입력 데이터 크기를 더한다

파라미터 인스턴스를 여러 개 넣어 각각 Q_I를 얻고 최종적으로 max를 취한다. 모든 Q_I가 유효한 하한이므로 최댓값도 유효하다.

실험 결과

C로 구현했고 ISL-0.13, barvinok-0.37, PET-0.05, GiNaC(기호식), PIP(선형계획)를 쓴다. PolyBench/C-4.2.1의 30개 커널 전부에 적용해 1초 미만에 결과가 나온다.

평가는 IOLB가 준 하한으로 만든 OI 상한(OI_up)과, 손으로 타일링 스케줄을 짜고 타일 크기를 최적화해서 얻은 OI 하한(OI_manual)을 비교한다. 항상 OI_up ≥ OI_manual이고, 둘이 같으면 양쪽 다 타이트하다는 증명이 된다.

jacobi-1d 예시. IOLB가 내놓는 식은

Q_low = 2 + N + max(0, TN/4S - N - T - (1/4S) - N/(3T·4S) - S + 5)

이고, 점근적으로 정리하면 TN/4S다. 연산 수가 6TN이므로 OI_up = 24S가 된다.

30개를 네 부류로 나눈다.

• (19개) 타일링이 가능하고 IOLB가 상수배 이내의 비자명한 상한을 준다. 8개는 점근적으로 정확히 일치, 6개는 2배 이내. gemm을 제외하면 전부 기존 발표 결과보다 개선이다
• (7개) #ops/#input이 상수 — 재사용 여지가 없는 커널. IOLB가 입력 크기를 하한으로 주고, 5개에서 점근적으로 타이트
• (2개) #ops/#input은 큰데 실제로는 타일링이 불가능한 경우 (adi, durbin). wavefront 덕분에 “최선의 OI가 상수”임을 증명한다
• (2개) IOLB가 너무 낙관적인 경우 (gramschmidt, nussinov). 모든 차원으로 타일링되지 않는 코드인데 IOLB가 이를 못 잡는다

기존 결과와의 비교가 흥미롭다. 행렬곱에 Loomis-Whitney를 처음 쓴 Irony et al.의 OI 상한이 4√2·√S인데 IOLB는 √S를 준다 — 4√2배 더 타이트하다. Ballard et al.이 6개 커널로 확장하며 얻은 8√S에 대해 IOLB는 4개에 √S, 2개에 2√S를 준다.

느슨한 케이스도 있다. stencil 계열에서 특히 그런데(fdtd-2d는 16√6배), 원인이 둘이다. (1) chain 의존성이 겹칠 수 있어 “사영 합산” 기법을 못 쓴다 — 이론적 한계다. (2) 배열이 여러 개일 때 차원당 하나만 고른다 — 구현 한계다.

PLuTo로 타일링한 스케줄을 만들고 Dinero 캐시 시뮬레이터로 실측 OI(OI_PLuTo)를 얻어, 단일 코어 / MB 8 word-per-cycle / 빠른 메모리 256kB(Skylake-X의 L2-L3 전송에 대응) 환경에서 비교한다.

• (18개) OI_PLuTo가 이미 MB 위 → compute-bound. 컴파일러가 충분히 했다
• (6개) OI_up이 MB 아래 → 알고리즘을 근본적으로 바꾸지 않는 한 영원히 bandwidth-bound다. atax, trisolv 등
• (6개) MB가 OI_PLuTo와 OI_up 사이 → 개선 여지가 있다는 신호. 실제로 확인해보니 ludcmp는 배열로 확장되지 않은 스칼라 때문에 PLuTo가 타일링을 못 하고 있었고, 손으로 확장해주니 OI가 MB 위로 올라갔다. Floyd-Warshall은 IOLB가 찾아낸 반복 공간 분해가 그대로 “코드를 이렇게 다시 쓰면 타일링된다”는 힌트가 되었다

마지막 항목이 이 도구의 실용적 가치다. 하한을 증명하는 과정에서 나온 분해 구조가 더 나은 스케줄을 찾는 단서가 된다.

성능이 아니라 에너지가 관심사라면 MB 비교는 무의미하고, OI_PLuTo와 OI_up의 간격 자체가 줄일 수 있는 데이터 이동(= 에너지)의 여지를 뜻한다.

선행 연구

• Christ et al. (2013) — Brascamp-Lieb을 쓴다는 아이디어의 출처. 하지만 연산과 데이터 원소의 연관만 보고 CDAG의 의존성을 못 잡아서 Jacobi 스텐실 같은 데서 매우 약한 하한이 나오고, 분해가 없어 루프 본문 전체를 원자적 statement로 보며(loop fission이 가능한 경우 틀린 결과가 된다), 2S-partition만 쓰고, 상수 없는 점근식만 나오며, 자동화되지 않았다
• Elango et al. (2014) — 재계산 없는 pebble game과 wavefront를 도입했지만 수동 적용이다
• Elango et al. (2015) — DFG 경로와 재사용 패턴을 연결해 처음으로 자동화했지만, O(…) 점근식만 나오고(상수가 없으면 roofline에 쓸 수 없다), 점근 하한은 겹쳐도 그냥 더할 수 있으니 분해 문제를 다루지 않았으며, loop parametrization과 같은 statement의 분해가 없고, wavefront를 쓰지 않아 adi·durbin에서 매우 약하다

정리

이 논문의 위치는 “affine 프로그램에 대한 데이터 이동 하한 증명을 자동화했다”는 것이고, 실질적 기여는 점근식이 아니라 상수까지 붙은 파라메트릭 식이라는 점이다. 상수가 있어야 특정 아키텍처의 machine balance와 직접 비교할 수 있고, 그래야 “이 코드는 최적화 여지가 남았는가 / 아니면 알고리즘을 바꿔야 하는가”를 판정할 수 있다.

과학계산과 기계학습 커널의 상당 부분이 폴리헤드럴 모델에 들어가는 affine 영역이므로, 캐시 크기·대역폭 같은 아키텍처 파라미터를 얼마나 잡아야 하는지 역으로 따지는 데에도 쓸 수 있다.