Paper deep dive
Analyzing Latency Hiding and Parallelism in an MLIR-based AI Kernel Compiler
Javed Absar, Samarth Narang, Muthu Baskaran
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 7/20/2026, 8:13:04 PM
Summary
This paper presents a benchmark methodology for analyzing latency hiding and parallelism in an MLIR-based AI kernel compiler for edge devices. It evaluates three compiler-controlled mechanisms: vectorization (Vec), multi-threading (MT), and double buffering (DB). Using an ablation ladder on vector-add and GELU kernels, the study finds that vectorization provides the primary performance gain for bandwidth-sensitive kernels, multi-threading offers substantial improvements once scheduling overhead is amortized, and double buffering adds incremental benefit by overlapping DMA transfers with compute.
Entities (8)
Relation Signals (7)
Multi-threading (MT) ā delivers ā substantial improvements
confidence 95% Ā· MT delivers substantial improvements once scheduling overhead is amortized
Vectorization (Vec) ā provides ā primary gain
confidence 95% Ā· The results show that vectorization provides the primary gain for bandwidth-sensitive kernels
Double Buffering (DB) ā provides ā additional benefit
confidence 95% Ā· DB provides additional benefit when transfers and compute can be overlapped
GELU ā usedtoevaluate ā Multi-threading (MT)
confidence 92% Ā· we quantify how MT speedup scales with problem size using GELU as a representative activation kernel
vec-add-2d ā usedtoevaluate ā Vectorization (Vec)
confidence 92% Ā· For the vector-add microbenchmark, we evaluate the ladder variants Scalar, Vec, Vec+MT, and Vec+MT+DB.
MLIR ā supports ā Double Buffering (DB)
confidence 90% Ā· We implement DB in two stages ... This is the composition used to realize the Vec+MT+DB rung of our ladder.
MLIR ā supports ā Multi-threading (MT)
confidence 90% Ā· Our implementation is built in MLIR [2] ... We implement MT as a two-stage pipeline
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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).
Tags
Links
- Source: https://arxiv.org/abs/2602.20204v1
- Canonical: https://arxiv.org/abs/2602.20204v1
Trouble viewing inline? Open PDF directly ā
Full Text
11,993 characters extracted from source content.
Expand or collapse full text
Analyzing Latency Hiding and Parallelism in an MLIR-based AI Kernel Compiler Javed Absar1, Samarth Narang2, and Muthu Baskaran2 1Qualcomm Technologies International, Ltd. 2Qualcomm Technologies, Inc. mabsar,samanara,muthub@qti.qualcomm.com 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, among other factors. 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). 1 Introduction Kernel optimization is difficult to automate: hand-written kernels remain hard to beat, yet production systems require portable and maintainable code generation across rapidly evolving architectures. On edge NPUs, performance is shaped by hierarchical memory, explicit DMA-managed transfers, and the need to schedule work so compute stays busy while transfers are in flight. End-to-end model results are essential, but they often obscure which mapping mechanisms improve performance and why. We therefore adopt a reproducible, kernel-oriented evaluation methodology aligned with Triton/Inductor-style code generation, benchmarking representative operator kernels under controlled problem sizes to separately quantify the impact of vectorization (Vec), multi-threading (MT), and double buffering (DB) [5]. Vec exploits data-level parallelism; MT exploits loop- and region-level parallelism by distributing independent tiles across hardware contexts; and DB reduces stall time by overlapping memory transfers with compute. To attribute gains to specific mechanisms, we report results using a simple ablation ladder: ScalarāVecāVec+MTāVec+MT+DB. Scalarā Vecā Vec+MTā Vec+MT+DB. In this ladder, Vec isolates SIMD-style lowering, Vec+MT quantifies incremental thread-level speedup, and Vec+MT+DB evaluates whether an explicit latency-hiding schedule provides additional improvement once Vec and MT are already in place. In addition to reporting results, we describe the concrete MLIR IR patterns and pass structure used to realize MT and DB, making the methodology easy to reproduce and extend. 2 Implementation Details: Multi-threading and Double Buffering Our implementation is built in MLIR [2] and follows a design principle: keep the intent expressed in structured IR for as long as possible, and lower to runtime constructs only after the compiler has imposed a schedule. This improves portability and makes transformations easier to validate. 2.1 Multi-threading (MT) Hardware model. The target NPU supports multi-threaded vector execution, where each hardware thread is associated with an independent vector context (e.g., vector register and predicate state). This enables concurrent execution of independent tiles of the same kernel, provided the tiled iteration space can be partitioned without cross-thread dependences. Two-stage MT lowering. We implement MT as a two-stage pipeline that preserves structured parallelism and then introduces an explicit forkājoin. First, Form-Virtual-Threads rewrites a tiled kernel (e.g., linalg.generic) into an explicitly parallel form using scf.forall. The pass uses a size-based profitability heuristic over the tile space and selects a distribution policy (block vs. block-cyclic) to balance work when ranges are uneven. Second, Form-Async-Threads lowers scf.forall to a forkājoin representation using MLIR ās Async [3]. Each tile becomes an async.execute region that produces a token; tokens are collected into an async group, and async.await_all forms a barrier before subsequent dependent computation. Canonical forkājoin skeleton in IR. The generated pattern is intentionally small and regular, which makes it straightforward to lower into a coroutine/task runtime [3]: %group = async.create_group %N scf.for %tile = ... %tok = async.execute /* tile body */ async.yield async.add_to_group %tok, %group async.await_all %group Why Async (instead of lowering MT directly). Keeping MT as a structured forkājoin in Async preserves parallel semantics in a declarative form until late lowering, while still enabling a straightforward translation to runtime scheduling [3]. 2.2 Double buffering (DB) Double buffering is a software-pipelining strategy that overlaps transfers and compute by alternating between two scratchpad buffers (ping and pong). The technique is closely related to modulo scheduling and pipelined loop execution [1, 4]. We implement DB in two stages, separating the construction of a pipelined schedule from the introduction of target-specific asynchronous transport primitives. Stage 1: structural pipelining. Stage 1 matches a single-buffered tiled-loop ānormal formā typically created by tiling: memref.subviewāmemref.allocāmemref.copy, memref.subviewā memref.allocā memref.copy, followed by compute and then a write-back sequence. From this structure, the pass builds an explicit pingāpong schedule. It emits a prologue that prefetches the first tile into ping buffers, then rebuilds the main loop into two alternating sub-kernels. Each sub-kernel (i) issues a prefetch of the next tile into the opposite buffer, (i) computes using the current buffer, and (i) reconstructs storeback by rematerializing subviews at the current induction variable. A boolean toggle selects ping vs. pong each iteration, and the pass attaches lightweight attributes as anchors so the next stage can reliably identify prefetch/compute/storeback regions. Stage 2: asynchronous DMA integration. Stage 2 replaces synchronous copies with explicit asynchronous DMA operations. Prefetch paths are rewritten to memref.dma_start with distinct ping/pong tags, and memref.dma_wait is inserted immediately before compute to ensure tile residency in TCM. Storeback copies can be handled similarly, and the pass emits balanced tag deallocations. By staging the transformation, we first establish a correct schedule and only then map it onto the targetās transport interface, mirroring classic compiler practice for pipelined loops [1, 4]. Composing DB with MT. DB composes naturally with MT. Once memref.dma_wait enforces that a tile is resident in TCM, the compute region can exploit scf.forall (virtual threads) and/or the lowered forkājoin representation to execute multiple sub-tiles concurrently. This is the composition used to realize the Vec+MT+DB rung of our ladder. 3 Benchmark Setup and Methodology We evaluate two representative kernels. The first is a bandwidth-centric 2D vector addition microbenchmark with shape [64,128Ć128][64,128Ć 128] elements, which is useful for understanding the relative impact of vectorization and latency-hiding in a memory-heavy regime. The second is GELU, a common activation kernel representative of transformer inference subgraphs; we use a Triton implementation of GELU as our concrete kernel instance. For the vector-add microbenchmark, we evaluate the ladder variants Scalar, Vec, Vec+MT, and Vec+MT+DB. For GELU, we report end-to-end latency across a problem-size sweep for both single-threaded and multi-threaded execution to expose MT overhead amortization and scaling. We report latency in microseconds (Ī¼ ) and, where applicable, speedup as the ratio of single-thread to multi-thread time. 4 Results 4.1 vec-add-2d Ablation Ladder Figure 1 summarizes the ladder for vec-add-2d. Vectorization accounts for the dominant improvement (132,479 Ī¼ ā 3,210 Ī¼ , ā¼ 41.3Ć), consistent with a bandwidth-oriented kernel that benefits immediately from SIMD-style lowering. MT and DB provide smaller but measurable incremental gains (3,210 Ī¼ ā 3,000 Ī¼ ā 2,689 Ī¼ ), suggesting that once vectorization is in place, remaining headroom comes from reducing synchronization overheads and partially overlapping transfer/compute effects rather than from increasing arithmetic throughput. Figure 1: vec-add-2d ([64,128Ć128][64,128Ć 128] elements) ablation ladder: Scalar (132,479 Ī¼ ), Vec (3,210 Ī¼ ), Vec+MT (3,000 Ī¼ ), Vec+MT+DB (2,689 Ī¼ ). 4.2 GELU: Single-thread vs Multi-thread Scaling Figure 2 reports GELU latency over a size sweep and Figure 3 shows the corresponding multi-thread speedup. MT improves performance across all tested sizes and the speedup grows with problem size, reaching ā¼ 3.91Ć at 1,048,576 elements (12,947 Ī¼ single-thread vs 3,313 Ī¼ multi-thread). This trend indicates that fixed overheads associated with forkājoin scheduling are amortized as per-thread work increases, while saturation at larger sizes suggests emerging shared bottlenecks such as memory bandwidth or barrier costs. 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. 4.3 Discussion Across these kernels, vectorization provides the key first-order win, especially in bandwidth-sensitive regimes. MT then yields substantial additional improvement when there is sufficient parallel slack and the kernel is large enough to amortize scheduling overhead. DB provides incremental benefit when transfers and compute are both significant; in the purely memory- bound limit, DB is constrained by transfer bandwidth, and in the purely compute-bound limit, overlap opportunities are limited. 5 Conclusion We presented a compact, reproducible methodology for analyzing vectorization, multi-threading, and double buffering in an MLIR-based kernel compiler. The proposed ladder makes it easier to attribute performance improvements to specific compiler mechanisms rather than reporting only end-to-end speedups. Our measurements show that vectorization is foundational, MT can provide large gains once overhead is amortized, and DB adds incremental benefit when transfer/compute overlap exists. Future work will broaden the benchmark set (e.g., RMSNorm, softmax) and connect these measurements to a predictive model that relates overlap efficiency to DMA throughput and scratchpad capacity. References [1] M. S. Lam (1988) Software pipelining: an effective scheduling technique for vliw machines. In Proceedings of the ACM SIGPLAN 1988 Conference on Programming Language Design and Implementation (PLDI), Cited by: Ā§2.2, Ā§2.2. [2] C. Lattner, M. Amini, U. Bondhugula, A. Cohen, A. Davis, J. Pienaar, R. Riddle, T. Shpeisman, N. Vasilache, and O. Zinenko (2020) MLIR: a compiler infrastructure for the end of mooreās law. External Links: 2002.11054 Cited by: Ā§2. [3] LLVM Project (2026) MLIR async dialect. Note: https://mlir.llvm.org/docs/Dialects/AsyncDialect/Accessed 2026-01-21 Cited by: Ā§2.1, Ā§2.1, Ā§2.1. [4] B. R. Rau (1994) Iterative modulo scheduling: an algorithm for software pipelining loops. Proceedings of the 27th Annual International Symposium on Microarchitecture (MICRO). Cited by: Ā§2.2, Ā§2.2. [5] P. Tillet, H. T. Kung, and D. Cox (2019) Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages (MAPL), Cited by: Ā§1.