In this paper, we present Hexagon-MLIR,an open-source compilation stack that targets Qualcomm Hexagon Neural Processing Unit (NPU) and provides unified support for lowering Triton kernels and PyTorch models . Built using the MLIR framework, our compiler applies a structured sequence of passes to exploit NPU architectural features to accelerate AI workloads. It enables faster deployment of new Triton kernels (hand-written or subgraphs from PyTorch 2.0), for our target by providing automated compilation from kernel to binary. By ingesting Triton kernels, we generate mega-kernels that maximize data locality in the NPU's Tightly Coupled Memory (TCM), reducing the bandwidth bottlenecks inherent in library-based approaches. This initiative complements our commercial toolchains by providing developers with an open-source MLIR-based compilation stack that gives them a path to advance AI compilation capabilities through a more flexible approach. Hexagon-MLIR is a work-in-progress, and we are continuing to add many more optimizations and capabilities in this effort.
Figures & tables
Figure 1: Hexagon-MLIR — AI Compilation Stack
Test
Speedup
Notes
Gaussian Error Linear Unit (GELU)
63.9x
float16, elements = 16384
Gaussian Error Linear Unit (GELU)
16.1x
float32, elements = 16384
Double Buffered GELU
15.4x
float32, element= 16 x 16384, double buffered
Flash Attention
4.7x
float32, N_CTX=1024, DIM_HEAD=64, BLOCK_N=64
Root Mean Square Norm (RMS-norm)
46.5x
float16, elements = 127 x 513,
Sigmoid Linear Unit (SiLU)
4.8
float16, elements = 16384
Table 1: Speedups by vectorization.
Figure 2: Execution times of single-threaded vs multi-threaded GELU kernel implementations.
Figure 3: Speedup of multi-threaded GELU over increasing problem-sizes.
Figure 4: Speedup for double buffering with perfect overlap as memory fraction varies.
Figure 5: Vector-Add 2D performance across optimization passes.
AI accelerator inference studies hardware and software for running neural networks faster and more efficiently across data-center, edge, and mobile devices.