HMX JIT

A just-in-time compiler that emits HVX and HMX machine code for GEMM on the Qualcomm Hexagon NPU. Kernels are generated at runtime and specialized to the matrix shape they're called with.

Target platform: Snapdragon 8 Gen 2 (cDSP v73), Hexagon SDK 6.0.0.2.

Building & running (each needs an attached device)

Requirements:

• the Hexagon SDK 6.0.0.2.
• adb-connected v73 device.

# set your HEXAGON_SDK & ANDROID_NDK & ANDROID_SERIAL
vim build_host.sh
vim build_device.sh
vim execute_script.sh

# The JIT (builds host + device, pushes, runs the on-device)
bash execute_script.sh

Results

[bench] JIT-ed FP16 HMX GEMM throughput (1000 iters)
  shape(MxNxK)  padded          us/run GFLOPS  max_abs  max_rel  
  32x32x32      -                 0.03 1927.5  1.7e-04  1.2e-02
  64x64x64      -                 0.10 5405.0  1.7e-04  1.2e-02
  128x128x128   -                 0.57 7410.4  2.0e-04  2.0e-03  
  96x96x96      -                 0.27 6529.4  1.7e-04  2.0e-03  
  512x512x512   -                28.77 9330.1  1.8e-04  2.4e-02  
  512x512x1024  -                56.44 9512.2  1.3e-04  1.0e-03  
  30x30x30      32x32x32          0.03 1636.4  1.5e-04  2.0e-03  
  48x48x48      64x64x64          0.09 2513.5  3.1e-04  2.5e-02  
  100x100x100   128x128x128       0.53 3752.3  1.5e-04  2.0e-03  
  50x80x40      64x96x64          0.13 2388.1  1.8e-04  2.5e-03  
  33x65x17      64x96x32          0.09 848.0   1.6e-04  1.9e-03  
  32x32x64      -                 0.03 4228.1  1.7e-04  1.2e-02  
  32x32x128     -                 0.04 7281.8  2.0e-04  2.0e-03  
  32x32x256     -                 0.06 9362.3  1.2e-04  1.2e-02  
  32x32x512     -                 0.11 9362.3  1.8e-04  2.4e-02  

References

• https://docs.qualcomm.com/doc/80-N2040-62/80-N2040-62_REV_AA_Qualcomm_Hexagon_V81_HMX_Programmers_Reference_Manual.pdf
• https://docs.qualcomm.com/doc/80-N2040-54/80-N2040-54.pdf
• https://scalable.uni-jena.de/opt/hpc/chapters/assignment_soc.html#cdsp
• https://github.com/haozixu/htp-ops-lib