JAX-QNN Architecture Specification

A comprehensive technical deep-dive into the JAX compiler pipeline, OpenXLA PJRT C API integration, and Qualcomm QNN native graph lowering.

            Core Principle: JAX-QNN integrates cleanly through OpenXLA's official GetPjRtApi C ABI. No custom JAX forks, no hacks, and zero Python execution overhead during compiled inference.
        

1. Compilation Pipeline Overview

JAX Python Code (@jax.jit)
       │
   JAXPR IR
       │
StableHLO MLIR Lowering
       │
PJRT_Client_Compile (C API)
       │
csrc/stablehlo_to_qnn.cc  ──▶  Direct AST walk & Op Mapping
       │
Qualcomm QNN Graph Construction (QnnGraph_create / QnnGraph_addNode)
       │
Qualcomm HTP Optimization (VTCM Allocation & Op Fusion)
       │
QnnGraph_finalize (Compiled Executable Handle)
       │
PJRT_LoadedExecutable_Execute (Direct Dispatch to Hexagon NPU 45 TOPS)
        

2. PJRT Plugin Dispatch (`pjrt_qnn.dll`)

The shared library exports the standard C ABI symbol GetPjRtApi() defined in csrc/pjrt/pjrt_c_api.h:

PJRT_Client_Create: Initializes the Qualcomm AI Engine Direct backend session (HTP / GPU / CPU).
PJRT_Client_Devices: Discovers addressable compute accelerators (e.g. Hexagon NPU ID 0).
PJRT_Client_Compile: Consumes serialized StableHLO bytecode and translates to a finalized QnnGraph_t.
PJRT_LoadedExecutable_Execute: Executes tensor buffers on Qualcomm unified memory with zero-copy where supported.

3. StableHLO to QNN Operator Matrix

StableHLO Primitive	Qualcomm QNN Operation	NPU Acceleration (HTP)
`stablehlo.dot_general`	`QNN_OP_MAT_MUL`	✔ Native Tensor Core
`stablehlo.add` / `subtract`	`QNN_OP_ELEMENT_WISE_ADD / SUB`	✔ Hexagon Vector Unit
`stablehlo.multiply` / `divide`	`QNN_OP_ELEMENT_WISE_MUL / DIV`	✔ Hexagon Vector Unit
`stablehlo.maximum` (ReLU)	`QNN_OP_RELU` / `QNN_OP_ELEMENT_WISE_MAX`	✔ Fused Activation
`stablehlo.convolution`	`QNN_OP_CONV_2D`	✔ 2D/3D HW Engine
`stablehlo.reduce`	`QNN_OP_REDUCE_SUM / MEAN`	✔ Vector Reduction
`stablehlo.reshape` / `transpose`	`QNN_OP_RESHAPE / TRANSPOSE`	✔ Zero-Copy View

4. Memory Management (VTCM vs DDR)

Qualcomm Snapdragon processors feature Vector Tightly-Coupled Memory (VTCM) directly embedded alongside the Hexagon NPU execution pipelines. During compilation, intermediate tensor allocations are placed in VTCM, completely eliminating main memory bus traffic during deep layer execution.