usearch 2.21.3 → 2.25.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (602) hide show
  1. package/README.md +24 -24
  2. package/binding.gyp +15 -37
  3. package/c/CMakeLists.txt +4 -5
  4. package/c/README.md +7 -7
  5. package/c/lib.cpp +32 -14
  6. package/c/test.c +83 -1
  7. package/c/usearch.h +24 -3
  8. package/cpp/CMakeLists.txt +4 -0
  9. package/cpp/README.md +5 -5
  10. package/cpp/bench.cpp +97 -116
  11. package/cpp/test.cpp +117 -20
  12. package/include/usearch/index.hpp +235 -52
  13. package/include/usearch/index_dense.hpp +105 -40
  14. package/include/usearch/index_plugins.hpp +1387 -204
  15. package/javascript/dist/cjs/usearch.d.ts +8 -3
  16. package/javascript/dist/cjs/usearch.js +12 -4
  17. package/javascript/dist/esm/usearch.d.ts +8 -3
  18. package/javascript/dist/esm/usearch.js +12 -4
  19. package/javascript/lib.cpp +17 -0
  20. package/javascript/usearch.ts +14 -6
  21. package/numkong/.clang-tidy +30 -0
  22. package/{simsimd → numkong}/.cmake-format.py +18 -0
  23. package/numkong/.editorconfig +12 -0
  24. package/numkong/.git-blame-ignore-revs +123 -0
  25. package/numkong/.mailmap +15 -0
  26. package/numkong/.spi.yml +11 -0
  27. package/numkong/.swift-format +14 -0
  28. package/numkong/CITATION.cff +81 -0
  29. package/numkong/CONTRIBUTING.md +500 -0
  30. package/numkong/Cargo.toml +54 -0
  31. package/numkong/MANIFEST.in +8 -0
  32. package/numkong/Package.swift +84 -0
  33. package/numkong/README.md +569 -0
  34. package/numkong/bench/README.md +217 -0
  35. package/numkong/bench/bench-browser.html +323 -0
  36. package/numkong/bench/bench.cpp +344 -0
  37. package/numkong/bench/bench.hpp +816 -0
  38. package/numkong/bench/bench.js +148 -0
  39. package/numkong/bench/bench.mjs +616 -0
  40. package/numkong/bench/bench_attention.cpp +113 -0
  41. package/numkong/bench/bench_cast.cpp +264 -0
  42. package/numkong/bench/bench_cross_amx.cpp +62 -0
  43. package/numkong/bench/bench_cross_arm.cpp +245 -0
  44. package/numkong/bench/bench_cross_blas.cpp +309 -0
  45. package/numkong/bench/bench_cross_loongarch.cpp +106 -0
  46. package/numkong/bench/bench_cross_power.cpp +84 -0
  47. package/numkong/bench/bench_cross_rvv.cpp +121 -0
  48. package/numkong/bench/bench_cross_serial.cpp +151 -0
  49. package/numkong/bench/bench_cross_sme.cpp +163 -0
  50. package/numkong/bench/bench_cross_wasm.cpp +142 -0
  51. package/numkong/bench/bench_cross_x86.cpp +455 -0
  52. package/numkong/bench/bench_curved.cpp +210 -0
  53. package/numkong/bench/bench_dot.cpp +315 -0
  54. package/numkong/bench/bench_each.cpp +503 -0
  55. package/numkong/bench/bench_geospatial.cpp +112 -0
  56. package/numkong/bench/bench_maxsim.cpp +154 -0
  57. package/numkong/bench/bench_mesh.cpp +148 -0
  58. package/numkong/bench/bench_probability.cpp +60 -0
  59. package/numkong/bench/bench_reduce.cpp +371 -0
  60. package/numkong/bench/bench_set.cpp +73 -0
  61. package/numkong/bench/bench_sparse.cpp +198 -0
  62. package/numkong/bench/bench_spatial.cpp +347 -0
  63. package/numkong/bench/bench_trigonometry.cpp +137 -0
  64. package/numkong/bench/swift/Bench.swift +506 -0
  65. package/numkong/binding.gyp +169 -0
  66. package/numkong/build.rs +484 -0
  67. package/numkong/c/dispatch.h +512 -0
  68. package/numkong/c/dispatch_bf16.c +447 -0
  69. package/numkong/c/dispatch_bf16c.c +52 -0
  70. package/numkong/c/dispatch_e2m3.c +296 -0
  71. package/numkong/c/dispatch_e3m2.c +284 -0
  72. package/numkong/c/dispatch_e4m3.c +318 -0
  73. package/numkong/c/dispatch_e5m2.c +316 -0
  74. package/numkong/c/dispatch_f16.c +452 -0
  75. package/numkong/c/dispatch_f16c.c +58 -0
  76. package/numkong/c/dispatch_f32.c +434 -0
  77. package/numkong/c/dispatch_f32c.c +99 -0
  78. package/numkong/c/dispatch_f64.c +348 -0
  79. package/numkong/c/dispatch_f64c.c +98 -0
  80. package/numkong/c/dispatch_i16.c +96 -0
  81. package/numkong/c/dispatch_i32.c +89 -0
  82. package/numkong/c/dispatch_i4.c +153 -0
  83. package/numkong/c/dispatch_i64.c +86 -0
  84. package/numkong/c/dispatch_i8.c +348 -0
  85. package/numkong/c/dispatch_other.c +348 -0
  86. package/numkong/c/dispatch_u1.c +193 -0
  87. package/numkong/c/dispatch_u16.c +124 -0
  88. package/numkong/c/dispatch_u32.c +118 -0
  89. package/numkong/c/dispatch_u4.c +153 -0
  90. package/numkong/c/dispatch_u64.c +102 -0
  91. package/numkong/c/dispatch_u8.c +364 -0
  92. package/numkong/c/numkong.c +940 -0
  93. package/numkong/cmake/nk_arm_isa_probes.cmake +61 -0
  94. package/numkong/cmake/nk_isa_probe.cmake +90 -0
  95. package/numkong/cmake/nk_loongarch_isa_probes.cmake +15 -0
  96. package/numkong/cmake/nk_math_lib.cmake +46 -0
  97. package/numkong/cmake/nk_power_isa_probes.cmake +15 -0
  98. package/numkong/cmake/nk_riscv_isa_probes.cmake +19 -0
  99. package/numkong/cmake/nk_x86_isa_probes.cmake +35 -0
  100. package/numkong/cmake/toolchain-aarch64-gnu.cmake +86 -0
  101. package/numkong/cmake/toolchain-android-arm64.cmake +25 -0
  102. package/numkong/cmake/toolchain-android-armv7.cmake +25 -0
  103. package/numkong/cmake/toolchain-loongarch64-gnu.cmake +68 -0
  104. package/numkong/cmake/toolchain-ppc64le-gnu.cmake +74 -0
  105. package/numkong/cmake/toolchain-riscv64-gnu.cmake +107 -0
  106. package/numkong/cmake/toolchain-riscv64-llvm.cmake +114 -0
  107. package/numkong/cmake/toolchain-wasi.cmake +77 -0
  108. package/numkong/cmake/toolchain-wasm.cmake +70 -0
  109. package/numkong/cmake/toolchain-wasm64.cmake +62 -0
  110. package/numkong/cmake/toolchain-x86_64-llvm.cmake +47 -0
  111. package/numkong/include/README.md +643 -0
  112. package/numkong/include/module.modulemap +129 -0
  113. package/numkong/include/numkong/attention/sapphireamx.h +1363 -0
  114. package/numkong/include/numkong/attention/sme.h +2075 -0
  115. package/numkong/include/numkong/attention.h +49 -0
  116. package/numkong/include/numkong/capabilities.h +930 -0
  117. package/numkong/include/numkong/cast/README.md +262 -0
  118. package/numkong/include/numkong/cast/diamond.h +64 -0
  119. package/numkong/include/numkong/cast/haswell.h +823 -0
  120. package/numkong/include/numkong/cast/icelake.h +475 -0
  121. package/numkong/include/numkong/cast/loongsonasx.h +252 -0
  122. package/numkong/include/numkong/cast/neon.h +1174 -0
  123. package/numkong/include/numkong/cast/powervsx.h +449 -0
  124. package/numkong/include/numkong/cast/rvv.h +970 -0
  125. package/numkong/include/numkong/cast/sapphire.h +262 -0
  126. package/numkong/include/numkong/cast/serial.h +2336 -0
  127. package/numkong/include/numkong/cast/skylake.h +915 -0
  128. package/numkong/include/numkong/cast/v128relaxed.h +537 -0
  129. package/numkong/include/numkong/cast.h +254 -0
  130. package/numkong/include/numkong/cast.hpp +44 -0
  131. package/numkong/include/numkong/curved/README.md +223 -0
  132. package/numkong/include/numkong/curved/genoa.h +182 -0
  133. package/numkong/include/numkong/curved/haswell.h +276 -0
  134. package/numkong/include/numkong/curved/neon.h +329 -0
  135. package/numkong/include/numkong/curved/neonbfdot.h +211 -0
  136. package/numkong/include/numkong/curved/rvv.h +305 -0
  137. package/numkong/include/numkong/curved/serial.h +207 -0
  138. package/numkong/include/numkong/curved/skylake.h +457 -0
  139. package/numkong/include/numkong/curved/smef64.h +530 -0
  140. package/numkong/include/numkong/curved.h +513 -0
  141. package/numkong/include/numkong/curved.hpp +144 -0
  142. package/numkong/include/numkong/dot/README.md +442 -0
  143. package/numkong/include/numkong/dot/alder.h +563 -0
  144. package/numkong/include/numkong/dot/diamond.h +158 -0
  145. package/numkong/include/numkong/dot/genoa.h +289 -0
  146. package/numkong/include/numkong/dot/haswell.h +1723 -0
  147. package/numkong/include/numkong/dot/icelake.h +997 -0
  148. package/numkong/include/numkong/dot/loongsonasx.h +671 -0
  149. package/numkong/include/numkong/dot/neon.h +869 -0
  150. package/numkong/include/numkong/dot/neonbfdot.h +243 -0
  151. package/numkong/include/numkong/dot/neonfhm.h +358 -0
  152. package/numkong/include/numkong/dot/neonfp8.h +323 -0
  153. package/numkong/include/numkong/dot/neonsdot.h +622 -0
  154. package/numkong/include/numkong/dot/powervsx.h +752 -0
  155. package/numkong/include/numkong/dot/rvv.h +722 -0
  156. package/numkong/include/numkong/dot/rvvbb.h +72 -0
  157. package/numkong/include/numkong/dot/rvvbf16.h +123 -0
  158. package/numkong/include/numkong/dot/rvvhalf.h +129 -0
  159. package/numkong/include/numkong/dot/sapphire.h +141 -0
  160. package/numkong/include/numkong/dot/serial.h +874 -0
  161. package/numkong/include/numkong/dot/sierra.h +405 -0
  162. package/numkong/include/numkong/dot/skylake.h +1116 -0
  163. package/numkong/include/numkong/dot/sve.h +420 -0
  164. package/numkong/include/numkong/dot/svebfdot.h +75 -0
  165. package/numkong/include/numkong/dot/svehalf.h +168 -0
  166. package/numkong/include/numkong/dot/svesdot.h +90 -0
  167. package/numkong/include/numkong/dot/v128relaxed.h +1293 -0
  168. package/numkong/include/numkong/dot.h +1136 -0
  169. package/numkong/include/numkong/dot.hpp +94 -0
  170. package/numkong/include/numkong/dots/README.md +496 -0
  171. package/numkong/include/numkong/dots/alder.h +117 -0
  172. package/numkong/include/numkong/dots/diamond.h +86 -0
  173. package/numkong/include/numkong/dots/genoa.h +100 -0
  174. package/numkong/include/numkong/dots/graniteamx.h +733 -0
  175. package/numkong/include/numkong/dots/haswell.h +310 -0
  176. package/numkong/include/numkong/dots/icelake.h +180 -0
  177. package/numkong/include/numkong/dots/loongsonasx.h +176 -0
  178. package/numkong/include/numkong/dots/neon.h +123 -0
  179. package/numkong/include/numkong/dots/neonbfdot.h +59 -0
  180. package/numkong/include/numkong/dots/neonfhm.h +96 -0
  181. package/numkong/include/numkong/dots/neonfp8.h +99 -0
  182. package/numkong/include/numkong/dots/neonsdot.h +144 -0
  183. package/numkong/include/numkong/dots/powervsx.h +194 -0
  184. package/numkong/include/numkong/dots/rvv.h +2593 -0
  185. package/numkong/include/numkong/dots/sapphireamx.h +4017 -0
  186. package/numkong/include/numkong/dots/serial.h +2867 -0
  187. package/numkong/include/numkong/dots/sierra.h +100 -0
  188. package/numkong/include/numkong/dots/skylake.h +207 -0
  189. package/numkong/include/numkong/dots/sme.h +5041 -0
  190. package/numkong/include/numkong/dots/smebi32.h +477 -0
  191. package/numkong/include/numkong/dots/smef64.h +1328 -0
  192. package/numkong/include/numkong/dots/smehalf.h +47 -0
  193. package/numkong/include/numkong/dots/v128relaxed.h +317 -0
  194. package/numkong/include/numkong/dots.h +3110 -0
  195. package/numkong/include/numkong/dots.hpp +641 -0
  196. package/numkong/include/numkong/each/README.md +465 -0
  197. package/numkong/include/numkong/each/haswell.h +1658 -0
  198. package/numkong/include/numkong/each/icelake.h +272 -0
  199. package/numkong/include/numkong/each/neon.h +1138 -0
  200. package/numkong/include/numkong/each/neonbfdot.h +211 -0
  201. package/numkong/include/numkong/each/neonhalf.h +318 -0
  202. package/numkong/include/numkong/each/rvv.h +1121 -0
  203. package/numkong/include/numkong/each/sapphire.h +343 -0
  204. package/numkong/include/numkong/each/serial.h +282 -0
  205. package/numkong/include/numkong/each/skylake.h +1562 -0
  206. package/numkong/include/numkong/each/v128relaxed.h +562 -0
  207. package/numkong/include/numkong/each.h +2232 -0
  208. package/numkong/include/numkong/each.hpp +434 -0
  209. package/numkong/include/numkong/geospatial/README.md +147 -0
  210. package/numkong/include/numkong/geospatial/haswell.h +633 -0
  211. package/numkong/include/numkong/geospatial/neon.h +615 -0
  212. package/numkong/include/numkong/geospatial/rvv.h +701 -0
  213. package/numkong/include/numkong/geospatial/serial.h +309 -0
  214. package/numkong/include/numkong/geospatial/skylake.h +613 -0
  215. package/numkong/include/numkong/geospatial/v128relaxed.h +652 -0
  216. package/numkong/include/numkong/geospatial.h +453 -0
  217. package/numkong/include/numkong/geospatial.hpp +235 -0
  218. package/numkong/include/numkong/matrix.hpp +336 -0
  219. package/numkong/include/numkong/maxsim/README.md +187 -0
  220. package/numkong/include/numkong/maxsim/alder.h +511 -0
  221. package/numkong/include/numkong/maxsim/genoa.h +115 -0
  222. package/numkong/include/numkong/maxsim/haswell.h +553 -0
  223. package/numkong/include/numkong/maxsim/icelake.h +480 -0
  224. package/numkong/include/numkong/maxsim/neonsdot.h +394 -0
  225. package/numkong/include/numkong/maxsim/sapphireamx.h +877 -0
  226. package/numkong/include/numkong/maxsim/serial.h +490 -0
  227. package/numkong/include/numkong/maxsim/sme.h +993 -0
  228. package/numkong/include/numkong/maxsim/v128relaxed.h +280 -0
  229. package/numkong/include/numkong/maxsim.h +571 -0
  230. package/numkong/include/numkong/maxsim.hpp +133 -0
  231. package/numkong/include/numkong/mesh/README.md +219 -0
  232. package/numkong/include/numkong/mesh/haswell.h +2013 -0
  233. package/numkong/include/numkong/mesh/neon.h +1821 -0
  234. package/numkong/include/numkong/mesh/neonbfdot.h +771 -0
  235. package/numkong/include/numkong/mesh/rvv.h +936 -0
  236. package/numkong/include/numkong/mesh/serial.h +721 -0
  237. package/numkong/include/numkong/mesh/skylake.h +2344 -0
  238. package/numkong/include/numkong/mesh/v128relaxed.h +1081 -0
  239. package/numkong/include/numkong/mesh.h +666 -0
  240. package/numkong/include/numkong/mesh.hpp +701 -0
  241. package/numkong/include/numkong/numkong.h +78 -0
  242. package/numkong/include/numkong/numkong.hpp +58 -0
  243. package/numkong/include/numkong/probability/README.md +177 -0
  244. package/numkong/include/numkong/probability/haswell.h +267 -0
  245. package/numkong/include/numkong/probability/neon.h +255 -0
  246. package/numkong/include/numkong/probability/rvv.h +414 -0
  247. package/numkong/include/numkong/probability/serial.h +181 -0
  248. package/numkong/include/numkong/probability/skylake.h +324 -0
  249. package/numkong/include/numkong/probability.h +380 -0
  250. package/numkong/include/numkong/probability.hpp +120 -0
  251. package/numkong/include/numkong/random.h +50 -0
  252. package/numkong/include/numkong/random.hpp +285 -0
  253. package/numkong/include/numkong/reduce/README.md +552 -0
  254. package/numkong/include/numkong/reduce/alder.h +636 -0
  255. package/numkong/include/numkong/reduce/genoa.h +201 -0
  256. package/numkong/include/numkong/reduce/haswell.h +3785 -0
  257. package/numkong/include/numkong/reduce/icelake.h +549 -0
  258. package/numkong/include/numkong/reduce/neon.h +3980 -0
  259. package/numkong/include/numkong/reduce/neonbfdot.h +142 -0
  260. package/numkong/include/numkong/reduce/neonfhm.h +293 -0
  261. package/numkong/include/numkong/reduce/neonsdot.h +357 -0
  262. package/numkong/include/numkong/reduce/rvv.h +3533 -0
  263. package/numkong/include/numkong/reduce/serial.h +753 -0
  264. package/numkong/include/numkong/reduce/sierra.h +338 -0
  265. package/numkong/include/numkong/reduce/skylake.h +3794 -0
  266. package/numkong/include/numkong/reduce/sve.h +52 -0
  267. package/numkong/include/numkong/reduce/v128relaxed.h +2302 -0
  268. package/numkong/include/numkong/reduce.h +1582 -0
  269. package/numkong/include/numkong/reduce.hpp +778 -0
  270. package/numkong/include/numkong/scalar/README.md +89 -0
  271. package/numkong/include/numkong/scalar/haswell.h +122 -0
  272. package/numkong/include/numkong/scalar/loongsonasx.h +74 -0
  273. package/numkong/include/numkong/scalar/neon.h +122 -0
  274. package/numkong/include/numkong/scalar/neonhalf.h +70 -0
  275. package/numkong/include/numkong/scalar/powervsx.h +96 -0
  276. package/numkong/include/numkong/scalar/rvv.h +211 -0
  277. package/numkong/include/numkong/scalar/sapphire.h +74 -0
  278. package/numkong/include/numkong/scalar/serial.h +332 -0
  279. package/numkong/include/numkong/scalar/v128relaxed.h +56 -0
  280. package/numkong/include/numkong/scalar.h +696 -0
  281. package/numkong/include/numkong/set/README.md +179 -0
  282. package/numkong/include/numkong/set/haswell.h +334 -0
  283. package/numkong/include/numkong/set/icelake.h +485 -0
  284. package/numkong/include/numkong/set/loongsonasx.h +181 -0
  285. package/numkong/include/numkong/set/neon.h +363 -0
  286. package/numkong/include/numkong/set/powervsx.h +326 -0
  287. package/numkong/include/numkong/set/rvv.h +226 -0
  288. package/numkong/include/numkong/set/rvvbb.h +117 -0
  289. package/numkong/include/numkong/set/serial.h +174 -0
  290. package/numkong/include/numkong/set/sve.h +187 -0
  291. package/numkong/include/numkong/set/v128relaxed.h +240 -0
  292. package/numkong/include/numkong/set.h +471 -0
  293. package/numkong/include/numkong/set.hpp +114 -0
  294. package/numkong/include/numkong/sets/README.md +149 -0
  295. package/numkong/include/numkong/sets/haswell.h +63 -0
  296. package/numkong/include/numkong/sets/icelake.h +66 -0
  297. package/numkong/include/numkong/sets/loongsonasx.h +52 -0
  298. package/numkong/include/numkong/sets/neon.h +61 -0
  299. package/numkong/include/numkong/sets/powervsx.h +65 -0
  300. package/numkong/include/numkong/sets/serial.h +43 -0
  301. package/numkong/include/numkong/sets/smebi32.h +1135 -0
  302. package/numkong/include/numkong/sets/v128relaxed.h +58 -0
  303. package/numkong/include/numkong/sets.h +382 -0
  304. package/numkong/include/numkong/sparse/README.md +156 -0
  305. package/numkong/include/numkong/sparse/icelake.h +463 -0
  306. package/numkong/include/numkong/sparse/neon.h +288 -0
  307. package/numkong/include/numkong/sparse/serial.h +117 -0
  308. package/numkong/include/numkong/sparse/sve2.h +504 -0
  309. package/numkong/include/numkong/sparse/turin.h +322 -0
  310. package/numkong/include/numkong/sparse.h +363 -0
  311. package/numkong/include/numkong/sparse.hpp +113 -0
  312. package/numkong/include/numkong/spatial/README.md +439 -0
  313. package/numkong/include/numkong/spatial/alder.h +604 -0
  314. package/numkong/include/numkong/spatial/diamond.h +240 -0
  315. package/numkong/include/numkong/spatial/genoa.h +222 -0
  316. package/numkong/include/numkong/spatial/haswell.h +979 -0
  317. package/numkong/include/numkong/spatial/icelake.h +1067 -0
  318. package/numkong/include/numkong/spatial/loongsonasx.h +483 -0
  319. package/numkong/include/numkong/spatial/neon.h +846 -0
  320. package/numkong/include/numkong/spatial/neonbfdot.h +164 -0
  321. package/numkong/include/numkong/spatial/neonfp8.h +258 -0
  322. package/numkong/include/numkong/spatial/neonsdot.h +429 -0
  323. package/numkong/include/numkong/spatial/powervsx.h +738 -0
  324. package/numkong/include/numkong/spatial/rvv.h +991 -0
  325. package/numkong/include/numkong/spatial/rvvbf16.h +128 -0
  326. package/numkong/include/numkong/spatial/rvvhalf.h +120 -0
  327. package/numkong/include/numkong/spatial/serial.h +347 -0
  328. package/numkong/include/numkong/spatial/sierra.h +516 -0
  329. package/numkong/include/numkong/spatial/skylake.h +605 -0
  330. package/numkong/include/numkong/spatial/sve.h +246 -0
  331. package/numkong/include/numkong/spatial/svebfdot.h +125 -0
  332. package/numkong/include/numkong/spatial/svehalf.h +136 -0
  333. package/numkong/include/numkong/spatial/svesdot.h +143 -0
  334. package/numkong/include/numkong/spatial/v128relaxed.h +969 -0
  335. package/numkong/include/numkong/spatial.h +1681 -0
  336. package/numkong/include/numkong/spatial.hpp +183 -0
  337. package/numkong/include/numkong/spatials/README.md +580 -0
  338. package/numkong/include/numkong/spatials/alder.h +94 -0
  339. package/numkong/include/numkong/spatials/diamond.h +82 -0
  340. package/numkong/include/numkong/spatials/genoa.h +94 -0
  341. package/numkong/include/numkong/spatials/graniteamx.h +173 -0
  342. package/numkong/include/numkong/spatials/haswell.h +219 -0
  343. package/numkong/include/numkong/spatials/icelake.h +113 -0
  344. package/numkong/include/numkong/spatials/loongsonasx.h +153 -0
  345. package/numkong/include/numkong/spatials/neon.h +109 -0
  346. package/numkong/include/numkong/spatials/neonbfdot.h +60 -0
  347. package/numkong/include/numkong/spatials/neonfhm.h +92 -0
  348. package/numkong/include/numkong/spatials/neonfp8.h +111 -0
  349. package/numkong/include/numkong/spatials/neonsdot.h +143 -0
  350. package/numkong/include/numkong/spatials/powervsx.h +153 -0
  351. package/numkong/include/numkong/spatials/rvv.h +1976 -0
  352. package/numkong/include/numkong/spatials/sapphireamx.h +1149 -0
  353. package/numkong/include/numkong/spatials/serial.h +248 -0
  354. package/numkong/include/numkong/spatials/sierra.h +96 -0
  355. package/numkong/include/numkong/spatials/skylake.h +184 -0
  356. package/numkong/include/numkong/spatials/sme.h +1927 -0
  357. package/numkong/include/numkong/spatials/smef64.h +483 -0
  358. package/numkong/include/numkong/spatials/v128relaxed.h +240 -0
  359. package/numkong/include/numkong/spatials.h +3374 -0
  360. package/numkong/include/numkong/spatials.hpp +509 -0
  361. package/numkong/include/numkong/tensor.hpp +1981 -0
  362. package/numkong/include/numkong/trigonometry/README.md +184 -0
  363. package/numkong/include/numkong/trigonometry/haswell.h +653 -0
  364. package/numkong/include/numkong/trigonometry/neon.h +638 -0
  365. package/numkong/include/numkong/trigonometry/rvv.h +700 -0
  366. package/numkong/include/numkong/trigonometry/serial.h +703 -0
  367. package/numkong/include/numkong/trigonometry/skylake.h +725 -0
  368. package/numkong/include/numkong/trigonometry/v128relaxed.h +667 -0
  369. package/numkong/include/numkong/trigonometry.h +467 -0
  370. package/numkong/include/numkong/trigonometry.hpp +166 -0
  371. package/numkong/include/numkong/types.h +1695 -0
  372. package/numkong/include/numkong/types.hpp +6027 -0
  373. package/numkong/include/numkong/vector.hpp +766 -0
  374. package/numkong/javascript/@numkong-darwin-arm64/package.json +26 -0
  375. package/numkong/javascript/@numkong-darwin-x64/package.json +26 -0
  376. package/numkong/javascript/@numkong-linux-arm64/package.json +25 -0
  377. package/numkong/javascript/@numkong-linux-x64/package.json +25 -0
  378. package/numkong/javascript/@numkong-win32-arm64/package.json +25 -0
  379. package/numkong/javascript/@numkong-win32-x64/package.json +25 -0
  380. package/numkong/javascript/README.md +254 -0
  381. package/numkong/javascript/node-gyp-build.d.ts +4 -0
  382. package/numkong/javascript/numkong-browser.ts +40 -0
  383. package/numkong/javascript/numkong-wasm.ts +750 -0
  384. package/numkong/javascript/numkong.c +712 -0
  385. package/numkong/javascript/numkong.ts +613 -0
  386. package/{simsimd → numkong}/javascript/tsconfig-base.json +6 -3
  387. package/numkong/javascript/tsconfig-cjs.json +13 -0
  388. package/{simsimd → numkong}/javascript/tsconfig-esm.json +1 -1
  389. package/numkong/javascript/types.ts +784 -0
  390. package/numkong/numkong.gypi +104 -0
  391. package/numkong/package.json +108 -0
  392. package/numkong/probes/arm_neon.c +8 -0
  393. package/numkong/probes/arm_neon_bfdot.c +9 -0
  394. package/numkong/probes/arm_neon_fhm.c +9 -0
  395. package/numkong/probes/arm_neon_half.c +8 -0
  396. package/numkong/probes/arm_neon_sdot.c +9 -0
  397. package/numkong/probes/arm_neonfp8.c +9 -0
  398. package/numkong/probes/arm_sme.c +16 -0
  399. package/numkong/probes/arm_sme2.c +16 -0
  400. package/numkong/probes/arm_sme2p1.c +16 -0
  401. package/numkong/probes/arm_sme_bf16.c +16 -0
  402. package/numkong/probes/arm_sme_bi32.c +16 -0
  403. package/numkong/probes/arm_sme_f64.c +16 -0
  404. package/numkong/probes/arm_sme_fa64.c +14 -0
  405. package/numkong/probes/arm_sme_half.c +16 -0
  406. package/numkong/probes/arm_sme_lut2.c +15 -0
  407. package/numkong/probes/arm_sve.c +18 -0
  408. package/numkong/probes/arm_sve2.c +20 -0
  409. package/numkong/probes/arm_sve2p1.c +18 -0
  410. package/numkong/probes/arm_sve_bfdot.c +20 -0
  411. package/numkong/probes/arm_sve_half.c +18 -0
  412. package/numkong/probes/arm_sve_sdot.c +21 -0
  413. package/numkong/probes/loongarch_lasx.c +12 -0
  414. package/numkong/probes/power_vsx.c +12 -0
  415. package/numkong/probes/probe.js +127 -0
  416. package/numkong/probes/riscv_rvv.c +14 -0
  417. package/numkong/probes/riscv_rvv_bb.c +15 -0
  418. package/numkong/probes/riscv_rvv_bf16.c +17 -0
  419. package/numkong/probes/riscv_rvv_half.c +14 -0
  420. package/numkong/probes/wasm_v128relaxed.c +11 -0
  421. package/numkong/probes/x86_alder.c +17 -0
  422. package/numkong/probes/x86_diamond.c +17 -0
  423. package/numkong/probes/x86_genoa.c +17 -0
  424. package/numkong/probes/x86_graniteamx.c +19 -0
  425. package/numkong/probes/x86_haswell.c +11 -0
  426. package/numkong/probes/x86_icelake.c +17 -0
  427. package/numkong/probes/x86_sapphire.c +16 -0
  428. package/numkong/probes/x86_sapphireamx.c +18 -0
  429. package/numkong/probes/x86_sierra.c +17 -0
  430. package/numkong/probes/x86_skylake.c +15 -0
  431. package/numkong/probes/x86_turin.c +17 -0
  432. package/numkong/rust/README.md +698 -0
  433. package/numkong/rust/capabilities.rs +96 -0
  434. package/numkong/rust/cast.rs +314 -0
  435. package/numkong/rust/curved.rs +391 -0
  436. package/numkong/rust/each.rs +2721 -0
  437. package/numkong/rust/geospatial.rs +343 -0
  438. package/numkong/rust/matrix.rs +5424 -0
  439. package/numkong/rust/maxsim.rs +540 -0
  440. package/numkong/rust/mesh.rs +837 -0
  441. package/numkong/rust/numkong.rs +409 -0
  442. package/numkong/rust/probability.rs +272 -0
  443. package/numkong/rust/reduce.rs +1249 -0
  444. package/numkong/rust/set.rs +194 -0
  445. package/numkong/rust/sparse.rs +513 -0
  446. package/numkong/rust/spatial.rs +1815 -0
  447. package/numkong/rust/tensor.rs +7565 -0
  448. package/numkong/rust/types.rs +3666 -0
  449. package/numkong/rust/vector.rs +1899 -0
  450. package/numkong/rustfmt.toml +2 -0
  451. package/numkong/swift/Matrix.swift +1388 -0
  452. package/numkong/swift/MaxSim.swift +151 -0
  453. package/numkong/swift/NumKong.swift +283 -0
  454. package/numkong/swift/Numerics.swift +329 -0
  455. package/numkong/swift/README.md +433 -0
  456. package/numkong/swift/Scalar.swift +741 -0
  457. package/numkong/swift/Tensor.swift +209 -0
  458. package/numkong/test/README.md +315 -0
  459. package/numkong/test/playwright.config.ts +16 -0
  460. package/numkong/test/swift/Test.swift +534 -0
  461. package/numkong/test/test-browser-bundle-runner.mjs +36 -0
  462. package/numkong/test/test-browser-bundle.html +188 -0
  463. package/numkong/test/test-browser-runner.mjs +43 -0
  464. package/numkong/test/test-browser.html +305 -0
  465. package/numkong/test/test-dtypes.mjs +453 -0
  466. package/numkong/test/test-wasm.mjs +267 -0
  467. package/numkong/test/test.cpp +442 -0
  468. package/numkong/test/test.hpp +580 -0
  469. package/numkong/test/test.mjs +363 -0
  470. package/numkong/test/test_base.py +1041 -0
  471. package/numkong/test/test_cast.cpp +242 -0
  472. package/numkong/test/test_cross.hpp +551 -0
  473. package/numkong/test/test_cross.py +467 -0
  474. package/numkong/test/test_cross_amx.cpp +50 -0
  475. package/numkong/test/test_cross_arm.cpp +225 -0
  476. package/numkong/test/test_cross_blas.cpp +354 -0
  477. package/numkong/test/test_cross_loongarch.cpp +90 -0
  478. package/numkong/test/test_cross_power.cpp +73 -0
  479. package/numkong/test/test_cross_rvv.cpp +109 -0
  480. package/numkong/test/test_cross_serial.cpp +181 -0
  481. package/numkong/test/test_cross_sme.cpp +148 -0
  482. package/numkong/test/test_cross_wasm.cpp +141 -0
  483. package/numkong/test/test_cross_x86.cpp +442 -0
  484. package/numkong/test/test_curved.cpp +182 -0
  485. package/numkong/test/test_curved.py +207 -0
  486. package/numkong/test/test_dot.cpp +353 -0
  487. package/numkong/test/test_dot.py +296 -0
  488. package/numkong/test/test_each.cpp +437 -0
  489. package/numkong/test/test_each.py +756 -0
  490. package/numkong/test/test_geospatial.cpp +155 -0
  491. package/numkong/test/test_geospatial.py +241 -0
  492. package/numkong/test/test_maxsim.cpp +138 -0
  493. package/numkong/test/test_maxsim.py +209 -0
  494. package/numkong/test/test_mesh.cpp +203 -0
  495. package/numkong/test/test_mesh.py +394 -0
  496. package/numkong/test/test_probability.cpp +136 -0
  497. package/numkong/test/test_probability.py +166 -0
  498. package/numkong/test/test_reduce.cpp +563 -0
  499. package/numkong/test/test_reduce.py +781 -0
  500. package/numkong/test/test_set.cpp +151 -0
  501. package/numkong/test/test_set.py +120 -0
  502. package/numkong/test/test_similarities.py +579 -0
  503. package/numkong/test/test_sparse.cpp +121 -0
  504. package/numkong/test/test_sparse.py +118 -0
  505. package/numkong/test/test_spatial.cpp +518 -0
  506. package/numkong/test/test_spatial.py +269 -0
  507. package/numkong/test/test_tensor.cpp +884 -0
  508. package/numkong/test/test_tensor.py +1417 -0
  509. package/numkong/test/test_trigonometry.cpp +156 -0
  510. package/numkong/test/test_trigonometry.py +180 -0
  511. package/package.json +7 -8
  512. package/prebuilds/darwin-arm64+x64/usearch.node +0 -0
  513. package/prebuilds/linux-arm64/usearch.node +0 -0
  514. package/prebuilds/linux-x64/usearch.node +0 -0
  515. package/fp16/CMakeLists.txt +0 -176
  516. package/fp16/LICENSE +0 -11
  517. package/fp16/README.md +0 -20
  518. package/fp16/bench/alt-element.cc +0 -89
  519. package/fp16/bench/from-alt-array.cc +0 -175
  520. package/fp16/bench/from-ieee-array.cc +0 -399
  521. package/fp16/bench/ieee-element.cc +0 -272
  522. package/fp16/bench/to-alt-array.cc +0 -91
  523. package/fp16/bench/to-ieee-array.cc +0 -303
  524. package/fp16/cmake/DownloadGoogleBenchmark.cmake +0 -15
  525. package/fp16/cmake/DownloadGoogleTest.cmake +0 -15
  526. package/fp16/cmake/DownloadPSimd.cmake +0 -15
  527. package/fp16/configure.py +0 -65
  528. package/fp16/confu.yaml +0 -8
  529. package/fp16/include/fp16/__init__.py +0 -0
  530. package/fp16/include/fp16/avx.py +0 -50
  531. package/fp16/include/fp16/avx2.py +0 -53
  532. package/fp16/include/fp16/bitcasts.h +0 -92
  533. package/fp16/include/fp16/fp16.h +0 -451
  534. package/fp16/include/fp16/psimd.h +0 -131
  535. package/fp16/include/fp16.h +0 -11
  536. package/fp16/jni/Android.mk +0 -6
  537. package/fp16/jni/Application.mk +0 -4
  538. package/fp16/test/alt-from-fp32-value.cc +0 -498
  539. package/fp16/test/alt-to-fp32-bits.cc +0 -263
  540. package/fp16/test/alt-to-fp32-psimd.cc +0 -145
  541. package/fp16/test/alt-to-fp32-value.cc +0 -348
  542. package/fp16/test/alt-to-fp32x2-psimd.cc +0 -245
  543. package/fp16/test/bitcasts.cc +0 -56
  544. package/fp16/test/ieee-from-fp32-value.cc +0 -523
  545. package/fp16/test/ieee-to-fp32-bits.cc +0 -333
  546. package/fp16/test/ieee-to-fp32-psimd.cc +0 -312
  547. package/fp16/test/ieee-to-fp32-value.cc +0 -420
  548. package/fp16/test/ieee-to-fp32x2-psimd.cc +0 -541
  549. package/fp16/test/peachpy/alt-xmm-to-fp32-xmm-avx.cc +0 -96
  550. package/fp16/test/peachpy/alt-xmm-to-fp32-ymm-avx2.cc +0 -96
  551. package/fp16/test/peachpy/stubs.py +0 -38
  552. package/fp16/test/tables.cc +0 -530
  553. package/fp16/test/tables.h +0 -13
  554. package/fp16/third-party/THHalf.h +0 -92
  555. package/fp16/third-party/eigen-half.h +0 -170
  556. package/fp16/third-party/float16-compressor.h +0 -80
  557. package/fp16/third-party/half.hpp +0 -2916
  558. package/fp16/third-party/npy-halffloat.h +0 -328
  559. package/simsimd/.git-blame-ignore-revs +0 -2
  560. package/simsimd/CONTRIBUTING.md +0 -321
  561. package/simsimd/Cargo.toml +0 -49
  562. package/simsimd/MANIFEST.in +0 -4
  563. package/simsimd/Package.swift +0 -28
  564. package/simsimd/README.md +0 -1493
  565. package/simsimd/binding.gyp +0 -35
  566. package/simsimd/build.rs +0 -70
  567. package/simsimd/c/lib.c +0 -374
  568. package/simsimd/include/module.modulemap +0 -4
  569. package/simsimd/include/simsimd/binary.h +0 -486
  570. package/simsimd/include/simsimd/curved.h +0 -1541
  571. package/simsimd/include/simsimd/dot.h +0 -1853
  572. package/simsimd/include/simsimd/elementwise.h +0 -2465
  573. package/simsimd/include/simsimd/geospatial.h +0 -43
  574. package/simsimd/include/simsimd/mesh.h +0 -69
  575. package/simsimd/include/simsimd/probability.h +0 -606
  576. package/simsimd/include/simsimd/simsimd.h +0 -2567
  577. package/simsimd/include/simsimd/sparse.h +0 -1384
  578. package/simsimd/include/simsimd/spatial.h +0 -2335
  579. package/simsimd/include/simsimd/types.h +0 -666
  580. package/simsimd/javascript/fallback.ts +0 -233
  581. package/simsimd/javascript/lib.c +0 -125
  582. package/simsimd/javascript/node-gyp-build.d.ts +0 -1
  583. package/simsimd/javascript/simsimd.ts +0 -179
  584. package/simsimd/javascript/tsconfig-cjs.json +0 -8
  585. package/simsimd/package-lock.json +0 -1761
  586. package/simsimd/package.json +0 -61
  587. package/simsimd/rust/lib.rs +0 -1821
  588. package/simsimd/scripts/bench.cxx +0 -1170
  589. package/simsimd/scripts/bench.ipynb +0 -187
  590. package/simsimd/scripts/bench.js +0 -127
  591. package/simsimd/scripts/bench.rs +0 -876
  592. package/simsimd/scripts/bench_vectors.py +0 -734
  593. package/simsimd/scripts/bench_vectors_live.py +0 -213
  594. package/simsimd/scripts/test.c +0 -187
  595. package/simsimd/scripts/test.mjs +0 -220
  596. package/simsimd/scripts/test.py +0 -1505
  597. package/simsimd/swift/SimSIMD.swift +0 -149
  598. package/simsimd/swift/Test.swift +0 -128
  599. /package/{simsimd → numkong}/CODE_OF_CONDUCT.md +0 -0
  600. /package/{simsimd → numkong}/LICENSE +0 -0
  601. /package/{simsimd → numkong}/javascript/dist-package-cjs.json +0 -0
  602. /package/{simsimd → numkong}/javascript/dist-package-esm.json +0 -0
