Hexagon-MLIR Compiler Stack

• Hexagon-MLIR is an open-source MLIR-based compilation stack that compiles Triton kernels and PyTorch models for Qualcomm Hexagon NPUs, emphasizing kernel fusion and memory locality.
• It employs a structured sequence of MLIR passes—including tiling, vectorization, and double buffering—to optimize data residency in tightly coupled memory and reduce DDR traffic.
• The system serves as both a high-performance compiler backend and a flexible research tool for rapid kernel innovation and advanced AI model deployment.

Hexagon-MLIR is an open-source compilation stack that targets the Qualcomm Hexagon Neural Processing Unit (NPU) and provides unified support for lowering Triton kernels and PyTorch models. Built using the MLIR framework, it is designed to accelerate AI workloads through a structured sequence of compiler passes that exploit NPU architectural features, while providing automated compilation from kernel to binary. Its central design objective is to ingest hand-written Triton kernels or subgraphs from PyTorch 2.0, fuse them into larger computation units, and generate binaries that maximize data locality in the NPU’s Tightly Coupled Memory (TCM), thereby reducing the bandwidth bottlenecks associated with library-based approaches. The system is presented as a complement to Qualcomm’s commercial toolchains and as a more flexible path for advancing AI compilation capabilities; the paper also characterizes Hexagon-MLIR as a work-in-progress with additional optimizations and capabilities under active development (Absar et al., 23 Feb 2026).

1. Motivation and design objectives

The motivating problem is the rapid evolution of AI kernels. Operator libraries in the style of cuDNN can provide very high single-operator performance, but they require painstaking manual assembly and scale poorly as new fused patterns or activation variants emerge, including attention variants and PolyCom activations (Absar et al., 23 Feb 2026). Hexagon-MLIR is positioned against that background not by rejecting highly tuned kernels, but by treating fusion as a first-class compilation concern.

A second motivation is memory-system pressure. In the library-based execution model described in the paper, successive operators round-trip intermediate data to DRAM (DDR), which exposes workloads to limited memory bandwidth. Hexagon-MLIR addresses this by ingesting Triton kernels and fusing them into “mega-kernels” that retain long chains of operations in TCM. The intended effect is to maximize data locality and drastically reduce DDR traffic. This suggests a compiler-centric alternative to deployment strategies that rely on assembling precompiled kernels around a fixed operator boundary.

The project is also framed as an open-source flexibility mechanism. Rather than replacing Qualcomm’s proprietary toolchains, it complements them with an MLIR-based stack that researchers and third-party developers can extend. The paper explicitly connects this flexibility to experimentation on new model classes such as LLMs, Mixture-of-Experts, and models with novel activations. A plausible implication is that Hexagon-MLIR is intended not only as a backend compiler, but also as a research vehicle for rapid kernel and lowering innovation.

2. System architecture and intermediate representations

Hexagon-MLIR sits atop LLVM’s MLIR framework and is organized around frontends, custom dialect and IR usage, and Hexagon-specific backends (Absar et al., 23 Feb 2026). On the frontend side, Torch-MLIR lowers PyTorch graphs, denoted $M$, into core MLIR constructs including Linalg, Standard, and SCF. In parallel, a triton-to-linalg path lowers Triton kernels, denoted $K$, into Linalg on Tensor operations. The architectural consequence is a unified mid-level representation for two distinct sources of AI computation: graph-oriented PyTorch programs and kernel-oriented Triton programs.

Within MLIR, the compilation stack uses several dialects and IR extensions with distinct roles. Linalg-on-Tensor serves as a uniform abstraction for linear algebra and element-wise operations through linalg.generic. SCF provides structured loop constructs such as scf.for and scf.forall for tiled execution. The Async dialect models fork-join HVX threads through async.execute, async.yield, and async.await_all. MemRef dialect extensions, including memref.copy, memref.dma_start, and memref.dma_wait, are used to represent explicit DDR↔TCM DMA. The stack also uses attributes and markers such as tiled_generic, db_prefetch, and db_prologue to anchor multistage transformations.

The backend lowers the transformed IR to LLVM IR, introduces calls into the Hexagon NPU runtime and the Qualcomm Hexagon Library (QHL) for math primitives, and then relies on LLVM’s Hexagon code generator for final HVX assembly generation. The paper describes the terminal artifact as a .hexagon_npu binary. Taken together, this architecture places structured MLIR rewrites at the center of the system while deferring ISA-specific code emission to the LLVM Hexagon toolchain.

