usearch 2.21.4 → 2.25.1

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 (604) hide show
  1. package/README.md +65 -69
  2. package/binding.gyp +15 -37
  3. package/c/CMakeLists.txt +4 -5
  4. package/c/README.md +8 -8
  5. package/c/lib.cpp +32 -14
  6. package/c/test.c +124 -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 +224 -48
  13. package/include/usearch/index_dense.hpp +105 -40
  14. package/include/usearch/index_plugins.hpp +1387 -204
  15. package/javascript/README.md +1 -1
  16. package/javascript/dist/cjs/usearch.d.ts +8 -3
  17. package/javascript/dist/cjs/usearch.js +12 -4
  18. package/javascript/dist/esm/usearch.d.ts +8 -3
  19. package/javascript/dist/esm/usearch.js +12 -4
  20. package/javascript/lib.cpp +17 -0
  21. package/javascript/usearch.test.js +117 -94
  22. package/javascript/usearch.ts +14 -6
  23. package/numkong/.clang-tidy +30 -0
  24. package/{simsimd → numkong}/.cmake-format.py +18 -0
  25. package/numkong/.editorconfig +12 -0
  26. package/numkong/.git-blame-ignore-revs +123 -0
  27. package/numkong/.mailmap +15 -0
  28. package/numkong/.spi.yml +11 -0
  29. package/numkong/.swift-format +14 -0
  30. package/numkong/CITATION.cff +81 -0
  31. package/numkong/CONTRIBUTING.md +500 -0
  32. package/numkong/Cargo.toml +54 -0
  33. package/numkong/MANIFEST.in +8 -0
  34. package/numkong/Package.swift +84 -0
  35. package/numkong/README.md +569 -0
  36. package/numkong/bench/README.md +217 -0
  37. package/numkong/bench/bench-browser.html +323 -0
  38. package/numkong/bench/bench.cpp +344 -0
  39. package/numkong/bench/bench.hpp +816 -0
  40. package/numkong/bench/bench.js +148 -0
  41. package/numkong/bench/bench.mjs +616 -0
  42. package/numkong/bench/bench_attention.cpp +113 -0
  43. package/numkong/bench/bench_cast.cpp +264 -0
  44. package/numkong/bench/bench_cross_amx.cpp +62 -0
  45. package/numkong/bench/bench_cross_arm.cpp +245 -0
  46. package/numkong/bench/bench_cross_blas.cpp +309 -0
  47. package/numkong/bench/bench_cross_loongarch.cpp +106 -0
  48. package/numkong/bench/bench_cross_power.cpp +84 -0
  49. package/numkong/bench/bench_cross_rvv.cpp +121 -0
  50. package/numkong/bench/bench_cross_serial.cpp +151 -0
  51. package/numkong/bench/bench_cross_sme.cpp +163 -0
  52. package/numkong/bench/bench_cross_wasm.cpp +142 -0
  53. package/numkong/bench/bench_cross_x86.cpp +455 -0
  54. package/numkong/bench/bench_curved.cpp +210 -0
  55. package/numkong/bench/bench_dot.cpp +315 -0
  56. package/numkong/bench/bench_each.cpp +503 -0
  57. package/numkong/bench/bench_geospatial.cpp +112 -0
  58. package/numkong/bench/bench_maxsim.cpp +154 -0
  59. package/numkong/bench/bench_mesh.cpp +148 -0
  60. package/numkong/bench/bench_probability.cpp +60 -0
  61. package/numkong/bench/bench_reduce.cpp +371 -0
  62. package/numkong/bench/bench_set.cpp +73 -0
  63. package/numkong/bench/bench_sparse.cpp +198 -0
  64. package/numkong/bench/bench_spatial.cpp +347 -0
  65. package/numkong/bench/bench_trigonometry.cpp +137 -0
  66. package/numkong/bench/swift/Bench.swift +506 -0
  67. package/numkong/binding.gyp +169 -0
  68. package/numkong/build.rs +484 -0
  69. package/numkong/c/dispatch.h +512 -0
  70. package/numkong/c/dispatch_bf16.c +447 -0
  71. package/numkong/c/dispatch_bf16c.c +52 -0
  72. package/numkong/c/dispatch_e2m3.c +296 -0
  73. package/numkong/c/dispatch_e3m2.c +284 -0
  74. package/numkong/c/dispatch_e4m3.c +318 -0
  75. package/numkong/c/dispatch_e5m2.c +316 -0
  76. package/numkong/c/dispatch_f16.c +452 -0
  77. package/numkong/c/dispatch_f16c.c +58 -0
  78. package/numkong/c/dispatch_f32.c +434 -0
  79. package/numkong/c/dispatch_f32c.c +99 -0
  80. package/numkong/c/dispatch_f64.c +348 -0
  81. package/numkong/c/dispatch_f64c.c +98 -0
  82. package/numkong/c/dispatch_i16.c +96 -0
  83. package/numkong/c/dispatch_i32.c +89 -0
  84. package/numkong/c/dispatch_i4.c +153 -0
  85. package/numkong/c/dispatch_i64.c +86 -0
  86. package/numkong/c/dispatch_i8.c +348 -0
  87. package/numkong/c/dispatch_other.c +348 -0
  88. package/numkong/c/dispatch_u1.c +193 -0
  89. package/numkong/c/dispatch_u16.c +124 -0
  90. package/numkong/c/dispatch_u32.c +118 -0
  91. package/numkong/c/dispatch_u4.c +153 -0
  92. package/numkong/c/dispatch_u64.c +102 -0
  93. package/numkong/c/dispatch_u8.c +364 -0
  94. package/numkong/c/numkong.c +940 -0
  95. package/numkong/cmake/nk_arm_isa_probes.cmake +61 -0
  96. package/numkong/cmake/nk_isa_probe.cmake +90 -0
  97. package/numkong/cmake/nk_loongarch_isa_probes.cmake +15 -0
  98. package/numkong/cmake/nk_math_lib.cmake +46 -0
  99. package/numkong/cmake/nk_power_isa_probes.cmake +15 -0
  100. package/numkong/cmake/nk_riscv_isa_probes.cmake +19 -0
  101. package/numkong/cmake/nk_x86_isa_probes.cmake +35 -0
  102. package/numkong/cmake/toolchain-aarch64-gnu.cmake +86 -0
  103. package/numkong/cmake/toolchain-android-arm64.cmake +25 -0
  104. package/numkong/cmake/toolchain-android-armv7.cmake +25 -0
  105. package/numkong/cmake/toolchain-loongarch64-gnu.cmake +68 -0
  106. package/numkong/cmake/toolchain-ppc64le-gnu.cmake +74 -0
  107. package/numkong/cmake/toolchain-riscv64-gnu.cmake +107 -0
  108. package/numkong/cmake/toolchain-riscv64-llvm.cmake +114 -0
  109. package/numkong/cmake/toolchain-wasi.cmake +77 -0
  110. package/numkong/cmake/toolchain-wasm.cmake +70 -0
  111. package/numkong/cmake/toolchain-wasm64.cmake +62 -0
  112. package/numkong/cmake/toolchain-x86_64-llvm.cmake +47 -0
  113. package/numkong/include/README.md +643 -0
  114. package/numkong/include/module.modulemap +129 -0
  115. package/numkong/include/numkong/attention/sapphireamx.h +1363 -0
  116. package/numkong/include/numkong/attention/sme.h +2075 -0
  117. package/numkong/include/numkong/attention.h +49 -0
  118. package/numkong/include/numkong/capabilities.h +930 -0
  119. package/numkong/include/numkong/cast/README.md +262 -0
  120. package/numkong/include/numkong/cast/diamond.h +64 -0
  121. package/numkong/include/numkong/cast/haswell.h +823 -0
  122. package/numkong/include/numkong/cast/icelake.h +475 -0
  123. package/numkong/include/numkong/cast/loongsonasx.h +252 -0
  124. package/numkong/include/numkong/cast/neon.h +1174 -0
  125. package/numkong/include/numkong/cast/powervsx.h +449 -0
  126. package/numkong/include/numkong/cast/rvv.h +970 -0
  127. package/numkong/include/numkong/cast/sapphire.h +262 -0
  128. package/numkong/include/numkong/cast/serial.h +2336 -0
  129. package/numkong/include/numkong/cast/skylake.h +915 -0
  130. package/numkong/include/numkong/cast/v128relaxed.h +537 -0
  131. package/numkong/include/numkong/cast.h +254 -0
  132. package/numkong/include/numkong/cast.hpp +44 -0
  133. package/numkong/include/numkong/curved/README.md +223 -0
  134. package/numkong/include/numkong/curved/genoa.h +182 -0
  135. package/numkong/include/numkong/curved/haswell.h +276 -0
  136. package/numkong/include/numkong/curved/neon.h +329 -0
  137. package/numkong/include/numkong/curved/neonbfdot.h +211 -0
  138. package/numkong/include/numkong/curved/rvv.h +305 -0
  139. package/numkong/include/numkong/curved/serial.h +207 -0
  140. package/numkong/include/numkong/curved/skylake.h +457 -0
  141. package/numkong/include/numkong/curved/smef64.h +530 -0
  142. package/numkong/include/numkong/curved.h +513 -0
  143. package/numkong/include/numkong/curved.hpp +144 -0
  144. package/numkong/include/numkong/dot/README.md +442 -0
  145. package/numkong/include/numkong/dot/alder.h +563 -0
  146. package/numkong/include/numkong/dot/diamond.h +158 -0
  147. package/numkong/include/numkong/dot/genoa.h +289 -0
  148. package/numkong/include/numkong/dot/haswell.h +1723 -0
  149. package/numkong/include/numkong/dot/icelake.h +997 -0
  150. package/numkong/include/numkong/dot/loongsonasx.h +671 -0
  151. package/numkong/include/numkong/dot/neon.h +869 -0
  152. package/numkong/include/numkong/dot/neonbfdot.h +243 -0
  153. package/numkong/include/numkong/dot/neonfhm.h +358 -0
  154. package/numkong/include/numkong/dot/neonfp8.h +323 -0
  155. package/numkong/include/numkong/dot/neonsdot.h +622 -0
  156. package/numkong/include/numkong/dot/powervsx.h +752 -0
  157. package/numkong/include/numkong/dot/rvv.h +722 -0
  158. package/numkong/include/numkong/dot/rvvbb.h +72 -0
  159. package/numkong/include/numkong/dot/rvvbf16.h +123 -0
  160. package/numkong/include/numkong/dot/rvvhalf.h +129 -0
  161. package/numkong/include/numkong/dot/sapphire.h +141 -0
  162. package/numkong/include/numkong/dot/serial.h +874 -0
  163. package/numkong/include/numkong/dot/sierra.h +405 -0
  164. package/numkong/include/numkong/dot/skylake.h +1116 -0
  165. package/numkong/include/numkong/dot/sve.h +420 -0
  166. package/numkong/include/numkong/dot/svebfdot.h +75 -0
  167. package/numkong/include/numkong/dot/svehalf.h +168 -0
  168. package/numkong/include/numkong/dot/svesdot.h +90 -0
  169. package/numkong/include/numkong/dot/v128relaxed.h +1293 -0
  170. package/numkong/include/numkong/dot.h +1136 -0
  171. package/numkong/include/numkong/dot.hpp +94 -0
  172. package/numkong/include/numkong/dots/README.md +496 -0
  173. package/numkong/include/numkong/dots/alder.h +117 -0
  174. package/numkong/include/numkong/dots/diamond.h +86 -0
  175. package/numkong/include/numkong/dots/genoa.h +100 -0
  176. package/numkong/include/numkong/dots/graniteamx.h +733 -0
  177. package/numkong/include/numkong/dots/haswell.h +310 -0
  178. package/numkong/include/numkong/dots/icelake.h +180 -0
  179. package/numkong/include/numkong/dots/loongsonasx.h +176 -0
  180. package/numkong/include/numkong/dots/neon.h +123 -0
  181. package/numkong/include/numkong/dots/neonbfdot.h +59 -0
  182. package/numkong/include/numkong/dots/neonfhm.h +96 -0
  183. package/numkong/include/numkong/dots/neonfp8.h +99 -0
  184. package/numkong/include/numkong/dots/neonsdot.h +144 -0
  185. package/numkong/include/numkong/dots/powervsx.h +194 -0
  186. package/numkong/include/numkong/dots/rvv.h +2593 -0
  187. package/numkong/include/numkong/dots/sapphireamx.h +4017 -0
  188. package/numkong/include/numkong/dots/serial.h +2867 -0
  189. package/numkong/include/numkong/dots/sierra.h +100 -0
  190. package/numkong/include/numkong/dots/skylake.h +207 -0
  191. package/numkong/include/numkong/dots/sme.h +5041 -0
  192. package/numkong/include/numkong/dots/smebi32.h +477 -0
  193. package/numkong/include/numkong/dots/smef64.h +1328 -0
  194. package/numkong/include/numkong/dots/smehalf.h +47 -0
  195. package/numkong/include/numkong/dots/v128relaxed.h +317 -0
  196. package/numkong/include/numkong/dots.h +3110 -0
  197. package/numkong/include/numkong/dots.hpp +641 -0
  198. package/numkong/include/numkong/each/README.md +465 -0
  199. package/numkong/include/numkong/each/haswell.h +1658 -0
  200. package/numkong/include/numkong/each/icelake.h +272 -0
  201. package/numkong/include/numkong/each/neon.h +1138 -0
  202. package/numkong/include/numkong/each/neonbfdot.h +211 -0
  203. package/numkong/include/numkong/each/neonhalf.h +318 -0
  204. package/numkong/include/numkong/each/rvv.h +1121 -0
  205. package/numkong/include/numkong/each/sapphire.h +343 -0
  206. package/numkong/include/numkong/each/serial.h +282 -0
  207. package/numkong/include/numkong/each/skylake.h +1562 -0
  208. package/numkong/include/numkong/each/v128relaxed.h +562 -0
  209. package/numkong/include/numkong/each.h +2232 -0
  210. package/numkong/include/numkong/each.hpp +434 -0
  211. package/numkong/include/numkong/geospatial/README.md +147 -0
  212. package/numkong/include/numkong/geospatial/haswell.h +633 -0
  213. package/numkong/include/numkong/geospatial/neon.h +615 -0
  214. package/numkong/include/numkong/geospatial/rvv.h +701 -0
  215. package/numkong/include/numkong/geospatial/serial.h +309 -0
  216. package/numkong/include/numkong/geospatial/skylake.h +613 -0
  217. package/numkong/include/numkong/geospatial/v128relaxed.h +652 -0
  218. package/numkong/include/numkong/geospatial.h +453 -0
  219. package/numkong/include/numkong/geospatial.hpp +235 -0
  220. package/numkong/include/numkong/matrix.hpp +336 -0
  221. package/numkong/include/numkong/maxsim/README.md +187 -0
  222. package/numkong/include/numkong/maxsim/alder.h +511 -0
  223. package/numkong/include/numkong/maxsim/genoa.h +115 -0
  224. package/numkong/include/numkong/maxsim/haswell.h +553 -0
  225. package/numkong/include/numkong/maxsim/icelake.h +480 -0
  226. package/numkong/include/numkong/maxsim/neonsdot.h +394 -0
  227. package/numkong/include/numkong/maxsim/sapphireamx.h +877 -0
  228. package/numkong/include/numkong/maxsim/serial.h +490 -0
  229. package/numkong/include/numkong/maxsim/sme.h +993 -0
  230. package/numkong/include/numkong/maxsim/v128relaxed.h +280 -0
  231. package/numkong/include/numkong/maxsim.h +571 -0
  232. package/numkong/include/numkong/maxsim.hpp +133 -0
  233. package/numkong/include/numkong/mesh/README.md +219 -0
  234. package/numkong/include/numkong/mesh/haswell.h +2013 -0
  235. package/numkong/include/numkong/mesh/neon.h +1821 -0
  236. package/numkong/include/numkong/mesh/neonbfdot.h +771 -0
  237. package/numkong/include/numkong/mesh/rvv.h +936 -0
  238. package/numkong/include/numkong/mesh/serial.h +721 -0
  239. package/numkong/include/numkong/mesh/skylake.h +2344 -0
  240. package/numkong/include/numkong/mesh/v128relaxed.h +1081 -0
  241. package/numkong/include/numkong/mesh.h +666 -0
  242. package/numkong/include/numkong/mesh.hpp +701 -0
  243. package/numkong/include/numkong/numkong.h +78 -0
  244. package/numkong/include/numkong/numkong.hpp +58 -0
  245. package/numkong/include/numkong/probability/README.md +177 -0
  246. package/numkong/include/numkong/probability/haswell.h +267 -0
  247. package/numkong/include/numkong/probability/neon.h +255 -0
  248. package/numkong/include/numkong/probability/rvv.h +414 -0
  249. package/numkong/include/numkong/probability/serial.h +181 -0
  250. package/numkong/include/numkong/probability/skylake.h +324 -0
  251. package/numkong/include/numkong/probability.h +380 -0
  252. package/numkong/include/numkong/probability.hpp +120 -0
  253. package/numkong/include/numkong/random.h +50 -0
  254. package/numkong/include/numkong/random.hpp +285 -0
  255. package/numkong/include/numkong/reduce/README.md +552 -0
  256. package/numkong/include/numkong/reduce/alder.h +636 -0
  257. package/numkong/include/numkong/reduce/genoa.h +201 -0
  258. package/numkong/include/numkong/reduce/haswell.h +3785 -0
  259. package/numkong/include/numkong/reduce/icelake.h +549 -0
  260. package/numkong/include/numkong/reduce/neon.h +3980 -0
  261. package/numkong/include/numkong/reduce/neonbfdot.h +142 -0
  262. package/numkong/include/numkong/reduce/neonfhm.h +293 -0
  263. package/numkong/include/numkong/reduce/neonsdot.h +357 -0
  264. package/numkong/include/numkong/reduce/rvv.h +3533 -0
  265. package/numkong/include/numkong/reduce/serial.h +753 -0
  266. package/numkong/include/numkong/reduce/sierra.h +338 -0
  267. package/numkong/include/numkong/reduce/skylake.h +3794 -0
  268. package/numkong/include/numkong/reduce/sve.h +52 -0
  269. package/numkong/include/numkong/reduce/v128relaxed.h +2302 -0
  270. package/numkong/include/numkong/reduce.h +1582 -0
  271. package/numkong/include/numkong/reduce.hpp +778 -0
  272. package/numkong/include/numkong/scalar/README.md +89 -0
  273. package/numkong/include/numkong/scalar/haswell.h +122 -0
  274. package/numkong/include/numkong/scalar/loongsonasx.h +74 -0
  275. package/numkong/include/numkong/scalar/neon.h +122 -0
  276. package/numkong/include/numkong/scalar/neonhalf.h +70 -0
  277. package/numkong/include/numkong/scalar/powervsx.h +96 -0
  278. package/numkong/include/numkong/scalar/rvv.h +211 -0
  279. package/numkong/include/numkong/scalar/sapphire.h +74 -0
  280. package/numkong/include/numkong/scalar/serial.h +332 -0
  281. package/numkong/include/numkong/scalar/v128relaxed.h +56 -0
  282. package/numkong/include/numkong/scalar.h +696 -0
  283. package/numkong/include/numkong/set/README.md +179 -0
  284. package/numkong/include/numkong/set/haswell.h +334 -0
  285. package/numkong/include/numkong/set/icelake.h +485 -0
  286. package/numkong/include/numkong/set/loongsonasx.h +181 -0
  287. package/numkong/include/numkong/set/neon.h +363 -0
  288. package/numkong/include/numkong/set/powervsx.h +326 -0
  289. package/numkong/include/numkong/set/rvv.h +226 -0
  290. package/numkong/include/numkong/set/rvvbb.h +117 -0
  291. package/numkong/include/numkong/set/serial.h +174 -0
  292. package/numkong/include/numkong/set/sve.h +187 -0
  293. package/numkong/include/numkong/set/v128relaxed.h +240 -0
  294. package/numkong/include/numkong/set.h +471 -0
  295. package/numkong/include/numkong/set.hpp +114 -0
  296. package/numkong/include/numkong/sets/README.md +149 -0
  297. package/numkong/include/numkong/sets/haswell.h +63 -0
  298. package/numkong/include/numkong/sets/icelake.h +66 -0
  299. package/numkong/include/numkong/sets/loongsonasx.h +52 -0
  300. package/numkong/include/numkong/sets/neon.h +61 -0
  301. package/numkong/include/numkong/sets/powervsx.h +65 -0
  302. package/numkong/include/numkong/sets/serial.h +43 -0
  303. package/numkong/include/numkong/sets/smebi32.h +1135 -0
  304. package/numkong/include/numkong/sets/v128relaxed.h +58 -0
  305. package/numkong/include/numkong/sets.h +382 -0
  306. package/numkong/include/numkong/sparse/README.md +156 -0
  307. package/numkong/include/numkong/sparse/icelake.h +463 -0
  308. package/numkong/include/numkong/sparse/neon.h +288 -0
  309. package/numkong/include/numkong/sparse/serial.h +117 -0
  310. package/numkong/include/numkong/sparse/sve2.h +504 -0
  311. package/numkong/include/numkong/sparse/turin.h +322 -0
  312. package/numkong/include/numkong/sparse.h +363 -0
  313. package/numkong/include/numkong/sparse.hpp +113 -0
  314. package/numkong/include/numkong/spatial/README.md +439 -0
  315. package/numkong/include/numkong/spatial/alder.h +604 -0
  316. package/numkong/include/numkong/spatial/diamond.h +240 -0
  317. package/numkong/include/numkong/spatial/genoa.h +222 -0
  318. package/numkong/include/numkong/spatial/haswell.h +979 -0
  319. package/numkong/include/numkong/spatial/icelake.h +1067 -0
  320. package/numkong/include/numkong/spatial/loongsonasx.h +483 -0
  321. package/numkong/include/numkong/spatial/neon.h +846 -0
  322. package/numkong/include/numkong/spatial/neonbfdot.h +164 -0
  323. package/numkong/include/numkong/spatial/neonfp8.h +258 -0
  324. package/numkong/include/numkong/spatial/neonsdot.h +429 -0
  325. package/numkong/include/numkong/spatial/powervsx.h +738 -0
  326. package/numkong/include/numkong/spatial/rvv.h +991 -0
  327. package/numkong/include/numkong/spatial/rvvbf16.h +128 -0
  328. package/numkong/include/numkong/spatial/rvvhalf.h +120 -0
  329. package/numkong/include/numkong/spatial/serial.h +347 -0
  330. package/numkong/include/numkong/spatial/sierra.h +516 -0
  331. package/numkong/include/numkong/spatial/skylake.h +605 -0
  332. package/numkong/include/numkong/spatial/sve.h +246 -0
  333. package/numkong/include/numkong/spatial/svebfdot.h +125 -0
  334. package/numkong/include/numkong/spatial/svehalf.h +136 -0
  335. package/numkong/include/numkong/spatial/svesdot.h +143 -0
  336. package/numkong/include/numkong/spatial/v128relaxed.h +969 -0
  337. package/numkong/include/numkong/spatial.h +1681 -0
  338. package/numkong/include/numkong/spatial.hpp +183 -0
  339. package/numkong/include/numkong/spatials/README.md +580 -0
  340. package/numkong/include/numkong/spatials/alder.h +94 -0
  341. package/numkong/include/numkong/spatials/diamond.h +82 -0
  342. package/numkong/include/numkong/spatials/genoa.h +94 -0
  343. package/numkong/include/numkong/spatials/graniteamx.h +173 -0
  344. package/numkong/include/numkong/spatials/haswell.h +219 -0
  345. package/numkong/include/numkong/spatials/icelake.h +113 -0
  346. package/numkong/include/numkong/spatials/loongsonasx.h +153 -0
  347. package/numkong/include/numkong/spatials/neon.h +109 -0
  348. package/numkong/include/numkong/spatials/neonbfdot.h +60 -0
  349. package/numkong/include/numkong/spatials/neonfhm.h +92 -0
  350. package/numkong/include/numkong/spatials/neonfp8.h +111 -0
  351. package/numkong/include/numkong/spatials/neonsdot.h +143 -0
  352. package/numkong/include/numkong/spatials/powervsx.h +153 -0
  353. package/numkong/include/numkong/spatials/rvv.h +1976 -0
  354. package/numkong/include/numkong/spatials/sapphireamx.h +1149 -0
  355. package/numkong/include/numkong/spatials/serial.h +248 -0
  356. package/numkong/include/numkong/spatials/sierra.h +96 -0
  357. package/numkong/include/numkong/spatials/skylake.h +184 -0
  358. package/numkong/include/numkong/spatials/sme.h +1927 -0
  359. package/numkong/include/numkong/spatials/smef64.h +483 -0
  360. package/numkong/include/numkong/spatials/v128relaxed.h +240 -0
  361. package/numkong/include/numkong/spatials.h +3374 -0
  362. package/numkong/include/numkong/spatials.hpp +509 -0
  363. package/numkong/include/numkong/tensor.hpp +1981 -0
  364. package/numkong/include/numkong/trigonometry/README.md +184 -0
  365. package/numkong/include/numkong/trigonometry/haswell.h +653 -0
  366. package/numkong/include/numkong/trigonometry/neon.h +638 -0
  367. package/numkong/include/numkong/trigonometry/rvv.h +700 -0
  368. package/numkong/include/numkong/trigonometry/serial.h +703 -0
  369. package/numkong/include/numkong/trigonometry/skylake.h +725 -0
  370. package/numkong/include/numkong/trigonometry/v128relaxed.h +667 -0
  371. package/numkong/include/numkong/trigonometry.h +467 -0
  372. package/numkong/include/numkong/trigonometry.hpp +166 -0
  373. package/numkong/include/numkong/types.h +1695 -0
  374. package/numkong/include/numkong/types.hpp +6027 -0
  375. package/numkong/include/numkong/vector.hpp +766 -0
  376. package/numkong/javascript/@numkong-darwin-arm64/package.json +26 -0
  377. package/numkong/javascript/@numkong-darwin-x64/package.json +26 -0
  378. package/numkong/javascript/@numkong-linux-arm64/package.json +25 -0
  379. package/numkong/javascript/@numkong-linux-x64/package.json +25 -0
  380. package/numkong/javascript/@numkong-win32-arm64/package.json +25 -0
  381. package/numkong/javascript/@numkong-win32-x64/package.json +25 -0
  382. package/numkong/javascript/README.md +254 -0
  383. package/numkong/javascript/node-gyp-build.d.ts +4 -0
  384. package/numkong/javascript/numkong-browser.ts +40 -0
  385. package/numkong/javascript/numkong-wasm.ts +750 -0
  386. package/numkong/javascript/numkong.c +712 -0
  387. package/numkong/javascript/numkong.ts +613 -0
  388. package/{simsimd → numkong}/javascript/tsconfig-base.json +6 -3
  389. package/numkong/javascript/tsconfig-cjs.json +13 -0
  390. package/{simsimd → numkong}/javascript/tsconfig-esm.json +1 -1
  391. package/numkong/javascript/types.ts +784 -0
  392. package/numkong/numkong.gypi +104 -0
  393. package/numkong/package.json +108 -0
  394. package/numkong/probes/arm_neon.c +8 -0
  395. package/numkong/probes/arm_neon_bfdot.c +9 -0
  396. package/numkong/probes/arm_neon_fhm.c +9 -0
  397. package/numkong/probes/arm_neon_half.c +8 -0
  398. package/numkong/probes/arm_neon_sdot.c +9 -0
  399. package/numkong/probes/arm_neonfp8.c +9 -0
  400. package/numkong/probes/arm_sme.c +16 -0
  401. package/numkong/probes/arm_sme2.c +16 -0
  402. package/numkong/probes/arm_sme2p1.c +16 -0
  403. package/numkong/probes/arm_sme_bf16.c +16 -0
  404. package/numkong/probes/arm_sme_bi32.c +16 -0
  405. package/numkong/probes/arm_sme_f64.c +16 -0
  406. package/numkong/probes/arm_sme_fa64.c +14 -0
  407. package/numkong/probes/arm_sme_half.c +16 -0
  408. package/numkong/probes/arm_sme_lut2.c +15 -0
  409. package/numkong/probes/arm_sve.c +18 -0
  410. package/numkong/probes/arm_sve2.c +20 -0
  411. package/numkong/probes/arm_sve2p1.c +18 -0
  412. package/numkong/probes/arm_sve_bfdot.c +20 -0
  413. package/numkong/probes/arm_sve_half.c +18 -0
  414. package/numkong/probes/arm_sve_sdot.c +21 -0
  415. package/numkong/probes/loongarch_lasx.c +12 -0
  416. package/numkong/probes/power_vsx.c +12 -0
  417. package/numkong/probes/probe.js +127 -0
  418. package/numkong/probes/riscv_rvv.c +14 -0
  419. package/numkong/probes/riscv_rvv_bb.c +15 -0
  420. package/numkong/probes/riscv_rvv_bf16.c +17 -0
  421. package/numkong/probes/riscv_rvv_half.c +14 -0
  422. package/numkong/probes/wasm_v128relaxed.c +11 -0
  423. package/numkong/probes/x86_alder.c +17 -0
  424. package/numkong/probes/x86_diamond.c +17 -0
  425. package/numkong/probes/x86_genoa.c +17 -0
  426. package/numkong/probes/x86_graniteamx.c +19 -0
  427. package/numkong/probes/x86_haswell.c +11 -0
  428. package/numkong/probes/x86_icelake.c +17 -0
  429. package/numkong/probes/x86_sapphire.c +16 -0
  430. package/numkong/probes/x86_sapphireamx.c +18 -0
  431. package/numkong/probes/x86_sierra.c +17 -0
  432. package/numkong/probes/x86_skylake.c +15 -0
  433. package/numkong/probes/x86_turin.c +17 -0
  434. package/numkong/rust/README.md +698 -0
  435. package/numkong/rust/capabilities.rs +96 -0
  436. package/numkong/rust/cast.rs +314 -0
  437. package/numkong/rust/curved.rs +391 -0
  438. package/numkong/rust/each.rs +2721 -0
  439. package/numkong/rust/geospatial.rs +343 -0
  440. package/numkong/rust/matrix.rs +5424 -0
  441. package/numkong/rust/maxsim.rs +540 -0
  442. package/numkong/rust/mesh.rs +837 -0
  443. package/numkong/rust/numkong.rs +409 -0
  444. package/numkong/rust/probability.rs +272 -0
  445. package/numkong/rust/reduce.rs +1249 -0
  446. package/numkong/rust/set.rs +194 -0
  447. package/numkong/rust/sparse.rs +513 -0
  448. package/numkong/rust/spatial.rs +1815 -0
  449. package/numkong/rust/tensor.rs +7565 -0
  450. package/numkong/rust/types.rs +3666 -0
  451. package/numkong/rust/vector.rs +1899 -0
  452. package/numkong/rustfmt.toml +2 -0
  453. package/numkong/swift/Matrix.swift +1388 -0
  454. package/numkong/swift/MaxSim.swift +151 -0
  455. package/numkong/swift/NumKong.swift +283 -0
  456. package/numkong/swift/Numerics.swift +329 -0
  457. package/numkong/swift/README.md +433 -0
  458. package/numkong/swift/Scalar.swift +741 -0
  459. package/numkong/swift/Tensor.swift +209 -0
  460. package/numkong/test/README.md +315 -0
  461. package/numkong/test/playwright.config.ts +16 -0
  462. package/numkong/test/swift/Test.swift +534 -0
  463. package/numkong/test/test-browser-bundle-runner.mjs +36 -0
  464. package/numkong/test/test-browser-bundle.html +188 -0
  465. package/numkong/test/test-browser-runner.mjs +43 -0
  466. package/numkong/test/test-browser.html +305 -0
  467. package/numkong/test/test-dtypes.mjs +453 -0
  468. package/numkong/test/test-wasm.mjs +267 -0
  469. package/numkong/test/test.cpp +442 -0
  470. package/numkong/test/test.hpp +580 -0
  471. package/numkong/test/test.mjs +363 -0
  472. package/numkong/test/test_base.py +1041 -0
  473. package/numkong/test/test_cast.cpp +242 -0
  474. package/numkong/test/test_cross.hpp +551 -0
  475. package/numkong/test/test_cross.py +467 -0
  476. package/numkong/test/test_cross_amx.cpp +50 -0
  477. package/numkong/test/test_cross_arm.cpp +225 -0
  478. package/numkong/test/test_cross_blas.cpp +354 -0
  479. package/numkong/test/test_cross_loongarch.cpp +90 -0
  480. package/numkong/test/test_cross_power.cpp +73 -0
  481. package/numkong/test/test_cross_rvv.cpp +109 -0
  482. package/numkong/test/test_cross_serial.cpp +181 -0
  483. package/numkong/test/test_cross_sme.cpp +148 -0
  484. package/numkong/test/test_cross_wasm.cpp +141 -0
  485. package/numkong/test/test_cross_x86.cpp +442 -0
  486. package/numkong/test/test_curved.cpp +182 -0
  487. package/numkong/test/test_curved.py +207 -0
  488. package/numkong/test/test_dot.cpp +353 -0
  489. package/numkong/test/test_dot.py +296 -0
  490. package/numkong/test/test_each.cpp +437 -0
  491. package/numkong/test/test_each.py +756 -0
  492. package/numkong/test/test_geospatial.cpp +155 -0
  493. package/numkong/test/test_geospatial.py +241 -0
  494. package/numkong/test/test_maxsim.cpp +138 -0
  495. package/numkong/test/test_maxsim.py +209 -0
  496. package/numkong/test/test_mesh.cpp +203 -0
  497. package/numkong/test/test_mesh.py +394 -0
  498. package/numkong/test/test_probability.cpp +136 -0
  499. package/numkong/test/test_probability.py +166 -0
  500. package/numkong/test/test_reduce.cpp +563 -0
  501. package/numkong/test/test_reduce.py +781 -0
  502. package/numkong/test/test_set.cpp +151 -0
  503. package/numkong/test/test_set.py +120 -0
  504. package/numkong/test/test_similarities.py +579 -0
  505. package/numkong/test/test_sparse.cpp +121 -0
  506. package/numkong/test/test_sparse.py +118 -0
  507. package/numkong/test/test_spatial.cpp +518 -0
  508. package/numkong/test/test_spatial.py +269 -0
  509. package/numkong/test/test_tensor.cpp +884 -0
  510. package/numkong/test/test_tensor.py +1417 -0
  511. package/numkong/test/test_trigonometry.cpp +156 -0
  512. package/numkong/test/test_trigonometry.py +180 -0
  513. package/package.json +7 -8
  514. package/prebuilds/darwin-arm64+x64/usearch.node +0 -0
  515. package/prebuilds/linux-arm64/usearch.node +0 -0
  516. package/prebuilds/linux-x64/usearch.node +0 -0
  517. package/fp16/CMakeLists.txt +0 -176
  518. package/fp16/LICENSE +0 -11
  519. package/fp16/README.md +0 -20
  520. package/fp16/bench/alt-element.cc +0 -89
  521. package/fp16/bench/from-alt-array.cc +0 -175
  522. package/fp16/bench/from-ieee-array.cc +0 -399
  523. package/fp16/bench/ieee-element.cc +0 -272
  524. package/fp16/bench/to-alt-array.cc +0 -91
  525. package/fp16/bench/to-ieee-array.cc +0 -303
  526. package/fp16/cmake/DownloadGoogleBenchmark.cmake +0 -15
  527. package/fp16/cmake/DownloadGoogleTest.cmake +0 -15
  528. package/fp16/cmake/DownloadPSimd.cmake +0 -15
  529. package/fp16/configure.py +0 -65
  530. package/fp16/confu.yaml +0 -8
  531. package/fp16/include/fp16/__init__.py +0 -0
  532. package/fp16/include/fp16/avx.py +0 -50
  533. package/fp16/include/fp16/avx2.py +0 -53
  534. package/fp16/include/fp16/bitcasts.h +0 -92
  535. package/fp16/include/fp16/fp16.h +0 -451
  536. package/fp16/include/fp16/psimd.h +0 -131
  537. package/fp16/include/fp16.h +0 -11
  538. package/fp16/jni/Android.mk +0 -6
  539. package/fp16/jni/Application.mk +0 -4
  540. package/fp16/test/alt-from-fp32-value.cc +0 -498
  541. package/fp16/test/alt-to-fp32-bits.cc +0 -263
  542. package/fp16/test/alt-to-fp32-psimd.cc +0 -145
  543. package/fp16/test/alt-to-fp32-value.cc +0 -348
  544. package/fp16/test/alt-to-fp32x2-psimd.cc +0 -245
  545. package/fp16/test/bitcasts.cc +0 -56
  546. package/fp16/test/ieee-from-fp32-value.cc +0 -523
  547. package/fp16/test/ieee-to-fp32-bits.cc +0 -333
  548. package/fp16/test/ieee-to-fp32-psimd.cc +0 -312
  549. package/fp16/test/ieee-to-fp32-value.cc +0 -420
  550. package/fp16/test/ieee-to-fp32x2-psimd.cc +0 -541
  551. package/fp16/test/peachpy/alt-xmm-to-fp32-xmm-avx.cc +0 -96
  552. package/fp16/test/peachpy/alt-xmm-to-fp32-ymm-avx2.cc +0 -96
  553. package/fp16/test/peachpy/stubs.py +0 -38
  554. package/fp16/test/tables.cc +0 -530
  555. package/fp16/test/tables.h +0 -13
  556. package/fp16/third-party/THHalf.h +0 -92
  557. package/fp16/third-party/eigen-half.h +0 -170
  558. package/fp16/third-party/float16-compressor.h +0 -80
  559. package/fp16/third-party/half.hpp +0 -2916
  560. package/fp16/third-party/npy-halffloat.h +0 -328
  561. package/simsimd/.git-blame-ignore-revs +0 -2
  562. package/simsimd/CONTRIBUTING.md +0 -321
  563. package/simsimd/Cargo.toml +0 -49
  564. package/simsimd/MANIFEST.in +0 -4
  565. package/simsimd/Package.swift +0 -28
  566. package/simsimd/README.md +0 -1493
  567. package/simsimd/binding.gyp +0 -35
  568. package/simsimd/build.rs +0 -70
  569. package/simsimd/c/lib.c +0 -374
  570. package/simsimd/include/module.modulemap +0 -4
  571. package/simsimd/include/simsimd/binary.h +0 -486
  572. package/simsimd/include/simsimd/curved.h +0 -1541
  573. package/simsimd/include/simsimd/dot.h +0 -1853
  574. package/simsimd/include/simsimd/elementwise.h +0 -2465
  575. package/simsimd/include/simsimd/geospatial.h +0 -43
  576. package/simsimd/include/simsimd/mesh.h +0 -69
  577. package/simsimd/include/simsimd/probability.h +0 -606
  578. package/simsimd/include/simsimd/simsimd.h +0 -2567
  579. package/simsimd/include/simsimd/sparse.h +0 -1384
  580. package/simsimd/include/simsimd/spatial.h +0 -2335
  581. package/simsimd/include/simsimd/types.h +0 -666
  582. package/simsimd/javascript/fallback.ts +0 -233
  583. package/simsimd/javascript/lib.c +0 -125
  584. package/simsimd/javascript/node-gyp-build.d.ts +0 -1
  585. package/simsimd/javascript/simsimd.ts +0 -179
  586. package/simsimd/javascript/tsconfig-cjs.json +0 -8
  587. package/simsimd/package-lock.json +0 -1761
  588. package/simsimd/package.json +0 -61
  589. package/simsimd/rust/lib.rs +0 -1821
  590. package/simsimd/scripts/bench.cxx +0 -1170
  591. package/simsimd/scripts/bench.ipynb +0 -187
  592. package/simsimd/scripts/bench.js +0 -127
  593. package/simsimd/scripts/bench.rs +0 -876
  594. package/simsimd/scripts/bench_vectors.py +0 -734
  595. package/simsimd/scripts/bench_vectors_live.py +0 -213
  596. package/simsimd/scripts/test.c +0 -187
  597. package/simsimd/scripts/test.mjs +0 -220
  598. package/simsimd/scripts/test.py +0 -1505
  599. package/simsimd/swift/SimSIMD.swift +0 -149
  600. package/simsimd/swift/Test.swift +0 -128
  601. /package/{simsimd → numkong}/CODE_OF_CONDUCT.md +0 -0
  602. /package/{simsimd → numkong}/LICENSE +0 -0
  603. /package/{simsimd → numkong}/javascript/dist-package-cjs.json +0 -0
  604. /package/{simsimd → numkong}/javascript/dist-package-esm.json +0 -0
