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