3. Compilation pipeline and lowering strategy

The paper denotes the overall compilation pipeline by $\Pi$. After frontend conversion, the compiler begins from an initial IR $\mathrm{IR}_0$, and the passes compose as

$$\Pi = p_n \circ \dots \circ \text{DB} \circ M \circ V \circ T \circ F \circ C \circ \dots \circ p_0.$$

Within that sequence, the key passes are described in a rough order (Absar et al., 23 Feb 2026). The pass $C$ performs canonicalization, common-subexpression elimination, and constant propagation. The operator-fusion pass $F$ merges chains of linalg.generic operations into single “mega-generics,” thereby eliminating temporaries. The paper gives Softmax as an illustrative case: separate subtraction and exponentiation stages can be fused so that a single expression computes exp(x[i] - max_j x[j]) without storing an intermediate tensor between them.

A layout transformation and propagation stage, denoted $L$, reduces pack/unpack operations around matmuls and convolutions. Tiling, denoted $T$, acts on a linalg.generic with iteration domain $\mathcal I \subset \mathbb Z^d$ and tile sizes $K$0. It splits $K$1 into outer loops expressed as scf.for and produces a smaller generic operating on a tile resident in TCM. The transformation inserts tensor.extract_slice, bufferization.alloc_tensor(memory_space=TCM), memref.copy for DDR→TCM transfer, a tiled linalg.generic on TCM buffers, and a memref.copy for TCM→DDR write-back. Loop interchange and the tiled_generic attribute are used so that vectorizable dimensions become innermost.

The pipeline also includes a quantization stage $K$2 for int8/4 targets. Vectorization, denoted $K$3, analyzes scalar loops inside the tiled generic and emits vector.transfer_read and vector.transfer_write, or lowers to HVX intrinsics for linalg.elementwise and linalg.contraction. Multi-threading, denoted $K$4, proceeds in two phases: Form-Virtual-Threads partitions tile work into scf.forall blocks according to a polytope-size heuristic, using either block or block-cyclic distribution, and Form-Async-Threads wraps each block in async.execute ... async.yield, groups tokens, and emits async.await_all at the end. The paper states that this lowers to LLVM coroutines.

Double buffering, denoted DB, is also two-stage. A structural stage locates single-buffer tiled loops and synthesizes a ping-pong prologue plus two subkernels, db_ping_kernel and db_pong_kernel, governed by a boolean toggle. An asynchronous DMA stage replaces memref.copy with memref.dma_start and memref.dma_wait, allocates unique tags for each prefetch and write-back, and emits waits immediately before computation. Math-library insertion then replaces polynomial-friendly operations such as exp, tanh, and sqrt either with QHL calls or with MLIR polynomial approximation expansions via Horner’s method. Final lowerings perform bufferization and deallocation, lower SCF and Async to LLVM IR and Hexagon runtime calls, and generate HVX assembly and the final binary.

4. Triton ingestion and the “mega-kernel” model

A defining capability of Hexagon-MLIR is the ingestion of Triton JIT kernels, whether hand-written or generated by PyTorch-2.0-Inductor, followed by lowering into a Linalg-based IR via triton-to-linalg (Absar et al., 23 Feb 2026). This representation makes the sequence of operations and affine indexing maps explicit, which in turn enables compiler-driven fusion across what would otherwise be distinct kernel stages.

The paper presents numerically stable Softmax as the canonical example:

$K$5

In the Linalg-generic form, the stages of subtraction, exponentiation, summation, and division are explicit. Fusion collapses that sequence into a single generic body computing the entire Softmax formula at each index. The resulting fused computation is the unit that is subsequently tiled into TCM. The paper’s terminology for this fused object is the “mega-kernel.”

The significance of this design lies in memory residency. By fusing across multiple Triton operations, Hexagon-MLIR keeps intermediate temporaries in TCM registers or local buffers rather than materializing them in DDR. This suggests that the compiler is optimizing not merely arithmetic throughput, but the granularity at which memory traffic is exposed to the hardware. In that sense, the “mega-kernel” model is both an IR-level fusion strategy and a memory-hierarchy strategy.

A common misconception would be to treat Triton support as a narrow import pathway for isolated custom kernels. The paper instead presents it as a mechanism for faster deployment of new Triton kernels, including hand-written kernels and subgraphs originating in PyTorch 2.0. The emphasis is therefore on end-to-end lowering from a higher-level kernel description to an NPU binary, rather than on pointwise translation alone.

