Futhark에서 일부 계산을 CPU에 남기기

AI 요약

Elias Smedegaard는 자동 미분(automatic differentiation)으로 희소 Jacobian 행렬을 효율적으로 계산하는 학사 논문을 작성했습니다. 이 과정에서 행렬의 희소 패턴에 해당하는 그래프를 거리 2 색칠(distance-2 colouring)해야 했습니다. 최적 그래프 색칠은 NP-hard 문제지만, 필요한 것은 최적해가 아니라 빠르게 구할 수 있는 적절한 색칠입니다. Smedegaard가 찾은 탐욕 알고리즘은 빠르지만 순차적으로 실행됩니다. 전체 작업에서 그래프 색칠은 순차 처리하고 Jacobian 계산은 GPU의 대규모 병렬성을 활용해야 합니다.

배열 메모리 위치가 만드는 문제

Futhark의 컴파일 모델은 단순합니다. GPU 백엔드를 사용하면 모든 배열을 GPU 메모리에 둡니다. 내부 표현은 CPU와 GPU 메모리를 함께 다룰 수 있지만, 구현을 간결하게 유지하려고 한쪽에 모읍니다. 이 방식은 CPU 코드가 GPU 메모리에 접근할 수 있다는 점에서 안전하지만, 속도는 보장하지 않습니다.

CPU 코드가 GPU 메모리 배열의 원소를 하나씩 읽으면 매번 통신을 설정하고 복사가 끝날 때까지 기다려야 합니다. 원소 하나를 옮기는 시간보다 통신 설정과 대기 비용이 훨씬 큽니다. 블로그 글은 원소마다 몇 마이크로초가 들 수 있다고 설명합니다. 따라서 GPU 배열을 CPU에서 반복문으로 순차 접근하면 성능이 크게 떨어집니다. 실제 그래프 색칠 코드는 그래프와 스택을 나타내는 배열을 여러 겹의 순차 반복문에서 다룹니다. GPU 백엔드에서는 순차 C 백엔드보다 성능이 세 자릿수 배 이상 나빠질 수 있습니다. 반면 C 백엔드는 색칠을 빠르게 처리하지만 Jacobian 계산까지 순차 실행합니다. 멀티코어 백엔드도 절충안이지만, 작성자는 Jacobian 계산에 GPU를 쓰고자 했습니다.

자동 메모리 배치 대신 함수 속성

근본적인 해법은 배열의 실제 사용 위치를 분석해 CPU나 GPU 메모리에 배치하는 것입니다. 필요한 경우 두 메모리 공간에 복사본을 둘 수도 있습니다. 다만 복사를 지나치게 자주 하거나 복사본을 오래 유지하면 통신 비용과 메모리 사용량이 늘어납니다. Futhark에는 단순한 순차 계산을 GPU로 옮기는 최적화가 이미 있지만, 반복문이 있는 그래프 색칠은 적용 대상이 아닙니다. 조건을 풀어 억지로 적용해도 단일 GPU 스레드의 순차 실행은 너무 느렸습니다.

작성자는 컴파일러가 데이터 이동을 자동으로 결정하는 최적화를 설계하는 대신, #[cpu_function] 속성을 추가했습니다. 이 속성을 붙인 함수는 CPU 코드로 컴파일하고 입력 배열과 중간 배열도 CPU 메모리에 둡니다. 호출부가 배열을 올바른 메모리 공간에 준비하도록 컴파일러가 복사 코드를 삽입합니다. 함수에는 보통 #[noinline]도 함께 붙여 인라인을 막습니다.

복사 시점을 조정한 구현

처음 설계에서는 함수의 출력도 CPU 메모리에 남았습니다. 그러자 그 출력 배열을 GPU 커널이 사용하는 코드가 문제가 됐습니다. 예를 들어 CPU에서 iota로 만든 배열을 반환하고, 호출부의 map이 GPU에서 실행되면 GPU 커널이 CPU 메모리 배열을 직접 참조하게 됩니다. 원칙적으로 컴파일러가 GPU 커널의 배열 접근을 분석해 필요한 배열을 복사해야 합니다. 그러나 커널이 반복문 안에서 호출될 때 반복마다 복사하지 않으면서, 복사본은 가능한 짧게 유지하는 처리가 필요합니다.

그래서 속성의 의미를 바꿨습니다. 입력과 중간 결과는 CPU 메모리에 두되, 함수 반환값은 GPU 메모리로 복사합니다. 그 결과 GPU 백엔드는 #[cpu_function] 함수 내부를 제외하면 배열이 GPU 메모리에 있다고 가정할 수 있습니다. 이 변경 뒤에도 GPU 백엔드가 iota와 replicate 결과를 항상 GPU 메모리에 둔다고 가정하는 문제가 드러났고, 관련 처리를 수정했습니다. 비슷한 버그가 다른 곳에도 남아 있을 가능성이 있다고 글은 덧붙입니다. #[cpu_function] 함수를 map에 적용할 때는 속성을 무시하고 일반 병렬 함수로 만드는 동작을 선택했습니다.

이 속성은 함수 호출 때 복사 비용이 드는 거친 해결책입니다. 그래도 원소별 복사 대신 배열 단위 복사를 하므로 오버헤드는 감당할 만했고, 그래프 색칠을 빠른 순차 코드로 처리하면서 같은 Futhark 프로그램에서 후속 수치 계산은 병렬 실행할 수 있게 됐습니다. 컴파일러 변경은 비교적 작았으며, 자동 메모리 배치 최적화가 나중에 구현돼도 이 속성을 유용한 힌트로 남길 가능성이 있다고 설명합니다.

원문: Futhark Blog / 번역·요약: Trawling