Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Javed Absar, Samarth Narang, Muthu Baskaran

TL;DR

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).

Abstract

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).

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

TL;DR

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).

Abstract

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).
Paper Structure (17 sections, 2 equations, 3 figures)

This paper contains 17 sections, 2 equations, 3 figures.

Table of Contents

  1. Introduction
  2. Implementation Details: Multi-threading and Double Buffering
  3. Multi-threading (MT)
  4. Hardware model.
  5. Two-stage MT lowering.
  6. Canonical fork--join skeleton in IR.
  7. Why Async (instead of lowering MT directly).
  8. Double buffering (DB)
  9. Stage 1: structural pipelining.
  10. Stage 2: asynchronous DMA integration.
  11. Composing DB with MT.
  12. Benchmark Setup and Methodology
  13. Results
  14. vec-add-2d Ablation Ladder
  15. GELU: Single-thread vs Multi-thread Scaling
  16. ...and 2 more sections

Figures (3)

  • Figure 1: vec-add-2d ($[64, 128 \times 128]$ elements) ablation ladder: Scalar (132,479 $\mu$s), Vec (3,210 $\mu$s), Vec+MT (3,000 $\mu$s), Vec+MT+DB (2,689 $\mu$s).
  • Figure 2: GELU latency vs problem size for single-threaded and multi-threaded execution.
  • Figure 3: GELU multi-thread speedup (Single/Multi) vs problem size.