5. Memory hierarchy management and scheduling

The paper’s memory and scheduling optimizations are structured around tiling, DMA orchestration, double buffering, and HVX thread formation (Absar et al., 23 Feb 2026). In the tiling model, the computation domain is an iteration domain $K$6, and tile sizes are selected as $K$7. Outer loops take the form of scf.for over dimension ranges with steps $K$8, while the inner linalg.generic operates on a tile of shape $K$9. This formulation makes tile residency in TCM a first-order compiler invariant.

DDR↔TCM movement is then expressed through memref.subview, alloc_tensor(memspace=TCM), memref.dma_start(srcDDR, dstTCM, size, tag), computation on the TCM buffer, and memref.dma_wait(tag). The explicit DMA representation is important because it allows the compiler to reason about overlap between memory transfer and computation, rather than treating data movement as an opaque side effect.

The double-buffering model is summarized in terms of compute time per tile $\Pi$0 and memory-transfer time per tile $\Pi$1. For an ideal two-stage pipeline with perfect overlap, the total time is approximated by

$\Pi$2

The paper also gives a speedup model,

$\Pi$3

and notes a peak of $\Pi$4 at $\Pi$5. The stated interpretation is that balanced tiles, for which compute and transfer costs are comparable, are the most favorable regime for double buffering. Conversely, memory-bound kernels exhibit diminishing returns.

Multi-thread scheduling is governed by a polytope-size heuristic threshold that determines when to spawn HVX contexts. Work partitioning is either block distribution when $\Pi$6 or block-cyclic distribution for load balancing. A plausible implication is that the scheduler is designed to remain regular when the problem size admits exact partitioning, while preserving utilization when it does not.

6. Empirical behavior, limitations, and projected directions

The paper reports vectorization speedups for several kernels and datatypes (Absar et al., 23 Feb 2026).

Kernel	Configuration	Reported speedup
GELU	float16, 16 K elements	63.9×
GELU	float32, 16 K	16.1×
RMSNorm	float16, 127×513	46.5×
SiLU	float32, 16 K	7.1×
FlashAttention	float32, $\Pi$7, DIM=64	4.7×

For single-thread versus multi-thread execution, the reported behavior is size-dependent. At small sizes, specifically below $\Pi$8 K, overhead dominates and speed is below $\Pi$9. Beyond $\mathrm{IR}_0$0 K, multi-threading yields $\mathrm{IR}_0$1–$\mathrm{IR}_0$2 speedup, with a dip to approximately $\mathrm{IR}_0$3 at $\mathrm{IR}_0$4 M elements, attributed to TCM/DMA pressure. For double buffering and multi-pass interactions, the configuration labeled “Vec + MT + DB” commonly yields a further $\mathrm{IR}_0$5–$\mathrm{IR}_0$6 improvement over vectorization plus multi-threading alone. The paper further states that memory-bound kernels see diminishing double-buffering returns, consistent with the overlap model above.

The reported resource-utilization picture is correspondingly specific: HVX contexts, stated as $\mathrm{IR}_0$7–$\mathrm{IR}_0$8 per core, are fully utilized for medium-to-large tile sizes, while DMA queue occupancy and toggle buffers enable nearly $\mathrm{IR}_0$9 overlap for balanced tiles. These claims place the performance discussion squarely at the interaction boundary between code generation, memory orchestration, and hardware concurrency.

The limitations and future directions are also explicit. The paper identifies broader operator coverage, including TopK, prefix-sum, and scatter/gather, as a next step, potentially via new dialects or composite Linalg patterns. It highlights autotuning of tile sizes $$\Pi = p_n \circ \dots \circ \text{DB} \circ M \circ V \circ T \circ F \circ C \circ \dots \circ p_0.$$0 and thread counts per kernel through small benchmarks or ML-guided heuristics; quantized int8/4 support with mixed-precision accumulation; extended layout propagation to reduce pack/unpack around MatMul engines; enhanced scheduling for multi-chip or multi-NPU parallelism, including pipelined Mixture-of-Experts; and integration of Helion-generated Triton kernels together with support for upcoming Triton features such as advanced memory reuse annotations. The paper’s repeated description of Hexagon-MLIR as a work-in-progress is significant here: the current system is presented as a functioning open-source stack, but not as a finalized endpoint.