MLIR 기반 AI 커널 컴파일러에서의 지연 시간 은폐 및 병렬성 분석

Analyzing Latency Hiding and Parallelism in an MLIR-based AI Kernel Compiler

에지 디바이스를 위한 AI 커널 컴파일은 계층적 메모리와 명시적 데이터 이동이 존재하는 환경에서 병렬성을 활용하고 메모리 지연 시간을 은폐하는 컴파일러의 역량에 좌우된다. 본 논문은 MLIR 기반 컴파일 파이프라인에서 컴파일러가 제어하는 세 가지 메커니즘인 벡터화(Vec), 하드웨어 컨텍스트 전반에 걸친 멀티스레딩(MT), 그리고 DMA 전송과 연산을 중첩하기 위해 핑퐁 스크래치패드 버퍼를 사용하는 이중 버퍼링(DB)에 대한 벤치마크 방법론과 그 결과를 보고한다. Triton/Inductor가 생성한 커널을 사용하여 Vec, MT 및 DB의 기여도를 분리하는 절제 단계(ablation ladder)를 제시하며, 대표적인 활성화 커널로 GELU를 사용하여 문제 크기에 따라 MT의 속도 향상이 어떻게 확장되는지 정량화한다. 결과에 따르면, 벡터화는 대역폭에 민감한 커널에서 주된 성능 이득을 제공하고, MT는 스케줄링 오버헤드가 상쇄된 후 상당한 개선을 가져오며, DB는 전송과 연산이 중첩될 수 있을 때(즉, 완전히 메모리 바운드되거나 완전히 연산 바운드되는 극단적인 상황을 벗어날 때) 추가적인 이점을 제공하는 것으로 나타났다.

AI kernel compilation for edge devices depends on the compiler's ability to exploit parallelism and hide memory latency in the presence of hierarchical memory and explicit data movement. This paper reports a benchmark methodology and corresponding results for three compiler-controlled mechanisms in an MLIR-based compilation pipeline: vectorization (Vec), multi-threading (MT) across hardware contexts, and double buffering (DB) using ping--pong scratchpad buffers to overlap DMA transfers with compute. Using Triton/Inductor-generated kernels, we present an ablation ladder that separates the contribution of Vec, MT, and DB, and we quantify how MT speedup scales with problem size using GELU as a representative activation kernel. The results show that vectorization provides the primary gain for bandwidth-sensitive kernels, MT delivers substantial improvements once scheduling overhead is amortized, and DB provides additional benefit when transfers and compute can be overlapped (i.e., outside the extremes of purely memory-bound or purely compute-bound behavior).