딥러닝 컴파일러 프레임워크의 중간 표현식(IR) 최적화와 메모리 계층 구조를 보조하는 오퍼레이터 퓨전

현업에서 대규모 트래픽을 다루다 보니 결국 기본기로 돌아가게 되더군요. 교과서적인 지식에 그치지 않고 실제 운영 환경에서 느낀 점을 공유합니다. 인공지능 모델의 규모가 폭발적으로 증가함에 따라 단순히 프레임워크 수준에서의 하드웨어 호출만으로는 연산 오버헤드와 메모리 대역폭의 병목을 감당할 수 없는 시대가 도래했습니다. Python 기반의 고수준 API로 작성된 PyTorch나 TensorFlow 모델은 그 자체로는 하드웨어 실행이 불가능한 추상적인 연산 그래프(Computation Graph)에 불과합니다. 이를 특정 하드웨어(CPU, GPU, TPU, 혹은 커스텀 NPU 등)에서 최적의 성능을 발휘하는 극도로 효율적인 기계어로 번역하기 위해, 현대 딥러닝 시스템 구조는 AI 컴파일러(Deep Learning Compiler)라는 거대하고 복잡한 중간 계층을 적극 도입하게 되었습니다. 구글의 XLA(Accelerated Linear Algebra)나 오픈소스인 Apache TVM 같은 딥러닝 컴파일러들은 기존의 정적 컴파일러인 GCC, LLVM의 아키텍처 철학을 차용하여 최적화 과정을 하드웨어 독립적인 프론트엔드 최적화와 하드웨어 의존적인 백엔드 커널 생성을 엄격하게 분리하여 수행합니다. 딥러닝 컴파일러의 첫 번째 핵심 역할은 고수준의 그래프 수준(Graph-level) 최적화입니다. 프레임워크로부터 전달받은 신경망 아키텍처를 고수준의 중간 표현식(High-level IR)으로 변환한 뒤, 연산 그래프 전체를 조망하며 구조적인 비효율성을 색출합니다. 안개처럼 흩어져 있는 불필요한 노드들을 제거하거나 상수로 덴스하게 만들어버리는 상수 폴딩(Constant Folding), 그리고 쓸모없는 브랜치를 삭제하는 데드 코드 제거(Dead Code Elimination)가 일어납니다. 이 단계의 하이라이트는 단연코 오퍼레이터 퓨전(Operator Fusion)입니다. 딥러닝 연산은 행렬 곱셈이나 합성곱 같이 연산 집약적인 작업 뒤에 으레 ReLU, 배치 정규화(Batch Normalization) 같은 메모리 집약적인 요소별(Element-wise) 연산이 뒤따릅니다. 만약 커널이 행렬 곱셈을 마친 뒤 결과를 VRAM(메인 비디오 메모리)에 저장하고, 다시 VRAM에서 값을 읽어와 ReLU 함수를 통과시킨 후 덮어쓴다면 막대한 메모리 입출력 대기시간이 발생하게 됩니다. 오퍼레이터 퓨전 기술은 이 분리된 연산 노드들을 하나의 거대한 단일 커널 파이프라인으로 강제 융합해 버립니다. 합성곱 연산 결과를 레지스터나 초고속 온칩 SRAM 캐시에 유지한 상태로 그 자리에서 즉시 활성화 함수까지 적용한 뒤에야 메인 메모리로 내려보냄으로써 캐시 히트율을 극대화하고 느린 글로벌 메모리 엑세스를 획기적으로 감축시킵니다. 이후 모델은 저수준의 중간 표현식(Low-level IR)으로 강하(Lowering)되며, 이때부터는 컴파일러 백엔드가 타겟 하드웨어의 미시적인 아키텍처 특성을 온전히 반영하기 시작합니다. 타겟 디바이스의 벡터 레지스터 크기, 스트리밍 멀티프로세서(SM)의 개수, 캐시 라인 크기 등의 하드웨어 스펙에 맞추어 루프 전개(Loop Unrolling), 타일링(Tiling), 텐서 메모리 레이아웃 변환(NCHW에서 NHWC로의 치환 등)을 수학적으로 계산하여 최적화합니다. 흥미로운 점은 이러한 백엔드 커널 최적화 과정에서 가능한 탐색 공간이 우주적 규모로 거대하다는 것입니다. 루프를 어느 크기 단위로 잘라서 처리해야 캐시 미스를 최소화할 수 있는지 인간 엔지니어가 수작업으로 산출하는 것은 불가능에 가깝습니다. 이를 해결하기 위해 TVM의 AutoTVM 같은 기술은 스스로 다양한 커널 형태를 생성하고 시뮬레이터나 실제 하드웨어 디바이스에 수만 번 직접 실행시켜보면서 소요된 마이크로초를 측정하는 기계 학습 탐색 방식(AutoTuning)을 활용합니다. 결과적으로 AI 스스로가 자신을 실행하기 위한 최적의 C++/CUDA/PTX 머신러닝 바이너리 커널 코드를 스스로 짜치고 컴파일하는 경지에 도달하였으며 이것이 거대 AI 인프라의 보이지 않는 중추 원리입니다.