usearch 2.21.4 → 2.25.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (602) hide show
  1. package/README.md +24 -24
  2. package/binding.gyp +15 -37
  3. package/c/CMakeLists.txt +4 -5
  4. package/c/README.md +7 -7
  5. package/c/lib.cpp +32 -14
  6. package/c/test.c +83 -1
  7. package/c/usearch.h +24 -3
  8. package/cpp/CMakeLists.txt +4 -0
  9. package/cpp/README.md +5 -5
  10. package/cpp/bench.cpp +97 -116
  11. package/cpp/test.cpp +117 -20
  12. package/include/usearch/index.hpp +224 -48
  13. package/include/usearch/index_dense.hpp +105 -40
  14. package/include/usearch/index_plugins.hpp +1387 -204
  15. package/javascript/dist/cjs/usearch.d.ts +8 -3
  16. package/javascript/dist/cjs/usearch.js +12 -4
  17. package/javascript/dist/esm/usearch.d.ts +8 -3
  18. package/javascript/dist/esm/usearch.js +12 -4
  19. package/javascript/lib.cpp +17 -0
  20. package/javascript/usearch.ts +14 -6
  21. package/numkong/.clang-tidy +30 -0
  22. package/{simsimd → numkong}/.cmake-format.py +18 -0
  23. package/numkong/.editorconfig +12 -0
  24. package/numkong/.git-blame-ignore-revs +123 -0
  25. package/numkong/.mailmap +15 -0
  26. package/numkong/.spi.yml +11 -0
  27. package/numkong/.swift-format +14 -0
  28. package/numkong/CITATION.cff +81 -0
  29. package/numkong/CONTRIBUTING.md +500 -0
  30. package/numkong/Cargo.toml +54 -0
  31. package/numkong/MANIFEST.in +8 -0
  32. package/numkong/Package.swift +84 -0
  33. package/numkong/README.md +569 -0
  34. package/numkong/bench/README.md +217 -0
  35. package/numkong/bench/bench-browser.html +323 -0
  36. package/numkong/bench/bench.cpp +344 -0
  37. package/numkong/bench/bench.hpp +816 -0
  38. package/numkong/bench/bench.js +148 -0
  39. package/numkong/bench/bench.mjs +616 -0
  40. package/numkong/bench/bench_attention.cpp +113 -0
  41. package/numkong/bench/bench_cast.cpp +264 -0
  42. package/numkong/bench/bench_cross_amx.cpp +62 -0
  43. package/numkong/bench/bench_cross_arm.cpp +245 -0
  44. package/numkong/bench/bench_cross_blas.cpp +309 -0
  45. package/numkong/bench/bench_cross_loongarch.cpp +106 -0
  46. package/numkong/bench/bench_cross_power.cpp +84 -0
  47. package/numkong/bench/bench_cross_rvv.cpp +121 -0
  48. package/numkong/bench/bench_cross_serial.cpp +151 -0
  49. package/numkong/bench/bench_cross_sme.cpp +163 -0
  50. package/numkong/bench/bench_cross_wasm.cpp +142 -0
  51. package/numkong/bench/bench_cross_x86.cpp +455 -0
  52. package/numkong/bench/bench_curved.cpp +210 -0
  53. package/numkong/bench/bench_dot.cpp +315 -0
  54. package/numkong/bench/bench_each.cpp +503 -0
  55. package/numkong/bench/bench_geospatial.cpp +112 -0
  56. package/numkong/bench/bench_maxsim.cpp +154 -0
  57. package/numkong/bench/bench_mesh.cpp +148 -0
  58. package/numkong/bench/bench_probability.cpp +60 -0
  59. package/numkong/bench/bench_reduce.cpp +371 -0
  60. package/numkong/bench/bench_set.cpp +73 -0
  61. package/numkong/bench/bench_sparse.cpp +198 -0
  62. package/numkong/bench/bench_spatial.cpp +347 -0
  63. package/numkong/bench/bench_trigonometry.cpp +137 -0
  64. package/numkong/bench/swift/Bench.swift +506 -0
  65. package/numkong/binding.gyp +169 -0
  66. package/numkong/build.rs +484 -0
  67. package/numkong/c/dispatch.h +512 -0
  68. package/numkong/c/dispatch_bf16.c +447 -0
  69. package/numkong/c/dispatch_bf16c.c +52 -0
  70. package/numkong/c/dispatch_e2m3.c +296 -0
  71. package/numkong/c/dispatch_e3m2.c +284 -0
  72. package/numkong/c/dispatch_e4m3.c +318 -0
  73. package/numkong/c/dispatch_e5m2.c +316 -0
  74. package/numkong/c/dispatch_f16.c +452 -0
  75. package/numkong/c/dispatch_f16c.c +58 -0
  76. package/numkong/c/dispatch_f32.c +434 -0
  77. package/numkong/c/dispatch_f32c.c +99 -0
  78. package/numkong/c/dispatch_f64.c +348 -0
  79. package/numkong/c/dispatch_f64c.c +98 -0
  80. package/numkong/c/dispatch_i16.c +96 -0
  81. package/numkong/c/dispatch_i32.c +89 -0
  82. package/numkong/c/dispatch_i4.c +153 -0
  83. package/numkong/c/dispatch_i64.c +86 -0
  84. package/numkong/c/dispatch_i8.c +348 -0
  85. package/numkong/c/dispatch_other.c +348 -0
  86. package/numkong/c/dispatch_u1.c +193 -0
  87. package/numkong/c/dispatch_u16.c +124 -0
  88. package/numkong/c/dispatch_u32.c +118 -0
  89. package/numkong/c/dispatch_u4.c +153 -0
  90. package/numkong/c/dispatch_u64.c +102 -0
  91. package/numkong/c/dispatch_u8.c +364 -0
  92. package/numkong/c/numkong.c +940 -0
  93. package/numkong/cmake/nk_arm_isa_probes.cmake +61 -0
  94. package/numkong/cmake/nk_isa_probe.cmake +90 -0
  95. package/numkong/cmake/nk_loongarch_isa_probes.cmake +15 -0
  96. package/numkong/cmake/nk_math_lib.cmake +46 -0
  97. package/numkong/cmake/nk_power_isa_probes.cmake +15 -0
  98. package/numkong/cmake/nk_riscv_isa_probes.cmake +19 -0
  99. package/numkong/cmake/nk_x86_isa_probes.cmake +35 -0
  100. package/numkong/cmake/toolchain-aarch64-gnu.cmake +86 -0
  101. package/numkong/cmake/toolchain-android-arm64.cmake +25 -0
  102. package/numkong/cmake/toolchain-android-armv7.cmake +25 -0
  103. package/numkong/cmake/toolchain-loongarch64-gnu.cmake +68 -0
  104. package/numkong/cmake/toolchain-ppc64le-gnu.cmake +74 -0
  105. package/numkong/cmake/toolchain-riscv64-gnu.cmake +107 -0
  106. package/numkong/cmake/toolchain-riscv64-llvm.cmake +114 -0
  107. package/numkong/cmake/toolchain-wasi.cmake +77 -0
  108. package/numkong/cmake/toolchain-wasm.cmake +70 -0
  109. package/numkong/cmake/toolchain-wasm64.cmake +62 -0
  110. package/numkong/cmake/toolchain-x86_64-llvm.cmake +47 -0
  111. package/numkong/include/README.md +643 -0
  112. package/numkong/include/module.modulemap +129 -0
  113. package/numkong/include/numkong/attention/sapphireamx.h +1363 -0
  114. package/numkong/include/numkong/attention/sme.h +2075 -0
  115. package/numkong/include/numkong/attention.h +49 -0
  116. package/numkong/include/numkong/capabilities.h +930 -0
  117. package/numkong/include/numkong/cast/README.md +262 -0
  118. package/numkong/include/numkong/cast/diamond.h +64 -0
  119. package/numkong/include/numkong/cast/haswell.h +823 -0
  120. package/numkong/include/numkong/cast/icelake.h +475 -0
  121. package/numkong/include/numkong/cast/loongsonasx.h +252 -0
  122. package/numkong/include/numkong/cast/neon.h +1174 -0
  123. package/numkong/include/numkong/cast/powervsx.h +449 -0
  124. package/numkong/include/numkong/cast/rvv.h +970 -0
  125. package/numkong/include/numkong/cast/sapphire.h +262 -0
  126. package/numkong/include/numkong/cast/serial.h +2336 -0
  127. package/numkong/include/numkong/cast/skylake.h +915 -0
  128. package/numkong/include/numkong/cast/v128relaxed.h +537 -0
  129. package/numkong/include/numkong/cast.h +254 -0
  130. package/numkong/include/numkong/cast.hpp +44 -0
  131. package/numkong/include/numkong/curved/README.md +223 -0
  132. package/numkong/include/numkong/curved/genoa.h +182 -0
  133. package/numkong/include/numkong/curved/haswell.h +276 -0
  134. package/numkong/include/numkong/curved/neon.h +329 -0
  135. package/numkong/include/numkong/curved/neonbfdot.h +211 -0
  136. package/numkong/include/numkong/curved/rvv.h +305 -0
  137. package/numkong/include/numkong/curved/serial.h +207 -0
  138. package/numkong/include/numkong/curved/skylake.h +457 -0
  139. package/numkong/include/numkong/curved/smef64.h +530 -0
  140. package/numkong/include/numkong/curved.h +513 -0
  141. package/numkong/include/numkong/curved.hpp +144 -0
  142. package/numkong/include/numkong/dot/README.md +442 -0
  143. package/numkong/include/numkong/dot/alder.h +563 -0
  144. package/numkong/include/numkong/dot/diamond.h +158 -0
  145. package/numkong/include/numkong/dot/genoa.h +289 -0
  146. package/numkong/include/numkong/dot/haswell.h +1723 -0
  147. package/numkong/include/numkong/dot/icelake.h +997 -0
  148. package/numkong/include/numkong/dot/loongsonasx.h +671 -0
  149. package/numkong/include/numkong/dot/neon.h +869 -0
  150. package/numkong/include/numkong/dot/neonbfdot.h +243 -0
  151. package/numkong/include/numkong/dot/neonfhm.h +358 -0
  152. package/numkong/include/numkong/dot/neonfp8.h +323 -0
  153. package/numkong/include/numkong/dot/neonsdot.h +622 -0
  154. package/numkong/include/numkong/dot/powervsx.h +752 -0
  155. package/numkong/include/numkong/dot/rvv.h +722 -0
  156. package/numkong/include/numkong/dot/rvvbb.h +72 -0
  157. package/numkong/include/numkong/dot/rvvbf16.h +123 -0
  158. package/numkong/include/numkong/dot/rvvhalf.h +129 -0
  159. package/numkong/include/numkong/dot/sapphire.h +141 -0
  160. package/numkong/include/numkong/dot/serial.h +874 -0
  161. package/numkong/include/numkong/dot/sierra.h +405 -0
  162. package/numkong/include/numkong/dot/skylake.h +1116 -0
  163. package/numkong/include/numkong/dot/sve.h +420 -0
  164. package/numkong/include/numkong/dot/svebfdot.h +75 -0
  165. package/numkong/include/numkong/dot/svehalf.h +168 -0
  166. package/numkong/include/numkong/dot/svesdot.h +90 -0
  167. package/numkong/include/numkong/dot/v128relaxed.h +1293 -0
  168. package/numkong/include/numkong/dot.h +1136 -0
  169. package/numkong/include/numkong/dot.hpp +94 -0
  170. package/numkong/include/numkong/dots/README.md +496 -0
  171. package/numkong/include/numkong/dots/alder.h +117 -0
  172. package/numkong/include/numkong/dots/diamond.h +86 -0
  173. package/numkong/include/numkong/dots/genoa.h +100 -0
  174. package/numkong/include/numkong/dots/graniteamx.h +733 -0
  175. package/numkong/include/numkong/dots/haswell.h +310 -0
  176. package/numkong/include/numkong/dots/icelake.h +180 -0
  177. package/numkong/include/numkong/dots/loongsonasx.h +176 -0
  178. package/numkong/include/numkong/dots/neon.h +123 -0
  179. package/numkong/include/numkong/dots/neonbfdot.h +59 -0
  180. package/numkong/include/numkong/dots/neonfhm.h +96 -0
  181. package/numkong/include/numkong/dots/neonfp8.h +99 -0
  182. package/numkong/include/numkong/dots/neonsdot.h +144 -0
  183. package/numkong/include/numkong/dots/powervsx.h +194 -0
  184. package/numkong/include/numkong/dots/rvv.h +2593 -0
  185. package/numkong/include/numkong/dots/sapphireamx.h +4017 -0
  186. package/numkong/include/numkong/dots/serial.h +2867 -0
  187. package/numkong/include/numkong/dots/sierra.h +100 -0
  188. package/numkong/include/numkong/dots/skylake.h +207 -0
  189. package/numkong/include/numkong/dots/sme.h +5041 -0
  190. package/numkong/include/numkong/dots/smebi32.h +477 -0
  191. package/numkong/include/numkong/dots/smef64.h +1328 -0
  192. package/numkong/include/numkong/dots/smehalf.h +47 -0
  193. package/numkong/include/numkong/dots/v128relaxed.h +317 -0
  194. package/numkong/include/numkong/dots.h +3110 -0
  195. package/numkong/include/numkong/dots.hpp +641 -0
  196. package/numkong/include/numkong/each/README.md +465 -0
  197. package/numkong/include/numkong/each/haswell.h +1658 -0
  198. package/numkong/include/numkong/each/icelake.h +272 -0
  199. package/numkong/include/numkong/each/neon.h +1138 -0
  200. package/numkong/include/numkong/each/neonbfdot.h +211 -0
  201. package/numkong/include/numkong/each/neonhalf.h +318 -0
  202. package/numkong/include/numkong/each/rvv.h +1121 -0
  203. package/numkong/include/numkong/each/sapphire.h +343 -0
  204. package/numkong/include/numkong/each/serial.h +282 -0
  205. package/numkong/include/numkong/each/skylake.h +1562 -0
  206. package/numkong/include/numkong/each/v128relaxed.h +562 -0
  207. package/numkong/include/numkong/each.h +2232 -0
  208. package/numkong/include/numkong/each.hpp +434 -0
  209. package/numkong/include/numkong/geospatial/README.md +147 -0
  210. package/numkong/include/numkong/geospatial/haswell.h +633 -0
  211. package/numkong/include/numkong/geospatial/neon.h +615 -0
  212. package/numkong/include/numkong/geospatial/rvv.h +701 -0
  213. package/numkong/include/numkong/geospatial/serial.h +309 -0
  214. package/numkong/include/numkong/geospatial/skylake.h +613 -0
  215. package/numkong/include/numkong/geospatial/v128relaxed.h +652 -0
  216. package/numkong/include/numkong/geospatial.h +453 -0
  217. package/numkong/include/numkong/geospatial.hpp +235 -0
  218. package/numkong/include/numkong/matrix.hpp +336 -0
  219. package/numkong/include/numkong/maxsim/README.md +187 -0
  220. package/numkong/include/numkong/maxsim/alder.h +511 -0
  221. package/numkong/include/numkong/maxsim/genoa.h +115 -0
  222. package/numkong/include/numkong/maxsim/haswell.h +553 -0
  223. package/numkong/include/numkong/maxsim/icelake.h +480 -0
  224. package/numkong/include/numkong/maxsim/neonsdot.h +394 -0
  225. package/numkong/include/numkong/maxsim/sapphireamx.h +877 -0
  226. package/numkong/include/numkong/maxsim/serial.h +490 -0
  227. package/numkong/include/numkong/maxsim/sme.h +993 -0
  228. package/numkong/include/numkong/maxsim/v128relaxed.h +280 -0
  229. package/numkong/include/numkong/maxsim.h +571 -0
  230. package/numkong/include/numkong/maxsim.hpp +133 -0
  231. package/numkong/include/numkong/mesh/README.md +219 -0
  232. package/numkong/include/numkong/mesh/haswell.h +2013 -0
  233. package/numkong/include/numkong/mesh/neon.h +1821 -0
  234. package/numkong/include/numkong/mesh/neonbfdot.h +771 -0
  235. package/numkong/include/numkong/mesh/rvv.h +936 -0
  236. package/numkong/include/numkong/mesh/serial.h +721 -0
  237. package/numkong/include/numkong/mesh/skylake.h +2344 -0
  238. package/numkong/include/numkong/mesh/v128relaxed.h +1081 -0
  239. package/numkong/include/numkong/mesh.h +666 -0
  240. package/numkong/include/numkong/mesh.hpp +701 -0
  241. package/numkong/include/numkong/numkong.h +78 -0
  242. package/numkong/include/numkong/numkong.hpp +58 -0
  243. package/numkong/include/numkong/probability/README.md +177 -0
  244. package/numkong/include/numkong/probability/haswell.h +267 -0
  245. package/numkong/include/numkong/probability/neon.h +255 -0
  246. package/numkong/include/numkong/probability/rvv.h +414 -0
  247. package/numkong/include/numkong/probability/serial.h +181 -0
  248. package/numkong/include/numkong/probability/skylake.h +324 -0
  249. package/numkong/include/numkong/probability.h +380 -0
  250. package/numkong/include/numkong/probability.hpp +120 -0
  251. package/numkong/include/numkong/random.h +50 -0
  252. package/numkong/include/numkong/random.hpp +285 -0
  253. package/numkong/include/numkong/reduce/README.md +552 -0
  254. package/numkong/include/numkong/reduce/alder.h +636 -0
  255. package/numkong/include/numkong/reduce/genoa.h +201 -0
  256. package/numkong/include/numkong/reduce/haswell.h +3785 -0
  257. package/numkong/include/numkong/reduce/icelake.h +549 -0
  258. package/numkong/include/numkong/reduce/neon.h +3980 -0
  259. package/numkong/include/numkong/reduce/neonbfdot.h +142 -0
  260. package/numkong/include/numkong/reduce/neonfhm.h +293 -0
  261. package/numkong/include/numkong/reduce/neonsdot.h +357 -0
  262. package/numkong/include/numkong/reduce/rvv.h +3533 -0
  263. package/numkong/include/numkong/reduce/serial.h +753 -0
  264. package/numkong/include/numkong/reduce/sierra.h +338 -0
  265. package/numkong/include/numkong/reduce/skylake.h +3794 -0
  266. package/numkong/include/numkong/reduce/sve.h +52 -0
  267. package/numkong/include/numkong/reduce/v128relaxed.h +2302 -0
  268. package/numkong/include/numkong/reduce.h +1582 -0
  269. package/numkong/include/numkong/reduce.hpp +778 -0
  270. package/numkong/include/numkong/scalar/README.md +89 -0
  271. package/numkong/include/numkong/scalar/haswell.h +122 -0
  272. package/numkong/include/numkong/scalar/loongsonasx.h +74 -0
  273. package/numkong/include/numkong/scalar/neon.h +122 -0
  274. package/numkong/include/numkong/scalar/neonhalf.h +70 -0
  275. package/numkong/include/numkong/scalar/powervsx.h +96 -0
  276. package/numkong/include/numkong/scalar/rvv.h +211 -0
  277. package/numkong/include/numkong/scalar/sapphire.h +74 -0
  278. package/numkong/include/numkong/scalar/serial.h +332 -0
  279. package/numkong/include/numkong/scalar/v128relaxed.h +56 -0
  280. package/numkong/include/numkong/scalar.h +696 -0
  281. package/numkong/include/numkong/set/README.md +179 -0
  282. package/numkong/include/numkong/set/haswell.h +334 -0
  283. package/numkong/include/numkong/set/icelake.h +485 -0
  284. package/numkong/include/numkong/set/loongsonasx.h +181 -0
  285. package/numkong/include/numkong/set/neon.h +363 -0
  286. package/numkong/include/numkong/set/powervsx.h +326 -0
  287. package/numkong/include/numkong/set/rvv.h +226 -0
  288. package/numkong/include/numkong/set/rvvbb.h +117 -0
  289. package/numkong/include/numkong/set/serial.h +174 -0
  290. package/numkong/include/numkong/set/sve.h +187 -0
  291. package/numkong/include/numkong/set/v128relaxed.h +240 -0
  292. package/numkong/include/numkong/set.h +471 -0
  293. package/numkong/include/numkong/set.hpp +114 -0
  294. package/numkong/include/numkong/sets/README.md +149 -0
  295. package/numkong/include/numkong/sets/haswell.h +63 -0
  296. package/numkong/include/numkong/sets/icelake.h +66 -0
  297. package/numkong/include/numkong/sets/loongsonasx.h +52 -0
  298. package/numkong/include/numkong/sets/neon.h +61 -0
  299. package/numkong/include/numkong/sets/powervsx.h +65 -0
  300. package/numkong/include/numkong/sets/serial.h +43 -0
  301. package/numkong/include/numkong/sets/smebi32.h +1135 -0
  302. package/numkong/include/numkong/sets/v128relaxed.h +58 -0
  303. package/numkong/include/numkong/sets.h +382 -0
  304. package/numkong/include/numkong/sparse/README.md +156 -0
  305. package/numkong/include/numkong/sparse/icelake.h +463 -0
  306. package/numkong/include/numkong/sparse/neon.h +288 -0
  307. package/numkong/include/numkong/sparse/serial.h +117 -0
  308. package/numkong/include/numkong/sparse/sve2.h +504 -0
  309. package/numkong/include/numkong/sparse/turin.h +322 -0
  310. package/numkong/include/numkong/sparse.h +363 -0
  311. package/numkong/include/numkong/sparse.hpp +113 -0
  312. package/numkong/include/numkong/spatial/README.md +439 -0
  313. package/numkong/include/numkong/spatial/alder.h +604 -0
  314. package/numkong/include/numkong/spatial/diamond.h +240 -0
  315. package/numkong/include/numkong/spatial/genoa.h +222 -0
  316. package/numkong/include/numkong/spatial/haswell.h +979 -0
  317. package/numkong/include/numkong/spatial/icelake.h +1067 -0
  318. package/numkong/include/numkong/spatial/loongsonasx.h +483 -0
  319. package/numkong/include/numkong/spatial/neon.h +846 -0
  320. package/numkong/include/numkong/spatial/neonbfdot.h +164 -0
  321. package/numkong/include/numkong/spatial/neonfp8.h +258 -0
  322. package/numkong/include/numkong/spatial/neonsdot.h +429 -0
  323. package/numkong/include/numkong/spatial/powervsx.h +738 -0
  324. package/numkong/include/numkong/spatial/rvv.h +991 -0
  325. package/numkong/include/numkong/spatial/rvvbf16.h +128 -0
  326. package/numkong/include/numkong/spatial/rvvhalf.h +120 -0
  327. package/numkong/include/numkong/spatial/serial.h +347 -0
  328. package/numkong/include/numkong/spatial/sierra.h +516 -0
  329. package/numkong/include/numkong/spatial/skylake.h +605 -0
  330. package/numkong/include/numkong/spatial/sve.h +246 -0
  331. package/numkong/include/numkong/spatial/svebfdot.h +125 -0
  332. package/numkong/include/numkong/spatial/svehalf.h +136 -0
  333. package/numkong/include/numkong/spatial/svesdot.h +143 -0
  334. package/numkong/include/numkong/spatial/v128relaxed.h +969 -0
  335. package/numkong/include/numkong/spatial.h +1681 -0
  336. package/numkong/include/numkong/spatial.hpp +183 -0
  337. package/numkong/include/numkong/spatials/README.md +580 -0
  338. package/numkong/include/numkong/spatials/alder.h +94 -0
  339. package/numkong/include/numkong/spatials/diamond.h +82 -0
  340. package/numkong/include/numkong/spatials/genoa.h +94 -0
  341. package/numkong/include/numkong/spatials/graniteamx.h +173 -0
  342. package/numkong/include/numkong/spatials/haswell.h +219 -0
  343. package/numkong/include/numkong/spatials/icelake.h +113 -0
  344. package/numkong/include/numkong/spatials/loongsonasx.h +153 -0
  345. package/numkong/include/numkong/spatials/neon.h +109 -0
  346. package/numkong/include/numkong/spatials/neonbfdot.h +60 -0
  347. package/numkong/include/numkong/spatials/neonfhm.h +92 -0
  348. package/numkong/include/numkong/spatials/neonfp8.h +111 -0
  349. package/numkong/include/numkong/spatials/neonsdot.h +143 -0
  350. package/numkong/include/numkong/spatials/powervsx.h +153 -0
  351. package/numkong/include/numkong/spatials/rvv.h +1976 -0
  352. package/numkong/include/numkong/spatials/sapphireamx.h +1149 -0
  353. package/numkong/include/numkong/spatials/serial.h +248 -0
  354. package/numkong/include/numkong/spatials/sierra.h +96 -0
  355. package/numkong/include/numkong/spatials/skylake.h +184 -0
  356. package/numkong/include/numkong/spatials/sme.h +1927 -0
  357. package/numkong/include/numkong/spatials/smef64.h +483 -0
  358. package/numkong/include/numkong/spatials/v128relaxed.h +240 -0
  359. package/numkong/include/numkong/spatials.h +3374 -0
  360. package/numkong/include/numkong/spatials.hpp +509 -0
  361. package/numkong/include/numkong/tensor.hpp +1981 -0
  362. package/numkong/include/numkong/trigonometry/README.md +184 -0
  363. package/numkong/include/numkong/trigonometry/haswell.h +653 -0
  364. package/numkong/include/numkong/trigonometry/neon.h +638 -0
  365. package/numkong/include/numkong/trigonometry/rvv.h +700 -0
  366. package/numkong/include/numkong/trigonometry/serial.h +703 -0
  367. package/numkong/include/numkong/trigonometry/skylake.h +725 -0
  368. package/numkong/include/numkong/trigonometry/v128relaxed.h +667 -0
  369. package/numkong/include/numkong/trigonometry.h +467 -0
  370. package/numkong/include/numkong/trigonometry.hpp +166 -0
  371. package/numkong/include/numkong/types.h +1695 -0
  372. package/numkong/include/numkong/types.hpp +6027 -0
  373. package/numkong/include/numkong/vector.hpp +766 -0
  374. package/numkong/javascript/@numkong-darwin-arm64/package.json +26 -0
  375. package/numkong/javascript/@numkong-darwin-x64/package.json +26 -0
  376. package/numkong/javascript/@numkong-linux-arm64/package.json +25 -0
  377. package/numkong/javascript/@numkong-linux-x64/package.json +25 -0
  378. package/numkong/javascript/@numkong-win32-arm64/package.json +25 -0
  379. package/numkong/javascript/@numkong-win32-x64/package.json +25 -0
  380. package/numkong/javascript/README.md +254 -0
  381. package/numkong/javascript/node-gyp-build.d.ts +4 -0
  382. package/numkong/javascript/numkong-browser.ts +40 -0
  383. package/numkong/javascript/numkong-wasm.ts +750 -0
  384. package/numkong/javascript/numkong.c +712 -0
  385. package/numkong/javascript/numkong.ts +613 -0
  386. package/{simsimd → numkong}/javascript/tsconfig-base.json +6 -3
  387. package/numkong/javascript/tsconfig-cjs.json +13 -0
  388. package/{simsimd → numkong}/javascript/tsconfig-esm.json +1 -1
  389. package/numkong/javascript/types.ts +784 -0
  390. package/numkong/numkong.gypi +104 -0
  391. package/numkong/package.json +108 -0
  392. package/numkong/probes/arm_neon.c +8 -0
  393. package/numkong/probes/arm_neon_bfdot.c +9 -0
  394. package/numkong/probes/arm_neon_fhm.c +9 -0
  395. package/numkong/probes/arm_neon_half.c +8 -0
  396. package/numkong/probes/arm_neon_sdot.c +9 -0
  397. package/numkong/probes/arm_neonfp8.c +9 -0
  398. package/numkong/probes/arm_sme.c +16 -0
  399. package/numkong/probes/arm_sme2.c +16 -0
  400. package/numkong/probes/arm_sme2p1.c +16 -0
  401. package/numkong/probes/arm_sme_bf16.c +16 -0
  402. package/numkong/probes/arm_sme_bi32.c +16 -0
  403. package/numkong/probes/arm_sme_f64.c +16 -0
  404. package/numkong/probes/arm_sme_fa64.c +14 -0
  405. package/numkong/probes/arm_sme_half.c +16 -0
  406. package/numkong/probes/arm_sme_lut2.c +15 -0
  407. package/numkong/probes/arm_sve.c +18 -0
  408. package/numkong/probes/arm_sve2.c +20 -0
  409. package/numkong/probes/arm_sve2p1.c +18 -0
  410. package/numkong/probes/arm_sve_bfdot.c +20 -0
  411. package/numkong/probes/arm_sve_half.c +18 -0
  412. package/numkong/probes/arm_sve_sdot.c +21 -0
  413. package/numkong/probes/loongarch_lasx.c +12 -0
  414. package/numkong/probes/power_vsx.c +12 -0
  415. package/numkong/probes/probe.js +127 -0
  416. package/numkong/probes/riscv_rvv.c +14 -0
  417. package/numkong/probes/riscv_rvv_bb.c +15 -0
  418. package/numkong/probes/riscv_rvv_bf16.c +17 -0
  419. package/numkong/probes/riscv_rvv_half.c +14 -0
  420. package/numkong/probes/wasm_v128relaxed.c +11 -0
  421. package/numkong/probes/x86_alder.c +17 -0
  422. package/numkong/probes/x86_diamond.c +17 -0
  423. package/numkong/probes/x86_genoa.c +17 -0
  424. package/numkong/probes/x86_graniteamx.c +19 -0
  425. package/numkong/probes/x86_haswell.c +11 -0
  426. package/numkong/probes/x86_icelake.c +17 -0
  427. package/numkong/probes/x86_sapphire.c +16 -0
  428. package/numkong/probes/x86_sapphireamx.c +18 -0
  429. package/numkong/probes/x86_sierra.c +17 -0
  430. package/numkong/probes/x86_skylake.c +15 -0
  431. package/numkong/probes/x86_turin.c +17 -0
  432. package/numkong/rust/README.md +698 -0
  433. package/numkong/rust/capabilities.rs +96 -0
  434. package/numkong/rust/cast.rs +314 -0
  435. package/numkong/rust/curved.rs +391 -0
  436. package/numkong/rust/each.rs +2721 -0
  437. package/numkong/rust/geospatial.rs +343 -0
  438. package/numkong/rust/matrix.rs +5424 -0
  439. package/numkong/rust/maxsim.rs +540 -0
  440. package/numkong/rust/mesh.rs +837 -0
  441. package/numkong/rust/numkong.rs +409 -0
  442. package/numkong/rust/probability.rs +272 -0
  443. package/numkong/rust/reduce.rs +1249 -0
  444. package/numkong/rust/set.rs +194 -0
  445. package/numkong/rust/sparse.rs +513 -0
  446. package/numkong/rust/spatial.rs +1815 -0
  447. package/numkong/rust/tensor.rs +7565 -0
  448. package/numkong/rust/types.rs +3666 -0
  449. package/numkong/rust/vector.rs +1899 -0
  450. package/numkong/rustfmt.toml +2 -0
  451. package/numkong/swift/Matrix.swift +1388 -0
  452. package/numkong/swift/MaxSim.swift +151 -0
  453. package/numkong/swift/NumKong.swift +283 -0
  454. package/numkong/swift/Numerics.swift +329 -0
  455. package/numkong/swift/README.md +433 -0
  456. package/numkong/swift/Scalar.swift +741 -0
  457. package/numkong/swift/Tensor.swift +209 -0
  458. package/numkong/test/README.md +315 -0
  459. package/numkong/test/playwright.config.ts +16 -0
  460. package/numkong/test/swift/Test.swift +534 -0
  461. package/numkong/test/test-browser-bundle-runner.mjs +36 -0
  462. package/numkong/test/test-browser-bundle.html +188 -0
  463. package/numkong/test/test-browser-runner.mjs +43 -0
  464. package/numkong/test/test-browser.html +305 -0
  465. package/numkong/test/test-dtypes.mjs +453 -0
  466. package/numkong/test/test-wasm.mjs +267 -0
  467. package/numkong/test/test.cpp +442 -0
  468. package/numkong/test/test.hpp +580 -0
  469. package/numkong/test/test.mjs +363 -0
  470. package/numkong/test/test_base.py +1041 -0
  471. package/numkong/test/test_cast.cpp +242 -0
  472. package/numkong/test/test_cross.hpp +551 -0
  473. package/numkong/test/test_cross.py +467 -0
  474. package/numkong/test/test_cross_amx.cpp +50 -0
  475. package/numkong/test/test_cross_arm.cpp +225 -0
  476. package/numkong/test/test_cross_blas.cpp +354 -0
  477. package/numkong/test/test_cross_loongarch.cpp +90 -0
  478. package/numkong/test/test_cross_power.cpp +73 -0
  479. package/numkong/test/test_cross_rvv.cpp +109 -0
  480. package/numkong/test/test_cross_serial.cpp +181 -0
  481. package/numkong/test/test_cross_sme.cpp +148 -0
  482. package/numkong/test/test_cross_wasm.cpp +141 -0
  483. package/numkong/test/test_cross_x86.cpp +442 -0
  484. package/numkong/test/test_curved.cpp +182 -0
  485. package/numkong/test/test_curved.py +207 -0
  486. package/numkong/test/test_dot.cpp +353 -0
  487. package/numkong/test/test_dot.py +296 -0
  488. package/numkong/test/test_each.cpp +437 -0
  489. package/numkong/test/test_each.py +756 -0
  490. package/numkong/test/test_geospatial.cpp +155 -0
  491. package/numkong/test/test_geospatial.py +241 -0
  492. package/numkong/test/test_maxsim.cpp +138 -0
  493. package/numkong/test/test_maxsim.py +209 -0
  494. package/numkong/test/test_mesh.cpp +203 -0
  495. package/numkong/test/test_mesh.py +394 -0
  496. package/numkong/test/test_probability.cpp +136 -0
  497. package/numkong/test/test_probability.py +166 -0
  498. package/numkong/test/test_reduce.cpp +563 -0
  499. package/numkong/test/test_reduce.py +781 -0
  500. package/numkong/test/test_set.cpp +151 -0
  501. package/numkong/test/test_set.py +120 -0
  502. package/numkong/test/test_similarities.py +579 -0
  503. package/numkong/test/test_sparse.cpp +121 -0
  504. package/numkong/test/test_sparse.py +118 -0
  505. package/numkong/test/test_spatial.cpp +518 -0
  506. package/numkong/test/test_spatial.py +269 -0
  507. package/numkong/test/test_tensor.cpp +884 -0
  508. package/numkong/test/test_tensor.py +1417 -0
  509. package/numkong/test/test_trigonometry.cpp +156 -0
  510. package/numkong/test/test_trigonometry.py +180 -0
  511. package/package.json +7 -8
  512. package/prebuilds/darwin-arm64+x64/usearch.node +0 -0
  513. package/prebuilds/linux-arm64/usearch.node +0 -0
  514. package/prebuilds/linux-x64/usearch.node +0 -0
  515. package/fp16/CMakeLists.txt +0 -176
  516. package/fp16/LICENSE +0 -11
  517. package/fp16/README.md +0 -20
  518. package/fp16/bench/alt-element.cc +0 -89
  519. package/fp16/bench/from-alt-array.cc +0 -175
  520. package/fp16/bench/from-ieee-array.cc +0 -399
  521. package/fp16/bench/ieee-element.cc +0 -272
  522. package/fp16/bench/to-alt-array.cc +0 -91
  523. package/fp16/bench/to-ieee-array.cc +0 -303
  524. package/fp16/cmake/DownloadGoogleBenchmark.cmake +0 -15
  525. package/fp16/cmake/DownloadGoogleTest.cmake +0 -15
  526. package/fp16/cmake/DownloadPSimd.cmake +0 -15
  527. package/fp16/configure.py +0 -65
  528. package/fp16/confu.yaml +0 -8
  529. package/fp16/include/fp16/__init__.py +0 -0
  530. package/fp16/include/fp16/avx.py +0 -50
  531. package/fp16/include/fp16/avx2.py +0 -53
  532. package/fp16/include/fp16/bitcasts.h +0 -92
  533. package/fp16/include/fp16/fp16.h +0 -451
  534. package/fp16/include/fp16/psimd.h +0 -131
  535. package/fp16/include/fp16.h +0 -11
  536. package/fp16/jni/Android.mk +0 -6
  537. package/fp16/jni/Application.mk +0 -4
  538. package/fp16/test/alt-from-fp32-value.cc +0 -498
  539. package/fp16/test/alt-to-fp32-bits.cc +0 -263
  540. package/fp16/test/alt-to-fp32-psimd.cc +0 -145
  541. package/fp16/test/alt-to-fp32-value.cc +0 -348
  542. package/fp16/test/alt-to-fp32x2-psimd.cc +0 -245
  543. package/fp16/test/bitcasts.cc +0 -56
  544. package/fp16/test/ieee-from-fp32-value.cc +0 -523
  545. package/fp16/test/ieee-to-fp32-bits.cc +0 -333
  546. package/fp16/test/ieee-to-fp32-psimd.cc +0 -312
  547. package/fp16/test/ieee-to-fp32-value.cc +0 -420
  548. package/fp16/test/ieee-to-fp32x2-psimd.cc +0 -541
  549. package/fp16/test/peachpy/alt-xmm-to-fp32-xmm-avx.cc +0 -96
  550. package/fp16/test/peachpy/alt-xmm-to-fp32-ymm-avx2.cc +0 -96
  551. package/fp16/test/peachpy/stubs.py +0 -38
  552. package/fp16/test/tables.cc +0 -530
  553. package/fp16/test/tables.h +0 -13
  554. package/fp16/third-party/THHalf.h +0 -92
  555. package/fp16/third-party/eigen-half.h +0 -170
  556. package/fp16/third-party/float16-compressor.h +0 -80
  557. package/fp16/third-party/half.hpp +0 -2916
  558. package/fp16/third-party/npy-halffloat.h +0 -328
  559. package/simsimd/.git-blame-ignore-revs +0 -2
  560. package/simsimd/CONTRIBUTING.md +0 -321
  561. package/simsimd/Cargo.toml +0 -49
  562. package/simsimd/MANIFEST.in +0 -4
  563. package/simsimd/Package.swift +0 -28
  564. package/simsimd/README.md +0 -1493
  565. package/simsimd/binding.gyp +0 -35
  566. package/simsimd/build.rs +0 -70
  567. package/simsimd/c/lib.c +0 -374
  568. package/simsimd/include/module.modulemap +0 -4
  569. package/simsimd/include/simsimd/binary.h +0 -486
  570. package/simsimd/include/simsimd/curved.h +0 -1541
  571. package/simsimd/include/simsimd/dot.h +0 -1853
  572. package/simsimd/include/simsimd/elementwise.h +0 -2465
  573. package/simsimd/include/simsimd/geospatial.h +0 -43
  574. package/simsimd/include/simsimd/mesh.h +0 -69
  575. package/simsimd/include/simsimd/probability.h +0 -606
  576. package/simsimd/include/simsimd/simsimd.h +0 -2567
  577. package/simsimd/include/simsimd/sparse.h +0 -1384
  578. package/simsimd/include/simsimd/spatial.h +0 -2335
  579. package/simsimd/include/simsimd/types.h +0 -666
  580. package/simsimd/javascript/fallback.ts +0 -233
  581. package/simsimd/javascript/lib.c +0 -125
  582. package/simsimd/javascript/node-gyp-build.d.ts +0 -1
  583. package/simsimd/javascript/simsimd.ts +0 -179
  584. package/simsimd/javascript/tsconfig-cjs.json +0 -8
  585. package/simsimd/package-lock.json +0 -1761
  586. package/simsimd/package.json +0 -61
  587. package/simsimd/rust/lib.rs +0 -1821
  588. package/simsimd/scripts/bench.cxx +0 -1170
  589. package/simsimd/scripts/bench.ipynb +0 -187
  590. package/simsimd/scripts/bench.js +0 -127
  591. package/simsimd/scripts/bench.rs +0 -876
  592. package/simsimd/scripts/bench_vectors.py +0 -734
  593. package/simsimd/scripts/bench_vectors_live.py +0 -213
  594. package/simsimd/scripts/test.c +0 -187
  595. package/simsimd/scripts/test.mjs +0 -220
  596. package/simsimd/scripts/test.py +0 -1505
  597. package/simsimd/swift/SimSIMD.swift +0 -149
  598. package/simsimd/swift/Test.swift +0 -128
  599. /package/{simsimd → numkong}/CODE_OF_CONDUCT.md +0 -0
  600. /package/{simsimd → numkong}/LICENSE +0 -0
  601. /package/{simsimd → numkong}/javascript/dist-package-cjs.json +0 -0
  602. /package/{simsimd → numkong}/javascript/dist-package-esm.json +0 -0
@@ -0,0 +1,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