Speed up NAX split-K by better tuning and routing and fix NAX addmm

[angeloskath]

Well the title says it all

• Split-K matmuls were being routed to the non NAX version
• Adds some tuning on the M5 Max to improve further
• AddMM was completely broken on NAX this fixes it
  • Also it is a bit unclear whether the const_for_loop helped with codegen cause there is some measurement variance but it is quite clean and comes with compile time bounds checking that already found a bug when I was fixing this.

Here are the first measurements on the Max

Shape (MxK)  Before  After   Speedup
-------------------------------------
128x128      0.009   0.01    0.9
128x256      0.012   0.01    1.2
128x512      0.014   0.012   1.17
128x1024     0.015   0.016   0.94
256x128      0.013   0.011   1.18
256x256      0.016   0.011   1.45
256x512      0.017   0.015   1.13
256x1024     0.025   0.016   1.56
256x2048     0.042   0.025   1.68
512x128      0.018   0.015   1.2
512x256      0.021   0.015   1.4
512x512      0.049   0.018   2.72
512x1024     0.057   0.029   1.97
512x2048     0.105   0.052   2.02
512x4096     0.195   0.086   2.27
1024x256     0.057   0.03    1.9
1024x512     0.105   0.048   2.19
1024x1024    0.09    0.09    1.0
1024x2048    0.164   0.164   1.0
1024x4096    0.313   0.301   1.04
1280x1280    0.164   0.164   1.0
1280x2048    0.238   0.238   1.0
1280x4096    0.47    0.459   1.02
2048x512     0.159   0.159   1.0
2048x1024    0.291   0.292   1.0
2048x2048    0.55    0.55    1.0
2048x4096    1.099   1.101   1.0
4096x1024    1.092   1.093   1.0
4096x2048    2.172   2.173   1.0
4096x4096    4.344   4.346   1.0
8192x2048    8.695   8.711   1.0
8192x4096    17.431  17.471  1.0
8192x8192    35.13   35.100  1.0

the benchmark is running (MxK @ KxM) @ MxK chained.

I will probably add some numbers from other M5s before merging.

[zcbenz]

👍