---
title: Hexagon-MLIR Compiler Stack
url: https://www.emergentmind.com/topics/hexagon-mlir
type: topic
---

# Hexagon-MLIR Compiler Stack

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 [2602.19762].

## 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 [2602.19762]. 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 [2602.19762]. 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 [2602.19762]. 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 \(\mathbf t=(t_1,\dots,t_d)\). It splits \(\mathcal I\) 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 \(Q\) for int8/4 targets. Vectorization, denoted \(V\), 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 \(M\), 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` [2602.19762]. 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:

$$
\sigma(z_i)
= \frac{\exp(z_i - \max_j z_j)}{\sum_k \exp(z_k - \max_j z_j)}.
$$

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 [2602.19762]. In the tiling model, the computation domain is an iteration domain \(\mathcal I \subset \mathbb Z^d\), and tile sizes are selected as \(\mathbf t = (t_1,\dots,t_d)\). Outer loops take the form of `scf.for` over dimension ranges with steps \(t_0, t_1,\dots\), while the inner `linalg.generic` operates on a tile of shape \(\mathbf t\). 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 \(C\) and memory-transfer time per tile \(M\). For an ideal two-stage pipeline with perfect overlap, the total time is approximated by

$$
\text{total} \approx N_{\rm tiles}\,\max(C,\,M).
$$

The paper also gives a speedup model,

$$
\text{Speedup} = \frac{1}{\max(m,\,1-m)},\quad m=\tfrac{M}{C+M},
$$

and notes a peak of \(2\times\) at \(m=0.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 \(N\%T=0\) 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 [2602.19762].

| 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, \(N_{nx}=1024\), DIM=64 | 4.7× |

For single-thread versus multi-thread execution, the reported behavior is size-dependent. At small sizes, specifically below \(16\) K, overhead dominates and speed is below \(1\). Beyond \(32\) K, multi-threading yields \(2.3\)–\(3.9\times\) speedup, with a dip to approximately \(3.4\times\) at \(1\) M elements, attributed to TCM/DMA pressure. For double buffering and multi-pass interactions, the configuration labeled “Vec + MT + DB” commonly yields a further \(1.1\)–\(1.3\times\) 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 \(4\)–\(8\) per core, are fully utilized for medium-to-large tile sizes, while DMA queue occupancy and toggle buffers enable nearly \(100\%\) 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 \(\mathbf t\) 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.

Source: https://www.emergentmind.com/topics/hexagon-mlir