@@ -0,0 +1,1976 @@
1
+ /**
2
+ * @brief Batched Spatial Distances for RISC-V Vector (RVV).
3
+ * @file include/numkong/spatials/rvv.h
4
+ * @author Ash Vardanian
5
+ * @date February 23, 2026
6
+ *
7
+ * @sa include/numkong/spatials.h
8
+ */
9
+ #ifndef NK_SPATIALS_RVV_H
10
+ #define NK_SPATIALS_RVV_H
11
+
12
+ #if NK_TARGET_RISCV64_
13
+ #if NK_TARGET_RVV
14
+
15
+ #include "numkong/dots/serial.h"
16
+ #include "numkong/dots/rvv.h"
17
+ #include "numkong/spatial/rvv.h"
18
+
19
+ #if defined(__cplusplus)
20
+ extern "C" {
21
+ #endif
22
+
23
+ #if defined(__clang__)
24
+ #pragma clang attribute push(__attribute__((target("arch=+v"))), apply_to = function)
25
+ #elif defined(__GNUC__)
26
+ #pragma GCC push_options
27
+ #pragma GCC target("arch=+v")
28
+ #endif
29
+
30
+ #pragma region F32 Floats
31
+
32
+ NK_INTERNAL void nk_angulars_packed_f32_rvv_finalize_(nk_f32_t const *a, void const *b_packed, nk_f64_t *c,
33
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
34
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
35
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
36
+ nk_f64_t const *target_norms = (nk_f64_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
37
+ header->column_count * header->depth_padded_values *
38
+ sizeof(nk_f32_t));
39
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
40
+ nk_f32_t const *a_row = a + row_index * a_stride_elements;
41
+ nk_f64_t *result_row = c + row_index * c_stride_elements;
42
+ nk_f64_t query_norm_sq_f64 = nk_dots_reduce_sumsq_f32_(a_row, depth);
43
+ nk_size_t count_columns = columns;
44
+ nk_f64_t *result_ptr = result_row;
45
+ nk_f64_t const *norms_ptr = target_norms;
46
+ while (count_columns > 0) {
47
+ size_t vector_length = __riscv_vsetvl_e64m1(count_columns);
48
+ vfloat64m1_t dots_f64m1 = __riscv_vle64_v_f64m1(result_ptr, vector_length);
49
+ vfloat64m1_t target_norms_sq_f64m1 = __riscv_vle64_v_f64m1(norms_ptr, vector_length);
50
+ vfloat64m1_t norms_product_f64m1 = __riscv_vfmul_vf_f64m1(target_norms_sq_f64m1, query_norm_sq_f64,
51
+ vector_length);
52
+ vfloat64m1_t rsqrt_f64m1 = nk_rsqrt_f64m1_rvv_(norms_product_f64m1, vector_length);
53
+ vfloat64m1_t normalized_dots_f64m1 = __riscv_vfmul_vv_f64m1(dots_f64m1, rsqrt_f64m1, vector_length);
54
+ vfloat64m1_t angular_f64m1 = __riscv_vfrsub_vf_f64m1(normalized_dots_f64m1, 1.0, vector_length);
55
+ angular_f64m1 = __riscv_vfmax_vf_f64m1(angular_f64m1, 0.0, vector_length);
56
+ __riscv_vse64_v_f64m1(result_ptr, angular_f64m1, vector_length);
57
+ result_ptr += vector_length;
58
+ norms_ptr += vector_length;
59
+ count_columns -= vector_length;
60
+ }
61
+ }
62
+ }
63
+
64
+ NK_PUBLIC void nk_angulars_packed_f32_rvv( //
65
+ nk_f32_t const *a, void const *b_packed, nk_f64_t *c, //
66
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
67
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
68
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f32_t);
69
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f64_t);
70
+ nk_dots_packed_f32_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
71
+ nk_angulars_packed_f32_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
72
+ }
73
+
74
+ NK_INTERNAL void nk_euclideans_packed_f32_rvv_finalize_(nk_f32_t const *a, void const *b_packed, nk_f64_t *c,
75
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
76
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
77
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
78
+ nk_f64_t const *target_norms = (nk_f64_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
79
+ header->column_count * header->depth_padded_values *
80
+ sizeof(nk_f32_t));
81
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
82
+ nk_f32_t const *a_row = a + row_index * a_stride_elements;
83
+ nk_f64_t *result_row = c + row_index * c_stride_elements;
84
+ nk_f64_t query_norm_sq_f64 = nk_dots_reduce_sumsq_f32_(a_row, depth);
85
+ nk_size_t count_columns = columns;
86
+ nk_f64_t *result_ptr = result_row;
87
+ nk_f64_t const *norms_ptr = target_norms;
88
+ while (count_columns > 0) {
89
+ size_t vector_length = __riscv_vsetvl_e64m1(count_columns);
90
+ vfloat64m1_t dots_f64m1 = __riscv_vle64_v_f64m1(result_ptr, vector_length);
91
+ vfloat64m1_t target_norms_sq_f64m1 = __riscv_vle64_v_f64m1(norms_ptr, vector_length);
92
+ vfloat64m1_t sum_sq_f64m1 = __riscv_vfadd_vf_f64m1(target_norms_sq_f64m1, query_norm_sq_f64, vector_length);
93
+ vfloat64m1_t dist_sq_f64m1 = __riscv_vfsub_vv_f64m1(
94
+ sum_sq_f64m1, __riscv_vfmul_vf_f64m1(dots_f64m1, 2.0, vector_length), vector_length);
95
+ dist_sq_f64m1 = __riscv_vfmax_vf_f64m1(dist_sq_f64m1, 0.0, vector_length);
96
+ __riscv_vse64_v_f64m1(result_ptr, __riscv_vfsqrt_v_f64m1(dist_sq_f64m1, vector_length), vector_length);
97
+ result_ptr += vector_length;
98
+ norms_ptr += vector_length;
99
+ count_columns -= vector_length;
100
+ }
101
+ }
102
+ }
103
+
104
+ NK_PUBLIC void nk_euclideans_packed_f32_rvv( //
105
+ nk_f32_t const *a, void const *b_packed, nk_f64_t *c, //
106
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
107
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
108
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f32_t);
109
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f64_t);
110
+ nk_dots_packed_f32_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
111
+ nk_euclideans_packed_f32_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
112
+ }
113
+
114
+ NK_INTERNAL void nk_angulars_symmetric_f32_rvv_finalize_(nk_f32_t const *vectors, nk_size_t vectors_count,
115
+ nk_size_t depth, nk_size_t stride_elements, nk_f64_t *result,
116
+ nk_size_t result_stride_elements, nk_size_t row_start,
117
+ nk_size_t row_count) {
118
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
119
+ nk_f64_t *result_row = result + row_index * result_stride_elements;
120
+ result_row[row_index] = nk_dots_reduce_sumsq_f32_(vectors + row_index * stride_elements, depth);
121
+ }
122
+ nk_f64_t norms_cache[256];
123
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
124
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
125
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
126
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_f32_(vectors + col * stride_elements, depth);
127
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
128
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
129
+ if (col_start >= chunk_end) continue;
130
+ nk_f64_t *result_row = result + row_index * result_stride_elements;
131
+ nk_f64_t query_norm_sq_f64 = result_row[row_index];
132
+ nk_size_t count_remaining = chunk_end - col_start;
133
+ nk_f64_t *result_ptr = result_row + col_start;
134
+ nk_f64_t const *norms_ptr = norms_cache + (col_start - chunk_start);
135
+ while (count_remaining > 0) {
136
+ size_t vector_length = __riscv_vsetvl_e64m1(count_remaining);
137
+ vfloat64m1_t dots_f64m1 = __riscv_vle64_v_f64m1(result_ptr, vector_length);
138
+ vfloat64m1_t target_norms_sq_f64m1 = __riscv_vle64_v_f64m1(norms_ptr, vector_length);
139
+ vfloat64m1_t norms_product_f64m1 = __riscv_vfmul_vf_f64m1(target_norms_sq_f64m1, query_norm_sq_f64,
140
+ vector_length);
141
+ vfloat64m1_t rsqrt_f64m1 = nk_rsqrt_f64m1_rvv_(norms_product_f64m1, vector_length);
142
+ vfloat64m1_t normalized_dots_f64m1 = __riscv_vfmul_vv_f64m1(dots_f64m1, rsqrt_f64m1, vector_length);
143
+ vfloat64m1_t angular_f64m1 = __riscv_vfrsub_vf_f64m1(normalized_dots_f64m1, 1.0, vector_length);
144
+ angular_f64m1 = __riscv_vfmax_vf_f64m1(angular_f64m1, 0.0, vector_length);
145
+ __riscv_vse64_v_f64m1(result_ptr, angular_f64m1, vector_length);
146
+ result_ptr += vector_length;
147
+ norms_ptr += vector_length;
148
+ count_remaining -= vector_length;
149
+ }
150
+ }
151
+ }
152
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
153
+ result[row_index * result_stride_elements + row_index] = 0;
154
+ }
155
+
156
+ NK_PUBLIC void nk_angulars_symmetric_f32_rvv( //
157
+ nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
158
+ nk_f64_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
159
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_f32_t);
160
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f64_t);
161
+ nk_dots_symmetric_f32_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes, row_start,
162
+ row_count);
163
+ nk_angulars_symmetric_f32_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
164
+ result_stride_elements, row_start, row_count);
165
+ }
166
+
167
+ NK_INTERNAL void nk_euclideans_symmetric_f32_rvv_finalize_(nk_f32_t const *vectors, nk_size_t vectors_count,
168
+ nk_size_t depth, nk_size_t stride_elements, nk_f64_t *result,
169
+ nk_size_t result_stride_elements, nk_size_t row_start,
170
+ nk_size_t row_count) {
171
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
172
+ nk_f64_t *result_row = result + row_index * result_stride_elements;
173
+ result_row[row_index] = nk_dots_reduce_sumsq_f32_(vectors + row_index * stride_elements, depth);
174
+ }
175
+ nk_f64_t norms_cache[256];
176
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
177
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
178
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
179
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_f32_(vectors + col * stride_elements, depth);
180
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
181
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
182
+ if (col_start >= chunk_end) continue;
183
+ nk_f64_t *result_row = result + row_index * result_stride_elements;
184
+ nk_f64_t query_norm_sq_f64 = result_row[row_index];
185
+ nk_size_t count_remaining = chunk_end - col_start;
186
+ nk_f64_t *result_ptr = result_row + col_start;
187
+ nk_f64_t const *norms_ptr = norms_cache + (col_start - chunk_start);
188
+ while (count_remaining > 0) {
189
+ size_t vector_length = __riscv_vsetvl_e64m1(count_remaining);
190
+ vfloat64m1_t dots_f64m1 = __riscv_vle64_v_f64m1(result_ptr, vector_length);
191
+ vfloat64m1_t target_norms_sq_f64m1 = __riscv_vle64_v_f64m1(norms_ptr, vector_length);
192
+ vfloat64m1_t sum_sq_f64m1 = __riscv_vfadd_vf_f64m1(target_norms_sq_f64m1, query_norm_sq_f64,
193
+ vector_length);
194
+ vfloat64m1_t dist_sq_f64m1 = __riscv_vfsub_vv_f64m1(
195
+ sum_sq_f64m1, __riscv_vfmul_vf_f64m1(dots_f64m1, 2.0, vector_length), vector_length);
196
+ dist_sq_f64m1 = __riscv_vfmax_vf_f64m1(dist_sq_f64m1, 0.0, vector_length);
197
+ __riscv_vse64_v_f64m1(result_ptr, __riscv_vfsqrt_v_f64m1(dist_sq_f64m1, vector_length), vector_length);
198
+ result_ptr += vector_length;
199
+ norms_ptr += vector_length;
200
+ count_remaining -= vector_length;
201
+ }
202
+ }
203
+ }
204
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
205
+ result[row_index * result_stride_elements + row_index] = 0;
206
+ }
207
+
208
+ NK_PUBLIC void nk_euclideans_symmetric_f32_rvv( //
209
+ nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
210
+ nk_f64_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
211
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_f32_t);
212
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f64_t);
213
+ nk_dots_symmetric_f32_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes, row_start,
214
+ row_count);
215
+ nk_euclideans_symmetric_f32_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
216
+ result_stride_elements, row_start, row_count);
217
+ }
218
+
219
+ #pragma endregion F32 Floats
220
+
221
+ #pragma region F64 Floats
222
+
223
+ NK_INTERNAL void nk_angulars_packed_f64_rvv_finalize_(nk_f64_t const *a, void const *b_packed, nk_f64_t *c,
224
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
225
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
226
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
227
+ nk_f64_t const *target_norms = (nk_f64_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
228
+ header->column_count * header->depth_padded_values *
229
+ sizeof(nk_f64_t));
230
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
231
+ nk_f64_t const *a_row = a + row_index * a_stride_elements;
232
+ nk_f64_t *result_row = c + row_index * c_stride_elements;
233
+ nk_f64_t query_norm_sq_f64 = nk_dots_reduce_sumsq_f64_(a_row, depth);
234
+ nk_size_t count_columns = columns;
235
+ nk_f64_t *result_ptr = result_row;
236
+ nk_f64_t const *norms_ptr = target_norms;
237
+ while (count_columns > 0) {
238
+ size_t vector_length = __riscv_vsetvl_e64m1(count_columns);
239
+ vfloat64m1_t dots_f64m1 = __riscv_vle64_v_f64m1(result_ptr, vector_length);
240
+ vfloat64m1_t target_norms_sq_f64m1 = __riscv_vle64_v_f64m1(norms_ptr, vector_length);
241
+ vfloat64m1_t norms_product_f64m1 = __riscv_vfmul_vf_f64m1(target_norms_sq_f64m1, query_norm_sq_f64,
242
+ vector_length);
243
+ vfloat64m1_t rsqrt_f64m1 = nk_rsqrt_f64m1_rvv_(norms_product_f64m1, vector_length);
244
+ vfloat64m1_t normalized_dots_f64m1 = __riscv_vfmul_vv_f64m1(dots_f64m1, rsqrt_f64m1, vector_length);
245
+ vfloat64m1_t angular_f64m1 = __riscv_vfrsub_vf_f64m1(normalized_dots_f64m1, 1.0, vector_length);
246
+ angular_f64m1 = __riscv_vfmax_vf_f64m1(angular_f64m1, 0.0, vector_length);
247
+ __riscv_vse64_v_f64m1(result_ptr, angular_f64m1, vector_length);
248
+ result_ptr += vector_length;
249
+ norms_ptr += vector_length;
250
+ count_columns -= vector_length;
251
+ }
252
+ }
253
+ }
254
+
255
+ NK_PUBLIC void nk_angulars_packed_f64_rvv( //
256
+ nk_f64_t const *a, void const *b_packed, nk_f64_t *c, //
257
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
258
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
259
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f64_t);
260
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f64_t);
261
+ nk_dots_packed_f64_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
262
+ nk_angulars_packed_f64_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
263
+ }
264
+
265
+ NK_INTERNAL void nk_euclideans_packed_f64_rvv_finalize_(nk_f64_t const *a, void const *b_packed, nk_f64_t *c,
266
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
267
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
268
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
269
+ nk_f64_t const *target_norms = (nk_f64_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
270
+ header->column_count * header->depth_padded_values *
271
+ sizeof(nk_f64_t));
272
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
273
+ nk_f64_t const *a_row = a + row_index * a_stride_elements;
274
+ nk_f64_t *result_row = c + row_index * c_stride_elements;
275
+ nk_f64_t query_norm_sq_f64 = nk_dots_reduce_sumsq_f64_(a_row, depth);
276
+ nk_size_t count_columns = columns;
277
+ nk_f64_t *result_ptr = result_row;
278
+ nk_f64_t const *norms_ptr = target_norms;
279
+ while (count_columns > 0) {
280
+ size_t vector_length = __riscv_vsetvl_e64m1(count_columns);
281
+ vfloat64m1_t dots_f64m1 = __riscv_vle64_v_f64m1(result_ptr, vector_length);
282
+ vfloat64m1_t target_norms_sq_f64m1 = __riscv_vle64_v_f64m1(norms_ptr, vector_length);
283
+ vfloat64m1_t sum_sq_f64m1 = __riscv_vfadd_vf_f64m1(target_norms_sq_f64m1, query_norm_sq_f64, vector_length);
284
+ vfloat64m1_t dist_sq_f64m1 = __riscv_vfsub_vv_f64m1(
285
+ sum_sq_f64m1, __riscv_vfmul_vf_f64m1(dots_f64m1, 2.0, vector_length), vector_length);
286
+ dist_sq_f64m1 = __riscv_vfmax_vf_f64m1(dist_sq_f64m1, 0.0, vector_length);
287
+ __riscv_vse64_v_f64m1(result_ptr, __riscv_vfsqrt_v_f64m1(dist_sq_f64m1, vector_length), vector_length);
288
+ result_ptr += vector_length;
289
+ norms_ptr += vector_length;
290
+ count_columns -= vector_length;
291
+ }
292
+ }
293
+ }
294
+
295
+ NK_PUBLIC void nk_euclideans_packed_f64_rvv( //
296
+ nk_f64_t const *a, void const *b_packed, nk_f64_t *c, //
297
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
298
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
299
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f64_t);
300
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f64_t);
301
+ nk_dots_packed_f64_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
302
+ nk_euclideans_packed_f64_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
303
+ }
304
+
305
+ NK_INTERNAL void nk_angulars_symmetric_f64_rvv_finalize_(nk_f64_t const *vectors, nk_size_t vectors_count,
306
+ nk_size_t depth, nk_size_t stride_elements, nk_f64_t *result,
307
+ nk_size_t result_stride_elements, nk_size_t row_start,
308
+ nk_size_t row_count) {
309
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
310
+ nk_f64_t *result_row = result + row_index * result_stride_elements;
311
+ result_row[row_index] = nk_dots_reduce_sumsq_f64_(vectors + row_index * stride_elements, depth);
312
+ }
313
+ nk_f64_t norms_cache[256];
314
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
315
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
316
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
317
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_f64_(vectors + col * stride_elements, depth);
318
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
319
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
320
+ if (col_start >= chunk_end) continue;
321
+ nk_f64_t *result_row = result + row_index * result_stride_elements;
322
+ nk_f64_t query_norm_sq_f64 = result_row[row_index];
323
+ nk_size_t count_remaining = chunk_end - col_start;
324
+ nk_f64_t *result_ptr = result_row + col_start;
325
+ nk_f64_t const *norms_ptr = norms_cache + (col_start - chunk_start);
326
+ while (count_remaining > 0) {
327
+ size_t vector_length = __riscv_vsetvl_e64m1(count_remaining);
328
+ vfloat64m1_t dots_f64m1 = __riscv_vle64_v_f64m1(result_ptr, vector_length);
329
+ vfloat64m1_t target_norms_sq_f64m1 = __riscv_vle64_v_f64m1(norms_ptr, vector_length);
330
+ vfloat64m1_t norms_product_f64m1 = __riscv_vfmul_vf_f64m1(target_norms_sq_f64m1, query_norm_sq_f64,
331
+ vector_length);
332
+ vfloat64m1_t rsqrt_f64m1 = nk_rsqrt_f64m1_rvv_(norms_product_f64m1, vector_length);
333
+ vfloat64m1_t normalized_dots_f64m1 = __riscv_vfmul_vv_f64m1(dots_f64m1, rsqrt_f64m1, vector_length);
334
+ vfloat64m1_t angular_f64m1 = __riscv_vfrsub_vf_f64m1(normalized_dots_f64m1, 1.0, vector_length);
335
+ angular_f64m1 = __riscv_vfmax_vf_f64m1(angular_f64m1, 0.0, vector_length);
336
+ __riscv_vse64_v_f64m1(result_ptr, angular_f64m1, vector_length);
337
+ result_ptr += vector_length;
338
+ norms_ptr += vector_length;
339
+ count_remaining -= vector_length;
340
+ }
341
+ }
342
+ }
343
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
344
+ result[row_index * result_stride_elements + row_index] = 0;
345
+ }
346
+
347
+ NK_PUBLIC void nk_angulars_symmetric_f64_rvv( //
348
+ nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
349
+ nk_f64_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
350
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_f64_t);
351
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f64_t);
352
+ nk_dots_symmetric_f64_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes, row_start,
353
+ row_count);
354
+ nk_angulars_symmetric_f64_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
355
+ result_stride_elements, row_start, row_count);
356
+ }
357
+
358
+ NK_INTERNAL void nk_euclideans_symmetric_f64_rvv_finalize_(nk_f64_t const *vectors, nk_size_t vectors_count,
359
+ nk_size_t depth, nk_size_t stride_elements, nk_f64_t *result,
360
+ nk_size_t result_stride_elements, nk_size_t row_start,
361
+ nk_size_t row_count) {
362
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
363
+ nk_f64_t *result_row = result + row_index * result_stride_elements;
364
+ result_row[row_index] = nk_dots_reduce_sumsq_f64_(vectors + row_index * stride_elements, depth);
365
+ }
366
+ nk_f64_t norms_cache[256];
367
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
368
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
369
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
370
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_f64_(vectors + col * stride_elements, depth);
371
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
372
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
373
+ if (col_start >= chunk_end) continue;
374
+ nk_f64_t *result_row = result + row_index * result_stride_elements;
375
+ nk_f64_t query_norm_sq_f64 = result_row[row_index];
376
+ nk_size_t count_remaining = chunk_end - col_start;
377
+ nk_f64_t *result_ptr = result_row + col_start;
378
+ nk_f64_t const *norms_ptr = norms_cache + (col_start - chunk_start);
379
+ while (count_remaining > 0) {
380
+ size_t vector_length = __riscv_vsetvl_e64m1(count_remaining);
381
+ vfloat64m1_t dots_f64m1 = __riscv_vle64_v_f64m1(result_ptr, vector_length);
382
+ vfloat64m1_t target_norms_sq_f64m1 = __riscv_vle64_v_f64m1(norms_ptr, vector_length);
383
+ vfloat64m1_t sum_sq_f64m1 = __riscv_vfadd_vf_f64m1(target_norms_sq_f64m1, query_norm_sq_f64,
384
+ vector_length);
385
+ vfloat64m1_t dist_sq_f64m1 = __riscv_vfsub_vv_f64m1(
386
+ sum_sq_f64m1, __riscv_vfmul_vf_f64m1(dots_f64m1, 2.0, vector_length), vector_length);
387
+ dist_sq_f64m1 = __riscv_vfmax_vf_f64m1(dist_sq_f64m1, 0.0, vector_length);
388
+ __riscv_vse64_v_f64m1(result_ptr, __riscv_vfsqrt_v_f64m1(dist_sq_f64m1, vector_length), vector_length);
389
+ result_ptr += vector_length;
390
+ norms_ptr += vector_length;
391
+ count_remaining -= vector_length;
392
+ }
393
+ }
394
+ }
395
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
396
+ result[row_index * result_stride_elements + row_index] = 0;
397
+ }
398
+
399
+ NK_PUBLIC void nk_euclideans_symmetric_f64_rvv( //
400
+ nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
401
+ nk_f64_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
402
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_f64_t);
403
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f64_t);
404
+ nk_dots_symmetric_f64_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes, row_start,
405
+ row_count);
406
+ nk_euclideans_symmetric_f64_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
407
+ result_stride_elements, row_start, row_count);
408
+ }
409
+
410
+ #pragma endregion F64 Floats
411
+
412
+ #pragma region F16 Floats
413
+
414
+ NK_INTERNAL void nk_angulars_packed_f16_rvv_finalize_(nk_f16_t const *a, void const *b_packed, nk_f32_t *c,
415
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
416
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
417
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
418
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
419
+ header->column_count * header->depth_padded_values *
420
+ sizeof(nk_f32_t));
421
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
422
+ nk_f16_t const *a_row = a + row_index * a_stride_elements;
423
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
424
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_f16_(a_row, depth);
425
+ nk_size_t count_columns = columns;
426
+ nk_f32_t *result_ptr = result_row;
427
+ nk_f32_t const *norms_ptr = target_norms;
428
+ while (count_columns > 0) {
429
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
430
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
431
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
432
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
433
+ vector_length);
434
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
435
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
436
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
437
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
438
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
439
+ result_ptr += vector_length;
440
+ norms_ptr += vector_length;
441
+ count_columns -= vector_length;
442
+ }
443
+ }
444
+ }
445
+
446
+ NK_PUBLIC void nk_angulars_packed_f16_rvv( //
447
+ nk_f16_t const *a, void const *b_packed, nk_f32_t *c, //
448
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
449
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
450
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f16_t);
451
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
452
+ nk_dots_packed_f16_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
453
+ nk_angulars_packed_f16_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
454
+ }
455
+
456
+ NK_INTERNAL void nk_euclideans_packed_f16_rvv_finalize_(nk_f16_t const *a, void const *b_packed, nk_f32_t *c,
457
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
458
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
459
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
460
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
461
+ header->column_count * header->depth_padded_values *
462
+ sizeof(nk_f32_t));
463
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
464
+ nk_f16_t const *a_row = a + row_index * a_stride_elements;
465
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
466
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_f16_(a_row, depth);
467
+ nk_size_t count_columns = columns;
468
+ nk_f32_t *result_ptr = result_row;
469
+ nk_f32_t const *norms_ptr = target_norms;
470
+ while (count_columns > 0) {
471
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
472
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
473
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
474
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32, vector_length);
475
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
476
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
477
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
478
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
479
+ result_ptr += vector_length;
480
+ norms_ptr += vector_length;
481
+ count_columns -= vector_length;
482
+ }
483
+ }
484
+ }
485
+
486
+ NK_PUBLIC void nk_euclideans_packed_f16_rvv( //
487
+ nk_f16_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_in_bytes, nk_size_t c_stride_in_bytes) {
490
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f16_t);
491
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
492
+ nk_dots_packed_f16_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
493
+ nk_euclideans_packed_f16_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
494
+ }
495
+
496
+ NK_INTERNAL void nk_angulars_symmetric_f16_rvv_finalize_(nk_f16_t const *vectors, nk_size_t vectors_count,
497
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
498
+ nk_size_t result_stride_elements, nk_size_t row_start,
499
+ nk_size_t row_count) {
500
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
501
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
502
+ result_row[row_index] = nk_dots_reduce_sumsq_f16_(vectors + row_index * stride_elements, depth);
503
+ }
504
+ nk_f32_t norms_cache[256];
505
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
506
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
507
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
508
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_f16_(vectors + col * stride_elements, depth);
509
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
510
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
511
+ if (col_start >= chunk_end) continue;
512
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
513
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
514
+ nk_size_t count_remaining = chunk_end - col_start;
515
+ nk_f32_t *result_ptr = result_row + col_start;
516
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
517
+ while (count_remaining > 0) {
518
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
519
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
520
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
521
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
522
+ vector_length);
523
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
524
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
525
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
526
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
527
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
528
+ result_ptr += vector_length;
529
+ norms_ptr += vector_length;
530
+ count_remaining -= vector_length;
531
+ }
532
+ }
533
+ }
534
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
535
+ result[row_index * result_stride_elements + row_index] = 0;
536
+ }
537
+
538
+ NK_PUBLIC void nk_angulars_symmetric_f16_rvv( //
539
+ nk_f16_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
540
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
541
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_f16_t);
542
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
543
+ nk_dots_symmetric_f16_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes, row_start,
544
+ row_count);
545
+ nk_angulars_symmetric_f16_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
546
+ result_stride_elements, row_start, row_count);
547
+ }
548
+
549
+ NK_INTERNAL void nk_euclideans_symmetric_f16_rvv_finalize_(nk_f16_t const *vectors, nk_size_t vectors_count,
550
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
551
+ nk_size_t result_stride_elements, nk_size_t row_start,
552
+ nk_size_t row_count) {
553
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
554
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
555
+ result_row[row_index] = nk_dots_reduce_sumsq_f16_(vectors + row_index * stride_elements, depth);
556
+ }
557
+ nk_f32_t norms_cache[256];
558
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
559
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
560
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
561
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_f16_(vectors + col * stride_elements, depth);
562
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
563
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
564
+ if (col_start >= chunk_end) continue;
565
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
566
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
567
+ nk_size_t count_remaining = chunk_end - col_start;
568
+ nk_f32_t *result_ptr = result_row + col_start;
569
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
570
+ while (count_remaining > 0) {
571
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
572
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
573
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
574
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
575
+ vector_length);
576
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
577
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
578
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
579
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
580
+ result_ptr += vector_length;
581
+ norms_ptr += vector_length;
582
+ count_remaining -= vector_length;
583
+ }
584
+ }
585
+ }
586
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
587
+ result[row_index * result_stride_elements + row_index] = 0;
588
+ }
589
+
590
+ NK_PUBLIC void nk_euclideans_symmetric_f16_rvv( //
591
+ nk_f16_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
592
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
593
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_f16_t);
594
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
595
+ nk_dots_symmetric_f16_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes, row_start,
596
+ row_count);
597
+ nk_euclideans_symmetric_f16_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
598
+ result_stride_elements, row_start, row_count);
599
+ }
600
+
601
+ #pragma endregion F16 Floats
602
+
603
+ #pragma region BF16 Floats
604
+
605
+ NK_INTERNAL void nk_angulars_packed_bf16_rvv_finalize_(nk_bf16_t const *a, void const *b_packed, nk_f32_t *c,
606
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
607
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
608
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
609
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
610
+ header->column_count * header->depth_padded_values *
611
+ sizeof(nk_f32_t));
612
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
613
+ nk_bf16_t const *a_row = a + row_index * a_stride_elements;
614
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
615
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_bf16_(a_row, depth);
616
+ nk_size_t count_columns = columns;
617
+ nk_f32_t *result_ptr = result_row;
618
+ nk_f32_t const *norms_ptr = target_norms;
619
+ while (count_columns > 0) {
620
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
621
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
622
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
623
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
624
+ vector_length);
625
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
626
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
627
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
628
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
629
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
630
+ result_ptr += vector_length;
631
+ norms_ptr += vector_length;
632
+ count_columns -= vector_length;
633
+ }
634
+ }
635
+ }
636
+
637
+ NK_PUBLIC void nk_angulars_packed_bf16_rvv( //
638
+ nk_bf16_t const *a, void const *b_packed, nk_f32_t *c, //
639
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
640
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
641
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_bf16_t);
642
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
643
+ nk_dots_packed_bf16_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
644
+ nk_angulars_packed_bf16_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
645
+ }
646
+
647
+ NK_INTERNAL void nk_euclideans_packed_bf16_rvv_finalize_(nk_bf16_t const *a, void const *b_packed, nk_f32_t *c,
648
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
649
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
650
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
651
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
652
+ header->column_count * header->depth_padded_values *
653
+ sizeof(nk_f32_t));
654
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
655
+ nk_bf16_t const *a_row = a + row_index * a_stride_elements;
656
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
657
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_bf16_(a_row, depth);
658
+ nk_size_t count_columns = columns;
659
+ nk_f32_t *result_ptr = result_row;
660
+ nk_f32_t const *norms_ptr = target_norms;
661
+ while (count_columns > 0) {
662
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
663
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
664
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
665
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32, vector_length);
666
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
667
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
668
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
669
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
670
+ result_ptr += vector_length;
671
+ norms_ptr += vector_length;
672
+ count_columns -= vector_length;
673
+ }
674
+ }
675
+ }
676
+
677
+ NK_PUBLIC void nk_euclideans_packed_bf16_rvv( //
678
+ nk_bf16_t const *a, void const *b_packed, nk_f32_t *c, //
679
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
680
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
681
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_bf16_t);
682
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
683
+ nk_dots_packed_bf16_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
684
+ nk_euclideans_packed_bf16_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
685
+ }
686
+
687
+ NK_INTERNAL void nk_angulars_symmetric_bf16_rvv_finalize_(nk_bf16_t const *vectors, nk_size_t vectors_count,
688
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
689
+ nk_size_t result_stride_elements, nk_size_t row_start,
690
+ nk_size_t row_count) {
691
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
692
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
693
+ result_row[row_index] = nk_dots_reduce_sumsq_bf16_(vectors + row_index * stride_elements, depth);
694
+ }
695
+ nk_f32_t norms_cache[256];
696
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
697
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
698
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
699
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_bf16_(vectors + col * stride_elements, depth);
700
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
701
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
702
+ if (col_start >= chunk_end) continue;
703
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
704
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
705
+ nk_size_t count_remaining = chunk_end - col_start;
706
+ nk_f32_t *result_ptr = result_row + col_start;
707
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
708
+ while (count_remaining > 0) {
709
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
710
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
711
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
712
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
713
+ vector_length);
714
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
715
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
716
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
717
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
718
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
719
+ result_ptr += vector_length;
720
+ norms_ptr += vector_length;
721
+ count_remaining -= vector_length;
722
+ }
723
+ }
724
+ }
725
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
726
+ result[row_index * result_stride_elements + row_index] = 0;
727
+ }
728
+
729
+ NK_PUBLIC void nk_angulars_symmetric_bf16_rvv( //
730
+ nk_bf16_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
731
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
732
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_bf16_t);
733
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
734
+ nk_dots_symmetric_bf16_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
735
+ row_start, row_count);
736
+ nk_angulars_symmetric_bf16_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
737
+ result_stride_elements, row_start, row_count);
738
+ }
739
+
740
+ NK_INTERNAL void nk_euclideans_symmetric_bf16_rvv_finalize_(nk_bf16_t const *vectors, nk_size_t vectors_count,
741
+ nk_size_t depth, nk_size_t stride_elements,
742
+ nk_f32_t *result, nk_size_t result_stride_elements,
743
+ nk_size_t row_start, nk_size_t row_count) {
744
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
745
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
746
+ result_row[row_index] = nk_dots_reduce_sumsq_bf16_(vectors + row_index * stride_elements, depth);
747
+ }
748
+ nk_f32_t norms_cache[256];
749
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
750
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
751
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
752
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_bf16_(vectors + col * stride_elements, depth);
753
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
754
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
755
+ if (col_start >= chunk_end) continue;
756
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
757
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
758
+ nk_size_t count_remaining = chunk_end - col_start;
759
+ nk_f32_t *result_ptr = result_row + col_start;
760
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
761
+ while (count_remaining > 0) {
762
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
763
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
764
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
765
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
766
+ vector_length);
767
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
768
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
769
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
770
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
771
+ result_ptr += vector_length;
772
+ norms_ptr += vector_length;
773
+ count_remaining -= vector_length;
774
+ }
775
+ }
776
+ }
777
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
778
+ result[row_index * result_stride_elements + row_index] = 0;
779
+ }
780
+
781
+ NK_PUBLIC void nk_euclideans_symmetric_bf16_rvv( //
782
+ nk_bf16_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
783
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
784
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_bf16_t);
785
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
786
+ nk_dots_symmetric_bf16_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
787
+ row_start, row_count);
788
+ nk_euclideans_symmetric_bf16_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
789
+ result_stride_elements, row_start, row_count);
790
+ }
791
+
792
+ #pragma endregion BF16 Floats
793
+
794
+ #pragma region E2M3 Floats
795
+
796
+ NK_INTERNAL void nk_angulars_packed_e2m3_rvv_finalize_(nk_e2m3_t const *a, void const *b_packed, nk_f32_t *c,
797
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
798
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
799
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
800
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
801
+ header->column_count * header->depth_padded_values *
802
+ sizeof(nk_e2m3_t));
803
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
804
+ nk_e2m3_t const *a_row = a + row_index * a_stride_elements;
805
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
806
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_e2m3_(a_row, depth);
807
+ nk_size_t count_columns = columns;
808
+ nk_f32_t *result_ptr = result_row;
809
+ nk_f32_t const *norms_ptr = target_norms;
810
+ while (count_columns > 0) {
811
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
812
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
813
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
814
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
815
+ vector_length);
816
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
817
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
818
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
819
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
820
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
821
+ result_ptr += vector_length;
822
+ norms_ptr += vector_length;
823
+ count_columns -= vector_length;
824
+ }
825
+ }
826
+ }
827
+
828
+ NK_PUBLIC void nk_angulars_packed_e2m3_rvv( //
829
+ nk_e2m3_t const *a, void const *b_packed, nk_f32_t *c, //
830
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
831
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
832
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e2m3_t);
833
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
834
+ nk_dots_packed_e2m3_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
835
+ nk_angulars_packed_e2m3_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
836
+ }
837
+
838
+ NK_INTERNAL void nk_euclideans_packed_e2m3_rvv_finalize_(nk_e2m3_t const *a, void const *b_packed, nk_f32_t *c,
839
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
840
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
841
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
842
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
843
+ header->column_count * header->depth_padded_values *
844
+ sizeof(nk_e2m3_t));
845
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
846
+ nk_e2m3_t const *a_row = a + row_index * a_stride_elements;
847
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
848
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_e2m3_(a_row, depth);
849
+ nk_size_t count_columns = columns;
850
+ nk_f32_t *result_ptr = result_row;
851
+ nk_f32_t const *norms_ptr = target_norms;
852
+ while (count_columns > 0) {
853
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
854
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
855
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
856
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32, vector_length);
857
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
858
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
859
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
860
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
861
+ result_ptr += vector_length;
862
+ norms_ptr += vector_length;
863
+ count_columns -= vector_length;
864
+ }
865
+ }
866
+ }
867
+
868
+ NK_PUBLIC void nk_euclideans_packed_e2m3_rvv( //
869
+ nk_e2m3_t const *a, void const *b_packed, nk_f32_t *c, //
870
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
871
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
872
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e2m3_t);
873
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
874
+ nk_dots_packed_e2m3_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
875
+ nk_euclideans_packed_e2m3_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
876
+ }
877
+
878
+ NK_INTERNAL void nk_angulars_symmetric_e2m3_rvv_finalize_(nk_e2m3_t const *vectors, nk_size_t vectors_count,
879
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
880
+ nk_size_t result_stride_elements, nk_size_t row_start,
881
+ nk_size_t row_count) {
882
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
883
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
884
+ result_row[row_index] = nk_dots_reduce_sumsq_e2m3_(vectors + row_index * stride_elements, depth);
885
+ }
886
+ nk_f32_t norms_cache[256];
887
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
888
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
889
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
890
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e2m3_(vectors + col * stride_elements, depth);
891
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
892
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
893
+ if (col_start >= chunk_end) continue;
894
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
895
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
896
+ nk_size_t count_remaining = chunk_end - col_start;
897
+ nk_f32_t *result_ptr = result_row + col_start;
898
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
899
+ while (count_remaining > 0) {
900
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
901
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
902
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
903
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
904
+ vector_length);
905
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
906
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
907
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
908
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
909
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
910
+ result_ptr += vector_length;
911
+ norms_ptr += vector_length;
912
+ count_remaining -= vector_length;
913
+ }
914
+ }
915
+ }
916
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
917
+ result[row_index * result_stride_elements + row_index] = 0;
918
+ }
919
+
920
+ NK_PUBLIC void nk_angulars_symmetric_e2m3_rvv( //
921
+ nk_e2m3_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
922
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
923
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e2m3_t);
924
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
925
+ nk_dots_symmetric_e2m3_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
926
+ row_start, row_count);
927
+ nk_angulars_symmetric_e2m3_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
928
+ result_stride_elements, row_start, row_count);
929
+ }
930
+
931
+ NK_INTERNAL void nk_euclideans_symmetric_e2m3_rvv_finalize_(nk_e2m3_t const *vectors, nk_size_t vectors_count,
932
+ nk_size_t depth, nk_size_t stride_elements,
933
+ nk_f32_t *result, nk_size_t result_stride_elements,
934
+ nk_size_t row_start, nk_size_t row_count) {
935
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
936
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
937
+ result_row[row_index] = nk_dots_reduce_sumsq_e2m3_(vectors + row_index * stride_elements, depth);
938
+ }
939
+ nk_f32_t norms_cache[256];
940
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
941
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
942
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
943
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e2m3_(vectors + col * stride_elements, depth);
944
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
945
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
946
+ if (col_start >= chunk_end) continue;
947
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
948
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
949
+ nk_size_t count_remaining = chunk_end - col_start;
950
+ nk_f32_t *result_ptr = result_row + col_start;
951
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
952
+ while (count_remaining > 0) {
953
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
954
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
955
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
956
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
957
+ vector_length);
958
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
959
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
960
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
961
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
962
+ result_ptr += vector_length;
963
+ norms_ptr += vector_length;
964
+ count_remaining -= vector_length;
965
+ }
966
+ }
967
+ }
968
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
969
+ result[row_index * result_stride_elements + row_index] = 0;
970
+ }
971
+
972
+ NK_PUBLIC void nk_euclideans_symmetric_e2m3_rvv( //
973
+ nk_e2m3_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
974
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
975
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e2m3_t);
976
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
977
+ nk_dots_symmetric_e2m3_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
978
+ row_start, row_count);
979
+ nk_euclideans_symmetric_e2m3_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
980
+ result_stride_elements, row_start, row_count);
981
+ }
982
+
983
+ #pragma endregion E2M3 Floats
984
+
985
+ #pragma region E3M2 Floats
986
+
987
+ NK_INTERNAL void nk_angulars_packed_e3m2_rvv_finalize_(nk_e3m2_t const *a, void const *b_packed, nk_f32_t *c,
988
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
989
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
990
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
991
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
992
+ header->column_count * header->depth_padded_values *
993
+ sizeof(nk_i16_t));
994
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
995
+ nk_e3m2_t const *a_row = a + row_index * a_stride_elements;
996
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
997
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_e3m2_(a_row, depth);
998
+ nk_size_t count_columns = columns;
999
+ nk_f32_t *result_ptr = result_row;
1000
+ nk_f32_t const *norms_ptr = target_norms;
1001
+ while (count_columns > 0) {
1002
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1003
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1004
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1005
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1006
+ vector_length);
1007
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1008
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1009
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1010
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1011
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1012
+ result_ptr += vector_length;
1013
+ norms_ptr += vector_length;
1014
+ count_columns -= vector_length;
1015
+ }
1016
+ }
1017
+ }
1018
+
1019
+ NK_PUBLIC void nk_angulars_packed_e3m2_rvv( //
1020
+ nk_e3m2_t const *a, void const *b_packed, nk_f32_t *c, //
1021
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1022
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1023
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e3m2_t);
1024
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1025
+ nk_dots_packed_e3m2_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1026
+ nk_angulars_packed_e3m2_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1027
+ }
1028
+
1029
+ NK_INTERNAL void nk_euclideans_packed_e3m2_rvv_finalize_(nk_e3m2_t const *a, void const *b_packed, nk_f32_t *c,
1030
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1031
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1032
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1033
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1034
+ header->column_count * header->depth_padded_values *
1035
+ sizeof(nk_i16_t));
1036
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1037
+ nk_e3m2_t const *a_row = a + row_index * a_stride_elements;
1038
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1039
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_e3m2_(a_row, depth);
1040
+ nk_size_t count_columns = columns;
1041
+ nk_f32_t *result_ptr = result_row;
1042
+ nk_f32_t const *norms_ptr = target_norms;
1043
+ while (count_columns > 0) {
1044
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1045
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1046
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1047
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32, vector_length);
1048
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1049
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1050
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1051
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1052
+ result_ptr += vector_length;
1053
+ norms_ptr += vector_length;
1054
+ count_columns -= vector_length;
1055
+ }
1056
+ }
1057
+ }
1058
+
1059
+ NK_PUBLIC void nk_euclideans_packed_e3m2_rvv( //
1060
+ nk_e3m2_t const *a, void const *b_packed, nk_f32_t *c, //
1061
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1062
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1063
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e3m2_t);
1064
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1065
+ nk_dots_packed_e3m2_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1066
+ nk_euclideans_packed_e3m2_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1067
+ }
1068
+
1069
+ NK_INTERNAL void nk_angulars_symmetric_e3m2_rvv_finalize_(nk_e3m2_t const *vectors, nk_size_t vectors_count,
1070
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
1071
+ nk_size_t result_stride_elements, nk_size_t row_start,
1072
+ nk_size_t row_count) {
1073
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1074
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1075
+ result_row[row_index] = nk_dots_reduce_sumsq_e3m2_(vectors + row_index * stride_elements, depth);
1076
+ }
1077
+ nk_f32_t norms_cache[256];
1078
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1079
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1080
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1081
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e3m2_(vectors + col * stride_elements, depth);
1082
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1083
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1084
+ if (col_start >= chunk_end) continue;
1085
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1086
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
1087
+ nk_size_t count_remaining = chunk_end - col_start;
1088
+ nk_f32_t *result_ptr = result_row + col_start;
1089
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1090
+ while (count_remaining > 0) {
1091
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1092
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1093
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1094
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1095
+ vector_length);
1096
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1097
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1098
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1099
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1100
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1101
+ result_ptr += vector_length;
1102
+ norms_ptr += vector_length;
1103
+ count_remaining -= vector_length;
1104
+ }
1105
+ }
1106
+ }
1107
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1108
+ result[row_index * result_stride_elements + row_index] = 0;
1109
+ }
1110
+
1111
+ NK_PUBLIC void nk_angulars_symmetric_e3m2_rvv( //
1112
+ nk_e3m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1113
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1114
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e3m2_t);
1115
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1116
+ nk_dots_symmetric_e3m2_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
1117
+ row_start, row_count);
1118
+ nk_angulars_symmetric_e3m2_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1119
+ result_stride_elements, row_start, row_count);
1120
+ }
1121
+
1122
+ NK_INTERNAL void nk_euclideans_symmetric_e3m2_rvv_finalize_(nk_e3m2_t const *vectors, nk_size_t vectors_count,
1123
+ nk_size_t depth, nk_size_t stride_elements,
1124
+ nk_f32_t *result, nk_size_t result_stride_elements,
1125
+ nk_size_t row_start, nk_size_t row_count) {
1126
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1127
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1128
+ result_row[row_index] = nk_dots_reduce_sumsq_e3m2_(vectors + row_index * stride_elements, depth);
1129
+ }
1130
+ nk_f32_t norms_cache[256];
1131
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1132
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1133
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1134
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e3m2_(vectors + col * stride_elements, depth);
1135
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1136
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1137
+ if (col_start >= chunk_end) continue;
1138
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1139
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
1140
+ nk_size_t count_remaining = chunk_end - col_start;
1141
+ nk_f32_t *result_ptr = result_row + col_start;
1142
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1143
+ while (count_remaining > 0) {
1144
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1145
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1146
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1147
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1148
+ vector_length);
1149
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1150
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1151
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1152
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1153
+ result_ptr += vector_length;
1154
+ norms_ptr += vector_length;
1155
+ count_remaining -= vector_length;
1156
+ }
1157
+ }
1158
+ }
1159
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1160
+ result[row_index * result_stride_elements + row_index] = 0;
1161
+ }
1162
+
1163
+ NK_PUBLIC void nk_euclideans_symmetric_e3m2_rvv( //
1164
+ nk_e3m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1165
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1166
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e3m2_t);
1167
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1168
+ nk_dots_symmetric_e3m2_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
1169
+ row_start, row_count);
1170
+ nk_euclideans_symmetric_e3m2_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1171
+ result_stride_elements, row_start, row_count);
1172
+ }
1173
+
1174
+ #pragma endregion E3M2 Floats
1175
+
1176
+ #pragma region E4M3 Floats
1177
+
1178
+ NK_INTERNAL void nk_angulars_packed_e4m3_rvv_finalize_(nk_e4m3_t const *a, void const *b_packed, nk_f32_t *c,
1179
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1180
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1181
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1182
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1183
+ header->column_count * header->depth_padded_values *
1184
+ sizeof(nk_f32_t));
1185
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1186
+ nk_e4m3_t const *a_row = a + row_index * a_stride_elements;
1187
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1188
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_e4m3_(a_row, depth);
1189
+ nk_size_t count_columns = columns;
1190
+ nk_f32_t *result_ptr = result_row;
1191
+ nk_f32_t const *norms_ptr = target_norms;
1192
+ while (count_columns > 0) {
1193
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1194
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1195
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1196
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1197
+ vector_length);
1198
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1199
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1200
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1201
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1202
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1203
+ result_ptr += vector_length;
1204
+ norms_ptr += vector_length;
1205
+ count_columns -= vector_length;
1206
+ }
1207
+ }
1208
+ }
1209
+
1210
+ NK_PUBLIC void nk_angulars_packed_e4m3_rvv( //
1211
+ nk_e4m3_t const *a, void const *b_packed, nk_f32_t *c, //
1212
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1213
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1214
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e4m3_t);
1215
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1216
+ nk_dots_packed_e4m3_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1217
+ nk_angulars_packed_e4m3_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1218
+ }
1219
+
1220
+ NK_INTERNAL void nk_euclideans_packed_e4m3_rvv_finalize_(nk_e4m3_t const *a, void const *b_packed, nk_f32_t *c,
1221
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1222
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1223
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1224
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1225
+ header->column_count * header->depth_padded_values *
1226
+ sizeof(nk_f32_t));
1227
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1228
+ nk_e4m3_t const *a_row = a + row_index * a_stride_elements;
1229
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1230
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_e4m3_(a_row, depth);
1231
+ nk_size_t count_columns = columns;
1232
+ nk_f32_t *result_ptr = result_row;
1233
+ nk_f32_t const *norms_ptr = target_norms;
1234
+ while (count_columns > 0) {
1235
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1236
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1237
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1238
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32, vector_length);
1239
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1240
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1241
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1242
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1243
+ result_ptr += vector_length;
1244
+ norms_ptr += vector_length;
1245
+ count_columns -= vector_length;
1246
+ }
1247
+ }
1248
+ }
1249
+
1250
+ NK_PUBLIC void nk_euclideans_packed_e4m3_rvv( //
1251
+ nk_e4m3_t const *a, void const *b_packed, nk_f32_t *c, //
1252
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1253
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1254
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e4m3_t);
1255
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1256
+ nk_dots_packed_e4m3_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1257
+ nk_euclideans_packed_e4m3_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1258
+ }
1259
+
1260
+ NK_INTERNAL void nk_angulars_symmetric_e4m3_rvv_finalize_(nk_e4m3_t const *vectors, nk_size_t vectors_count,
1261
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
1262
+ nk_size_t result_stride_elements, nk_size_t row_start,
1263
+ nk_size_t row_count) {
1264
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1265
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1266
+ result_row[row_index] = nk_dots_reduce_sumsq_e4m3_(vectors + row_index * stride_elements, depth);
1267
+ }
1268
+ nk_f32_t norms_cache[256];
1269
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1270
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1271
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1272
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e4m3_(vectors + col * stride_elements, depth);
1273
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1274
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1275
+ if (col_start >= chunk_end) continue;
1276
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1277
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
1278
+ nk_size_t count_remaining = chunk_end - col_start;
1279
+ nk_f32_t *result_ptr = result_row + col_start;
1280
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1281
+ while (count_remaining > 0) {
1282
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1283
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1284
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1285
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1286
+ vector_length);
1287
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1288
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1289
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1290
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1291
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1292
+ result_ptr += vector_length;
1293
+ norms_ptr += vector_length;
1294
+ count_remaining -= vector_length;
1295
+ }
1296
+ }
1297
+ }
1298
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1299
+ result[row_index * result_stride_elements + row_index] = 0;
1300
+ }
1301
+
1302
+ NK_PUBLIC void nk_angulars_symmetric_e4m3_rvv( //
1303
+ nk_e4m3_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1304
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1305
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e4m3_t);
1306
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1307
+ nk_dots_symmetric_e4m3_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
1308
+ row_start, row_count);
1309
+ nk_angulars_symmetric_e4m3_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1310
+ result_stride_elements, row_start, row_count);
1311
+ }
1312
+
1313
+ NK_INTERNAL void nk_euclideans_symmetric_e4m3_rvv_finalize_(nk_e4m3_t const *vectors, nk_size_t vectors_count,
1314
+ nk_size_t depth, nk_size_t stride_elements,
1315
+ nk_f32_t *result, nk_size_t result_stride_elements,
1316
+ nk_size_t row_start, nk_size_t row_count) {
1317
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1318
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1319
+ result_row[row_index] = nk_dots_reduce_sumsq_e4m3_(vectors + row_index * stride_elements, depth);
1320
+ }
1321
+ nk_f32_t norms_cache[256];
1322
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1323
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1324
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1325
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e4m3_(vectors + col * stride_elements, depth);
1326
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1327
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1328
+ if (col_start >= chunk_end) continue;
1329
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1330
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
1331
+ nk_size_t count_remaining = chunk_end - col_start;
1332
+ nk_f32_t *result_ptr = result_row + col_start;
1333
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1334
+ while (count_remaining > 0) {
1335
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1336
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1337
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1338
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1339
+ vector_length);
1340
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1341
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1342
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1343
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1344
+ result_ptr += vector_length;
1345
+ norms_ptr += vector_length;
1346
+ count_remaining -= vector_length;
1347
+ }
1348
+ }
1349
+ }
1350
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1351
+ result[row_index * result_stride_elements + row_index] = 0;
1352
+ }
1353
+
1354
+ NK_PUBLIC void nk_euclideans_symmetric_e4m3_rvv( //
1355
+ nk_e4m3_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1356
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1357
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e4m3_t);
1358
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1359
+ nk_dots_symmetric_e4m3_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
1360
+ row_start, row_count);
1361
+ nk_euclideans_symmetric_e4m3_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1362
+ result_stride_elements, row_start, row_count);
1363
+ }
1364
+
1365
+ #pragma endregion E4M3 Floats
1366
+
1367
+ #pragma region E5M2 Floats
1368
+
1369
+ NK_INTERNAL void nk_angulars_packed_e5m2_rvv_finalize_(nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c,
1370
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1371
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1372
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1373
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1374
+ header->column_count * header->depth_padded_values *
1375
+ sizeof(nk_f32_t));
1376
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1377
+ nk_e5m2_t const *a_row = a + row_index * a_stride_elements;
1378
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1379
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_e5m2_(a_row, depth);
1380
+ nk_size_t count_columns = columns;
1381
+ nk_f32_t *result_ptr = result_row;
1382
+ nk_f32_t const *norms_ptr = target_norms;
1383
+ while (count_columns > 0) {
1384
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1385
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1386
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1387
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1388
+ vector_length);
1389
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1390
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1391
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1392
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1393
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1394
+ result_ptr += vector_length;
1395
+ norms_ptr += vector_length;
1396
+ count_columns -= vector_length;
1397
+ }
1398
+ }
1399
+ }
1400
+
1401
+ NK_PUBLIC void nk_angulars_packed_e5m2_rvv( //
1402
+ nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c, //
1403
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1404
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1405
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e5m2_t);
1406
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1407
+ nk_dots_packed_e5m2_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1408
+ nk_angulars_packed_e5m2_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1409
+ }
1410
+
1411
+ NK_INTERNAL void nk_euclideans_packed_e5m2_rvv_finalize_(nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c,
1412
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1413
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1414
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1415
+ nk_f32_t const *target_norms = (nk_f32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1416
+ header->column_count * header->depth_padded_values *
1417
+ sizeof(nk_f32_t));
1418
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1419
+ nk_e5m2_t const *a_row = a + row_index * a_stride_elements;
1420
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1421
+ nk_f32_t query_norm_sq_f32 = nk_dots_reduce_sumsq_e5m2_(a_row, depth);
1422
+ nk_size_t count_columns = columns;
1423
+ nk_f32_t *result_ptr = result_row;
1424
+ nk_f32_t const *norms_ptr = target_norms;
1425
+ while (count_columns > 0) {
1426
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1427
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1428
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1429
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32, vector_length);
1430
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1431
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1432
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1433
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1434
+ result_ptr += vector_length;
1435
+ norms_ptr += vector_length;
1436
+ count_columns -= vector_length;
1437
+ }
1438
+ }
1439
+ }
1440
+
1441
+ NK_PUBLIC void nk_euclideans_packed_e5m2_rvv( //
1442
+ nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c, //
1443
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1444
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1445
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_e5m2_t);
1446
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1447
+ nk_dots_packed_e5m2_rvv(a, b_packed, c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1448
+ nk_euclideans_packed_e5m2_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1449
+ }
1450
+
1451
+ NK_INTERNAL void nk_angulars_symmetric_e5m2_rvv_finalize_(nk_e5m2_t const *vectors, nk_size_t vectors_count,
1452
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
1453
+ nk_size_t result_stride_elements, nk_size_t row_start,
1454
+ nk_size_t row_count) {
1455
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1456
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1457
+ result_row[row_index] = nk_dots_reduce_sumsq_e5m2_(vectors + row_index * stride_elements, depth);
1458
+ }
1459
+ nk_f32_t norms_cache[256];
1460
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1461
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1462
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1463
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e5m2_(vectors + col * stride_elements, depth);
1464
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1465
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1466
+ if (col_start >= chunk_end) continue;
1467
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1468
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
1469
+ nk_size_t count_remaining = chunk_end - col_start;
1470
+ nk_f32_t *result_ptr = result_row + col_start;
1471
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1472
+ while (count_remaining > 0) {
1473
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1474
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1475
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1476
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1477
+ vector_length);
1478
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1479
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1480
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1481
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1482
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1483
+ result_ptr += vector_length;
1484
+ norms_ptr += vector_length;
1485
+ count_remaining -= vector_length;
1486
+ }
1487
+ }
1488
+ }
1489
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1490
+ result[row_index * result_stride_elements + row_index] = 0;
1491
+ }
1492
+
1493
+ NK_PUBLIC void nk_angulars_symmetric_e5m2_rvv( //
1494
+ nk_e5m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1495
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1496
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e5m2_t);
1497
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1498
+ nk_dots_symmetric_e5m2_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
1499
+ row_start, row_count);
1500
+ nk_angulars_symmetric_e5m2_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1501
+ result_stride_elements, row_start, row_count);
1502
+ }
1503
+
1504
+ NK_INTERNAL void nk_euclideans_symmetric_e5m2_rvv_finalize_(nk_e5m2_t const *vectors, nk_size_t vectors_count,
1505
+ nk_size_t depth, nk_size_t stride_elements,
1506
+ nk_f32_t *result, nk_size_t result_stride_elements,
1507
+ nk_size_t row_start, nk_size_t row_count) {
1508
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1509
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1510
+ result_row[row_index] = nk_dots_reduce_sumsq_e5m2_(vectors + row_index * stride_elements, depth);
1511
+ }
1512
+ nk_f32_t norms_cache[256];
1513
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1514
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1515
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1516
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_e5m2_(vectors + col * stride_elements, depth);
1517
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1518
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1519
+ if (col_start >= chunk_end) continue;
1520
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1521
+ nk_f32_t query_norm_sq_f32 = result_row[row_index];
1522
+ nk_size_t count_remaining = chunk_end - col_start;
1523
+ nk_f32_t *result_ptr = result_row + col_start;
1524
+ nk_f32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1525
+ while (count_remaining > 0) {
1526
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1527
+ vfloat32m1_t dots_f32m1 = __riscv_vle32_v_f32m1(result_ptr, vector_length);
1528
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vle32_v_f32m1(norms_ptr, vector_length);
1529
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1530
+ vector_length);
1531
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1532
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1533
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1534
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1535
+ result_ptr += vector_length;
1536
+ norms_ptr += vector_length;
1537
+ count_remaining -= vector_length;
1538
+ }
1539
+ }
1540
+ }
1541
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1542
+ result[row_index * result_stride_elements + row_index] = 0;
1543
+ }
1544
+
1545
+ NK_PUBLIC void nk_euclideans_symmetric_e5m2_rvv( //
1546
+ nk_e5m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1547
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1548
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_e5m2_t);
1549
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1550
+ nk_dots_symmetric_e5m2_rvv(vectors, vectors_count, depth, stride_in_bytes, result, result_stride_in_bytes,
1551
+ row_start, row_count);
1552
+ nk_euclideans_symmetric_e5m2_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1553
+ result_stride_elements, row_start, row_count);
1554
+ }
1555
+
1556
+ #pragma endregion E5M2 Floats
1557
+
1558
+ #pragma region I8 Integers
1559
+
1560
+ NK_INTERNAL void nk_angulars_packed_i8_rvv_finalize_(nk_i8_t const *a, void const *b_packed, nk_f32_t *c,
1561
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1562
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1563
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1564
+ nk_u32_t const *target_norms = (nk_u32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1565
+ header->column_count * header->depth_padded_values *
1566
+ sizeof(nk_i8_t));
1567
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1568
+ nk_i8_t const *a_row = a + row_index * a_stride_elements;
1569
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1570
+ nk_u32_t query_norm_sq = nk_dots_reduce_sumsq_i8_(a_row, depth);
1571
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)query_norm_sq;
1572
+ nk_size_t count_columns = columns;
1573
+ nk_f32_t *result_ptr = result_row;
1574
+ nk_u32_t const *norms_ptr = target_norms;
1575
+ while (count_columns > 0) {
1576
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1577
+ vfloat32m1_t dots_f32m1 = __riscv_vfcvt_f_x_v_f32m1(
1578
+ __riscv_vle32_v_i32m1((nk_i32_t const *)result_ptr, vector_length), vector_length);
1579
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1580
+ __riscv_vle32_v_u32m1((nk_u32_t const *)norms_ptr, vector_length), vector_length);
1581
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1582
+ vector_length);
1583
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1584
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1585
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1586
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1587
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1588
+ result_ptr += vector_length;
1589
+ norms_ptr += vector_length;
1590
+ count_columns -= vector_length;
1591
+ }
1592
+ }
1593
+ }
1594
+
1595
+ NK_PUBLIC void nk_angulars_packed_i8_rvv( //
1596
+ nk_i8_t const *a, void const *b_packed, nk_f32_t *c, //
1597
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1598
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1599
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_i8_t);
1600
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1601
+ nk_dots_packed_i8_rvv(a, b_packed, (nk_i32_t *)c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1602
+ nk_angulars_packed_i8_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1603
+ }
1604
+
1605
+ NK_INTERNAL void nk_euclideans_packed_i8_rvv_finalize_(nk_i8_t const *a, void const *b_packed, nk_f32_t *c,
1606
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1607
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1608
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1609
+ nk_u32_t const *target_norms = (nk_u32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1610
+ header->column_count * header->depth_padded_values *
1611
+ sizeof(nk_i8_t));
1612
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1613
+ nk_i8_t const *a_row = a + row_index * a_stride_elements;
1614
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1615
+ nk_u32_t query_norm_sq = nk_dots_reduce_sumsq_i8_(a_row, depth);
1616
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)query_norm_sq;
1617
+ nk_size_t count_columns = columns;
1618
+ nk_f32_t *result_ptr = result_row;
1619
+ nk_u32_t const *norms_ptr = target_norms;
1620
+ while (count_columns > 0) {
1621
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1622
+ vfloat32m1_t dots_f32m1 = __riscv_vfcvt_f_x_v_f32m1(
1623
+ __riscv_vle32_v_i32m1((nk_i32_t const *)result_ptr, vector_length), vector_length);
1624
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1625
+ __riscv_vle32_v_u32m1((nk_u32_t const *)norms_ptr, vector_length), vector_length);
1626
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32, vector_length);
1627
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1628
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1629
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1630
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1631
+ result_ptr += vector_length;
1632
+ norms_ptr += vector_length;
1633
+ count_columns -= vector_length;
1634
+ }
1635
+ }
1636
+ }
1637
+
1638
+ NK_PUBLIC void nk_euclideans_packed_i8_rvv( //
1639
+ nk_i8_t const *a, void const *b_packed, nk_f32_t *c, //
1640
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1641
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1642
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_i8_t);
1643
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1644
+ nk_dots_packed_i8_rvv(a, b_packed, (nk_i32_t *)c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1645
+ nk_euclideans_packed_i8_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1646
+ }
1647
+
1648
+ NK_INTERNAL void nk_angulars_symmetric_i8_rvv_finalize_(nk_i8_t const *vectors, nk_size_t vectors_count,
1649
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
1650
+ nk_size_t result_stride_elements, nk_size_t row_start,
1651
+ nk_size_t row_count) {
1652
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1653
+ nk_u32_t norm = nk_dots_reduce_sumsq_i8_(vectors + row_index * stride_elements, depth);
1654
+ ((nk_u32_t *)(result + row_index * result_stride_elements))[row_index] = norm;
1655
+ }
1656
+ nk_u32_t norms_cache[256];
1657
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1658
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1659
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1660
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_i8_(vectors + col * stride_elements, depth);
1661
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1662
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1663
+ if (col_start >= chunk_end) continue;
1664
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1665
+ nk_u32_t query_norm_sq = ((nk_u32_t *)result_row)[row_index];
1666
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)query_norm_sq;
1667
+ nk_size_t count_remaining = chunk_end - col_start;
1668
+ nk_f32_t *result_ptr = result_row + col_start;
1669
+ nk_u32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1670
+ while (count_remaining > 0) {
1671
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1672
+ vfloat32m1_t dots_f32m1 = __riscv_vfcvt_f_x_v_f32m1(
1673
+ __riscv_vle32_v_i32m1((nk_i32_t const *)result_ptr, vector_length), vector_length);
1674
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1675
+ __riscv_vle32_v_u32m1((nk_u32_t const *)norms_ptr, vector_length), vector_length);
1676
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1677
+ vector_length);
1678
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1679
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1680
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1681
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1682
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1683
+ result_ptr += vector_length;
1684
+ norms_ptr += vector_length;
1685
+ count_remaining -= vector_length;
1686
+ }
1687
+ }
1688
+ }
1689
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1690
+ result[row_index * result_stride_elements + row_index] = 0;
1691
+ }
1692
+
1693
+ NK_PUBLIC void nk_angulars_symmetric_i8_rvv( //
1694
+ nk_i8_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1695
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1696
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_i8_t);
1697
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1698
+ nk_dots_symmetric_i8_rvv(vectors, vectors_count, depth, stride_in_bytes, (nk_i32_t *)result, result_stride_in_bytes,
1699
+ row_start, row_count);
1700
+ nk_angulars_symmetric_i8_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1701
+ result_stride_elements, row_start, row_count);
1702
+ }
1703
+
1704
+ NK_INTERNAL void nk_euclideans_symmetric_i8_rvv_finalize_(nk_i8_t const *vectors, nk_size_t vectors_count,
1705
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
1706
+ nk_size_t result_stride_elements, nk_size_t row_start,
1707
+ nk_size_t row_count) {
1708
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1709
+ nk_u32_t norm = nk_dots_reduce_sumsq_i8_(vectors + row_index * stride_elements, depth);
1710
+ ((nk_u32_t *)(result + row_index * result_stride_elements))[row_index] = norm;
1711
+ }
1712
+ nk_u32_t norms_cache[256];
1713
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1714
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1715
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1716
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_i8_(vectors + col * stride_elements, depth);
1717
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1718
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1719
+ if (col_start >= chunk_end) continue;
1720
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1721
+ nk_u32_t query_norm_sq = ((nk_u32_t *)result_row)[row_index];
1722
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)query_norm_sq;
1723
+ nk_size_t count_remaining = chunk_end - col_start;
1724
+ nk_f32_t *result_ptr = result_row + col_start;
1725
+ nk_u32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1726
+ while (count_remaining > 0) {
1727
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1728
+ vfloat32m1_t dots_f32m1 = __riscv_vfcvt_f_x_v_f32m1(
1729
+ __riscv_vle32_v_i32m1((nk_i32_t const *)result_ptr, vector_length), vector_length);
1730
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1731
+ __riscv_vle32_v_u32m1((nk_u32_t const *)norms_ptr, vector_length), vector_length);
1732
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1733
+ vector_length);
1734
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1735
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1736
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1737
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1738
+ result_ptr += vector_length;
1739
+ norms_ptr += vector_length;
1740
+ count_remaining -= vector_length;
1741
+ }
1742
+ }
1743
+ }
1744
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1745
+ result[row_index * result_stride_elements + row_index] = 0;
1746
+ }
1747
+
1748
+ NK_PUBLIC void nk_euclideans_symmetric_i8_rvv( //
1749
+ nk_i8_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1750
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1751
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_i8_t);
1752
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1753
+ nk_dots_symmetric_i8_rvv(vectors, vectors_count, depth, stride_in_bytes, (nk_i32_t *)result, result_stride_in_bytes,
1754
+ row_start, row_count);
1755
+ nk_euclideans_symmetric_i8_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1756
+ result_stride_elements, row_start, row_count);
1757
+ }
1758
+
1759
+ #pragma endregion I8 Integers
1760
+
1761
+ #pragma region U8 Integers
1762
+
1763
+ NK_INTERNAL void nk_angulars_packed_u8_rvv_finalize_(nk_u8_t const *a, void const *b_packed, nk_f32_t *c,
1764
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1765
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1766
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1767
+ nk_u32_t const *target_norms = (nk_u32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1768
+ header->column_count * header->depth_padded_values *
1769
+ sizeof(nk_u8_t));
1770
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1771
+ nk_u8_t const *a_row = a + row_index * a_stride_elements;
1772
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1773
+ nk_u32_t query_norm_sq = nk_dots_reduce_sumsq_u8_(a_row, depth);
1774
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)query_norm_sq;
1775
+ nk_size_t count_columns = columns;
1776
+ nk_f32_t *result_ptr = result_row;
1777
+ nk_u32_t const *norms_ptr = target_norms;
1778
+ while (count_columns > 0) {
1779
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1780
+ vfloat32m1_t dots_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1781
+ __riscv_vle32_v_u32m1((nk_u32_t const *)result_ptr, vector_length), vector_length);
1782
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1783
+ __riscv_vle32_v_u32m1((nk_u32_t const *)norms_ptr, vector_length), vector_length);
1784
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1785
+ vector_length);
1786
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1787
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1788
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1789
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1790
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1791
+ result_ptr += vector_length;
1792
+ norms_ptr += vector_length;
1793
+ count_columns -= vector_length;
1794
+ }
1795
+ }
1796
+ }
1797
+
1798
+ NK_PUBLIC void nk_angulars_packed_u8_rvv( //
1799
+ nk_u8_t const *a, void const *b_packed, nk_f32_t *c, //
1800
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1801
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1802
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_u8_t);
1803
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1804
+ nk_dots_packed_u8_rvv(a, b_packed, (nk_u32_t *)c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1805
+ nk_angulars_packed_u8_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1806
+ }
1807
+
1808
+ NK_INTERNAL void nk_euclideans_packed_u8_rvv_finalize_(nk_u8_t const *a, void const *b_packed, nk_f32_t *c,
1809
+ nk_size_t rows, nk_size_t columns, nk_size_t depth,
1810
+ nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
1811
+ nk_cross_packed_buffer_header_t const *header = (nk_cross_packed_buffer_header_t const *)b_packed;
1812
+ nk_u32_t const *target_norms = (nk_u32_t const *)((char const *)b_packed + sizeof(nk_cross_packed_buffer_header_t) +
1813
+ header->column_count * header->depth_padded_values *
1814
+ sizeof(nk_u8_t));
1815
+ for (nk_size_t row_index = 0; row_index < rows; row_index++) {
1816
+ nk_u8_t const *a_row = a + row_index * a_stride_elements;
1817
+ nk_f32_t *result_row = c + row_index * c_stride_elements;
1818
+ nk_u32_t query_norm_sq = nk_dots_reduce_sumsq_u8_(a_row, depth);
1819
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)query_norm_sq;
1820
+ nk_size_t count_columns = columns;
1821
+ nk_f32_t *result_ptr = result_row;
1822
+ nk_u32_t const *norms_ptr = target_norms;
1823
+ while (count_columns > 0) {
1824
+ size_t vector_length = __riscv_vsetvl_e32m1(count_columns);
1825
+ vfloat32m1_t dots_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1826
+ __riscv_vle32_v_u32m1((nk_u32_t const *)result_ptr, vector_length), vector_length);
1827
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1828
+ __riscv_vle32_v_u32m1((nk_u32_t const *)norms_ptr, vector_length), vector_length);
1829
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32, vector_length);
1830
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1831
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1832
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1833
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1834
+ result_ptr += vector_length;
1835
+ norms_ptr += vector_length;
1836
+ count_columns -= vector_length;
1837
+ }
1838
+ }
1839
+ }
1840
+
1841
+ NK_PUBLIC void nk_euclideans_packed_u8_rvv( //
1842
+ nk_u8_t const *a, void const *b_packed, nk_f32_t *c, //
1843
+ nk_size_t rows, nk_size_t columns, nk_size_t depth, //
1844
+ nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
1845
+ nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_u8_t);
1846
+ nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f32_t);
1847
+ nk_dots_packed_u8_rvv(a, b_packed, (nk_u32_t *)c, rows, columns, depth, a_stride_in_bytes, c_stride_in_bytes);
1848
+ nk_euclideans_packed_u8_rvv_finalize_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
1849
+ }
1850
+
1851
+ NK_INTERNAL void nk_angulars_symmetric_u8_rvv_finalize_(nk_u8_t const *vectors, nk_size_t vectors_count,
1852
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
1853
+ nk_size_t result_stride_elements, nk_size_t row_start,
1854
+ nk_size_t row_count) {
1855
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1856
+ nk_u32_t norm = nk_dots_reduce_sumsq_u8_(vectors + row_index * stride_elements, depth);
1857
+ ((nk_u32_t *)(result + row_index * result_stride_elements))[row_index] = norm;
1858
+ }
1859
+ nk_u32_t norms_cache[256];
1860
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1861
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1862
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1863
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_u8_(vectors + col * stride_elements, depth);
1864
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1865
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1866
+ if (col_start >= chunk_end) continue;
1867
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1868
+ nk_u32_t query_norm_sq = ((nk_u32_t *)result_row)[row_index];
1869
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)query_norm_sq;
1870
+ nk_size_t count_remaining = chunk_end - col_start;
1871
+ nk_f32_t *result_ptr = result_row + col_start;
1872
+ nk_u32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1873
+ while (count_remaining > 0) {
1874
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1875
+ vfloat32m1_t dots_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1876
+ __riscv_vle32_v_u32m1((nk_u32_t const *)result_ptr, vector_length), vector_length);
1877
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1878
+ __riscv_vle32_v_u32m1((nk_u32_t const *)norms_ptr, vector_length), vector_length);
1879
+ vfloat32m1_t norms_product_f32m1 = __riscv_vfmul_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1880
+ vector_length);
1881
+ vfloat32m1_t rsqrt_f32m1 = nk_rsqrt_f32m1_rvv_(norms_product_f32m1, vector_length);
1882
+ vfloat32m1_t normalized_dots_f32m1 = __riscv_vfmul_vv_f32m1(dots_f32m1, rsqrt_f32m1, vector_length);
1883
+ vfloat32m1_t angular_f32m1 = __riscv_vfrsub_vf_f32m1(normalized_dots_f32m1, 1.0f, vector_length);
1884
+ angular_f32m1 = __riscv_vfmax_vf_f32m1(angular_f32m1, 0.0f, vector_length);
1885
+ __riscv_vse32_v_f32m1(result_ptr, angular_f32m1, vector_length);
1886
+ result_ptr += vector_length;
1887
+ norms_ptr += vector_length;
1888
+ count_remaining -= vector_length;
1889
+ }
1890
+ }
1891
+ }
1892
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1893
+ result[row_index * result_stride_elements + row_index] = 0;
1894
+ }
1895
+
1896
+ NK_PUBLIC void nk_angulars_symmetric_u8_rvv( //
1897
+ nk_u8_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1898
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1899
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_u8_t);
1900
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1901
+ nk_dots_symmetric_u8_rvv(vectors, vectors_count, depth, stride_in_bytes, (nk_u32_t *)result, result_stride_in_bytes,
1902
+ row_start, row_count);
1903
+ nk_angulars_symmetric_u8_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1904
+ result_stride_elements, row_start, row_count);
1905
+ }
1906
+
1907
+ NK_INTERNAL void nk_euclideans_symmetric_u8_rvv_finalize_(nk_u8_t const *vectors, nk_size_t vectors_count,
1908
+ nk_size_t depth, nk_size_t stride_elements, nk_f32_t *result,
1909
+ nk_size_t result_stride_elements, nk_size_t row_start,
1910
+ nk_size_t row_count) {
1911
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1912
+ nk_u32_t norm = nk_dots_reduce_sumsq_u8_(vectors + row_index * stride_elements, depth);
1913
+ ((nk_u32_t *)(result + row_index * result_stride_elements))[row_index] = norm;
1914
+ }
1915
+ nk_u32_t norms_cache[256];
1916
+ for (nk_size_t chunk_start = 0; chunk_start < vectors_count; chunk_start += 256) {
1917
+ nk_size_t chunk_end = chunk_start + 256 < vectors_count ? chunk_start + 256 : vectors_count;
1918
+ for (nk_size_t col = chunk_start; col < chunk_end; ++col)
1919
+ norms_cache[col - chunk_start] = nk_dots_reduce_sumsq_u8_(vectors + col * stride_elements, depth);
1920
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
1921
+ nk_size_t col_start = row_index + 1 > chunk_start ? row_index + 1 : chunk_start;
1922
+ if (col_start >= chunk_end) continue;
1923
+ nk_f32_t *result_row = result + row_index * result_stride_elements;
1924
+ nk_u32_t query_norm_sq = ((nk_u32_t *)result_row)[row_index];
1925
+ nk_f32_t query_norm_sq_f32 = (nk_f32_t)query_norm_sq;
1926
+ nk_size_t count_remaining = chunk_end - col_start;
1927
+ nk_f32_t *result_ptr = result_row + col_start;
1928
+ nk_u32_t const *norms_ptr = norms_cache + (col_start - chunk_start);
1929
+ while (count_remaining > 0) {
1930
+ size_t vector_length = __riscv_vsetvl_e32m1(count_remaining);
1931
+ vfloat32m1_t dots_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1932
+ __riscv_vle32_v_u32m1((nk_u32_t const *)result_ptr, vector_length), vector_length);
1933
+ vfloat32m1_t target_norms_sq_f32m1 = __riscv_vfcvt_f_xu_v_f32m1(
1934
+ __riscv_vle32_v_u32m1((nk_u32_t const *)norms_ptr, vector_length), vector_length);
1935
+ vfloat32m1_t sum_sq_f32m1 = __riscv_vfadd_vf_f32m1(target_norms_sq_f32m1, query_norm_sq_f32,
1936
+ vector_length);
1937
+ vfloat32m1_t dist_sq_f32m1 = __riscv_vfsub_vv_f32m1(
1938
+ sum_sq_f32m1, __riscv_vfmul_vf_f32m1(dots_f32m1, 2.0f, vector_length), vector_length);
1939
+ dist_sq_f32m1 = __riscv_vfmax_vf_f32m1(dist_sq_f32m1, 0.0f, vector_length);
1940
+ __riscv_vse32_v_f32m1(result_ptr, __riscv_vfsqrt_v_f32m1(dist_sq_f32m1, vector_length), vector_length);
1941
+ result_ptr += vector_length;
1942
+ norms_ptr += vector_length;
1943
+ count_remaining -= vector_length;
1944
+ }
1945
+ }
1946
+ }
1947
+ for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index)
1948
+ result[row_index * result_stride_elements + row_index] = 0;
1949
+ }
1950
+
1951
+ NK_PUBLIC void nk_euclideans_symmetric_u8_rvv( //
1952
+ nk_u8_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
1953
+ nk_f32_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
1954
+ nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_u8_t);
1955
+ nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);
1956
+ nk_dots_symmetric_u8_rvv(vectors, vectors_count, depth, stride_in_bytes, (nk_u32_t *)result, result_stride_in_bytes,
1957
+ row_start, row_count);
1958
+ nk_euclideans_symmetric_u8_rvv_finalize_(vectors, vectors_count, depth, stride_elements, result,
1959
+ result_stride_elements, row_start, row_count);
1960
+ }
1961
+
1962
+ #pragma endregion U8 Integers
1963
+
1964
+ #if defined(__clang__)
1965
+ #pragma clang attribute pop
1966
+ #elif defined(__GNUC__)
1967
+ #pragma GCC pop_options
1968
+ #endif
1969
+
1970
+ #if defined(__cplusplus)
1971
+ } // extern "C"
1972
+ #endif
1973
+
1974
+ #endif // NK_TARGET_RVV
1975
+ #endif // NK_TARGET_RISCV64_
1976
+ #endif // NK_SPATIALS_RVV_H