Compiler Engineering • Machine Learning

Accelerating JAX on Qualcomm Snapdragon: Building a Native QNN Backend with OpenXLA PJRT

Published by carrycooldude • August 2026 • 6 min read

Introduction

JAX has emerged as the weapon of choice for machine learning researchers and systems engineers thanks to its functional purity, composable transformations (jit, grad, vmap), and first-class compiler pipeline. Historically, however, JAX’s high-performance hardware execution has been constrained primarily to NVIDIA GPUs and Google TPUs, with CPU as a fallback.

With the advent of the Qualcomm Snapdragon® X Elite (45 TOPS Hexagon NPU) and Snapdragon mobile SoCs, a massive amount of high-efficiency tensor compute became available on edge and client devices. In this deep-dive, we explore JAX-QNN — an open-source, native backend for JAX that maps StableHLO operations directly into Qualcomm execution graphs using OpenXLA's PJRT C API standard.

Why Not Just Wrap ONNX?

A naive approach to supporting Qualcomm hardware is converting Python tensors to ONNX and invoking ONNX Runtime through Python. While functional for quick prototypes, this introduces several fundamental flaws:

Host-Device Synchronization Bottlenecks: Round-tripping NumPy arrays through Python wrappers breaks JAX’s asynchronous dispatch model and buffer lifecycle management.
Loss of Transformation Semantics: JAX functional transforms (vmap, custom_vjp, lax.scan) cannot be natively reasoned about across a static ONNX boundary.
Memory Overhead: Python-level runtime orchestration incurs CPU cache thrashing, bypassing on-chip Vector Tightly-Coupled Memory (VTCM).

System Architecture Diagram

The compiler architecture interfaces directly at the OpenXLA PJRT layer:

flowchart TD
    subgraph JAX_Frontend["1. Python / JAX Frontend"]
        A["Python Program
jax.jit(backend='qnn')"] --> B["JAX Core Engine
(Tracing & JAXPR)"]
        B --> C["MLIR Lowering
(StableHLO Bytecode)"]
    end

    subgraph PJRT_Plugin["2. Native C++ PJRT Plugin (pjrt_qnn.dll)"]
        D["PJRT_Client_Compile()"] --> E["AST Parser & Op Mapper
(csrc/stablehlo_to_qnn.cc)"]
        E --> F["Qualcomm QNN Graph Construction
(QnnGraph_create / addNode)"]
        F --> G["HTP Hardware Optimizer
VTCM Allocation & Op Fusion"]
        G --> H["QnnGraph_finalize()
(Compiled Executable Handle)"]
    end

    subgraph Hardware_Execution["3. Qualcomm Hardware Layer"]
        I["PJRT_LoadedExecutable_Execute()"] --> J["Qualcomm Hexagon NPU
(45 TOPS HTP Core)"]
        I --> K["Qualcomm Adreno GPU
(OpenCL / Direct3D)"]
        I --> L["Qualcomm Oryon CPU
(12-Core Reference)"]
    end

    C --> D
    H --> I

    classDef jaxNode fill:#0284c7,stroke:#38bdf8,stroke-width:2px,color:#ffffff;
    classDef pjrtNode fill:#7c3aed,stroke:#c084fc,stroke-width:2px,color:#ffffff;
    classDef hwNode fill:#059669,stroke:#34d399,stroke-width:2px,color:#ffffff;

    class A,B,C jaxNode;
    class D,E,F,G,H pjrtNode;
    class I,J,K,L hwNode;

Compiler Sequence & Hardware Dispatch

During execution, the native PJRT plugin compiles StableHLO IR into Qualcomm microcode and dispatches tensors directly to the Hexagon NPU:

sequenceDiagram
    autonumber
    participant JAX as JAX Runtime
    participant PJRT as PJRT Plugin (pjrt_qnn.dll)
    participant QNN as Qualcomm QNN SDK
    participant NPU as Hexagon NPU (HTP)

    JAX->>PJRT: PJRT_Client_Compile(StableHLO Module)
    PJRT->>QNN: QnnGraph_create(context, "jax_model")
    PJRT->>QNN: QnnGraph_addNode(QNN_OP_MAT_MUL, inputs, outputs)
    PJRT->>QNN: QnnGraph_addNode(QNN_OP_ELEMENT_WISE_ADD, ...)
    PJRT->>QNN: QnnGraph_addNode(QNN_OP_RELU, ...)
    PJRT->>QNN: QnnGraph_finalize()
    QNN->>NPU: Allocate VTCM Vector Memory & Generate Microcode
    QNN-->>PJRT: Executable Handle
    PJRT-->>JAX: PJRT_LoadedExecutable

    Note over JAX,NPU: Execution Phase (Sub-millisecond latency)
    JAX->>PJRT: PJRT_LoadedExecutable_Execute(input_buffers)
    PJRT->>NPU: Direct DMA / Vector Register Dispatch
    NPU-->>PJRT: Result Ready
    PJRT-->>JAX: Output PJRT_Buffer

Performance Benchmarks (Snapdragon® X Elite)

Evaluated on a Snapdragon® X Elite (X1E80100) 12-Core Oryon CPU & 45 TOPS Hexagon NPU running Windows 11 ARM64:

Workload	Host CPU	Adreno GPU	Hexagon NPU (HTP)	Speedup vs CPU
512×512 Dense Layer (GEMM + ReLU)	2.67 ms	1.12 ms	0.48 ms (2,104 FPS)	5.6x
1024×1024 Dense Layer	6.30 ms	2.95 ms	1.35 ms (740 FPS)	4.7x
2D Convolution (64x64x32, 3x3)	4.15 ms	1.45 ms	0.62 ms (1,613 FPS)	6.7x
Transformer Self-Attention (Seq 128)	3.85 ms	1.60 ms	0.71 ms (1,408 FPS)	5.4x

Quickstart: Try It in 30 Seconds

import jax
import jax.numpy as jnp
import jax_qnn

# 1. Inspect Devices
print("Devices:", jax.devices("qnn"))

# 2. Define Model
@jax.jit
def model(x, w, b):
    return jax.nn.relu(jnp.matmul(x, w) + b)

# 3. Execute on Qualcomm Hexagon NPU
x = jnp.ones((128, 512), dtype=jnp.float32)
w = jnp.ones((512, 512), dtype=jnp.float32)
b = jnp.zeros((512,), dtype=jnp.float32)

y = jax.jit(model, backend="qnn")(x, w, b)
print("Output Shape:", y.shape)
        

Summary & Resources

JAX-QNN unlocks high-efficiency, sub-millisecond JAX execution on Qualcomm Snapdragon hardware without sacrificing JAX's developer ergonomics.

GitHub Repository: github.com/carrycooldude/JAX-QNN
Documentation: carrycooldude.github.io/JAX-QNN/