@@ -0,0 +1,1149 @@
1
+ /**
2
+ * @brief Batched Spatial Distances for Sapphire Rapids (AMX) with AVX-512 Finalization.
3
+ * @file include/numkong/spatials/sapphireamx.h
4
+ * @author Ash Vardanian
5
+ * @date February 23, 2026
6
+ *
7
+ * @sa include/numkong/spatials.h
8
+ */
9
+ #ifndef NK_SPATIALS_SAPPHIREAMX_H
10
+ #define NK_SPATIALS_SAPPHIREAMX_H
11
+
12
+ #if NK_TARGET_X8664_
13
+ #if NK_TARGET_SAPPHIREAMX
14
+
15
+ #include "numkong/spatial/skylake.h"
16
+ #include "numkong/spatial/serial.h"
17
+ #include "numkong/dots/sapphireamx.h"
18
+
19
+ #if defined(__cplusplus)
20
+ extern "C" {
21
+ #endif
22
+
23
+ #if defined(__clang__)
24
+ #pragma clang attribute push( \
25
+ __attribute__((target( \
26
+ "avx2,avx512f,avx512vl,avx512bw,avx512dq,avx512fp16,avx512vbmi,f16c,fma,bmi,bmi2,amx-tile,amx-bf16,amx-int8"))), \
27
+ apply_to = function)
28
+ #elif defined(__GNUC__)
29
+ #pragma GCC push_options
30
+ #pragma GCC target("avx2", "avx512f", "avx512vl", "avx512bw", "avx512dq", "avx512fp16", "avx512vbmi", "f16c", "fma", \
31
+ "bmi", "bmi2", "amx-tile", "amx-bf16", "amx-int8")
32
+ #endif
33
+
34
+ #pragma region Row Finalize Helpers
35
+
36
+ NK_INTERNAL void nk_angulars_row_f32dots_sapphireamx_(nk_f32_t *results, nk_f32_t const *norms, nk_f32_t query_norm_sq,
37
+ nk_size_t count) {
38
+ __m512 query_norm_sq_f32x16 = _mm512_set1_ps(query_norm_sq);
39
+ nk_size_t i = 0;
40
+ for (; i + 16 <= count; i += 16) {
41
+ __m512 dots_f32x16 = _mm512_loadu_ps(results + i);
42
+ __m512 norms_f32x16 = _mm512_loadu_ps(norms + i);
43
+ __m512 products_f32x16 = _mm512_mul_ps(query_norm_sq_f32x16, norms_f32x16);
44
+ __m512 rsqrt_f32x16 = nk_rsqrt_f32x16_skylake_(products_f32x16);
45
+ __m512 normalized_f32x16 = _mm512_mul_ps(dots_f32x16, rsqrt_f32x16);
46
+ __m512 angular_f32x16 = _mm512_sub_ps(_mm512_set1_ps(1.0f), normalized_f32x16);
47
+ _mm512_storeu_ps(results + i, _mm512_max_ps(angular_f32x16, _mm512_setzero_ps()));
48
+ }
49
+ if (i < count) {
50
+ __mmask16 tail = (__mmask16)((1u << (count - i)) - 1);
51
+ __m512 dots_f32x16 = _mm512_maskz_loadu_ps(tail, results + i);
52
+ __m512 norms_f32x16 = _mm512_maskz_loadu_ps(tail, norms + i);
53
+ __m512 products_f32x16 = _mm512_mul_ps(query_norm_sq_f32x16, norms_f32x16);
54
+ __m512 rsqrt_f32x16 = nk_rsqrt_f32x16_skylake_(products_f32x16);
55
+ __m512 normalized_f32x16 = _mm512_mul_ps(dots_f32x16, rsqrt_f32x16);
56
+ __m512 angular_f32x16 = _mm512_sub_ps(_mm512_set1_ps(1.0f), normalized_f32x16);
57
+ _mm512_mask_storeu_ps(results + i, tail, _mm512_max_ps(angular_f32x16, _mm512_setzero_ps()));
58
+ }
59
+ }
60
+
61
+ NK_INTERNAL void nk_euclideans_row_f32dots_sapphireamx_(nk_f32_t *results, nk_f32_t const *norms,
62
+ nk_f32_t query_norm_sq, nk_size_t count) {
63
+ __m512 query_norm_sq_f32x16 = _mm512_set1_ps(query_norm_sq);
64
+ __m512 two_f32x16 = _mm512_set1_ps(2.0f);
65
+ nk_size_t i = 0;
66
+ for (; i + 16 <= count; i += 16) {
67
+ __m512 dots_f32x16 = _mm512_loadu_ps(results + i);
68
+ __m512 norms_f32x16 = _mm512_loadu_ps(norms + i);
69
+ __m512 sum_norms_f32x16 = _mm512_add_ps(query_norm_sq_f32x16, norms_f32x16);
70
+ __m512 dist_sq_f32x16 = _mm512_fnmadd_ps(two_f32x16, dots_f32x16, sum_norms_f32x16);
71
+ dist_sq_f32x16 = _mm512_max_ps(dist_sq_f32x16, _mm512_setzero_ps());
72
+ _mm512_storeu_ps(results + i, _mm512_sqrt_ps(dist_sq_f32x16));
73
+ }
74
+ if (i < count) {
75
+ __mmask16 tail = (__mmask16)((1u << (count - i)) - 1);
76
+ __m512 dots_f32x16 = _mm512_maskz_loadu_ps(tail, results + i);
77
+ __m512 norms_f32x16 = _mm512_maskz_loadu_ps(tail, norms + i);
78
+ __m512 sum_norms_f32x16 = _mm512_add_ps(query_norm_sq_f32x16, norms_f32x16);
79
+ __m512 dist_sq_f32x16 = _mm512_fnmadd_ps(two_f32x16, dots_f32x16, sum_norms_f32x16);
80
+ dist_sq_f32x16 = _mm512_max_ps(dist_sq_f32x16, _mm512_setzero_ps());
81
+ _mm512_mask_storeu_ps(results + i, tail, _mm512_sqrt_ps(dist_sq_f32x16));
82
+ }
83
+ }
84
+
85
+ NK_INTERNAL void nk_angulars_row_i32dots_sapphireamx_(nk_f32_t *results, nk_u32_t const *norms, nk_f32_t query_norm_sq,
86
+ nk_size_t count) {
87
+ nk_i32_t *results_i32 = (nk_i32_t *)results;
88
+ __m512 query_norm_sq_f32x16 = _mm512_set1_ps(query_norm_sq);
89
+ nk_size_t i = 0;
90
+ for (; i + 16 <= count; i += 16) {
91
+ __m512 dots_f32x16 = _mm512_cvtepi32_ps(_mm512_loadu_si512(results_i32 + i));
92
+ __m512 norms_f32x16 = _mm512_cvtepu32_ps(_mm512_loadu_si512((__m512i const *)(norms + i)));
93
+ __m512 products_f32x16 = _mm512_mul_ps(query_norm_sq_f32x16, norms_f32x16);
94
+ __m512 rsqrt_f32x16 = nk_rsqrt_f32x16_skylake_(products_f32x16);
95
+ __m512 normalized_f32x16 = _mm512_mul_ps(dots_f32x16, rsqrt_f32x16);
96
+ __m512 angular_f32x16 = _mm512_sub_ps(_mm512_set1_ps(1.0f), normalized_f32x16);
97
+ _mm512_storeu_ps(results + i, _mm512_max_ps(angular_f32x16, _mm512_setzero_ps()));
98
+ }
99
+ if (i < count) {
100
+ __mmask16 tail = (__mmask16)((1u << (count - i)) - 1);
101
+ __m512 dots_f32x16 = _mm512_cvtepi32_ps(_mm512_maskz_loadu_epi32(tail, results_i32 + i));
102
+ __m512 norms_f32x16 = _mm512_cvtepu32_ps(_mm512_maskz_loadu_epi32(tail, norms + i));
103
+ __m512 products_f32x16 = _mm512_mul_ps(query_norm_sq_f32x16, norms_f32x16);
104
+ __m512 rsqrt_f32x16 = nk_rsqrt_f32x16_skylake_(products_f32x16);
105
+ __m512 normalized_f32x16 = _mm512_mul_ps(dots_f32x16, rsqrt_f32x16);
106
+ __m512 angular_f32x16 = _mm512_sub_ps(_mm512_set1_ps(1.0f), normalized_f32x16);
107
+ _mm512_mask_storeu_ps(results + i, tail, _mm512_max_ps(angular_f32x16, _mm512_setzero_ps()));
108
+ }
109
+ }
110
+
111
+ NK_INTERNAL void nk_euclideans_row_i32dots_sapphireamx_(nk_f32_t *results, nk_u32_t const *norms,
112
+ nk_f32_t query_norm_sq, nk_size_t count) {
113
+ nk_i32_t *results_i32 = (nk_i32_t *)results;
114
+ __m512 query_norm_sq_f32x16 = _mm512_set1_ps(query_norm_sq);
115
+ __m512 two_f32x16 = _mm512_set1_ps(2.0f);
116
+ nk_size_t i = 0;
117
+ for (; i + 16 <= count; i += 16) {
118
+ __m512 dots_f32x16 = _mm512_cvtepi32_ps(_mm512_loadu_si512(results_i32 + i));
119
+ __m512 norms_f32x16 = _mm512_cvtepu32_ps(_mm512_loadu_si512((__m512i const *)(norms + i)));
120
+ __m512 sum_norms_f32x16 = _mm512_add_ps(query_norm_sq_f32x16, norms_f32x16);
121
+ __m512 dist_sq_f32x16 = _mm512_fnmadd_ps(two_f32x16, dots_f32x16, sum_norms_f32x16);
122
+ dist_sq_f32x16 = _mm512_max_ps(dist_sq_f32x16, _mm512_setzero_ps());
123
+ _mm512_storeu_ps(results + i, _mm512_sqrt_ps(dist_sq_f32x16));
124
+ }
125
+ if (i < count) {
126
+ __mmask16 tail = (__mmask16)((1u << (count - i)) - 1);
127
+ __m512 dots_f32x16 = _mm512_cvtepi32_ps(_mm512_maskz_loadu_epi32(tail, results_i32 + i));
128
+ __m512 norms_f32x16 = _mm512_cvtepu32_ps(_mm512_maskz_loadu_epi32(tail, norms + i));
129
+ __m512 sum_norms_f32x16 = _mm512_add_ps(query_norm_sq_f32x16, norms_f32x16);
130
+ __m512 dist_sq_f32x16 = _mm512_fnmadd_ps(two_f32x16, dots_f32x16, sum_norms_f32x16);
131
+ dist_sq_f32x16 = _mm512_max_ps(dist_sq_f32x16, _mm512_setzero_ps());
132
+ _mm512_mask_storeu_ps(results + i, tail, _mm512_sqrt_ps(dist_sq_f32x16));
133
+ }
134
+ }
135
+
136
+ NK_INTERNAL void nk_angulars_row_u32dots_sapphireamx_(nk_f32_t *results, nk_u32_t const *norms, nk_f32_t query_norm_sq,
137
+ nk_size_t count) {
138
+ nk_u32_t *results_u32 = (nk_u32_t *)results;
139
+ __m512 query_norm_sq_f32x16 = _mm512_set1_ps(query_norm_sq);
140
+ nk_size_t i = 0;
141
+ for (; i + 16 <= count; i += 16) {
142
+ __m512 dots_f32x16 = _mm512_cvtepu32_ps(_mm512_loadu_si512((__m512i const *)(results_u32 + i)));
143
+ __m512 norms_f32x16 = _mm512_cvtepu32_ps(_mm512_loadu_si512((__m512i const *)(norms + i)));
144
+ __m512 products_f32x16 = _mm512_mul_ps(query_norm_sq_f32x16, norms_f32x16);
145
+ __m512 rsqrt_f32x16 = nk_rsqrt_f32x16_skylake_(products_f32x16);
146
+ __m512 normalized_f32x16 = _mm512_mul_ps(dots_f32x16, rsqrt_f32x16);
147
+ __m512 angular_f32x16 = _mm512_sub_ps(_mm512_set1_ps(1.0f), normalized_f32x16);
148
+ _mm512_storeu_ps(results + i, _mm512_max_ps(angular_f32x16, _mm512_setzero_ps()));
149
+ }
150
+ if (i < count) {
151
+ __mmask16 tail = (__mmask16)((1u << (count - i)) - 1);
152
+ __m512 dots_f32x16 = _mm512_cvtepu32_ps(_mm512_maskz_loadu_epi32(tail, results_u32 + i));
153
+ __m512 norms_f32x16 = _mm512_cvtepu32_ps(_mm512_maskz_loadu_epi32(tail, norms + i));
154
+ __m512 products_f32x16 = _mm512_mul_ps(query_norm_sq_f32x16, norms_f32x16);
155
+ __m512 rsqrt_f32x16 = nk_rsqrt_f32x16_skylake_(products_f32x16);
156
+ __m512 normalized_f32x16 = _mm512_mul_ps(dots_f32x16, rsqrt_f32x16);
157
+ __m512 angular_f32x16 = _mm512_sub_ps(_mm512_set1_ps(1.0f), normalized_f32x16);
158
+ _mm512_mask_storeu_ps(results + i, tail, _mm512_max_ps(angular_f32x16, _mm512_setzero_ps()));
159
+ }
160
+ }
161
+
162
+ NK_INTERNAL void nk_euclideans_row_u32dots_sapphireamx_(nk_f32_t *results, nk_u32_t const *norms,
163
+ nk_f32_t query_norm_sq, nk_size_t count) {
164
+ nk_u32_t *results_u32 = (nk_u32_t *)results;
165
+ __m512 query_norm_sq_f32x16 = _mm512_set1_ps(query_norm_sq);
166
+ __m512 two_f32x16 = _mm512_set1_ps(2.0f);
167
+ nk_size_t i = 0;
168
+ for (; i + 16 <= count; i += 16) {
169
+ __m512 dots_f32x16 = _mm512_cvtepu32_ps(_mm512_loadu_si512((__m512i const *)(results_u32 + i)));
170
+ __m512 norms_f32x16 = _mm512_cvtepu32_ps(_mm512_loadu_si512((__m512i const *)(norms + i)));
171
+ __m512 sum_norms_f32x16 = _mm512_add_ps(query_norm_sq_f32x16, norms_f32x16);
172
+ __m512 dist_sq_f32x16 = _mm512_fnmadd_ps(two_f32x16, dots_f32x16, sum_norms_f32x16);
173
+ dist_sq_f32x16 = _mm512_max_ps(dist_sq_f32x16, _mm512_setzero_ps());
174
+ _mm512_storeu_ps(results + i, _mm512_sqrt_ps(dist_sq_f32x16));
175
+ }
176
+ if (i < count) {
177
+ __mmask16 tail = (__mmask16)((1u << (count - i)) - 1);
178
+ __m512 dots_f32x16 = _mm512_cvtepu32_ps(_mm512_maskz_loadu_epi32(tail, results_u32 + i));
179
+ __m512 norms_f32x16 = _mm512_cvtepu32_ps(_mm512_maskz_loadu_epi32(tail, norms + i));
180
+ __m512 sum_norms_f32x16 = _mm512_add_ps(query_norm_sq_f32x16, norms_f32x16);
181
+ __m512 dist_sq_f32x16 = _mm512_fnmadd_ps(two_f32x16, dots_f32x16, sum_norms_f32x16);
182
+ dist_sq_f32x16 = _mm512_max_ps(dist_sq_f32x16, _mm512_setzero_ps());
183
+ _mm512_mask_storeu_ps(results + i, tail, _mm512_sqrt_ps(dist_sq_f32x16));
184
+ }
185
+ }
186
+
187
+ #pragma endregion Row Finalize Helpers
188
+
189
+ #pragma region BF16 Packed
190
+
191
+ NK_INTERNAL void nk_angulars_packed_bf16_sapphireamx_finalize_(nk_bf16_t const *a, void const *b_packed, nk_f32_t *c,
192
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
193
+ nk_size_t a_stride_elements,
194
+ nk_size_t c_stride_elements) {
195
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
196
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
197
+ for (nk_size_t row = 0; row < rows; row++) {
198
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_bf16_(a + row * a_stride_elements, depth);
199
+ nk_angulars_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
200
+ }
201
+ }
202
+
203
+ NK_PUBLIC void nk_angulars_packed_bf16_sapphireamx( //
204
+ nk_bf16_t const *a, void const *b_packed, nk_f32_t *c, //
205
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
206
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
207
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_bf16_t);
208
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
209
+ nk_dots_packed_bf16_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
210
+ nk_angulars_packed_bf16_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
211
+ c_stride_elements);
212
+ }
213
+
214
+ NK_INTERNAL void nk_euclideans_packed_bf16_sapphireamx_finalize_(nk_bf16_t const *a, void const *b_packed, nk_f32_t *c,
215
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
216
+ nk_size_t a_stride_elements,
217
+ nk_size_t c_stride_elements) {
218
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
219
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
220
+ for (nk_size_t row = 0; row < rows; row++) {
221
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_bf16_(a + row * a_stride_elements, depth);
222
+ nk_euclideans_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
223
+ }
224
+ }
225
+
226
+ NK_PUBLIC void nk_euclideans_packed_bf16_sapphireamx( //
227
+ nk_bf16_t const *a, void const *b_packed, nk_f32_t *c, //
228
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
229
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
230
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_bf16_t);
231
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
232
+ nk_dots_packed_bf16_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
233
+ nk_euclideans_packed_bf16_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
234
+ c_stride_elements);
235
+ }
236
+
237
+ #pragma endregion BF16 Packed
238
+
239
+ #pragma region BF16 Symmetric
240
+
241
+ NK_INTERNAL void nk_angulars_symmetric_bf16_sapphireamx_finalize_(nk_bf16_t const *vectors, nk_size_t vectors_count,
242
+ nk_size_t depth, nk_size_t stride_elements,
243
+ nk_f32_t *result, nk_size_t result_stride_elements,
244
+ nk_size_t row_start, nk_size_t row_count) {
245
+
246
+ // Phase 1: Cache row norms on diagonal
247
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
248
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_bf16_(vectors + row * stride_elements, depth);
249
+
250
+ // Phase 2: 256-column chunks with cached norms
251
+ nk_f32_t column_norms_cache[256];
252
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
253
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
254
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
255
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_bf16_(vectors + col * stride_elements, depth);
256
+
257
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
258
+ nk_f32_t *r_row = result + row * result_stride_elements;
259
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
260
+ if (col_start >= chunk_end) continue;
261
+ nk_angulars_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
262
+ r_row[row], chunk_end - col_start);
263
+ }
264
+ }
265
+
266
+ // Phase 3: Zero diagonal
267
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
268
+ }
269
+
270
+ NK_PUBLIC void nk_angulars_symmetric_bf16_sapphireamx( //
271
+ nk_bf16_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
272
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
273
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_bf16_t);
274
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
275
+ nk_dots_symmetric_bf16_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
276
+ result_stride_in_bytes, row_start, row_count);
277
+ nk_angulars_symmetric_bf16_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
278
+ result_stride_elements, row_start, row_count);
279
+ }
280
+
281
+ NK_INTERNAL void nk_euclideans_symmetric_bf16_sapphireamx_finalize_(nk_bf16_t const *vectors, nk_size_t vectors_count,
282
+ nk_size_t depth, nk_size_t stride_elements,
283
+ nk_f32_t *result, nk_size_t result_stride_elements,
284
+ nk_size_t row_start, nk_size_t row_count) {
285
+
286
+ // Phase 1: Cache row norms on diagonal
287
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
288
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_bf16_(vectors + row * stride_elements, depth);
289
+
290
+ // Phase 2: 256-column chunks with cached norms
291
+ nk_f32_t column_norms_cache[256];
292
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
293
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
294
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
295
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_bf16_(vectors + col * stride_elements, depth);
296
+
297
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
298
+ nk_f32_t *r_row = result + row * result_stride_elements;
299
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
300
+ if (col_start >= chunk_end) continue;
301
+ nk_euclideans_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
302
+ r_row[row], chunk_end - col_start);
303
+ }
304
+ }
305
+
306
+ // Phase 3: Zero diagonal
307
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
308
+ }
309
+
310
+ NK_PUBLIC void nk_euclideans_symmetric_bf16_sapphireamx( //
311
+ nk_bf16_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
312
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
313
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_bf16_t);
314
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
315
+ nk_dots_symmetric_bf16_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
316
+ result_stride_in_bytes, row_start, row_count);
317
+ nk_euclideans_symmetric_bf16_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
318
+ result_stride_elements, row_start, row_count);
319
+ }
320
+
321
+ #pragma endregion BF16 Symmetric
322
+
323
+ #pragma region I8 Packed
324
+
325
+ NK_INTERNAL void nk_angulars_packed_i8_sapphireamx_finalize_(nk_i8_t const *a, void const *b_packed, nk_f32_t *c,
326
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
327
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
328
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
329
+ nk_u32_t const *b_norms = (nk_u32_t const *)((char const *)b_packed + header->norms_byte_offset);
330
+ for (nk_size_t row = 0; row < rows; row++) {
331
+ nk_f32_t query_norm_sq = (nk_f32_t)nk_dots_reduce_sumsq_i8_(a + row * a_stride_elements, depth);
332
+ nk_angulars_row_i32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
333
+ }
334
+ }
335
+
336
+ NK_PUBLIC void nk_angulars_packed_i8_sapphireamx( //
337
+ nk_i8_t const *a, void const *b_packed, nk_f32_t *c, //
338
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
339
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
340
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_i8_t);
341
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
342
+ nk_dots_packed_i8_sapphireamx(a, b_packed, (nk_i32_t *)c, rows, columns, depth, a_stride_in_bytes,
343
+ c_stride_in_bytes);
344
+ nk_angulars_packed_i8_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
345
+ c_stride_elements);
346
+ }
347
+
348
+ NK_INTERNAL void nk_euclideans_packed_i8_sapphireamx_finalize_(nk_i8_t const *a, void const *b_packed, nk_f32_t *c,
349
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
350
+ nk_size_t a_stride_elements,
351
+ nk_size_t c_stride_elements) {
352
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
353
+ nk_u32_t const *b_norms = (nk_u32_t const *)((char const *)b_packed + header->norms_byte_offset);
354
+ for (nk_size_t row = 0; row < rows; row++) {
355
+ nk_f32_t query_norm_sq = (nk_f32_t)nk_dots_reduce_sumsq_i8_(a + row * a_stride_elements, depth);
356
+ nk_euclideans_row_i32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
357
+ }
358
+ }
359
+
360
+ NK_PUBLIC void nk_euclideans_packed_i8_sapphireamx( //
361
+ nk_i8_t const *a, void const *b_packed, nk_f32_t *c, //
362
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
363
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
364
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_i8_t);
365
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
366
+ nk_dots_packed_i8_sapphireamx(a, b_packed, (nk_i32_t *)c, rows, columns, depth, a_stride_in_bytes,
367
+ c_stride_in_bytes);
368
+ nk_euclideans_packed_i8_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
369
+ c_stride_elements);
370
+ }
371
+
372
+ #pragma endregion I8 Packed
373
+
374
+ #pragma region I8 Symmetric
375
+
376
+ NK_INTERNAL void nk_angulars_symmetric_i8_sapphireamx_finalize_(nk_i8_t const *vectors, nk_size_t vectors_count,
377
+ nk_size_t depth, nk_size_t stride_elements,
378
+ nk_f32_t *result, nk_size_t result_stride_elements,
379
+ nk_size_t row_start, nk_size_t row_count) {
380
+
381
+ // Phase 1: Cache row norms on diagonal (stored as u32 reinterpreted in f32 slot)
382
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
383
+ ((nk_u32_t *)(result + row * result_stride_elements))[row] = nk_dots_reduce_sumsq_i8_(
384
+ vectors + row * stride_elements, depth);
385
+
386
+ // Phase 2: 256-column chunks with cached norms
387
+ nk_u32_t column_norms_cache[256];
388
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
389
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
390
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
391
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_i8_(vectors + col * stride_elements, depth);
392
+
393
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
394
+ nk_f32_t *r_row = result + row * result_stride_elements;
395
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
396
+ if (col_start >= chunk_end) continue;
397
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)((nk_u32_t *)r_row)[row];
398
+ nk_angulars_row_i32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
399
+ query_norm_sq_f32, chunk_end - col_start);
400
+ }
401
+ }
402
+
403
+ // Phase 3: Zero diagonal
404
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
405
+ }
406
+
407
+ NK_PUBLIC void nk_angulars_symmetric_i8_sapphireamx( //
408
+ nk_i8_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
409
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
410
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_i8_t);
411
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
412
+ nk_dots_symmetric_i8_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_i32_t *)result,
413
+ result_stride_in_bytes, row_start, row_count);
414
+ nk_angulars_symmetric_i8_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
415
+ result_stride_elements, row_start, row_count);
416
+ }
417
+
418
+ NK_INTERNAL void nk_euclideans_symmetric_i8_sapphireamx_finalize_(nk_i8_t const *vectors, nk_size_t vectors_count,
419
+ nk_size_t depth, nk_size_t stride_elements,
420
+ nk_f32_t *result, nk_size_t result_stride_elements,
421
+ nk_size_t row_start, nk_size_t row_count) {
422
+
423
+ // Phase 1: Cache row norms on diagonal (stored as u32 reinterpreted in f32 slot)
424
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
425
+ ((nk_u32_t *)(result + row * result_stride_elements))[row] = nk_dots_reduce_sumsq_i8_(
426
+ vectors + row * stride_elements, depth);
427
+
428
+ // Phase 2: 256-column chunks with cached norms
429
+ nk_u32_t column_norms_cache[256];
430
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
431
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
432
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
433
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_i8_(vectors + col * stride_elements, depth);
434
+
435
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
436
+ nk_f32_t *r_row = result + row * result_stride_elements;
437
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
438
+ if (col_start >= chunk_end) continue;
439
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)((nk_u32_t *)r_row)[row];
440
+ nk_euclideans_row_i32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
441
+ query_norm_sq_f32, chunk_end - col_start);
442
+ }
443
+ }
444
+
445
+ // Phase 3: Zero diagonal
446
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
447
+ }
448
+
449
+ NK_PUBLIC void nk_euclideans_symmetric_i8_sapphireamx( //
450
+ nk_i8_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
451
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
452
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_i8_t);
453
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
454
+ nk_dots_symmetric_i8_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_i32_t *)result,
455
+ result_stride_in_bytes, row_start, row_count);
456
+ nk_euclideans_symmetric_i8_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
457
+ result_stride_elements, row_start, row_count);
458
+ }
459
+
460
+ #pragma endregion I8 Symmetric
461
+
462
+ #pragma region U8 Packed
463
+
464
+ NK_INTERNAL void nk_angulars_packed_u8_sapphireamx_finalize_(nk_u8_t const *a, void const *b_packed, nk_f32_t *c,
465
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
466
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
467
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
468
+ nk_u32_t const *b_norms = (nk_u32_t const *)((char const *)b_packed + header->norms_byte_offset);
469
+ for (nk_size_t row = 0; row < rows; row++) {
470
+ nk_f32_t query_norm_sq = (nk_f32_t)nk_dots_reduce_sumsq_u8_(a + row * a_stride_elements, depth);
471
+ nk_angulars_row_u32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
472
+ }
473
+ }
474
+
475
+ NK_PUBLIC void nk_angulars_packed_u8_sapphireamx( //
476
+ nk_u8_t const *a, void const *b_packed, nk_f32_t *c, //
477
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
478
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
479
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_u8_t);
480
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
481
+ nk_dots_packed_u8_sapphireamx(a, b_packed, (nk_u32_t *)c, rows, columns, depth, a_stride_in_bytes,
482
+ c_stride_in_bytes);
483
+ nk_angulars_packed_u8_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
484
+ c_stride_elements);
485
+ }
486
+
487
+ NK_INTERNAL void nk_euclideans_packed_u8_sapphireamx_finalize_(nk_u8_t const *a, void const *b_packed, nk_f32_t *c,
488
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
489
+ nk_size_t a_stride_elements,
490
+ nk_size_t c_stride_elements) {
491
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
492
+ nk_u32_t const *b_norms = (nk_u32_t const *)((char const *)b_packed + header->norms_byte_offset);
493
+ for (nk_size_t row = 0; row < rows; row++) {
494
+ nk_f32_t query_norm_sq = (nk_f32_t)nk_dots_reduce_sumsq_u8_(a + row * a_stride_elements, depth);
495
+ nk_euclideans_row_u32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
496
+ }
497
+ }
498
+
499
+ NK_PUBLIC void nk_euclideans_packed_u8_sapphireamx( //
500
+ nk_u8_t const *a, void const *b_packed, nk_f32_t *c, //
501
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
502
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
503
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_u8_t);
504
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
505
+ nk_dots_packed_u8_sapphireamx(a, b_packed, (nk_u32_t *)c, rows, columns, depth, a_stride_in_bytes,
506
+ c_stride_in_bytes);
507
+ nk_euclideans_packed_u8_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
508
+ c_stride_elements);
509
+ }
510
+
511
+ #pragma endregion U8 Packed
512
+
513
+ #pragma region U8 Symmetric
514
+
515
+ NK_INTERNAL void nk_angulars_symmetric_u8_sapphireamx_finalize_(nk_u8_t const *vectors, nk_size_t vectors_count,
516
+ nk_size_t depth, nk_size_t stride_elements,
517
+ nk_f32_t *result, nk_size_t result_stride_elements,
518
+ nk_size_t row_start, nk_size_t row_count) {
519
+
520
+ // Phase 1: Cache row norms on diagonal (stored as u32 reinterpreted in f32 slot)
521
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
522
+ ((nk_u32_t *)(result + row * result_stride_elements))[row] = nk_dots_reduce_sumsq_u8_(
523
+ vectors + row * stride_elements, depth);
524
+
525
+ // Phase 2: 256-column chunks with cached norms
526
+ nk_u32_t column_norms_cache[256];
527
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
528
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
529
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
530
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_u8_(vectors + col * stride_elements, depth);
531
+
532
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
533
+ nk_f32_t *r_row = result + row * result_stride_elements;
534
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
535
+ if (col_start >= chunk_end) continue;
536
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)((nk_u32_t *)r_row)[row];
537
+ nk_angulars_row_u32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
538
+ query_norm_sq_f32, chunk_end - col_start);
539
+ }
540
+ }
541
+
542
+ // Phase 3: Zero diagonal
543
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
544
+ }
545
+
546
+ NK_PUBLIC void nk_angulars_symmetric_u8_sapphireamx( //
547
+ nk_u8_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
548
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
549
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_u8_t);
550
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
551
+ nk_dots_symmetric_u8_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_u32_t *)result,
552
+ result_stride_in_bytes, row_start, row_count);
553
+ nk_angulars_symmetric_u8_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
554
+ result_stride_elements, row_start, row_count);
555
+ }
556
+
557
+ NK_INTERNAL void nk_euclideans_symmetric_u8_sapphireamx_finalize_(nk_u8_t const *vectors, nk_size_t vectors_count,
558
+ nk_size_t depth, nk_size_t stride_elements,
559
+ nk_f32_t *result, nk_size_t result_stride_elements,
560
+ nk_size_t row_start, nk_size_t row_count) {
561
+
562
+ // Phase 1: Cache row norms on diagonal (stored as u32 reinterpreted in f32 slot)
563
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
564
+ ((nk_u32_t *)(result + row * result_stride_elements))[row] = nk_dots_reduce_sumsq_u8_(
565
+ vectors + row * stride_elements, depth);
566
+
567
+ // Phase 2: 256-column chunks with cached norms
568
+ nk_u32_t column_norms_cache[256];
569
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
570
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
571
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
572
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_u8_(vectors + col * stride_elements, depth);
573
+
574
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
575
+ nk_f32_t *r_row = result + row * result_stride_elements;
576
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
577
+ if (col_start >= chunk_end) continue;
578
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)((nk_u32_t *)r_row)[row];
579
+ nk_euclideans_row_u32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
580
+ query_norm_sq_f32, chunk_end - col_start);
581
+ }
582
+ }
583
+
584
+ // Phase 3: Zero diagonal
585
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
586
+ }
587
+
588
+ NK_PUBLIC void nk_euclideans_symmetric_u8_sapphireamx( //
589
+ nk_u8_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
590
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
591
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_u8_t);
592
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
593
+ nk_dots_symmetric_u8_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_u32_t *)result,
594
+ result_stride_in_bytes, row_start, row_count);
595
+ nk_euclideans_symmetric_u8_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
596
+ result_stride_elements, row_start, row_count);
597
+ }
598
+
599
+ #pragma endregion U8 Symmetric
600
+
601
+ #pragma region E4M3 Packed
602
+
603
+ NK_INTERNAL void nk_angulars_packed_e4m3_sapphireamx_finalize_(nk_e4m3_t const *a, void const *b_packed, nk_f32_t *c,
604
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
605
+ nk_size_t a_stride_elements,
606
+ nk_size_t c_stride_elements) {
607
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
608
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
609
+ for (nk_size_t row = 0; row < rows; row++) {
610
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_e4m3_(a + row * a_stride_elements, depth);
611
+ nk_angulars_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
612
+ }
613
+ }
614
+
615
+ NK_PUBLIC void nk_angulars_packed_e4m3_sapphireamx( //
616
+ nk_e4m3_t const *a, void const *b_packed, nk_f32_t *c, //
617
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
618
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
619
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e4m3_t);
620
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
621
+ nk_dots_packed_e4m3_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
622
+ nk_angulars_packed_e4m3_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
623
+ c_stride_elements);
624
+ }
625
+
626
+ NK_INTERNAL void nk_euclideans_packed_e4m3_sapphireamx_finalize_(nk_e4m3_t const *a, void const *b_packed, nk_f32_t *c,
627
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
628
+ nk_size_t a_stride_elements,
629
+ nk_size_t c_stride_elements) {
630
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
631
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
632
+ for (nk_size_t row = 0; row < rows; row++) {
633
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_e4m3_(a + row * a_stride_elements, depth);
634
+ nk_euclideans_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
635
+ }
636
+ }
637
+
638
+ NK_PUBLIC void nk_euclideans_packed_e4m3_sapphireamx( //
639
+ nk_e4m3_t const *a, void const *b_packed, nk_f32_t *c, //
640
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
641
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
642
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e4m3_t);
643
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
644
+ nk_dots_packed_e4m3_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
645
+ nk_euclideans_packed_e4m3_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
646
+ c_stride_elements);
647
+ }
648
+
649
+ #pragma endregion E4M3 Packed
650
+
651
+ #pragma region E5M2 Packed
652
+
653
+ NK_INTERNAL void nk_angulars_packed_e5m2_sapphireamx_finalize_(nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c,
654
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
655
+ nk_size_t a_stride_elements,
656
+ nk_size_t c_stride_elements) {
657
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
658
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
659
+ for (nk_size_t row = 0; row < rows; row++) {
660
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_e5m2_(a + row * a_stride_elements, depth);
661
+ nk_angulars_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
662
+ }
663
+ }
664
+
665
+ NK_PUBLIC void nk_angulars_packed_e5m2_sapphireamx( //
666
+ nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c, //
667
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
668
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
669
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e5m2_t);
670
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
671
+ nk_dots_packed_e5m2_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
672
+ nk_angulars_packed_e5m2_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
673
+ c_stride_elements);
674
+ }
675
+
676
+ NK_INTERNAL void nk_euclideans_packed_e5m2_sapphireamx_finalize_(nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c,
677
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
678
+ nk_size_t a_stride_elements,
679
+ nk_size_t c_stride_elements) {
680
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
681
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
682
+ for (nk_size_t row = 0; row < rows; row++) {
683
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_e5m2_(a + row * a_stride_elements, depth);
684
+ nk_euclideans_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
685
+ }
686
+ }
687
+
688
+ NK_PUBLIC void nk_euclideans_packed_e5m2_sapphireamx( //
689
+ nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c, //
690
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
691
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
692
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e5m2_t);
693
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
694
+ nk_dots_packed_e5m2_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
695
+ nk_euclideans_packed_e5m2_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
696
+ c_stride_elements);
697
+ }
698
+
699
+ #pragma endregion E5M2 Packed
700
+
701
+ #pragma region E5M2 Symmetric
702
+
703
+ NK_INTERNAL void nk_angulars_symmetric_e5m2_sapphireamx_finalize_(nk_e5m2_t const *vectors, nk_size_t vectors_count,
704
+ nk_size_t depth, nk_size_t stride_elements,
705
+ nk_f32_t *result, nk_size_t result_stride_elements,
706
+ nk_size_t row_start, nk_size_t row_count) {
707
+
708
+ // Phase 1: Cache row norms on diagonal
709
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
710
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_e5m2_(vectors + row * stride_elements, depth);
711
+
712
+ // Phase 2: 256-column chunks with cached norms
713
+ nk_f32_t column_norms_cache[256];
714
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
715
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
716
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
717
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e5m2_(vectors + col * stride_elements, depth);
718
+
719
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
720
+ nk_f32_t *r_row = result + row * result_stride_elements;
721
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
722
+ if (col_start >= chunk_end) continue;
723
+ nk_angulars_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
724
+ r_row[row], chunk_end - col_start);
725
+ }
726
+ }
727
+
728
+ // Phase 3: Zero diagonal
729
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
730
+ }
731
+
732
+ NK_PUBLIC void nk_angulars_symmetric_e5m2_sapphireamx( //
733
+ nk_e5m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
734
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
735
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e5m2_t);
736
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
737
+ nk_dots_symmetric_e5m2_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
738
+ result_stride_in_bytes, row_start, row_count);
739
+ nk_angulars_symmetric_e5m2_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
740
+ result_stride_elements, row_start, row_count);
741
+ }
742
+
743
+ NK_INTERNAL void nk_euclideans_symmetric_e5m2_sapphireamx_finalize_(nk_e5m2_t const *vectors, nk_size_t vectors_count,
744
+ nk_size_t depth, nk_size_t stride_elements,
745
+ nk_f32_t *result, nk_size_t result_stride_elements,
746
+ nk_size_t row_start, nk_size_t row_count) {
747
+
748
+ // Phase 1: Cache row norms on diagonal
749
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
750
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_e5m2_(vectors + row * stride_elements, depth);
751
+
752
+ // Phase 2: 256-column chunks with cached norms
753
+ nk_f32_t column_norms_cache[256];
754
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
755
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
756
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
757
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e5m2_(vectors + col * stride_elements, depth);
758
+
759
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
760
+ nk_f32_t *r_row = result + row * result_stride_elements;
761
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
762
+ if (col_start >= chunk_end) continue;
763
+ nk_euclideans_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
764
+ r_row[row], chunk_end - col_start);
765
+ }
766
+ }
767
+
768
+ // Phase 3: Zero diagonal
769
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
770
+ }
771
+
772
+ NK_PUBLIC void nk_euclideans_symmetric_e5m2_sapphireamx( //
773
+ nk_e5m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
774
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
775
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e5m2_t);
776
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
777
+ nk_dots_symmetric_e5m2_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
778
+ result_stride_in_bytes, row_start, row_count);
779
+ nk_euclideans_symmetric_e5m2_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
780
+ result_stride_elements, row_start, row_count);
781
+ }
782
+
783
+ #pragma endregion E5M2 Symmetric
784
+
785
+ #pragma region E4M3 Symmetric
786
+
787
+ NK_INTERNAL void nk_angulars_symmetric_e4m3_sapphireamx_finalize_(nk_e4m3_t const *vectors, nk_size_t vectors_count,
788
+ nk_size_t depth, nk_size_t stride_elements,
789
+ nk_f32_t *result, nk_size_t result_stride_elements,
790
+ nk_size_t row_start, nk_size_t row_count) {
791
+
792
+ // Phase 1: Cache row norms on diagonal
793
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
794
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_e4m3_(vectors + row * stride_elements, depth);
795
+
796
+ // Phase 2: 256-column chunks with cached norms
797
+ nk_f32_t column_norms_cache[256];
798
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
799
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
800
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
801
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e4m3_(vectors + col * stride_elements, depth);
802
+
803
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
804
+ nk_f32_t *r_row = result + row * result_stride_elements;
805
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
806
+ if (col_start >= chunk_end) continue;
807
+ nk_angulars_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
808
+ r_row[row], chunk_end - col_start);
809
+ }
810
+ }
811
+
812
+ // Phase 3: Zero diagonal
813
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
814
+ }
815
+
816
+ NK_PUBLIC void nk_angulars_symmetric_e4m3_sapphireamx( //
817
+ nk_e4m3_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
818
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
819
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e4m3_t);
820
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
821
+ nk_dots_symmetric_e4m3_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
822
+ result_stride_in_bytes, row_start, row_count);
823
+ nk_angulars_symmetric_e4m3_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
824
+ result_stride_elements, row_start, row_count);
825
+ }
826
+
827
+ NK_INTERNAL void nk_euclideans_symmetric_e4m3_sapphireamx_finalize_(nk_e4m3_t const *vectors, nk_size_t vectors_count,
828
+ nk_size_t depth, nk_size_t stride_elements,
829
+ nk_f32_t *result, nk_size_t result_stride_elements,
830
+ nk_size_t row_start, nk_size_t row_count) {
831
+
832
+ // Phase 1: Cache row norms on diagonal
833
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
834
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_e4m3_(vectors + row * stride_elements, depth);
835
+
836
+ // Phase 2: 256-column chunks with cached norms
837
+ nk_f32_t column_norms_cache[256];
838
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
839
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
840
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
841
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e4m3_(vectors + col * stride_elements, depth);
842
+
843
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
844
+ nk_f32_t *r_row = result + row * result_stride_elements;
845
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
846
+ if (col_start >= chunk_end) continue;
847
+ nk_euclideans_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
848
+ r_row[row], chunk_end - col_start);
849
+ }
850
+ }
851
+
852
+ // Phase 3: Zero diagonal
853
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
854
+ }
855
+
856
+ NK_PUBLIC void nk_euclideans_symmetric_e4m3_sapphireamx( //
857
+ nk_e4m3_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
858
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
859
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e4m3_t);
860
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
861
+ nk_dots_symmetric_e4m3_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
862
+ result_stride_in_bytes, row_start, row_count);
863
+ nk_euclideans_symmetric_e4m3_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
864
+ result_stride_elements, row_start, row_count);
865
+ }
866
+
867
+ #pragma endregion E4M3 Symmetric
868
+
869
+ #pragma region E2M3 Packed
870
+
871
+ NK_INTERNAL void nk_angulars_packed_e2m3_sapphireamx_finalize_(nk_e2m3_t const *a, void const *b_packed, nk_f32_t *c,
872
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
873
+ nk_size_t a_stride_elements,
874
+ nk_size_t c_stride_elements) {
875
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
876
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
877
+ for (nk_size_t row = 0; row < rows; row++) {
878
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_e2m3_(a + row * a_stride_elements, depth);
879
+ nk_angulars_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
880
+ }
881
+ }
882
+
883
+ NK_PUBLIC void nk_angulars_packed_e2m3_sapphireamx( //
884
+ nk_e2m3_t const *a, void const *b_packed, nk_f32_t *c, //
885
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
886
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
887
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e2m3_t);
888
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
889
+ nk_dots_packed_e2m3_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
890
+ nk_angulars_packed_e2m3_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
891
+ c_stride_elements);
892
+ }
893
+
894
+ NK_INTERNAL void nk_euclideans_packed_e2m3_sapphireamx_finalize_(nk_e2m3_t const *a, void const *b_packed, nk_f32_t *c,
895
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
896
+ nk_size_t a_stride_elements,
897
+ nk_size_t c_stride_elements) {
898
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
899
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
900
+ for (nk_size_t row = 0; row < rows; row++) {
901
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_e2m3_(a + row * a_stride_elements, depth);
902
+ nk_euclideans_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
903
+ }
904
+ }
905
+
906
+ NK_PUBLIC void nk_euclideans_packed_e2m3_sapphireamx( //
907
+ nk_e2m3_t const *a, void const *b_packed, nk_f32_t *c, //
908
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
909
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
910
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e2m3_t);
911
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
912
+ nk_dots_packed_e2m3_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
913
+ nk_euclideans_packed_e2m3_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
914
+ c_stride_elements);
915
+ }
916
+
917
+ #pragma endregion E2M3 Packed
918
+
919
+ #pragma region E2M3 Symmetric
920
+
921
+ NK_INTERNAL void nk_angulars_symmetric_e2m3_sapphireamx_finalize_(nk_e2m3_t const *vectors, nk_size_t vectors_count,
922
+ nk_size_t depth, nk_size_t stride_elements,
923
+ nk_f32_t *result, nk_size_t result_stride_elements,
924
+ nk_size_t row_start, nk_size_t row_count) {
925
+
926
+ // Phase 1: Cache row norms on diagonal
927
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
928
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_e2m3_(vectors + row * stride_elements, depth);
929
+
930
+ // Phase 2: 256-column chunks with cached norms
931
+ nk_f32_t column_norms_cache[256];
932
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
933
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
934
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
935
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e2m3_(vectors + col * stride_elements, depth);
936
+
937
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
938
+ nk_f32_t *r_row = result + row * result_stride_elements;
939
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
940
+ if (col_start >= chunk_end) continue;
941
+ nk_angulars_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
942
+ r_row[row], chunk_end - col_start);
943
+ }
944
+ }
945
+
946
+ // Phase 3: Zero diagonal
947
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
948
+ }
949
+
950
+ NK_PUBLIC void nk_angulars_symmetric_e2m3_sapphireamx( //
951
+ nk_e2m3_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
952
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
953
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e2m3_t);
954
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
955
+ nk_dots_symmetric_e2m3_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
956
+ result_stride_in_bytes, row_start, row_count);
957
+ nk_angulars_symmetric_e2m3_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
958
+ result_stride_elements, row_start, row_count);
959
+ }
960
+
961
+ NK_INTERNAL void nk_euclideans_symmetric_e2m3_sapphireamx_finalize_(nk_e2m3_t const *vectors, nk_size_t vectors_count,
962
+ nk_size_t depth, nk_size_t stride_elements,
963
+ nk_f32_t *result, nk_size_t result_stride_elements,
964
+ nk_size_t row_start, nk_size_t row_count) {
965
+
966
+ // Phase 1: Cache row norms on diagonal
967
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
968
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_e2m3_(vectors + row * stride_elements, depth);
969
+
970
+ // Phase 2: 256-column chunks with cached norms
971
+ nk_f32_t column_norms_cache[256];
972
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
973
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
974
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
975
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e2m3_(vectors + col * stride_elements, depth);
976
+
977
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
978
+ nk_f32_t *r_row = result + row * result_stride_elements;
979
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
980
+ if (col_start >= chunk_end) continue;
981
+ nk_euclideans_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
982
+ r_row[row], chunk_end - col_start);
983
+ }
984
+ }
985
+
986
+ // Phase 3: Zero diagonal
987
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
988
+ }
989
+
990
+ NK_PUBLIC void nk_euclideans_symmetric_e2m3_sapphireamx( //
991
+ nk_e2m3_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
992
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
993
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e2m3_t);
994
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
995
+ nk_dots_symmetric_e2m3_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
996
+ result_stride_in_bytes, row_start, row_count);
997
+ nk_euclideans_symmetric_e2m3_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
998
+ result_stride_elements, row_start, row_count);
999
+ }
1000
+
1001
+ #pragma endregion E2M3 Symmetric
1002
+
1003
+ #pragma region E3M2 Packed
1004
+
1005
+ NK_INTERNAL void nk_angulars_packed_e3m2_sapphireamx_finalize_(nk_e3m2_t const *a, void const *b_packed, nk_f32_t *c,
1006
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1007
+ nk_size_t a_stride_elements,
1008
+ nk_size_t c_stride_elements) {
1009
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
1010
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
1011
+ for (nk_size_t row = 0; row < rows; row++) {
1012
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_e3m2_(a + row * a_stride_elements, depth);
1013
+ nk_angulars_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
1014
+ }
1015
+ }
1016
+
1017
+ NK_PUBLIC void nk_angulars_packed_e3m2_sapphireamx( //
1018
+ nk_e3m2_t const *a, void const *b_packed, nk_f32_t *c, //
1019
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1020
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1021
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e3m2_t);
1022
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1023
+ nk_dots_packed_e3m2_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1024
+ nk_angulars_packed_e3m2_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
1025
+ c_stride_elements);
1026
+ }
1027
+
1028
+ NK_INTERNAL void nk_euclideans_packed_e3m2_sapphireamx_finalize_(nk_e3m2_t const *a, void const *b_packed, nk_f32_t *c,
1029
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1030
+ nk_size_t a_stride_elements,
1031
+ nk_size_t c_stride_elements) {
1032
+ nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
1033
+ nk_f32_t const *b_norms = (nk_f32_t const *)((char const *)b_packed + header->norms_byte_offset);
1034
+ for (nk_size_t row = 0; row < rows; row++) {
1035
+ nk_f32_t query_norm_sq = nk_dots_reduce_sumsq_e3m2_(a + row * a_stride_elements, depth);
1036
+ nk_euclideans_row_f32dots_sapphireamx_(c + row * c_stride_elements, b_norms, query_norm_sq, columns);
1037
+ }
1038
+ }
1039
+
1040
+ NK_PUBLIC void nk_euclideans_packed_e3m2_sapphireamx( //
1041
+ nk_e3m2_t const *a, void const *b_packed, nk_f32_t *c, //
1042
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1043
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1044
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e3m2_t);
1045
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1046
+ nk_dots_packed_e3m2_sapphireamx(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1047
+ nk_euclideans_packed_e3m2_sapphireamx_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements,
1048
+ c_stride_elements);
1049
+ }
1050
+
1051
+ #pragma endregion E3M2 Packed
1052
+
1053
+ #pragma region E3M2 Symmetric
1054
+
1055
+ NK_INTERNAL void nk_angulars_symmetric_e3m2_sapphireamx_finalize_(nk_e3m2_t const *vectors, nk_size_t vectors_count,
1056
+ nk_size_t depth, nk_size_t stride_elements,
1057
+ nk_f32_t *result, nk_size_t result_stride_elements,
1058
+ nk_size_t row_start, nk_size_t row_count) {
1059
+
1060
+ // Phase 1: Cache row norms on diagonal
1061
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
1062
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_e3m2_(vectors + row * stride_elements, depth);
1063
+
1064
+ // Phase 2: 256-column chunks with cached norms
1065
+ nk_f32_t column_norms_cache[256];
1066
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1067
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1068
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
1069
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e3m2_(vectors + col * stride_elements, depth);
1070
+
1071
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
1072
+ nk_f32_t *r_row = result + row * result_stride_elements;
1073
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
1074
+ if (col_start >= chunk_end) continue;
1075
+ nk_angulars_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
1076
+ r_row[row], chunk_end - col_start);
1077
+ }
1078
+ }
1079
+
1080
+ // Phase 3: Zero diagonal
1081
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
1082
+ }
1083
+
1084
+ NK_PUBLIC void nk_angulars_symmetric_e3m2_sapphireamx( //
1085
+ nk_e3m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1086
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1087
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e3m2_t);
1088
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1089
+ nk_dots_symmetric_e3m2_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
1090
+ result_stride_in_bytes, row_start, row_count);
1091
+ nk_angulars_symmetric_e3m2_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
1092
+ result_stride_elements, row_start, row_count);
1093
+ }
1094
+
1095
+ NK_INTERNAL void nk_euclideans_symmetric_e3m2_sapphireamx_finalize_(nk_e3m2_t const *vectors, nk_size_t vectors_count,
1096
+ nk_size_t depth, nk_size_t stride_elements,
1097
+ nk_f32_t *result, nk_size_t result_stride_elements,
1098
+ nk_size_t row_start, nk_size_t row_count) {
1099
+
1100
+ // Phase 1: Cache row norms on diagonal
1101
+ for (nk_size_t row = row_start; row < row_start + row_count; row++)
1102
+ result[row * result_stride_elements + row] = nk_dots_reduce_sumsq_e3m2_(vectors + row * stride_elements, depth);
1103
+
1104
+ // Phase 2: 256-column chunks with cached norms
1105
+ nk_f32_t column_norms_cache[256];
1106
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1107
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1108
+ for (nk_size_t col = chunk_start; col < chunk_end; col++)
1109
+ column_norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e3m2_(vectors + col * stride_elements, depth);
1110
+
1111
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) {
1112
+ nk_f32_t *r_row = result + row * result_stride_elements;
1113
+ nk_size_t col_start = chunk_start > row + 1 ? chunk_start : row + 1;
1114
+ if (col_start >= chunk_end) continue;
1115
+ nk_euclideans_row_f32dots_sapphireamx_(r_row + col_start, column_norms_cache + col_start - chunk_start,
1116
+ r_row[row], chunk_end - col_start);
1117
+ }
1118
+ }
1119
+
1120
+ // Phase 3: Zero diagonal
1121
+ for (nk_size_t row = row_start; row < row_start + row_count; row++) result[row * result_stride_elements + row] = 0;
1122
+ }
1123
+
1124
+ NK_PUBLIC void nk_euclideans_symmetric_e3m2_sapphireamx( //
1125
+ nk_e3m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1126
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1127
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e3m2_t);
1128
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1129
+ nk_dots_symmetric_e3m2_sapphireamx(vectors, vectors_count, depth, stride_in_bytes, (nk_f32_t *)result,
1130
+ result_stride_in_bytes, row_start, row_count);
1131
+ nk_euclideans_symmetric_e3m2_sapphireamx_finalize_(vectors, vectors_count, depth, stride_elements, result,
1132
+ result_stride_elements, row_start, row_count);
1133
+ }
1134
+
1135
+ #pragma endregion E3M2 Symmetric
1136
+
1137
+ #if defined(__clang__)
1138
+ #pragma clang attribute pop
1139
+ #elif defined(__GNUC__)
1140
+ #pragma GCC pop_options
1141
+ #endif
1142
+
1143
+ #if defined(__cplusplus)
1144
+ } // extern "C"
1145
+ #endif
1146
+
1147
+ #endif // NK_TARGET_SAPPHIREAMX
1148
+ #endif // NK_TARGET_X8664_
1149
+ #endif // NK_SPATIALS_SAPPHIREAMX_H