Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Hexagon-MLIR: An AI Compilation Stack For Qualcomm's Neural Processing Units (NPUs)

Mohammed Javed Absar, Muthu Baskaran, Abhikrant Sharma, Abhilash Bhandari, Ankit Aggarwal, Arun Rangasamy, Dibyendu Das, Fateme Hosseini, Franck Slama, Iulian Brumar, Jyotsna Verma, Krishnaprasad Bindumadhavan, Mitesh Kothari, Mohit Gupta, Ravishankar Kolachana, Richard Lethin, Samarth Narang, Sanjay Motilal Ladwa, Shalini Jain, Snigdha Suresh Dalvi, Tasmia Rahman, Venkat Rasagna Reddy Komatireddy, Vivek Vasudevbhai Pandya, Xiyue Shi, Zachary Zipper

TL;DR

Hexagon-MLIR is 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.

Abstract

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.

Hexagon-MLIR: An AI Compilation Stack For Qualcomm's Neural Processing Units (NPUs)

TL;DR

Hexagon-MLIR is 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.

Abstract

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.
Paper Structure (31 sections, 17 equations, 7 figures, 1 table)

This paper contains 31 sections, 17 equations, 7 figures, 1 table.

Table of Contents

  1. Introduction
  2. Organization
  3. Hexagon-MLIR: Overview
  4. Domains.
  5. Semantics preservation.
  6. Hexagon-MLIR Pass Pipeline.
  7. Hexagon-MLIR: Details
  8. Triton to Linalg
  9. Softmax
  10. Triton Kernel.
  11. Linalg-Generic
  12. Indexing maps:
  13. Iterator types:
  14. Payload region:
  15. Operator Fusion
  16. ...and 16 more sections

Figures (7)

  • Figure 1: Hexagon-MLIR — AI Compilation Stack
  • 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.
  • ...and 2 more figures