Question about achievable XDNA2 NPU throughput with IRON

[yichao-yuan-99]

Hi, I’m experimenting with IRON on an AI Max+ 395 and noticed that I’m only seeing ~700 GFLOP/s, like below, when running the example gemm pytest with 8 columns.
I am using the devel branch.

operators/gemm/test.py::test_gemm[iter0-gemm_2048x2048x2048_64x64x64_8_cols_0_bcolmaj_0_ccolmaj_0_0] 
Latency (us): 24679.3
Effective Bandwidth: 1.019712e+00 GB/s
Throughput: 6.961237e+02 GFLOP/s

PASSED

Here is the output of xrt-smi validate on my machine

Validate Device           : [0000:c6:00.1]
    Platform              : NPU Strix Halo
    Power Mode            : Default
-------------------------------------------------------------------------------
Test 1 [0000:c6:00.1]     : gemm                                                
    Details               : TOPS: 51.0
    Test Status           : [PASSED]
-------------------------------------------------------------------------------
Test 2 [0000:c6:00.1]     : latency                                             
    Details               : Average latency: 52.0 us
    Test Status           : [PASSED]
-------------------------------------------------------------------------------
Test 3 [0000:c6:00.1]     : throughput                                          
    Details               : Average throughput: 78771.0 op/s
    Test Status           : [PASSED]
-------------------------------------------------------------------------------

Based on the spec of AIE2p cores, which shows the throughput of BF16 is half of INT8, I am expecting a throughput that is close to the order of magnitude of 50/2=25TOPS.

0.7TOPS looks far away from my expectations, thus I wonder if this is due to some configuration problems, or wrong understanding, or something else?

Thanks!

[AngryLoki]

Things that you are missing:

1. You probably encountered a bug after pytest rewrite. Looks like you edited hardcoded num_aie_columns in the test, then rerun it, and it used 2 columns from cached results. Either clean build directory, or change

   IRON/operators/gemm/op.py

   Line 103 in 87e12ba

   file_name_total_base = f"{prefix}{M}x{K}x{N}_{tile_m}x{tile_k}x{tile_n}_{int(b_col_maj)}_{int(c_col_maj)}"

   to

file_name_total_base = f"{prefix}{M}x{K}x{N}_{tile_m}x{tile_k}x{tile_n}_{num_aie_columns}cols_{int(b_col_maj)}_{int(c_col_maj)}"

With 8 columns instead of 2, it reaches 2.75 TFLOPS (3.92 times faster).

2. "the throughput of BF16 is half of INT8" - this is somewhat correct. BF16 is 256 MACs/cycle, BFP16 is 512 MACs/cycle.
   If you enable emulate_bf16_mmul_with_bfp16=True, prio_accuracy=False, you will see 4.5 TFLOPS (with less accurate result).

3. Given clock 1800 MHz, the theoretical peak per core is exactly 512 MACs/cycle × 2 ops/MAC × 1.8 GHz = 1.84 TFLOPS. Per whole array = 58.8 TFLOPS. So IRON gemm (split from mlir-aie gemm) design for that size is around 8% (4.5/58.8) from theoretical peak.

4. There is a paper https://arxiv.org/abs/2511.16041v1 which used 8192×4224×4096 (microkernel 64×88×64) as a baseline with 4.8 TFLOPS. Note that they switched from 64×64×64 size because it it not optimal.

A buffer: m × k × 2 bytes × 2 buffers = 64 × 88 × 2 × 2 = 22 KB
B buffer: k × n × 2 bytes × 2 buffers = 88 × 64 × 2 × 2 = 22 KB
C buffer: m × n × 2 bytes × 2 buffers = 64 × 64 × 2 × 2 = 16 KB
Stack:    3.25 KB
────────────────────────────────────────────
Total:    63.25 KB (fits in 64 KB per-core limit)

Then they improved from the baseline, reaching 24.6 TFLOPS (42% of theoretical peak) for selected configuration. You can check that version in https://github.com/ChengyueWang/mlir-aie/tree/atb-gemm/programming_examples/ml/block_datatypes/atb_mm