usearch 2.21.3 → 2.25.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (602) hide show
  1. package/README.md +24 -24
  2. package/binding.gyp +15 -37
  3. package/c/CMakeLists.txt +4 -5
  4. package/c/README.md +7 -7
  5. package/c/lib.cpp +32 -14
  6. package/c/test.c +83 -1
  7. package/c/usearch.h +24 -3
  8. package/cpp/CMakeLists.txt +4 -0
  9. package/cpp/README.md +5 -5
  10. package/cpp/bench.cpp +97 -116
  11. package/cpp/test.cpp +117 -20
  12. package/include/usearch/index.hpp +235 -52
  13. package/include/usearch/index_dense.hpp +105 -40
  14. package/include/usearch/index_plugins.hpp +1387 -204
  15. package/javascript/dist/cjs/usearch.d.ts +8 -3
  16. package/javascript/dist/cjs/usearch.js +12 -4
  17. package/javascript/dist/esm/usearch.d.ts +8 -3
  18. package/javascript/dist/esm/usearch.js +12 -4
  19. package/javascript/lib.cpp +17 -0
  20. package/javascript/usearch.ts +14 -6
  21. package/numkong/.clang-tidy +30 -0
  22. package/{simsimd → numkong}/.cmake-format.py +18 -0
  23. package/numkong/.editorconfig +12 -0
  24. package/numkong/.git-blame-ignore-revs +123 -0
  25. package/numkong/.mailmap +15 -0
  26. package/numkong/.spi.yml +11 -0
  27. package/numkong/.swift-format +14 -0
  28. package/numkong/CITATION.cff +81 -0
  29. package/numkong/CONTRIBUTING.md +500 -0
  30. package/numkong/Cargo.toml +54 -0
  31. package/numkong/MANIFEST.in +8 -0
  32. package/numkong/Package.swift +84 -0
  33. package/numkong/README.md +569 -0
  34. package/numkong/bench/README.md +217 -0
  35. package/numkong/bench/bench-browser.html +323 -0
  36. package/numkong/bench/bench.cpp +344 -0
  37. package/numkong/bench/bench.hpp +816 -0
  38. package/numkong/bench/bench.js +148 -0
  39. package/numkong/bench/bench.mjs +616 -0
  40. package/numkong/bench/bench_attention.cpp +113 -0
  41. package/numkong/bench/bench_cast.cpp +264 -0
  42. package/numkong/bench/bench_cross_amx.cpp +62 -0
  43. package/numkong/bench/bench_cross_arm.cpp +245 -0
  44. package/numkong/bench/bench_cross_blas.cpp +309 -0
  45. package/numkong/bench/bench_cross_loongarch.cpp +106 -0
  46. package/numkong/bench/bench_cross_power.cpp +84 -0
  47. package/numkong/bench/bench_cross_rvv.cpp +121 -0
  48. package/numkong/bench/bench_cross_serial.cpp +151 -0
  49. package/numkong/bench/bench_cross_sme.cpp +163 -0
  50. package/numkong/bench/bench_cross_wasm.cpp +142 -0
  51. package/numkong/bench/bench_cross_x86.cpp +455 -0
  52. package/numkong/bench/bench_curved.cpp +210 -0
  53. package/numkong/bench/bench_dot.cpp +315 -0
  54. package/numkong/bench/bench_each.cpp +503 -0
  55. package/numkong/bench/bench_geospatial.cpp +112 -0
  56. package/numkong/bench/bench_maxsim.cpp +154 -0
  57. package/numkong/bench/bench_mesh.cpp +148 -0
  58. package/numkong/bench/bench_probability.cpp +60 -0
  59. package/numkong/bench/bench_reduce.cpp +371 -0
  60. package/numkong/bench/bench_set.cpp +73 -0
  61. package/numkong/bench/bench_sparse.cpp +198 -0
  62. package/numkong/bench/bench_spatial.cpp +347 -0
  63. package/numkong/bench/bench_trigonometry.cpp +137 -0
  64. package/numkong/bench/swift/Bench.swift +506 -0
  65. package/numkong/binding.gyp +169 -0
  66. package/numkong/build.rs +484 -0
  67. package/numkong/c/dispatch.h +512 -0
  68. package/numkong/c/dispatch_bf16.c +447 -0
  69. package/numkong/c/dispatch_bf16c.c +52 -0
  70. package/numkong/c/dispatch_e2m3.c +296 -0
  71. package/numkong/c/dispatch_e3m2.c +284 -0
  72. package/numkong/c/dispatch_e4m3.c +318 -0
  73. package/numkong/c/dispatch_e5m2.c +316 -0
  74. package/numkong/c/dispatch_f16.c +452 -0
  75. package/numkong/c/dispatch_f16c.c +58 -0
  76. package/numkong/c/dispatch_f32.c +434 -0
  77. package/numkong/c/dispatch_f32c.c +99 -0
  78. package/numkong/c/dispatch_f64.c +348 -0
  79. package/numkong/c/dispatch_f64c.c +98 -0
  80. package/numkong/c/dispatch_i16.c +96 -0
  81. package/numkong/c/dispatch_i32.c +89 -0
  82. package/numkong/c/dispatch_i4.c +153 -0
  83. package/numkong/c/dispatch_i64.c +86 -0
  84. package/numkong/c/dispatch_i8.c +348 -0
  85. package/numkong/c/dispatch_other.c +348 -0
  86. package/numkong/c/dispatch_u1.c +193 -0
  87. package/numkong/c/dispatch_u16.c +124 -0
  88. package/numkong/c/dispatch_u32.c +118 -0
  89. package/numkong/c/dispatch_u4.c +153 -0
  90. package/numkong/c/dispatch_u64.c +102 -0
  91. package/numkong/c/dispatch_u8.c +364 -0
  92. package/numkong/c/numkong.c +940 -0
  93. package/numkong/cmake/nk_arm_isa_probes.cmake +61 -0
  94. package/numkong/cmake/nk_isa_probe.cmake +90 -0
  95. package/numkong/cmake/nk_loongarch_isa_probes.cmake +15 -0
  96. package/numkong/cmake/nk_math_lib.cmake +46 -0
  97. package/numkong/cmake/nk_power_isa_probes.cmake +15 -0
  98. package/numkong/cmake/nk_riscv_isa_probes.cmake +19 -0
  99. package/numkong/cmake/nk_x86_isa_probes.cmake +35 -0
  100. package/numkong/cmake/toolchain-aarch64-gnu.cmake +86 -0
  101. package/numkong/cmake/toolchain-android-arm64.cmake +25 -0
  102. package/numkong/cmake/toolchain-android-armv7.cmake +25 -0
  103. package/numkong/cmake/toolchain-loongarch64-gnu.cmake +68 -0
  104. package/numkong/cmake/toolchain-ppc64le-gnu.cmake +74 -0
  105. package/numkong/cmake/toolchain-riscv64-gnu.cmake +107 -0
  106. package/numkong/cmake/toolchain-riscv64-llvm.cmake +114 -0
  107. package/numkong/cmake/toolchain-wasi.cmake +77 -0
  108. package/numkong/cmake/toolchain-wasm.cmake +70 -0
  109. package/numkong/cmake/toolchain-wasm64.cmake +62 -0
  110. package/numkong/cmake/toolchain-x86_64-llvm.cmake +47 -0
  111. package/numkong/include/README.md +643 -0
  112. package/numkong/include/module.modulemap +129 -0
  113. package/numkong/include/numkong/attention/sapphireamx.h +1363 -0
  114. package/numkong/include/numkong/attention/sme.h +2075 -0
  115. package/numkong/include/numkong/attention.h +49 -0
  116. package/numkong/include/numkong/capabilities.h +930 -0
  117. package/numkong/include/numkong/cast/README.md +262 -0
  118. package/numkong/include/numkong/cast/diamond.h +64 -0
  119. package/numkong/include/numkong/cast/haswell.h +823 -0
  120. package/numkong/include/numkong/cast/icelake.h +475 -0
  121. package/numkong/include/numkong/cast/loongsonasx.h +252 -0
  122. package/numkong/include/numkong/cast/neon.h +1174 -0
  123. package/numkong/include/numkong/cast/powervsx.h +449 -0
  124. package/numkong/include/numkong/cast/rvv.h +970 -0
  125. package/numkong/include/numkong/cast/sapphire.h +262 -0
  126. package/numkong/include/numkong/cast/serial.h +2336 -0
  127. package/numkong/include/numkong/cast/skylake.h +915 -0
  128. package/numkong/include/numkong/cast/v128relaxed.h +537 -0
  129. package/numkong/include/numkong/cast.h +254 -0
  130. package/numkong/include/numkong/cast.hpp +44 -0
  131. package/numkong/include/numkong/curved/README.md +223 -0
  132. package/numkong/include/numkong/curved/genoa.h +182 -0
  133. package/numkong/include/numkong/curved/haswell.h +276 -0
  134. package/numkong/include/numkong/curved/neon.h +329 -0
  135. package/numkong/include/numkong/curved/neonbfdot.h +211 -0
  136. package/numkong/include/numkong/curved/rvv.h +305 -0
  137. package/numkong/include/numkong/curved/serial.h +207 -0
  138. package/numkong/include/numkong/curved/skylake.h +457 -0
  139. package/numkong/include/numkong/curved/smef64.h +530 -0
  140. package/numkong/include/numkong/curved.h +513 -0
  141. package/numkong/include/numkong/curved.hpp +144 -0
  142. package/numkong/include/numkong/dot/README.md +442 -0
  143. package/numkong/include/numkong/dot/alder.h +563 -0
  144. package/numkong/include/numkong/dot/diamond.h +158 -0
  145. package/numkong/include/numkong/dot/genoa.h +289 -0
  146. package/numkong/include/numkong/dot/haswell.h +1723 -0
  147. package/numkong/include/numkong/dot/icelake.h +997 -0
  148. package/numkong/include/numkong/dot/loongsonasx.h +671 -0
  149. package/numkong/include/numkong/dot/neon.h +869 -0
  150. package/numkong/include/numkong/dot/neonbfdot.h +243 -0
  151. package/numkong/include/numkong/dot/neonfhm.h +358 -0
  152. package/numkong/include/numkong/dot/neonfp8.h +323 -0
  153. package/numkong/include/numkong/dot/neonsdot.h +622 -0
  154. package/numkong/include/numkong/dot/powervsx.h +752 -0
  155. package/numkong/include/numkong/dot/rvv.h +722 -0
  156. package/numkong/include/numkong/dot/rvvbb.h +72 -0
  157. package/numkong/include/numkong/dot/rvvbf16.h +123 -0
  158. package/numkong/include/numkong/dot/rvvhalf.h +129 -0
  159. package/numkong/include/numkong/dot/sapphire.h +141 -0
  160. package/numkong/include/numkong/dot/serial.h +874 -0
  161. package/numkong/include/numkong/dot/sierra.h +405 -0
  162. package/numkong/include/numkong/dot/skylake.h +1116 -0
  163. package/numkong/include/numkong/dot/sve.h +420 -0
  164. package/numkong/include/numkong/dot/svebfdot.h +75 -0
  165. package/numkong/include/numkong/dot/svehalf.h +168 -0
  166. package/numkong/include/numkong/dot/svesdot.h +90 -0
  167. package/numkong/include/numkong/dot/v128relaxed.h +1293 -0
  168. package/numkong/include/numkong/dot.h +1136 -0
  169. package/numkong/include/numkong/dot.hpp +94 -0
  170. package/numkong/include/numkong/dots/README.md +496 -0
  171. package/numkong/include/numkong/dots/alder.h +117 -0
  172. package/numkong/include/numkong/dots/diamond.h +86 -0
  173. package/numkong/include/numkong/dots/genoa.h +100 -0
  174. package/numkong/include/numkong/dots/graniteamx.h +733 -0
  175. package/numkong/include/numkong/dots/haswell.h +310 -0
  176. package/numkong/include/numkong/dots/icelake.h +180 -0
  177. package/numkong/include/numkong/dots/loongsonasx.h +176 -0
  178. package/numkong/include/numkong/dots/neon.h +123 -0
  179. package/numkong/include/numkong/dots/neonbfdot.h +59 -0
  180. package/numkong/include/numkong/dots/neonfhm.h +96 -0
  181. package/numkong/include/numkong/dots/neonfp8.h +99 -0
  182. package/numkong/include/numkong/dots/neonsdot.h +144 -0
  183. package/numkong/include/numkong/dots/powervsx.h +194 -0
  184. package/numkong/include/numkong/dots/rvv.h +2593 -0
  185. package/numkong/include/numkong/dots/sapphireamx.h +4017 -0
  186. package/numkong/include/numkong/dots/serial.h +2867 -0
  187. package/numkong/include/numkong/dots/sierra.h +100 -0
  188. package/numkong/include/numkong/dots/skylake.h +207 -0
  189. package/numkong/include/numkong/dots/sme.h +5041 -0
  190. package/numkong/include/numkong/dots/smebi32.h +477 -0
  191. package/numkong/include/numkong/dots/smef64.h +1328 -0
  192. package/numkong/include/numkong/dots/smehalf.h +47 -0
  193. package/numkong/include/numkong/dots/v128relaxed.h +317 -0
  194. package/numkong/include/numkong/dots.h +3110 -0
  195. package/numkong/include/numkong/dots.hpp +641 -0
  196. package/numkong/include/numkong/each/README.md +465 -0
  197. package/numkong/include/numkong/each/haswell.h +1658 -0
  198. package/numkong/include/numkong/each/icelake.h +272 -0
  199. package/numkong/include/numkong/each/neon.h +1138 -0
  200. package/numkong/include/numkong/each/neonbfdot.h +211 -0
  201. package/numkong/include/numkong/each/neonhalf.h +318 -0
  202. package/numkong/include/numkong/each/rvv.h +1121 -0
  203. package/numkong/include/numkong/each/sapphire.h +343 -0
  204. package/numkong/include/numkong/each/serial.h +282 -0
  205. package/numkong/include/numkong/each/skylake.h +1562 -0
  206. package/numkong/include/numkong/each/v128relaxed.h +562 -0
  207. package/numkong/include/numkong/each.h +2232 -0
  208. package/numkong/include/numkong/each.hpp +434 -0
  209. package/numkong/include/numkong/geospatial/README.md +147 -0
  210. package/numkong/include/numkong/geospatial/haswell.h +633 -0
  211. package/numkong/include/numkong/geospatial/neon.h +615 -0
  212. package/numkong/include/numkong/geospatial/rvv.h +701 -0
  213. package/numkong/include/numkong/geospatial/serial.h +309 -0
  214. package/numkong/include/numkong/geospatial/skylake.h +613 -0
  215. package/numkong/include/numkong/geospatial/v128relaxed.h +652 -0
  216. package/numkong/include/numkong/geospatial.h +453 -0
  217. package/numkong/include/numkong/geospatial.hpp +235 -0
  218. package/numkong/include/numkong/matrix.hpp +336 -0
  219. package/numkong/include/numkong/maxsim/README.md +187 -0
  220. package/numkong/include/numkong/maxsim/alder.h +511 -0
  221. package/numkong/include/numkong/maxsim/genoa.h +115 -0
  222. package/numkong/include/numkong/maxsim/haswell.h +553 -0
  223. package/numkong/include/numkong/maxsim/icelake.h +480 -0
  224. package/numkong/include/numkong/maxsim/neonsdot.h +394 -0
  225. package/numkong/include/numkong/maxsim/sapphireamx.h +877 -0
  226. package/numkong/include/numkong/maxsim/serial.h +490 -0
  227. package/numkong/include/numkong/maxsim/sme.h +993 -0
  228. package/numkong/include/numkong/maxsim/v128relaxed.h +280 -0
  229. package/numkong/include/numkong/maxsim.h +571 -0
  230. package/numkong/include/numkong/maxsim.hpp +133 -0
  231. package/numkong/include/numkong/mesh/README.md +219 -0
  232. package/numkong/include/numkong/mesh/haswell.h +2013 -0
  233. package/numkong/include/numkong/mesh/neon.h +1821 -0
  234. package/numkong/include/numkong/mesh/neonbfdot.h +771 -0
  235. package/numkong/include/numkong/mesh/rvv.h +936 -0
  236. package/numkong/include/numkong/mesh/serial.h +721 -0
  237. package/numkong/include/numkong/mesh/skylake.h +2344 -0
  238. package/numkong/include/numkong/mesh/v128relaxed.h +1081 -0
  239. package/numkong/include/numkong/mesh.h +666 -0
  240. package/numkong/include/numkong/mesh.hpp +701 -0
  241. package/numkong/include/numkong/numkong.h +78 -0
  242. package/numkong/include/numkong/numkong.hpp +58 -0
  243. package/numkong/include/numkong/probability/README.md +177 -0
  244. package/numkong/include/numkong/probability/haswell.h +267 -0
  245. package/numkong/include/numkong/probability/neon.h +255 -0
  246. package/numkong/include/numkong/probability/rvv.h +414 -0
  247. package/numkong/include/numkong/probability/serial.h +181 -0
  248. package/numkong/include/numkong/probability/skylake.h +324 -0
  249. package/numkong/include/numkong/probability.h +380 -0
  250. package/numkong/include/numkong/probability.hpp +120 -0
  251. package/numkong/include/numkong/random.h +50 -0
  252. package/numkong/include/numkong/random.hpp +285 -0
  253. package/numkong/include/numkong/reduce/README.md +552 -0
  254. package/numkong/include/numkong/reduce/alder.h +636 -0
  255. package/numkong/include/numkong/reduce/genoa.h +201 -0
  256. package/numkong/include/numkong/reduce/haswell.h +3785 -0
  257. package/numkong/include/numkong/reduce/icelake.h +549 -0
  258. package/numkong/include/numkong/reduce/neon.h +3980 -0
  259. package/numkong/include/numkong/reduce/neonbfdot.h +142 -0
  260. package/numkong/include/numkong/reduce/neonfhm.h +293 -0
  261. package/numkong/include/numkong/reduce/neonsdot.h +357 -0
  262. package/numkong/include/numkong/reduce/rvv.h +3533 -0
  263. package/numkong/include/numkong/reduce/serial.h +753 -0
  264. package/numkong/include/numkong/reduce/sierra.h +338 -0
  265. package/numkong/include/numkong/reduce/skylake.h +3794 -0
  266. package/numkong/include/numkong/reduce/sve.h +52 -0
  267. package/numkong/include/numkong/reduce/v128relaxed.h +2302 -0
  268. package/numkong/include/numkong/reduce.h +1582 -0
  269. package/numkong/include/numkong/reduce.hpp +778 -0
  270. package/numkong/include/numkong/scalar/README.md +89 -0
  271. package/numkong/include/numkong/scalar/haswell.h +122 -0
  272. package/numkong/include/numkong/scalar/loongsonasx.h +74 -0
  273. package/numkong/include/numkong/scalar/neon.h +122 -0
  274. package/numkong/include/numkong/scalar/neonhalf.h +70 -0
  275. package/numkong/include/numkong/scalar/powervsx.h +96 -0
  276. package/numkong/include/numkong/scalar/rvv.h +211 -0
  277. package/numkong/include/numkong/scalar/sapphire.h +74 -0
  278. package/numkong/include/numkong/scalar/serial.h +332 -0
  279. package/numkong/include/numkong/scalar/v128relaxed.h +56 -0
  280. package/numkong/include/numkong/scalar.h +696 -0
  281. package/numkong/include/numkong/set/README.md +179 -0
  282. package/numkong/include/numkong/set/haswell.h +334 -0
  283. package/numkong/include/numkong/set/icelake.h +485 -0
  284. package/numkong/include/numkong/set/loongsonasx.h +181 -0
  285. package/numkong/include/numkong/set/neon.h +363 -0
  286. package/numkong/include/numkong/set/powervsx.h +326 -0
  287. package/numkong/include/numkong/set/rvv.h +226 -0
  288. package/numkong/include/numkong/set/rvvbb.h +117 -0
  289. package/numkong/include/numkong/set/serial.h +174 -0
  290. package/numkong/include/numkong/set/sve.h +187 -0
  291. package/numkong/include/numkong/set/v128relaxed.h +240 -0
  292. package/numkong/include/numkong/set.h +471 -0
  293. package/numkong/include/numkong/set.hpp +114 -0
  294. package/numkong/include/numkong/sets/README.md +149 -0
  295. package/numkong/include/numkong/sets/haswell.h +63 -0
  296. package/numkong/include/numkong/sets/icelake.h +66 -0
  297. package/numkong/include/numkong/sets/loongsonasx.h +52 -0
  298. package/numkong/include/numkong/sets/neon.h +61 -0
  299. package/numkong/include/numkong/sets/powervsx.h +65 -0
  300. package/numkong/include/numkong/sets/serial.h +43 -0
  301. package/numkong/include/numkong/sets/smebi32.h +1135 -0
  302. package/numkong/include/numkong/sets/v128relaxed.h +58 -0
  303. package/numkong/include/numkong/sets.h +382 -0
  304. package/numkong/include/numkong/sparse/README.md +156 -0
  305. package/numkong/include/numkong/sparse/icelake.h +463 -0
  306. package/numkong/include/numkong/sparse/neon.h +288 -0
  307. package/numkong/include/numkong/sparse/serial.h +117 -0
  308. package/numkong/include/numkong/sparse/sve2.h +504 -0
  309. package/numkong/include/numkong/sparse/turin.h +322 -0
  310. package/numkong/include/numkong/sparse.h +363 -0
  311. package/numkong/include/numkong/sparse.hpp +113 -0
  312. package/numkong/include/numkong/spatial/README.md +439 -0
  313. package/numkong/include/numkong/spatial/alder.h +604 -0
  314. package/numkong/include/numkong/spatial/diamond.h +240 -0
  315. package/numkong/include/numkong/spatial/genoa.h +222 -0
  316. package/numkong/include/numkong/spatial/haswell.h +979 -0
  317. package/numkong/include/numkong/spatial/icelake.h +1067 -0
  318. package/numkong/include/numkong/spatial/loongsonasx.h +483 -0
  319. package/numkong/include/numkong/spatial/neon.h +846 -0
  320. package/numkong/include/numkong/spatial/neonbfdot.h +164 -0
  321. package/numkong/include/numkong/spatial/neonfp8.h +258 -0
  322. package/numkong/include/numkong/spatial/neonsdot.h +429 -0
  323. package/numkong/include/numkong/spatial/powervsx.h +738 -0
  324. package/numkong/include/numkong/spatial/rvv.h +991 -0
  325. package/numkong/include/numkong/spatial/rvvbf16.h +128 -0
  326. package/numkong/include/numkong/spatial/rvvhalf.h +120 -0
  327. package/numkong/include/numkong/spatial/serial.h +347 -0
  328. package/numkong/include/numkong/spatial/sierra.h +516 -0
  329. package/numkong/include/numkong/spatial/skylake.h +605 -0
  330. package/numkong/include/numkong/spatial/sve.h +246 -0
  331. package/numkong/include/numkong/spatial/svebfdot.h +125 -0
  332. package/numkong/include/numkong/spatial/svehalf.h +136 -0
  333. package/numkong/include/numkong/spatial/svesdot.h +143 -0
  334. package/numkong/include/numkong/spatial/v128relaxed.h +969 -0
  335. package/numkong/include/numkong/spatial.h +1681 -0
  336. package/numkong/include/numkong/spatial.hpp +183 -0
  337. package/numkong/include/numkong/spatials/README.md +580 -0
  338. package/numkong/include/numkong/spatials/alder.h +94 -0
  339. package/numkong/include/numkong/spatials/diamond.h +82 -0
  340. package/numkong/include/numkong/spatials/genoa.h +94 -0
  341. package/numkong/include/numkong/spatials/graniteamx.h +173 -0
  342. package/numkong/include/numkong/spatials/haswell.h +219 -0
  343. package/numkong/include/numkong/spatials/icelake.h +113 -0
  344. package/numkong/include/numkong/spatials/loongsonasx.h +153 -0
  345. package/numkong/include/numkong/spatials/neon.h +109 -0
  346. package/numkong/include/numkong/spatials/neonbfdot.h +60 -0
  347. package/numkong/include/numkong/spatials/neonfhm.h +92 -0
  348. package/numkong/include/numkong/spatials/neonfp8.h +111 -0
  349. package/numkong/include/numkong/spatials/neonsdot.h +143 -0
  350. package/numkong/include/numkong/spatials/powervsx.h +153 -0
  351. package/numkong/include/numkong/spatials/rvv.h +1976 -0
  352. package/numkong/include/numkong/spatials/sapphireamx.h +1149 -0
  353. package/numkong/include/numkong/spatials/serial.h +248 -0
  354. package/numkong/include/numkong/spatials/sierra.h +96 -0
  355. package/numkong/include/numkong/spatials/skylake.h +184 -0
  356. package/numkong/include/numkong/spatials/sme.h +1927 -0
  357. package/numkong/include/numkong/spatials/smef64.h +483 -0
  358. package/numkong/include/numkong/spatials/v128relaxed.h +240 -0
  359. package/numkong/include/numkong/spatials.h +3374 -0
  360. package/numkong/include/numkong/spatials.hpp +509 -0
  361. package/numkong/include/numkong/tensor.hpp +1981 -0
  362. package/numkong/include/numkong/trigonometry/README.md +184 -0
  363. package/numkong/include/numkong/trigonometry/haswell.h +653 -0
  364. package/numkong/include/numkong/trigonometry/neon.h +638 -0
  365. package/numkong/include/numkong/trigonometry/rvv.h +700 -0
  366. package/numkong/include/numkong/trigonometry/serial.h +703 -0
  367. package/numkong/include/numkong/trigonometry/skylake.h +725 -0
  368. package/numkong/include/numkong/trigonometry/v128relaxed.h +667 -0
  369. package/numkong/include/numkong/trigonometry.h +467 -0
  370. package/numkong/include/numkong/trigonometry.hpp +166 -0
  371. package/numkong/include/numkong/types.h +1695 -0
  372. package/numkong/include/numkong/types.hpp +6027 -0
  373. package/numkong/include/numkong/vector.hpp +766 -0
  374. package/numkong/javascript/@numkong-darwin-arm64/package.json +26 -0
  375. package/numkong/javascript/@numkong-darwin-x64/package.json +26 -0
  376. package/numkong/javascript/@numkong-linux-arm64/package.json +25 -0
  377. package/numkong/javascript/@numkong-linux-x64/package.json +25 -0
  378. package/numkong/javascript/@numkong-win32-arm64/package.json +25 -0
  379. package/numkong/javascript/@numkong-win32-x64/package.json +25 -0
  380. package/numkong/javascript/README.md +254 -0
  381. package/numkong/javascript/node-gyp-build.d.ts +4 -0
  382. package/numkong/javascript/numkong-browser.ts +40 -0
  383. package/numkong/javascript/numkong-wasm.ts +750 -0
  384. package/numkong/javascript/numkong.c +712 -0
  385. package/numkong/javascript/numkong.ts +613 -0
  386. package/{simsimd → numkong}/javascript/tsconfig-base.json +6 -3
  387. package/numkong/javascript/tsconfig-cjs.json +13 -0
  388. package/{simsimd → numkong}/javascript/tsconfig-esm.json +1 -1
  389. package/numkong/javascript/types.ts +784 -0
  390. package/numkong/numkong.gypi +104 -0
  391. package/numkong/package.json +108 -0
  392. package/numkong/probes/arm_neon.c +8 -0
  393. package/numkong/probes/arm_neon_bfdot.c +9 -0
  394. package/numkong/probes/arm_neon_fhm.c +9 -0
  395. package/numkong/probes/arm_neon_half.c +8 -0
  396. package/numkong/probes/arm_neon_sdot.c +9 -0
  397. package/numkong/probes/arm_neonfp8.c +9 -0
  398. package/numkong/probes/arm_sme.c +16 -0
  399. package/numkong/probes/arm_sme2.c +16 -0
  400. package/numkong/probes/arm_sme2p1.c +16 -0
  401. package/numkong/probes/arm_sme_bf16.c +16 -0
  402. package/numkong/probes/arm_sme_bi32.c +16 -0
  403. package/numkong/probes/arm_sme_f64.c +16 -0
  404. package/numkong/probes/arm_sme_fa64.c +14 -0
  405. package/numkong/probes/arm_sme_half.c +16 -0
  406. package/numkong/probes/arm_sme_lut2.c +15 -0
  407. package/numkong/probes/arm_sve.c +18 -0
  408. package/numkong/probes/arm_sve2.c +20 -0
  409. package/numkong/probes/arm_sve2p1.c +18 -0
  410. package/numkong/probes/arm_sve_bfdot.c +20 -0
  411. package/numkong/probes/arm_sve_half.c +18 -0
  412. package/numkong/probes/arm_sve_sdot.c +21 -0
  413. package/numkong/probes/loongarch_lasx.c +12 -0
  414. package/numkong/probes/power_vsx.c +12 -0
  415. package/numkong/probes/probe.js +127 -0
  416. package/numkong/probes/riscv_rvv.c +14 -0
  417. package/numkong/probes/riscv_rvv_bb.c +15 -0
  418. package/numkong/probes/riscv_rvv_bf16.c +17 -0
  419. package/numkong/probes/riscv_rvv_half.c +14 -0
  420. package/numkong/probes/wasm_v128relaxed.c +11 -0
  421. package/numkong/probes/x86_alder.c +17 -0
  422. package/numkong/probes/x86_diamond.c +17 -0
  423. package/numkong/probes/x86_genoa.c +17 -0
  424. package/numkong/probes/x86_graniteamx.c +19 -0
  425. package/numkong/probes/x86_haswell.c +11 -0
  426. package/numkong/probes/x86_icelake.c +17 -0
  427. package/numkong/probes/x86_sapphire.c +16 -0
  428. package/numkong/probes/x86_sapphireamx.c +18 -0
  429. package/numkong/probes/x86_sierra.c +17 -0
  430. package/numkong/probes/x86_skylake.c +15 -0
  431. package/numkong/probes/x86_turin.c +17 -0
  432. package/numkong/rust/README.md +698 -0
  433. package/numkong/rust/capabilities.rs +96 -0
  434. package/numkong/rust/cast.rs +314 -0
  435. package/numkong/rust/curved.rs +391 -0
  436. package/numkong/rust/each.rs +2721 -0
  437. package/numkong/rust/geospatial.rs +343 -0
  438. package/numkong/rust/matrix.rs +5424 -0
  439. package/numkong/rust/maxsim.rs +540 -0
  440. package/numkong/rust/mesh.rs +837 -0
  441. package/numkong/rust/numkong.rs +409 -0
  442. package/numkong/rust/probability.rs +272 -0
  443. package/numkong/rust/reduce.rs +1249 -0
  444. package/numkong/rust/set.rs +194 -0
  445. package/numkong/rust/sparse.rs +513 -0
  446. package/numkong/rust/spatial.rs +1815 -0
  447. package/numkong/rust/tensor.rs +7565 -0
  448. package/numkong/rust/types.rs +3666 -0
  449. package/numkong/rust/vector.rs +1899 -0
  450. package/numkong/rustfmt.toml +2 -0
  451. package/numkong/swift/Matrix.swift +1388 -0
  452. package/numkong/swift/MaxSim.swift +151 -0
  453. package/numkong/swift/NumKong.swift +283 -0
  454. package/numkong/swift/Numerics.swift +329 -0
  455. package/numkong/swift/README.md +433 -0
  456. package/numkong/swift/Scalar.swift +741 -0
  457. package/numkong/swift/Tensor.swift +209 -0
  458. package/numkong/test/README.md +315 -0
  459. package/numkong/test/playwright.config.ts +16 -0
  460. package/numkong/test/swift/Test.swift +534 -0
  461. package/numkong/test/test-browser-bundle-runner.mjs +36 -0
  462. package/numkong/test/test-browser-bundle.html +188 -0
  463. package/numkong/test/test-browser-runner.mjs +43 -0
  464. package/numkong/test/test-browser.html +305 -0
  465. package/numkong/test/test-dtypes.mjs +453 -0
  466. package/numkong/test/test-wasm.mjs +267 -0
  467. package/numkong/test/test.cpp +442 -0
  468. package/numkong/test/test.hpp +580 -0
  469. package/numkong/test/test.mjs +363 -0
  470. package/numkong/test/test_base.py +1041 -0
  471. package/numkong/test/test_cast.cpp +242 -0
  472. package/numkong/test/test_cross.hpp +551 -0
  473. package/numkong/test/test_cross.py +467 -0
  474. package/numkong/test/test_cross_amx.cpp +50 -0
  475. package/numkong/test/test_cross_arm.cpp +225 -0
  476. package/numkong/test/test_cross_blas.cpp +354 -0
  477. package/numkong/test/test_cross_loongarch.cpp +90 -0
  478. package/numkong/test/test_cross_power.cpp +73 -0
  479. package/numkong/test/test_cross_rvv.cpp +109 -0
  480. package/numkong/test/test_cross_serial.cpp +181 -0
  481. package/numkong/test/test_cross_sme.cpp +148 -0
  482. package/numkong/test/test_cross_wasm.cpp +141 -0
  483. package/numkong/test/test_cross_x86.cpp +442 -0
  484. package/numkong/test/test_curved.cpp +182 -0
  485. package/numkong/test/test_curved.py +207 -0
  486. package/numkong/test/test_dot.cpp +353 -0
  487. package/numkong/test/test_dot.py +296 -0
  488. package/numkong/test/test_each.cpp +437 -0
  489. package/numkong/test/test_each.py +756 -0
  490. package/numkong/test/test_geospatial.cpp +155 -0
  491. package/numkong/test/test_geospatial.py +241 -0
  492. package/numkong/test/test_maxsim.cpp +138 -0
  493. package/numkong/test/test_maxsim.py +209 -0
  494. package/numkong/test/test_mesh.cpp +203 -0
  495. package/numkong/test/test_mesh.py +394 -0
  496. package/numkong/test/test_probability.cpp +136 -0
  497. package/numkong/test/test_probability.py +166 -0
  498. package/numkong/test/test_reduce.cpp +563 -0
  499. package/numkong/test/test_reduce.py +781 -0
  500. package/numkong/test/test_set.cpp +151 -0
  501. package/numkong/test/test_set.py +120 -0
  502. package/numkong/test/test_similarities.py +579 -0
  503. package/numkong/test/test_sparse.cpp +121 -0
  504. package/numkong/test/test_sparse.py +118 -0
  505. package/numkong/test/test_spatial.cpp +518 -0
  506. package/numkong/test/test_spatial.py +269 -0
  507. package/numkong/test/test_tensor.cpp +884 -0
  508. package/numkong/test/test_tensor.py +1417 -0
  509. package/numkong/test/test_trigonometry.cpp +156 -0
  510. package/numkong/test/test_trigonometry.py +180 -0
  511. package/package.json +7 -8
  512. package/prebuilds/darwin-arm64+x64/usearch.node +0 -0
  513. package/prebuilds/linux-arm64/usearch.node +0 -0
  514. package/prebuilds/linux-x64/usearch.node +0 -0
  515. package/fp16/CMakeLists.txt +0 -176
  516. package/fp16/LICENSE +0 -11
  517. package/fp16/README.md +0 -20
  518. package/fp16/bench/alt-element.cc +0 -89
  519. package/fp16/bench/from-alt-array.cc +0 -175
  520. package/fp16/bench/from-ieee-array.cc +0 -399
  521. package/fp16/bench/ieee-element.cc +0 -272
  522. package/fp16/bench/to-alt-array.cc +0 -91
  523. package/fp16/bench/to-ieee-array.cc +0 -303
  524. package/fp16/cmake/DownloadGoogleBenchmark.cmake +0 -15
  525. package/fp16/cmake/DownloadGoogleTest.cmake +0 -15
  526. package/fp16/cmake/DownloadPSimd.cmake +0 -15
  527. package/fp16/configure.py +0 -65
  528. package/fp16/confu.yaml +0 -8
  529. package/fp16/include/fp16/__init__.py +0 -0
  530. package/fp16/include/fp16/avx.py +0 -50
  531. package/fp16/include/fp16/avx2.py +0 -53
  532. package/fp16/include/fp16/bitcasts.h +0 -92
  533. package/fp16/include/fp16/fp16.h +0 -451
  534. package/fp16/include/fp16/psimd.h +0 -131
  535. package/fp16/include/fp16.h +0 -11
  536. package/fp16/jni/Android.mk +0 -6
  537. package/fp16/jni/Application.mk +0 -4
  538. package/fp16/test/alt-from-fp32-value.cc +0 -498
  539. package/fp16/test/alt-to-fp32-bits.cc +0 -263
  540. package/fp16/test/alt-to-fp32-psimd.cc +0 -145
  541. package/fp16/test/alt-to-fp32-value.cc +0 -348
  542. package/fp16/test/alt-to-fp32x2-psimd.cc +0 -245
  543. package/fp16/test/bitcasts.cc +0 -56
  544. package/fp16/test/ieee-from-fp32-value.cc +0 -523
  545. package/fp16/test/ieee-to-fp32-bits.cc +0 -333
  546. package/fp16/test/ieee-to-fp32-psimd.cc +0 -312
  547. package/fp16/test/ieee-to-fp32-value.cc +0 -420
  548. package/fp16/test/ieee-to-fp32x2-psimd.cc +0 -541
  549. package/fp16/test/peachpy/alt-xmm-to-fp32-xmm-avx.cc +0 -96
  550. package/fp16/test/peachpy/alt-xmm-to-fp32-ymm-avx2.cc +0 -96
  551. package/fp16/test/peachpy/stubs.py +0 -38
  552. package/fp16/test/tables.cc +0 -530
  553. package/fp16/test/tables.h +0 -13
  554. package/fp16/third-party/THHalf.h +0 -92
  555. package/fp16/third-party/eigen-half.h +0 -170
  556. package/fp16/third-party/float16-compressor.h +0 -80
  557. package/fp16/third-party/half.hpp +0 -2916
  558. package/fp16/third-party/npy-halffloat.h +0 -328
  559. package/simsimd/.git-blame-ignore-revs +0 -2
  560. package/simsimd/CONTRIBUTING.md +0 -321
  561. package/simsimd/Cargo.toml +0 -49
  562. package/simsimd/MANIFEST.in +0 -4
  563. package/simsimd/Package.swift +0 -28
  564. package/simsimd/README.md +0 -1493
  565. package/simsimd/binding.gyp +0 -35
  566. package/simsimd/build.rs +0 -70
  567. package/simsimd/c/lib.c +0 -374
  568. package/simsimd/include/module.modulemap +0 -4
  569. package/simsimd/include/simsimd/binary.h +0 -486
  570. package/simsimd/include/simsimd/curved.h +0 -1541
  571. package/simsimd/include/simsimd/dot.h +0 -1853
  572. package/simsimd/include/simsimd/elementwise.h +0 -2465
  573. package/simsimd/include/simsimd/geospatial.h +0 -43
  574. package/simsimd/include/simsimd/mesh.h +0 -69
  575. package/simsimd/include/simsimd/probability.h +0 -606
  576. package/simsimd/include/simsimd/simsimd.h +0 -2567
  577. package/simsimd/include/simsimd/sparse.h +0 -1384
  578. package/simsimd/include/simsimd/spatial.h +0 -2335
  579. package/simsimd/include/simsimd/types.h +0 -666
  580. package/simsimd/javascript/fallback.ts +0 -233
  581. package/simsimd/javascript/lib.c +0 -125
  582. package/simsimd/javascript/node-gyp-build.d.ts +0 -1
  583. package/simsimd/javascript/simsimd.ts +0 -179
  584. package/simsimd/javascript/tsconfig-cjs.json +0 -8
  585. package/simsimd/package-lock.json +0 -1761
  586. package/simsimd/package.json +0 -61
  587. package/simsimd/rust/lib.rs +0 -1821
  588. package/simsimd/scripts/bench.cxx +0 -1170
  589. package/simsimd/scripts/bench.ipynb +0 -187
  590. package/simsimd/scripts/bench.js +0 -127
  591. package/simsimd/scripts/bench.rs +0 -876
  592. package/simsimd/scripts/bench_vectors.py +0 -734
  593. package/simsimd/scripts/bench_vectors_live.py +0 -213
  594. package/simsimd/scripts/test.c +0 -187
  595. package/simsimd/scripts/test.mjs +0 -220
  596. package/simsimd/scripts/test.py +0 -1505
  597. package/simsimd/swift/SimSIMD.swift +0 -149
  598. package/simsimd/swift/Test.swift +0 -128
  599. /package/{simsimd → numkong}/CODE_OF_CONDUCT.md +0 -0
  600. /package/{simsimd → numkong}/LICENSE +0 -0
  601. /package/{simsimd → numkong}/javascript/dist-package-cjs.json +0 -0
  602. /package/{simsimd → numkong}/javascript/dist-package-esm.json +0 -0
@@ -0,0 +1,3980 @@
1
+ /**
2
+ * @brief Base NEON (ARMv8-A) implementations for the redesigned reduction API (moments + minmax).
3
+ * @file include/numkong/reduce/neon.h
4
+ * @author Ash Vardanian
5
+ * @date February 13, 2026
6
+ *
7
+ * @sa include/numkong/reduce.h
8
+ */
9
+ #ifndef NK_REDUCE_NEON_H
10
+ #define NK_REDUCE_NEON_H
11
+
12
+ #if NK_TARGET_ARM64_
13
+ #if NK_TARGET_NEON
14
+
15
+ #include "numkong/types.h" // `nk_size_t`
16
+ #include "numkong/cast/neon.h" // `nk_e4m3x16_to_f16x8x2_neon_`
17
+ #include "numkong/cast/serial.h" // `nk_e4m3_to_f16_serial`
18
+
19
+ #if defined(__cplusplus)
20
+ extern "C" {
21
+ #endif
22
+
23
+ #if defined(__clang__)
24
+ #pragma clang attribute push(__attribute__((target("arch=armv8-a+simd"))), apply_to = function)
25
+ #elif defined(__GNUC__)
26
+ #pragma GCC push_options
27
+ #pragma GCC target("arch=armv8-a+simd")
28
+ #endif
29
+
30
+ NK_INTERNAL nk_u64_t nk_reduce_sadd_u64x2_neon_(uint64x2_t v_u64x2) {
31
+ uint64x2_t swapped_u64x2 = vextq_u64(v_u64x2, v_u64x2, 1);
32
+ return vgetq_lane_u64(vqaddq_u64(v_u64x2, swapped_u64x2), 0);
33
+ }
34
+
35
+ /** @brief Saturating square of each i64 lane → u64. If |a| >= 2^32, a² overflows u64 → saturate. */
36
+ NK_INTERNAL uint64x2_t nk_i64_smul_sq_i64x2_neon_(int64x2_t val_i64x2) {
37
+ uint64x2_t absolute_u64x2 = vreinterpretq_u64_s64(vabsq_s64(val_i64x2));
38
+ uint32x2_t low_halves_u32x2 = vmovn_u64(absolute_u64x2);
39
+ uint64x2_t high_bits_u64x2 = vshrq_n_u64(absolute_u64x2, 32);
40
+ uint64x2_t low_squared_u64x2 = vmull_u32(low_halves_u32x2, low_halves_u32x2);
41
+ uint64x2_t is_small_u64x2 = vceqq_u64(high_bits_u64x2, vdupq_n_u64(0));
42
+ return vbslq_u64(is_small_u64x2, low_squared_u64x2, vdupq_n_u64(NK_U64_MAX));
43
+ }
44
+
45
+ /** @brief Saturating square of each u64 lane → u64. If a >= 2^32, a² overflows u64 → saturate. */
46
+ NK_INTERNAL uint64x2_t nk_u64_smul_sq_u64x2_neon_(uint64x2_t val_u64x2) {
47
+ uint32x2_t low_halves_u32x2 = vmovn_u64(val_u64x2);
48
+ uint64x2_t high_bits_u64x2 = vshrq_n_u64(val_u64x2, 32);
49
+ uint64x2_t low_squared_u64x2 = vmull_u32(low_halves_u32x2, low_halves_u32x2);
50
+ uint64x2_t is_small_u64x2 = vceqq_u64(high_bits_u64x2, vdupq_n_u64(0));
51
+ return vbslq_u64(is_small_u64x2, low_squared_u64x2, vdupq_n_u64(NK_U64_MAX));
52
+ }
53
+
54
+ NK_INTERNAL void nk_reduce_moments_f32_neon_contiguous_( //
55
+ nk_f32_t const *data_ptr, nk_size_t count, //
56
+ nk_f64_t *sum_ptr, nk_f64_t *sumsq_ptr) {
57
+ float64x2_t sum_f64x2 = vdupq_n_f64(0), sumsq_f64x2 = vdupq_n_f64(0);
58
+ nk_size_t idx = 0;
59
+ for (; idx + 4 <= count; idx += 4) {
60
+ float32x4_t data_f32x4 = vld1q_f32(data_ptr + idx);
61
+ float64x2_t data_low_f64x2 = vcvt_f64_f32(vget_low_f32(data_f32x4));
62
+ float64x2_t data_high_f64x2 = vcvt_high_f64_f32(data_f32x4);
63
+ sum_f64x2 = vaddq_f64(sum_f64x2, data_low_f64x2);
64
+ sum_f64x2 = vaddq_f64(sum_f64x2, data_high_f64x2);
65
+ sumsq_f64x2 = vfmaq_f64(sumsq_f64x2, data_low_f64x2, data_low_f64x2);
66
+ sumsq_f64x2 = vfmaq_f64(sumsq_f64x2, data_high_f64x2, data_high_f64x2);
67
+ }
68
+ nk_f64_t sum = vaddvq_f64(sum_f64x2), sumsq = vaddvq_f64(sumsq_f64x2);
69
+ for (; idx < count; ++idx) {
70
+ nk_f64_t value = (nk_f64_t)data_ptr[idx];
71
+ sum += value, sumsq += value * value;
72
+ }
73
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
74
+ }
75
+
76
+ NK_INTERNAL void nk_reduce_moments_f32_neon_strided_( //
77
+ nk_f32_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
78
+ nk_f64_t *sum_ptr, nk_f64_t *sumsq_ptr) {
79
+ float64x2_t sum_f64x2 = vdupq_n_f64(0), sumsq_f64x2 = vdupq_n_f64(0);
80
+ nk_size_t idx = 0;
81
+ if (stride_elements == 2) {
82
+ for (; idx + 4 < count; idx += 4) {
83
+ float32x4x2_t loaded_f32x4x2 = vld2q_f32(data_ptr + idx * 2);
84
+ float64x2_t data_low_f64x2 = vcvt_f64_f32(vget_low_f32(loaded_f32x4x2.val[0]));
85
+ float64x2_t data_high_f64x2 = vcvt_high_f64_f32(loaded_f32x4x2.val[0]);
86
+ sum_f64x2 = vaddq_f64(sum_f64x2, data_low_f64x2);
87
+ sum_f64x2 = vaddq_f64(sum_f64x2, data_high_f64x2);
88
+ sumsq_f64x2 = vfmaq_f64(sumsq_f64x2, data_low_f64x2, data_low_f64x2);
89
+ sumsq_f64x2 = vfmaq_f64(sumsq_f64x2, data_high_f64x2, data_high_f64x2);
90
+ }
91
+ }
92
+ else if (stride_elements == 3) {
93
+ for (; idx + 4 < count; idx += 4) {
94
+ float32x4x3_t loaded_f32x4x3 = vld3q_f32(data_ptr + idx * 3);
95
+ float64x2_t data_low_f64x2 = vcvt_f64_f32(vget_low_f32(loaded_f32x4x3.val[0]));
96
+ float64x2_t data_high_f64x2 = vcvt_high_f64_f32(loaded_f32x4x3.val[0]);
97
+ sum_f64x2 = vaddq_f64(sum_f64x2, data_low_f64x2);
98
+ sum_f64x2 = vaddq_f64(sum_f64x2, data_high_f64x2);
99
+ sumsq_f64x2 = vfmaq_f64(sumsq_f64x2, data_low_f64x2, data_low_f64x2);
100
+ sumsq_f64x2 = vfmaq_f64(sumsq_f64x2, data_high_f64x2, data_high_f64x2);
101
+ }
102
+ }
103
+ else {
104
+ for (; idx + 4 < count; idx += 4) {
105
+ float32x4x4_t loaded_f32x4x4 = vld4q_f32(data_ptr + idx * 4);
106
+ float64x2_t data_low_f64x2 = vcvt_f64_f32(vget_low_f32(loaded_f32x4x4.val[0]));
107
+ float64x2_t data_high_f64x2 = vcvt_high_f64_f32(loaded_f32x4x4.val[0]);
108
+ sum_f64x2 = vaddq_f64(sum_f64x2, data_low_f64x2);
109
+ sum_f64x2 = vaddq_f64(sum_f64x2, data_high_f64x2);
110
+ sumsq_f64x2 = vfmaq_f64(sumsq_f64x2, data_low_f64x2, data_low_f64x2);
111
+ sumsq_f64x2 = vfmaq_f64(sumsq_f64x2, data_high_f64x2, data_high_f64x2);
112
+ }
113
+ }
114
+ nk_f64_t sum = vaddvq_f64(sum_f64x2), sumsq = vaddvq_f64(sumsq_f64x2);
115
+ nk_f32_t const *current_ptr = data_ptr + idx * stride_elements;
116
+ for (; idx < count; ++idx, current_ptr += stride_elements) {
117
+ nk_f64_t value = (nk_f64_t)(*current_ptr);
118
+ sum += value, sumsq += value * value;
119
+ }
120
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
121
+ }
122
+
123
+ NK_PUBLIC void nk_reduce_moments_f32_neon( //
124
+ nk_f32_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
125
+ nk_f64_t *sum_ptr, nk_f64_t *sumsq_ptr) {
126
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_f32_t);
127
+ int aligned = (stride_bytes % sizeof(nk_f32_t) == 0);
128
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
129
+ else if (!aligned) nk_reduce_moments_f32_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
130
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 4) {
131
+ nk_size_t left_count = count / 2;
132
+ nk_f64_t left_sum, left_sumsq, right_sum, right_sumsq;
133
+ nk_reduce_moments_f32_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
134
+ nk_reduce_moments_f32_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
135
+ &right_sum, &right_sumsq);
136
+ *sum_ptr = left_sum + right_sum, *sumsq_ptr = left_sumsq + right_sumsq;
137
+ }
138
+ else if (stride_elements == 1) nk_reduce_moments_f32_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
139
+ else if (stride_elements <= 4)
140
+ nk_reduce_moments_f32_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
141
+ else nk_reduce_moments_f32_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
142
+ }
143
+
144
+ NK_INTERNAL void nk_reduce_minmax_f32_neon_contiguous_( //
145
+ nk_f32_t const *data_ptr, nk_size_t count, //
146
+ nk_f32_t *min_value_ptr, nk_size_t *min_index_ptr, //
147
+ nk_f32_t *max_value_ptr, nk_size_t *max_index_ptr) {
148
+ float32x4_t min_f32x4 = vdupq_n_f32(NK_F32_MAX), max_f32x4 = vdupq_n_f32(NK_F32_MIN);
149
+ uint32x4_t min_iter_u32x4 = vdupq_n_u32(0), max_iter_u32x4 = vdupq_n_u32(0);
150
+ uint32x4_t iter_u32x4 = vdupq_n_u32(0), one_u32x4 = vdupq_n_u32(1);
151
+ nk_size_t idx = 0;
152
+ for (; idx + 4 <= count; idx += 4) {
153
+ float32x4_t data_f32x4 = vld1q_f32(data_ptr + idx);
154
+ uint32x4_t less_u32x4 = vcltq_f32(data_f32x4, min_f32x4);
155
+ uint32x4_t greater_u32x4 = vcgtq_f32(data_f32x4, max_f32x4);
156
+ min_f32x4 = vbslq_f32(less_u32x4, data_f32x4, min_f32x4);
157
+ max_f32x4 = vbslq_f32(greater_u32x4, data_f32x4, max_f32x4);
158
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
159
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
160
+ iter_u32x4 = vaddq_u32(iter_u32x4, one_u32x4);
161
+ }
162
+ nk_size_t remaining = count - idx;
163
+ if (remaining > 0) {
164
+ nk_b128_vec_t tail_vec;
165
+ nk_partial_load_b32x4_serial_(data_ptr + idx, &tail_vec, remaining);
166
+ uint32x4_t lane_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
167
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
168
+ uint32x4_t valid_u32x4 = vcltq_u32(lane_u32x4, vdupq_n_u32((nk_u32_t)remaining));
169
+ float32x4_t data_min_f32x4 = vbslq_f32(valid_u32x4, tail_vec.f32x4, min_f32x4);
170
+ float32x4_t data_max_f32x4 = vbslq_f32(valid_u32x4, tail_vec.f32x4, max_f32x4);
171
+ uint32x4_t less_u32x4 = vcltq_f32(data_min_f32x4, min_f32x4);
172
+ uint32x4_t greater_u32x4 = vcgtq_f32(data_max_f32x4, max_f32x4);
173
+ min_f32x4 = vbslq_f32(less_u32x4, data_min_f32x4, min_f32x4);
174
+ max_f32x4 = vbslq_f32(greater_u32x4, data_max_f32x4, max_f32x4);
175
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
176
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
177
+ }
178
+ nk_f32_t min_value = vminvq_f32(min_f32x4), max_value = vmaxvq_f32(max_f32x4);
179
+
180
+ // All-NaN / sentinel check: sentinels remain unchanged when all data is NaN.
181
+ if (min_value == NK_F32_MAX && max_value == NK_F32_MIN) {
182
+ *min_value_ptr = NK_F32_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_F32_MIN,
183
+ *max_index_ptr = NK_SIZE_MAX;
184
+ return;
185
+ }
186
+
187
+ uint32x4_t min_value_match_u32x4 = vceqq_f32(min_f32x4, vdupq_n_f32(min_value));
188
+ uint32x4_t masked_min_iter_u32x4 = vbslq_u32(min_value_match_u32x4, min_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
189
+ nk_u32_t earliest_min_cycle = vminvq_u32(masked_min_iter_u32x4);
190
+ uint32x4_t max_value_match_u32x4 = vceqq_f32(max_f32x4, vdupq_n_f32(max_value));
191
+ uint32x4_t masked_max_iter_u32x4 = vbslq_u32(max_value_match_u32x4, max_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
192
+ nk_u32_t earliest_max_cycle = vminvq_u32(masked_max_iter_u32x4);
193
+ uint32x4_t lane_indices_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
194
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
195
+ uint32x4_t min_cycle_match_u32x4 = vceqq_u32(min_iter_u32x4, vdupq_n_u32(earliest_min_cycle));
196
+ uint32x4_t min_both_match_u32x4 = vandq_u32(min_value_match_u32x4, min_cycle_match_u32x4);
197
+ uint32x4_t min_masked_lanes_u32x4 = vbslq_u32(min_both_match_u32x4, lane_indices_u32x4, vdupq_n_u32(NK_U32_MAX));
198
+ nk_u32_t min_lane_offset = vminvq_u32(min_masked_lanes_u32x4);
199
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 4 + (nk_size_t)min_lane_offset;
200
+ uint32x4_t max_cycle_match_u32x4 = vceqq_u32(max_iter_u32x4, vdupq_n_u32(earliest_max_cycle));
201
+ uint32x4_t max_both_match_u32x4 = vandq_u32(max_value_match_u32x4, max_cycle_match_u32x4);
202
+ uint32x4_t max_masked_lanes_u32x4 = vbslq_u32(max_both_match_u32x4, lane_indices_u32x4, vdupq_n_u32(NK_U32_MAX));
203
+ nk_u32_t max_lane_offset = vminvq_u32(max_masked_lanes_u32x4);
204
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 4 + (nk_size_t)max_lane_offset;
205
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
206
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
207
+ }
208
+
209
+ NK_INTERNAL void nk_reduce_minmax_f32_neon_strided_( //
210
+ nk_f32_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
211
+ nk_f32_t *min_value_ptr, nk_size_t *min_index_ptr, //
212
+ nk_f32_t *max_value_ptr, nk_size_t *max_index_ptr) {
213
+ float32x4_t min_f32x4 = vdupq_n_f32(NK_F32_MAX), max_f32x4 = vdupq_n_f32(NK_F32_MIN);
214
+ uint32x4_t min_iter_u32x4 = vdupq_n_u32(0), max_iter_u32x4 = vdupq_n_u32(0);
215
+ uint32x4_t iter_u32x4 = vdupq_n_u32(0), one_u32x4 = vdupq_n_u32(1);
216
+ uint32x4_t lane_indices_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
217
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
218
+ nk_size_t idx = 0;
219
+ float32x4_t data_for_min_f32x4, data_for_max_f32x4;
220
+
221
+ nk_reduce_minmax_f32_neon_cycle:
222
+ if (stride_elements == 2 && idx + 4 < count) {
223
+ float32x4x2_t loaded = vld2q_f32(data_ptr + idx * 2);
224
+ data_for_min_f32x4 = loaded.val[0];
225
+ data_for_max_f32x4 = loaded.val[0];
226
+ idx += 4;
227
+ }
228
+ else if (stride_elements == 3 && idx + 4 < count) {
229
+ float32x4x3_t loaded = vld3q_f32(data_ptr + idx * 3);
230
+ data_for_min_f32x4 = loaded.val[0];
231
+ data_for_max_f32x4 = loaded.val[0];
232
+ idx += 4;
233
+ }
234
+ else if (stride_elements == 4 && idx + 4 < count) {
235
+ float32x4x4_t loaded = vld4q_f32(data_ptr + idx * 4);
236
+ data_for_min_f32x4 = loaded.val[0];
237
+ data_for_max_f32x4 = loaded.val[0];
238
+ idx += 4;
239
+ }
240
+ else if (idx < count) {
241
+ nk_b128_vec_t tail_vec;
242
+ nk_strided_load_b32x4_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
243
+ uint32x4_t valid_u32x4 = vcltq_u32(lane_indices_u32x4, vdupq_n_u32((nk_u32_t)(count - idx)));
244
+ data_for_min_f32x4 = vbslq_f32(valid_u32x4, tail_vec.f32x4, min_f32x4);
245
+ data_for_max_f32x4 = vbslq_f32(valid_u32x4, tail_vec.f32x4, max_f32x4);
246
+ idx = count;
247
+ }
248
+ else {
249
+ nk_f32_t min_value = vminvq_f32(min_f32x4), max_value = vmaxvq_f32(max_f32x4);
250
+ if (min_value == NK_F32_MAX && max_value == NK_F32_MIN) {
251
+ *min_value_ptr = NK_F32_MAX, *min_index_ptr = NK_SIZE_MAX;
252
+ *max_value_ptr = NK_F32_MIN, *max_index_ptr = NK_SIZE_MAX;
253
+ return;
254
+ }
255
+ uint32x4_t min_value_match_u32x4 = vceqq_f32(min_f32x4, vdupq_n_f32(min_value));
256
+ uint32x4_t masked_min_iter_u32x4 = vbslq_u32(min_value_match_u32x4, min_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
257
+ nk_u32_t earliest_min_cycle = vminvq_u32(masked_min_iter_u32x4);
258
+ uint32x4_t max_value_match_u32x4 = vceqq_f32(max_f32x4, vdupq_n_f32(max_value));
259
+ uint32x4_t masked_max_iter_u32x4 = vbslq_u32(max_value_match_u32x4, max_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
260
+ nk_u32_t earliest_max_cycle = vminvq_u32(masked_max_iter_u32x4);
261
+ uint32x4_t min_cycle_match_u32x4 = vceqq_u32(min_iter_u32x4, vdupq_n_u32(earliest_min_cycle));
262
+ uint32x4_t min_both_match_u32x4 = vandq_u32(min_value_match_u32x4, min_cycle_match_u32x4);
263
+ uint32x4_t min_masked_lanes_u32x4 = vbslq_u32(min_both_match_u32x4, lane_indices_u32x4,
264
+ vdupq_n_u32(NK_U32_MAX));
265
+ nk_u32_t min_lane_offset = vminvq_u32(min_masked_lanes_u32x4);
266
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 4 + (nk_size_t)min_lane_offset;
267
+ uint32x4_t max_cycle_match_u32x4 = vceqq_u32(max_iter_u32x4, vdupq_n_u32(earliest_max_cycle));
268
+ uint32x4_t max_both_match_u32x4 = vandq_u32(max_value_match_u32x4, max_cycle_match_u32x4);
269
+ uint32x4_t max_masked_lanes_u32x4 = vbslq_u32(max_both_match_u32x4, lane_indices_u32x4,
270
+ vdupq_n_u32(NK_U32_MAX));
271
+ nk_u32_t max_lane_offset = vminvq_u32(max_masked_lanes_u32x4);
272
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 4 + (nk_size_t)max_lane_offset;
273
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
274
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
275
+ return;
276
+ }
277
+
278
+ // Shared update body
279
+ uint32x4_t less_u32x4 = vcltq_f32(data_for_min_f32x4, min_f32x4);
280
+ uint32x4_t greater_u32x4 = vcgtq_f32(data_for_max_f32x4, max_f32x4);
281
+ min_f32x4 = vbslq_f32(less_u32x4, data_for_min_f32x4, min_f32x4);
282
+ max_f32x4 = vbslq_f32(greater_u32x4, data_for_max_f32x4, max_f32x4);
283
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
284
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
285
+ iter_u32x4 = vaddq_u32(iter_u32x4, one_u32x4);
286
+ goto nk_reduce_minmax_f32_neon_cycle;
287
+ }
288
+
289
+ NK_PUBLIC void nk_reduce_minmax_f32_neon( //
290
+ nk_f32_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
291
+ nk_f32_t *min_value_ptr, nk_size_t *min_index_ptr, //
292
+ nk_f32_t *max_value_ptr, nk_size_t *max_index_ptr) {
293
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_f32_t);
294
+ int aligned = (stride_bytes % sizeof(nk_f32_t) == 0);
295
+ if (count == 0)
296
+ *min_value_ptr = NK_F32_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_F32_MIN,
297
+ *max_index_ptr = NK_SIZE_MAX;
298
+ else if (!aligned)
299
+ nk_reduce_minmax_f32_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
300
+ max_index_ptr);
301
+ else if (count > (nk_size_t)NK_U32_MAX * 4) {
302
+ nk_size_t left_count = count / 2;
303
+ nk_f32_t left_min_value, right_min_value, left_max_value, right_max_value;
304
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
305
+ nk_reduce_minmax_f32_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index, &left_max_value,
306
+ &left_max_index);
307
+ nk_reduce_minmax_f32_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
308
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
309
+ if (right_min_value < left_min_value)
310
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
311
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
312
+ if (right_max_value > left_max_value)
313
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
314
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
315
+ }
316
+ else if (stride_elements == 1)
317
+ nk_reduce_minmax_f32_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
318
+ max_index_ptr);
319
+ else if (stride_elements <= 4)
320
+ nk_reduce_minmax_f32_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
321
+ max_value_ptr, max_index_ptr);
322
+ else
323
+ nk_reduce_minmax_f32_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
324
+ max_index_ptr);
325
+ }
326
+
327
+ NK_INTERNAL void nk_reduce_moments_f64_neon_contiguous_( //
328
+ nk_f64_t const *data_ptr, nk_size_t count, //
329
+ nk_f64_t *sum_ptr, nk_f64_t *sumsq_ptr) {
330
+ float64x2_t sum_f64x2 = vdupq_n_f64(0), sum_compensation_f64x2 = vdupq_n_f64(0);
331
+ float64x2_t sumsq_f64x2 = vdupq_n_f64(0), sumsq_compensation_f64x2 = vdupq_n_f64(0);
332
+ nk_size_t idx = 0;
333
+ for (; idx + 2 <= count; idx += 2) {
334
+ float64x2_t data_f64x2 = vld1q_f64(data_ptr + idx);
335
+ float64x2_t temp_sum_f64x2 = vaddq_f64(sum_f64x2, data_f64x2);
336
+ float64x2_t residual_f64x2 = vsubq_f64(temp_sum_f64x2, sum_f64x2);
337
+ sum_compensation_f64x2 = vaddq_f64(sum_compensation_f64x2,
338
+ vaddq_f64(vsubq_f64(sum_f64x2, vsubq_f64(temp_sum_f64x2, residual_f64x2)),
339
+ vsubq_f64(data_f64x2, residual_f64x2)));
340
+ sum_f64x2 = temp_sum_f64x2;
341
+ float64x2_t data_squared_f64x2 = vmulq_f64(data_f64x2, data_f64x2);
342
+ float64x2_t temp_sumsq_f64x2 = vaddq_f64(sumsq_f64x2, data_squared_f64x2);
343
+ float64x2_t residual_sumsq_f64x2 = vsubq_f64(temp_sumsq_f64x2, sumsq_f64x2);
344
+ sumsq_compensation_f64x2 = vaddq_f64(
345
+ sumsq_compensation_f64x2,
346
+ vaddq_f64(vsubq_f64(sumsq_f64x2, vsubq_f64(temp_sumsq_f64x2, residual_sumsq_f64x2)),
347
+ vsubq_f64(data_squared_f64x2, residual_sumsq_f64x2)));
348
+ sumsq_f64x2 = temp_sumsq_f64x2;
349
+ }
350
+ nk_size_t remaining = count - idx;
351
+ if (remaining > 0) {
352
+ nk_b128_vec_t tail_vec;
353
+ nk_partial_load_b64x2_serial_(data_ptr + idx, &tail_vec, remaining);
354
+ float64x2_t data_f64x2 = tail_vec.f64x2;
355
+ float64x2_t temp_sum_f64x2 = vaddq_f64(sum_f64x2, data_f64x2);
356
+ float64x2_t residual_f64x2 = vsubq_f64(temp_sum_f64x2, sum_f64x2);
357
+ sum_compensation_f64x2 = vaddq_f64(sum_compensation_f64x2,
358
+ vaddq_f64(vsubq_f64(sum_f64x2, vsubq_f64(temp_sum_f64x2, residual_f64x2)),
359
+ vsubq_f64(data_f64x2, residual_f64x2)));
360
+ sum_f64x2 = temp_sum_f64x2;
361
+ float64x2_t data_squared_f64x2 = vmulq_f64(data_f64x2, data_f64x2);
362
+ float64x2_t temp_sumsq_f64x2 = vaddq_f64(sumsq_f64x2, data_squared_f64x2);
363
+ float64x2_t residual_sumsq_f64x2 = vsubq_f64(temp_sumsq_f64x2, sumsq_f64x2);
364
+ sumsq_compensation_f64x2 = vaddq_f64(
365
+ sumsq_compensation_f64x2,
366
+ vaddq_f64(vsubq_f64(sumsq_f64x2, vsubq_f64(temp_sumsq_f64x2, residual_sumsq_f64x2)),
367
+ vsubq_f64(data_squared_f64x2, residual_sumsq_f64x2)));
368
+ sumsq_f64x2 = temp_sumsq_f64x2;
369
+ }
370
+ *sum_ptr = vaddvq_f64(vaddq_f64(sum_f64x2, sum_compensation_f64x2));
371
+ *sumsq_ptr = vaddvq_f64(vaddq_f64(sumsq_f64x2, sumsq_compensation_f64x2));
372
+ }
373
+
374
+ NK_PUBLIC void nk_reduce_moments_f64_neon( //
375
+ nk_f64_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
376
+ nk_f64_t *sum_ptr, nk_f64_t *sumsq_ptr) {
377
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_f64_t);
378
+ int aligned = (stride_bytes % sizeof(nk_f64_t) == 0);
379
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
380
+ else if (!aligned) nk_reduce_moments_f64_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
381
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 2) {
382
+ nk_size_t left_count = count / 2;
383
+ nk_f64_t left_sum, left_sumsq, right_sum, right_sumsq;
384
+ nk_reduce_moments_f64_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
385
+ nk_reduce_moments_f64_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
386
+ &right_sum, &right_sumsq);
387
+ *sum_ptr = left_sum + right_sum, *sumsq_ptr = left_sumsq + right_sumsq;
388
+ }
389
+ else if (stride_elements == 1) nk_reduce_moments_f64_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
390
+ else nk_reduce_moments_f64_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
391
+ }
392
+
393
+ NK_INTERNAL void nk_reduce_minmax_f64_neon_contiguous_( //
394
+ nk_f64_t const *data_ptr, nk_size_t count, //
395
+ nk_f64_t *min_value_ptr, nk_size_t *min_index_ptr, //
396
+ nk_f64_t *max_value_ptr, nk_size_t *max_index_ptr) {
397
+ float64x2_t min_f64x2 = vdupq_n_f64(NK_F64_MAX), max_f64x2 = vdupq_n_f64(NK_F64_MIN);
398
+ uint64x2_t min_iter_u64x2 = vdupq_n_u64(0), max_iter_u64x2 = vdupq_n_u64(0);
399
+ uint64x2_t iter_u64x2 = vdupq_n_u64(0), one_u64x2 = vdupq_n_u64(1);
400
+ nk_size_t idx = 0;
401
+ for (; idx + 2 <= count; idx += 2) {
402
+ float64x2_t data_f64x2 = vld1q_f64(data_ptr + idx);
403
+ uint64x2_t less_u64x2 = vcltq_f64(data_f64x2, min_f64x2);
404
+ uint64x2_t greater_u64x2 = vcgtq_f64(data_f64x2, max_f64x2);
405
+ min_f64x2 = vbslq_f64(less_u64x2, data_f64x2, min_f64x2);
406
+ max_f64x2 = vbslq_f64(greater_u64x2, data_f64x2, max_f64x2);
407
+ min_iter_u64x2 = vbslq_u64(less_u64x2, iter_u64x2, min_iter_u64x2);
408
+ max_iter_u64x2 = vbslq_u64(greater_u64x2, iter_u64x2, max_iter_u64x2);
409
+ iter_u64x2 = vaddq_u64(iter_u64x2, one_u64x2);
410
+ }
411
+ nk_b128_vec_t min_values_vec, max_values_vec, min_indices_vec, max_indices_vec;
412
+ min_values_vec.f64x2 = min_f64x2;
413
+ min_indices_vec.u64x2 = min_iter_u64x2;
414
+ max_values_vec.f64x2 = max_f64x2;
415
+ max_indices_vec.u64x2 = max_iter_u64x2;
416
+ nk_f64_t min_value, max_value;
417
+ nk_size_t min_index, max_index;
418
+ if (min_values_vec.f64s[0] <= min_values_vec.f64s[1])
419
+ min_value = min_values_vec.f64s[0], min_index = (nk_size_t)min_indices_vec.u64s[0] * 2;
420
+ else min_value = min_values_vec.f64s[1], min_index = (nk_size_t)min_indices_vec.u64s[1] * 2 + 1;
421
+ if (max_values_vec.f64s[0] >= max_values_vec.f64s[1])
422
+ max_value = max_values_vec.f64s[0], max_index = (nk_size_t)max_indices_vec.u64s[0] * 2;
423
+ else max_value = max_values_vec.f64s[1], max_index = (nk_size_t)max_indices_vec.u64s[1] * 2 + 1;
424
+ for (; idx < count; ++idx) {
425
+ nk_f64_t val = data_ptr[idx];
426
+ if (val < min_value) min_value = val, min_index = idx;
427
+ if (val > max_value) max_value = val, max_index = idx;
428
+ }
429
+ // All-NaN / sentinel check: sentinels remain unchanged when all data is NaN.
430
+ if (min_value == NK_F64_MAX && max_value == NK_F64_MIN) {
431
+ *min_value_ptr = NK_F64_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_F64_MIN,
432
+ *max_index_ptr = NK_SIZE_MAX;
433
+ return;
434
+ }
435
+ *min_value_ptr = min_value, *min_index_ptr = min_index;
436
+ *max_value_ptr = max_value, *max_index_ptr = max_index;
437
+ }
438
+
439
+ NK_PUBLIC void nk_reduce_minmax_f64_neon( //
440
+ nk_f64_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
441
+ nk_f64_t *min_value_ptr, nk_size_t *min_index_ptr, //
442
+ nk_f64_t *max_value_ptr, nk_size_t *max_index_ptr) {
443
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_f64_t);
444
+ int aligned = (stride_bytes % sizeof(nk_f64_t) == 0);
445
+ if (count == 0)
446
+ *min_value_ptr = NK_F64_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_F64_MIN,
447
+ *max_index_ptr = NK_SIZE_MAX;
448
+ else if (!aligned)
449
+ nk_reduce_minmax_f64_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
450
+ max_index_ptr);
451
+ else if (stride_elements == 1)
452
+ nk_reduce_minmax_f64_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
453
+ max_index_ptr);
454
+ else
455
+ nk_reduce_minmax_f64_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
456
+ max_index_ptr);
457
+ }
458
+
459
+ NK_INTERNAL void nk_reduce_moments_i8_neon_contiguous_( //
460
+ nk_i8_t const *data_ptr, nk_size_t count, //
461
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
462
+ int32x4_t sum_i32x4 = vdupq_n_s32(0);
463
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
464
+ nk_size_t idx = 0;
465
+ for (; idx + 16 <= count; idx += 16) {
466
+ int8x16_t data_i8x16 = vld1q_s8(data_ptr + idx);
467
+ int16x8_t pairwise_i16x8 = vpaddlq_s8(data_i8x16);
468
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(pairwise_i16x8));
469
+ int16x8_t squares_low_i16x8 = vmull_s8(vget_low_s8(data_i8x16), vget_low_s8(data_i8x16));
470
+ int16x8_t squares_high_i16x8 = vmull_high_s8(data_i8x16, data_i8x16);
471
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_low_i16x8))));
472
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_high_i16x8))));
473
+ }
474
+ nk_i64_t sum = vaddlvq_s32(sum_i32x4);
475
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
476
+ for (; idx < count; ++idx) {
477
+ nk_i64_t value_i64 = (nk_i64_t)data_ptr[idx];
478
+ sum += value_i64, sumsq += (nk_u64_t)(value_i64 * value_i64);
479
+ }
480
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
481
+ }
482
+
483
+ NK_INTERNAL void nk_reduce_moments_i8_neon_strided_( //
484
+ nk_i8_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
485
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
486
+ int32x4_t sum_i32x4 = vdupq_n_s32(0);
487
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
488
+ nk_size_t idx = 0;
489
+ if (stride_elements == 2) {
490
+ for (; idx + 16 < count; idx += 16) {
491
+ int8x16x2_t loaded_i8x16x2 = vld2q_s8(data_ptr + idx * 2);
492
+ int8x16_t data_i8x16 = loaded_i8x16x2.val[0];
493
+ int16x8_t pairwise_i16x8 = vpaddlq_s8(data_i8x16);
494
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(pairwise_i16x8));
495
+ int16x8_t squares_low_i16x8 = vmull_s8(vget_low_s8(data_i8x16), vget_low_s8(data_i8x16));
496
+ int16x8_t squares_high_i16x8 = vmull_high_s8(data_i8x16, data_i8x16);
497
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_low_i16x8))));
498
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_high_i16x8))));
499
+ }
500
+ }
501
+ else if (stride_elements == 3) {
502
+ for (; idx + 16 < count; idx += 16) {
503
+ int8x16x3_t loaded_i8x16x3 = vld3q_s8(data_ptr + idx * 3);
504
+ int8x16_t data_i8x16 = loaded_i8x16x3.val[0];
505
+ int16x8_t pairwise_i16x8 = vpaddlq_s8(data_i8x16);
506
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(pairwise_i16x8));
507
+ int16x8_t squares_low_i16x8 = vmull_s8(vget_low_s8(data_i8x16), vget_low_s8(data_i8x16));
508
+ int16x8_t squares_high_i16x8 = vmull_high_s8(data_i8x16, data_i8x16);
509
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_low_i16x8))));
510
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_high_i16x8))));
511
+ }
512
+ }
513
+ else {
514
+ for (; idx + 16 < count; idx += 16) {
515
+ int8x16x4_t loaded_i8x16x4 = vld4q_s8(data_ptr + idx * 4);
516
+ int8x16_t data_i8x16 = loaded_i8x16x4.val[0];
517
+ int16x8_t pairwise_i16x8 = vpaddlq_s8(data_i8x16);
518
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(pairwise_i16x8));
519
+ int16x8_t squares_low_i16x8 = vmull_s8(vget_low_s8(data_i8x16), vget_low_s8(data_i8x16));
520
+ int16x8_t squares_high_i16x8 = vmull_high_s8(data_i8x16, data_i8x16);
521
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_low_i16x8))));
522
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_high_i16x8))));
523
+ }
524
+ }
525
+ nk_i64_t sum = vaddlvq_s32(sum_i32x4);
526
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
527
+ for (; idx < count; ++idx) {
528
+ nk_i64_t value_i64 = (nk_i64_t)data_ptr[idx * stride_elements];
529
+ sum += value_i64, sumsq += (nk_u64_t)(value_i64 * value_i64);
530
+ }
531
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
532
+ }
533
+
534
+ NK_PUBLIC void nk_reduce_moments_i8_neon( //
535
+ nk_i8_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
536
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
537
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_i8_t);
538
+ int aligned = (stride_bytes % sizeof(nk_i8_t) == 0);
539
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
540
+ else if (!aligned) nk_reduce_moments_i8_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
541
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 16) {
542
+ nk_size_t left_count = count / 2;
543
+ nk_i64_t left_sum, right_sum;
544
+ nk_u64_t left_sumsq, right_sumsq;
545
+ nk_reduce_moments_i8_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
546
+ nk_reduce_moments_i8_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes, &right_sum,
547
+ &right_sumsq);
548
+ *sum_ptr = nk_i64_saturating_add_serial(left_sum, right_sum);
549
+ *sumsq_ptr = nk_u64_saturating_add_serial(left_sumsq, right_sumsq);
550
+ }
551
+ else if (stride_elements == 1) nk_reduce_moments_i8_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
552
+ else if (stride_elements <= 4)
553
+ nk_reduce_moments_i8_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
554
+ else nk_reduce_moments_i8_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
555
+ }
556
+
557
+ NK_INTERNAL void nk_reduce_minmax_i8_neon_contiguous_( //
558
+ nk_i8_t const *data_ptr, nk_size_t count, //
559
+ nk_i8_t *min_value_ptr, nk_size_t *min_index_ptr, //
560
+ nk_i8_t *max_value_ptr, nk_size_t *max_index_ptr) {
561
+ int8x16_t min_i8x16 = vdupq_n_s8(NK_I8_MAX), max_i8x16 = vdupq_n_s8(NK_I8_MIN);
562
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
563
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
564
+ nk_size_t idx = 0;
565
+ for (; idx + 16 <= count; idx += 16) {
566
+ int8x16_t data_i8x16 = vld1q_s8(data_ptr + idx);
567
+ uint8x16_t less_u8x16 = vcltq_s8(data_i8x16, min_i8x16);
568
+ uint8x16_t greater_u8x16 = vcgtq_s8(data_i8x16, max_i8x16);
569
+ min_i8x16 = vbslq_s8(less_u8x16, data_i8x16, min_i8x16);
570
+ max_i8x16 = vbslq_s8(greater_u8x16, data_i8x16, max_i8x16);
571
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
572
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
573
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
574
+ }
575
+ nk_size_t remaining = count - idx;
576
+ if (remaining > 0) {
577
+ nk_b128_vec_t tail_vec;
578
+ nk_partial_load_b8x16_serial_(data_ptr + idx, &tail_vec, remaining);
579
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
580
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
581
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)remaining));
582
+ int8x16_t data_for_min_i8x16 = vbslq_s8(valid_u8x16, tail_vec.i8x16, vdupq_n_s8(NK_I8_MAX));
583
+ int8x16_t data_for_max_i8x16 = vbslq_s8(valid_u8x16, tail_vec.i8x16, vdupq_n_s8(NK_I8_MIN));
584
+ uint8x16_t less_u8x16 = vcltq_s8(data_for_min_i8x16, min_i8x16);
585
+ uint8x16_t greater_u8x16 = vcgtq_s8(data_for_max_i8x16, max_i8x16);
586
+ min_i8x16 = vbslq_s8(less_u8x16, data_for_min_i8x16, min_i8x16);
587
+ max_i8x16 = vbslq_s8(greater_u8x16, data_for_max_i8x16, max_i8x16);
588
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
589
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
590
+ }
591
+ nk_i8_t min_value = vminvq_s8(min_i8x16), max_value = vmaxvq_s8(max_i8x16);
592
+ uint8x16_t min_value_match_u8x16 = vceqq_s8(min_i8x16, vdupq_n_s8(min_value));
593
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
594
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
595
+ uint8x16_t max_value_match_u8x16 = vceqq_s8(max_i8x16, vdupq_n_s8(max_value));
596
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
597
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
598
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
599
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
600
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
601
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
602
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
603
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
604
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
605
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
606
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
607
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
608
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
609
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
610
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
611
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
612
+ }
613
+
614
+ NK_INTERNAL void nk_reduce_minmax_i8_neon_strided_( //
615
+ nk_i8_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
616
+ nk_i8_t *min_value_ptr, nk_size_t *min_index_ptr, //
617
+ nk_i8_t *max_value_ptr, nk_size_t *max_index_ptr) {
618
+ int8x16_t min_i8x16 = vdupq_n_s8(NK_I8_MAX), max_i8x16 = vdupq_n_s8(NK_I8_MIN);
619
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
620
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
621
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
622
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
623
+ nk_size_t idx = 0;
624
+ int8x16_t data_for_min_i8x16, data_for_max_i8x16;
625
+
626
+ nk_reduce_minmax_i8_neon_cycle:
627
+ if (stride_elements == 2 && idx + 16 < count) {
628
+ int8x16x2_t loaded = vld2q_s8(data_ptr + idx * 2);
629
+ data_for_min_i8x16 = loaded.val[0];
630
+ data_for_max_i8x16 = loaded.val[0];
631
+ idx += 16;
632
+ }
633
+ else if (stride_elements == 3 && idx + 16 < count) {
634
+ int8x16x3_t loaded = vld3q_s8(data_ptr + idx * 3);
635
+ data_for_min_i8x16 = loaded.val[0];
636
+ data_for_max_i8x16 = loaded.val[0];
637
+ idx += 16;
638
+ }
639
+ else if (stride_elements == 4 && idx + 16 < count) {
640
+ int8x16x4_t loaded_i8x16x4 = vld4q_s8(data_ptr + idx * 4);
641
+ data_for_min_i8x16 = loaded_i8x16x4.val[0];
642
+ data_for_max_i8x16 = loaded_i8x16x4.val[0];
643
+ idx += 16;
644
+ }
645
+ else if (idx < count) {
646
+ nk_b128_vec_t tail_vec;
647
+ nk_strided_load_b8x16_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
648
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)(count - idx)));
649
+ data_for_min_i8x16 = vbslq_s8(valid_u8x16, tail_vec.i8x16, min_i8x16);
650
+ data_for_max_i8x16 = vbslq_s8(valid_u8x16, tail_vec.i8x16, max_i8x16);
651
+ idx = count;
652
+ }
653
+ else {
654
+ nk_i8_t min_value = vminvq_s8(min_i8x16), max_value = vmaxvq_s8(max_i8x16);
655
+ uint8x16_t min_value_match_u8x16 = vceqq_s8(min_i8x16, vdupq_n_s8(min_value));
656
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
657
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
658
+ uint8x16_t max_value_match_u8x16 = vceqq_s8(max_i8x16, vdupq_n_s8(max_value));
659
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
660
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
661
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
662
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
663
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
664
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
665
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
666
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
667
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
668
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
669
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
670
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
671
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
672
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
673
+ return;
674
+ }
675
+
676
+ // Shared update body
677
+ uint8x16_t less_u8x16 = vcltq_s8(data_for_min_i8x16, min_i8x16);
678
+ uint8x16_t greater_u8x16 = vcgtq_s8(data_for_max_i8x16, max_i8x16);
679
+ min_i8x16 = vbslq_s8(less_u8x16, data_for_min_i8x16, min_i8x16);
680
+ max_i8x16 = vbslq_s8(greater_u8x16, data_for_max_i8x16, max_i8x16);
681
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
682
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
683
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
684
+ goto nk_reduce_minmax_i8_neon_cycle;
685
+ }
686
+
687
+ NK_PUBLIC void nk_reduce_minmax_i8_neon( //
688
+ nk_i8_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
689
+ nk_i8_t *min_value_ptr, nk_size_t *min_index_ptr, //
690
+ nk_i8_t *max_value_ptr, nk_size_t *max_index_ptr) {
691
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_i8_t);
692
+ int aligned = (stride_bytes % sizeof(nk_i8_t) == 0);
693
+ if (count == 0)
694
+ *min_value_ptr = NK_I8_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_I8_MIN,
695
+ *max_index_ptr = NK_SIZE_MAX;
696
+ else if (!aligned)
697
+ nk_reduce_minmax_i8_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
698
+ max_index_ptr);
699
+ else if (count > (nk_size_t)(NK_U8_MAX + 1) * 16) {
700
+ nk_size_t left_count = count / 2;
701
+ nk_i8_t left_min_value, right_min_value, left_max_value, right_max_value;
702
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
703
+ nk_reduce_minmax_i8_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index, &left_max_value,
704
+ &left_max_index);
705
+ nk_reduce_minmax_i8_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
706
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
707
+ if (right_min_value < left_min_value)
708
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
709
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
710
+ if (right_max_value > left_max_value)
711
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
712
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
713
+ }
714
+ else if (stride_elements == 1)
715
+ nk_reduce_minmax_i8_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
716
+ max_index_ptr);
717
+ else if (stride_elements <= 4)
718
+ nk_reduce_minmax_i8_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr, max_value_ptr,
719
+ max_index_ptr);
720
+ else
721
+ nk_reduce_minmax_i8_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
722
+ max_index_ptr);
723
+ }
724
+
725
+ NK_INTERNAL void nk_reduce_moments_u8_neon_contiguous_( //
726
+ nk_u8_t const *data_ptr, nk_size_t count, //
727
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
728
+ uint32x4_t sum_u32x4 = vdupq_n_u32(0);
729
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
730
+ nk_size_t idx = 0;
731
+ for (; idx + 16 <= count; idx += 16) {
732
+ uint8x16_t data_u8x16 = vld1q_u8(data_ptr + idx);
733
+ uint16x8_t pairwise_sum_u16x8 = vpaddlq_u8(data_u8x16);
734
+ sum_u32x4 = vaddq_u32(sum_u32x4, vpaddlq_u16(pairwise_sum_u16x8));
735
+ uint16x8_t squares_low_u16x8 = vmull_u8(vget_low_u8(data_u8x16), vget_low_u8(data_u8x16));
736
+ uint16x8_t squares_high_u16x8 = vmull_high_u8(data_u8x16, data_u8x16);
737
+ uint32x4_t squares_low_u32x4 = vpaddlq_u16(squares_low_u16x8);
738
+ uint32x4_t squares_high_u32x4 = vpaddlq_u16(squares_high_u16x8);
739
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(squares_low_u32x4));
740
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(squares_high_u32x4));
741
+ }
742
+ nk_u64_t sum = vaddlvq_u32(sum_u32x4);
743
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
744
+ for (; idx < count; ++idx) {
745
+ nk_u64_t value = (nk_u64_t)data_ptr[idx];
746
+ sum += value, sumsq += value * value;
747
+ }
748
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
749
+ }
750
+
751
+ NK_INTERNAL void nk_reduce_moments_u8_neon_strided_( //
752
+ nk_u8_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
753
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
754
+ uint32x4_t sum_u32x4 = vdupq_n_u32(0);
755
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
756
+ nk_size_t idx = 0;
757
+ if (stride_elements == 2) {
758
+ for (; idx + 16 < count; idx += 16) {
759
+ uint8x16x2_t loaded_u8x16x2 = vld2q_u8(data_ptr + idx * 2);
760
+ uint8x16_t data_u8x16 = loaded_u8x16x2.val[0];
761
+ uint16x8_t pairwise_u16x8 = vpaddlq_u8(data_u8x16);
762
+ sum_u32x4 = vaddq_u32(sum_u32x4, vpaddlq_u16(pairwise_u16x8));
763
+ uint16x8_t squares_low_u16x8 = vmull_u8(vget_low_u8(data_u8x16), vget_low_u8(data_u8x16));
764
+ uint16x8_t squares_high_u16x8 = vmull_high_u8(data_u8x16, data_u8x16);
765
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(squares_low_u16x8)));
766
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(squares_high_u16x8)));
767
+ }
768
+ }
769
+ else if (stride_elements == 3) {
770
+ for (; idx + 16 < count; idx += 16) {
771
+ uint8x16x3_t loaded_u8x16x3 = vld3q_u8(data_ptr + idx * 3);
772
+ uint8x16_t data_u8x16 = loaded_u8x16x3.val[0];
773
+ uint16x8_t pairwise_u16x8 = vpaddlq_u8(data_u8x16);
774
+ sum_u32x4 = vaddq_u32(sum_u32x4, vpaddlq_u16(pairwise_u16x8));
775
+ uint16x8_t squares_low_u16x8 = vmull_u8(vget_low_u8(data_u8x16), vget_low_u8(data_u8x16));
776
+ uint16x8_t squares_high_u16x8 = vmull_high_u8(data_u8x16, data_u8x16);
777
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(squares_low_u16x8)));
778
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(squares_high_u16x8)));
779
+ }
780
+ }
781
+ else {
782
+ for (; idx + 16 < count; idx += 16) {
783
+ uint8x16x4_t loaded_u8x16x4 = vld4q_u8(data_ptr + idx * 4);
784
+ uint8x16_t data_u8x16 = loaded_u8x16x4.val[0];
785
+ uint16x8_t pairwise_u16x8 = vpaddlq_u8(data_u8x16);
786
+ sum_u32x4 = vaddq_u32(sum_u32x4, vpaddlq_u16(pairwise_u16x8));
787
+ uint16x8_t squares_low_u16x8 = vmull_u8(vget_low_u8(data_u8x16), vget_low_u8(data_u8x16));
788
+ uint16x8_t squares_high_u16x8 = vmull_high_u8(data_u8x16, data_u8x16);
789
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(squares_low_u16x8)));
790
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(squares_high_u16x8)));
791
+ }
792
+ }
793
+ nk_u64_t sum = vaddlvq_u32(sum_u32x4);
794
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
795
+ for (; idx < count; ++idx) {
796
+ nk_u64_t value = (nk_u64_t)data_ptr[idx * stride_elements];
797
+ sum += value, sumsq += value * value;
798
+ }
799
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
800
+ }
801
+
802
+ NK_PUBLIC void nk_reduce_moments_u8_neon( //
803
+ nk_u8_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
804
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
805
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_u8_t);
806
+ int aligned = (stride_bytes % sizeof(nk_u8_t) == 0);
807
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
808
+ else if (!aligned) nk_reduce_moments_u8_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
809
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 16) {
810
+ nk_size_t left_count = count / 2;
811
+ nk_u64_t left_sum, left_sumsq, right_sum, right_sumsq;
812
+ nk_reduce_moments_u8_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
813
+ nk_reduce_moments_u8_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes, &right_sum,
814
+ &right_sumsq);
815
+ *sum_ptr = nk_u64_saturating_add_serial(left_sum, right_sum);
816
+ *sumsq_ptr = nk_u64_saturating_add_serial(left_sumsq, right_sumsq);
817
+ }
818
+ else if (stride_elements == 1) nk_reduce_moments_u8_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
819
+ else if (stride_elements <= 4)
820
+ nk_reduce_moments_u8_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
821
+ else nk_reduce_moments_u8_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
822
+ }
823
+
824
+ NK_INTERNAL void nk_reduce_minmax_u8_neon_contiguous_( //
825
+ nk_u8_t const *data_ptr, nk_size_t count, //
826
+ nk_u8_t *min_value_ptr, nk_size_t *min_index_ptr, //
827
+ nk_u8_t *max_value_ptr, nk_size_t *max_index_ptr) {
828
+ uint8x16_t min_u8x16 = vdupq_n_u8(NK_U8_MAX), max_u8x16 = vdupq_n_u8(0);
829
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
830
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
831
+ nk_size_t idx = 0;
832
+ for (; idx + 16 <= count; idx += 16) {
833
+ uint8x16_t data_u8x16 = vld1q_u8(data_ptr + idx);
834
+ uint8x16_t less_u8x16 = vcltq_u8(data_u8x16, min_u8x16);
835
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_u8x16, max_u8x16);
836
+ min_u8x16 = vbslq_u8(less_u8x16, data_u8x16, min_u8x16);
837
+ max_u8x16 = vbslq_u8(greater_u8x16, data_u8x16, max_u8x16);
838
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
839
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
840
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
841
+ }
842
+ nk_size_t remaining = count - idx;
843
+ if (remaining > 0) {
844
+ nk_b128_vec_t tail_vec;
845
+ nk_partial_load_b8x16_serial_(data_ptr + idx, &tail_vec, remaining);
846
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
847
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
848
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)remaining));
849
+ uint8x16_t data_for_min_u8x16 = vbslq_u8(valid_u8x16, tail_vec.u8x16, vdupq_n_u8(NK_U8_MAX));
850
+ uint8x16_t data_for_max_u8x16 = vbslq_u8(valid_u8x16, tail_vec.u8x16, vdupq_n_u8(0));
851
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
852
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
853
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
854
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
855
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
856
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
857
+ }
858
+ nk_u8_t min_value = vminvq_u8(min_u8x16), max_value = vmaxvq_u8(max_u8x16);
859
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_value));
860
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
861
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
862
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_value));
863
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
864
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
865
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
866
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
867
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
868
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
869
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
870
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
871
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
872
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
873
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
874
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
875
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
876
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
877
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
878
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
879
+ }
880
+
881
+ NK_INTERNAL void nk_reduce_minmax_u8_neon_strided_( //
882
+ nk_u8_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
883
+ nk_u8_t *min_value_ptr, nk_size_t *min_index_ptr, //
884
+ nk_u8_t *max_value_ptr, nk_size_t *max_index_ptr) {
885
+ uint8x16_t min_u8x16 = vdupq_n_u8(NK_U8_MAX), max_u8x16 = vdupq_n_u8(0);
886
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
887
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
888
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
889
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
890
+ nk_size_t idx = 0;
891
+ uint8x16_t data_for_min_u8x16, data_for_max_u8x16;
892
+
893
+ nk_reduce_minmax_u8_neon_cycle:
894
+ if (stride_elements == 2 && idx + 16 < count) {
895
+ uint8x16x2_t loaded = vld2q_u8((nk_u8_t const *)data_ptr + idx * 2);
896
+ data_for_min_u8x16 = loaded.val[0];
897
+ data_for_max_u8x16 = loaded.val[0];
898
+ idx += 16;
899
+ }
900
+ else if (stride_elements == 3 && idx + 16 < count) {
901
+ uint8x16x3_t loaded = vld3q_u8((nk_u8_t const *)data_ptr + idx * 3);
902
+ data_for_min_u8x16 = loaded.val[0];
903
+ data_for_max_u8x16 = loaded.val[0];
904
+ idx += 16;
905
+ }
906
+ else if (stride_elements == 4 && idx + 16 < count) {
907
+ uint8x16x4_t loaded_u8x16x4 = vld4q_u8((nk_u8_t const *)data_ptr + idx * 4);
908
+ data_for_min_u8x16 = loaded_u8x16x4.val[0];
909
+ data_for_max_u8x16 = loaded_u8x16x4.val[0];
910
+ idx += 16;
911
+ }
912
+ else if (idx < count) {
913
+ nk_b128_vec_t tail_vec;
914
+ nk_strided_load_b8x16_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
915
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)(count - idx)));
916
+ data_for_min_u8x16 = vbslq_u8(valid_u8x16, tail_vec.u8x16, min_u8x16);
917
+ data_for_max_u8x16 = vbslq_u8(valid_u8x16, tail_vec.u8x16, max_u8x16);
918
+ idx = count;
919
+ }
920
+ else {
921
+ nk_u8_t min_value = vminvq_u8(min_u8x16), max_value = vmaxvq_u8(max_u8x16);
922
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_value));
923
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
924
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
925
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_value));
926
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
927
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
928
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
929
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
930
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
931
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
932
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
933
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
934
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
935
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
936
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
937
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
938
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
939
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
940
+ return;
941
+ }
942
+
943
+ // Shared update body
944
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
945
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
946
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
947
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
948
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
949
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
950
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
951
+ goto nk_reduce_minmax_u8_neon_cycle;
952
+ }
953
+
954
+ NK_PUBLIC void nk_reduce_minmax_u8_neon( //
955
+ nk_u8_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
956
+ nk_u8_t *min_value_ptr, nk_size_t *min_index_ptr, //
957
+ nk_u8_t *max_value_ptr, nk_size_t *max_index_ptr) {
958
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_u8_t);
959
+ int aligned = (stride_bytes % sizeof(nk_u8_t) == 0);
960
+ if (count == 0)
961
+ *min_value_ptr = NK_U8_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = 0, *max_index_ptr = NK_SIZE_MAX;
962
+ else if (!aligned)
963
+ nk_reduce_minmax_u8_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
964
+ max_index_ptr);
965
+ else if (count > (nk_size_t)(NK_U8_MAX + 1) * 16) {
966
+ nk_size_t left_count = count / 2;
967
+ nk_u8_t left_min_value, right_min_value, left_max_value, right_max_value;
968
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
969
+ nk_reduce_minmax_u8_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index, &left_max_value,
970
+ &left_max_index);
971
+ nk_reduce_minmax_u8_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
972
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
973
+ if (right_min_value < left_min_value)
974
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
975
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
976
+ if (right_max_value > left_max_value)
977
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
978
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
979
+ }
980
+ else if (stride_elements == 1)
981
+ nk_reduce_minmax_u8_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
982
+ max_index_ptr);
983
+ else if (stride_elements <= 4)
984
+ nk_reduce_minmax_u8_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr, max_value_ptr,
985
+ max_index_ptr);
986
+ else
987
+ nk_reduce_minmax_u8_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
988
+ max_index_ptr);
989
+ }
990
+
991
+ NK_INTERNAL void nk_reduce_moments_i16_neon_contiguous_( //
992
+ nk_i16_t const *data_ptr, nk_size_t count, //
993
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
994
+ int64x2_t sum_i64x2 = vdupq_n_s64(0);
995
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
996
+ nk_size_t idx = 0;
997
+ for (; idx + 8 <= count; idx += 8) {
998
+ int16x8_t data_i16x8 = vld1q_s16(data_ptr + idx);
999
+ int32x4_t sum32_i32x4 = vpaddlq_s16(data_i16x8);
1000
+ sum_i64x2 = vaddq_s64(sum_i64x2, vpaddlq_s32(sum32_i32x4));
1001
+ // sumsq: widening multiply i16*i16 -> i32, then widen to u64
1002
+ int32x4_t sq_low_i32x4 = vmull_s16(vget_low_s16(data_i16x8), vget_low_s16(data_i16x8));
1003
+ int32x4_t sq_high_i32x4 = vmull_high_s16(data_i16x8, data_i16x8);
1004
+ // i16*i16 squares are always non-negative, safe to reinterpret as u32
1005
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(sq_low_i32x4)));
1006
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(sq_high_i32x4)));
1007
+ }
1008
+ nk_i64_t sum = vgetq_lane_s64(sum_i64x2, 0) + vgetq_lane_s64(sum_i64x2, 1);
1009
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
1010
+ for (; idx < count; ++idx) {
1011
+ nk_i64_t value_i64 = (nk_i64_t)data_ptr[idx];
1012
+ sum += value_i64;
1013
+ sumsq += (nk_u64_t)(value_i64 * value_i64);
1014
+ }
1015
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
1016
+ }
1017
+
1018
+ NK_INTERNAL void nk_reduce_moments_i16_neon_strided_( //
1019
+ nk_i16_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
1020
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1021
+ int64x2_t sum_i64x2 = vdupq_n_s64(0);
1022
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
1023
+ nk_size_t idx = 0;
1024
+ if (stride_elements == 2) {
1025
+ for (; idx + 8 < count; idx += 8) {
1026
+ int16x8x2_t loaded_i16x8x2 = vld2q_s16(data_ptr + idx * 2);
1027
+ int16x8_t data_i16x8 = loaded_i16x8x2.val[0];
1028
+ int32x4_t pairwise_i32x4 = vpaddlq_s16(data_i16x8);
1029
+ sum_i64x2 = vaddq_s64(sum_i64x2, vpaddlq_s32(pairwise_i32x4));
1030
+ int32x4_t squares_low_i32x4 = vmull_s16(vget_low_s16(data_i16x8), vget_low_s16(data_i16x8));
1031
+ int32x4_t squares_high_i32x4 = vmull_high_s16(data_i16x8, data_i16x8);
1032
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_i32x4)));
1033
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_i32x4)));
1034
+ }
1035
+ }
1036
+ else if (stride_elements == 3) {
1037
+ for (; idx + 8 < count; idx += 8) {
1038
+ int16x8x3_t loaded_i16x8x3 = vld3q_s16(data_ptr + idx * 3);
1039
+ int16x8_t data_i16x8 = loaded_i16x8x3.val[0];
1040
+ int32x4_t pairwise_i32x4 = vpaddlq_s16(data_i16x8);
1041
+ sum_i64x2 = vaddq_s64(sum_i64x2, vpaddlq_s32(pairwise_i32x4));
1042
+ int32x4_t squares_low_i32x4 = vmull_s16(vget_low_s16(data_i16x8), vget_low_s16(data_i16x8));
1043
+ int32x4_t squares_high_i32x4 = vmull_high_s16(data_i16x8, data_i16x8);
1044
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_i32x4)));
1045
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_i32x4)));
1046
+ }
1047
+ }
1048
+ else {
1049
+ for (; idx + 8 < count; idx += 8) {
1050
+ int16x8x4_t loaded_i16x8x4 = vld4q_s16(data_ptr + idx * 4);
1051
+ int16x8_t data_i16x8 = loaded_i16x8x4.val[0];
1052
+ int32x4_t pairwise_i32x4 = vpaddlq_s16(data_i16x8);
1053
+ sum_i64x2 = vaddq_s64(sum_i64x2, vpaddlq_s32(pairwise_i32x4));
1054
+ int32x4_t squares_low_i32x4 = vmull_s16(vget_low_s16(data_i16x8), vget_low_s16(data_i16x8));
1055
+ int32x4_t squares_high_i32x4 = vmull_high_s16(data_i16x8, data_i16x8);
1056
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_i32x4)));
1057
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_i32x4)));
1058
+ }
1059
+ }
1060
+ nk_i64_t sum = vgetq_lane_s64(sum_i64x2, 0) + vgetq_lane_s64(sum_i64x2, 1);
1061
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
1062
+ for (; idx < count; ++idx) {
1063
+ nk_i64_t value_i64 = (nk_i64_t)data_ptr[idx * stride_elements];
1064
+ sum += value_i64, sumsq += (nk_u64_t)(value_i64 * value_i64);
1065
+ }
1066
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
1067
+ }
1068
+
1069
+ NK_PUBLIC void nk_reduce_moments_i16_neon( //
1070
+ nk_i16_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
1071
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1072
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_i16_t);
1073
+ int aligned = (stride_bytes % sizeof(nk_i16_t) == 0);
1074
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
1075
+ else if (!aligned) nk_reduce_moments_i16_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
1076
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 8) {
1077
+ nk_size_t left_count = count / 2;
1078
+ nk_i64_t left_sum, right_sum;
1079
+ nk_u64_t left_sumsq, right_sumsq;
1080
+ nk_reduce_moments_i16_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
1081
+ nk_reduce_moments_i16_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
1082
+ &right_sum, &right_sumsq);
1083
+ *sum_ptr = nk_i64_saturating_add_serial(left_sum, right_sum);
1084
+ *sumsq_ptr = nk_u64_saturating_add_serial(left_sumsq, right_sumsq);
1085
+ }
1086
+ else if (stride_elements == 1) nk_reduce_moments_i16_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
1087
+ else if (stride_elements <= 4)
1088
+ nk_reduce_moments_i16_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
1089
+ else nk_reduce_moments_i16_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
1090
+ }
1091
+
1092
+ NK_INTERNAL void nk_reduce_minmax_i16_neon_contiguous_( //
1093
+ nk_i16_t const *data_ptr, nk_size_t count, //
1094
+ nk_i16_t *min_value_ptr, nk_size_t *min_index_ptr, //
1095
+ nk_i16_t *max_value_ptr, nk_size_t *max_index_ptr) {
1096
+ int16x8_t min_i16x8 = vdupq_n_s16(NK_I16_MAX), max_i16x8 = vdupq_n_s16(NK_I16_MIN);
1097
+ uint16x8_t min_iter_u16x8 = vdupq_n_u16(0), max_iter_u16x8 = vdupq_n_u16(0);
1098
+ uint16x8_t iter_u16x8 = vdupq_n_u16(0), one_u16x8 = vdupq_n_u16(1);
1099
+ nk_size_t idx = 0;
1100
+ for (; idx + 8 <= count; idx += 8) {
1101
+ int16x8_t data_i16x8 = vld1q_s16(data_ptr + idx);
1102
+ uint16x8_t less_u16x8 = vcltq_s16(data_i16x8, min_i16x8);
1103
+ uint16x8_t greater_u16x8 = vcgtq_s16(data_i16x8, max_i16x8);
1104
+ min_i16x8 = vbslq_s16(less_u16x8, data_i16x8, min_i16x8);
1105
+ max_i16x8 = vbslq_s16(greater_u16x8, data_i16x8, max_i16x8);
1106
+ min_iter_u16x8 = vbslq_u16(less_u16x8, iter_u16x8, min_iter_u16x8);
1107
+ max_iter_u16x8 = vbslq_u16(greater_u16x8, iter_u16x8, max_iter_u16x8);
1108
+ iter_u16x8 = vaddq_u16(iter_u16x8, one_u16x8);
1109
+ }
1110
+ nk_size_t remaining = count - idx;
1111
+ if (remaining > 0) {
1112
+ nk_b128_vec_t tail_vec;
1113
+ nk_partial_load_b16x8_serial_(data_ptr + idx, &tail_vec, remaining);
1114
+ uint16x8_t lane_indices_u16x8 = vcombine_u16(vreinterpret_u16_u64(vcreate_u64(0x0003000200010000ULL)),
1115
+ vreinterpret_u16_u64(vcreate_u64(0x0007000600050004ULL)));
1116
+ uint16x8_t valid_u16x8 = vcltq_u16(lane_indices_u16x8, vdupq_n_u16((nk_u16_t)remaining));
1117
+ int16x8_t data_for_min_i16x8 = vbslq_s16(valid_u16x8, tail_vec.i16x8, vdupq_n_s16(NK_I16_MAX));
1118
+ int16x8_t data_for_max_i16x8 = vbslq_s16(valid_u16x8, tail_vec.i16x8, vdupq_n_s16(NK_I16_MIN));
1119
+ uint16x8_t less_u16x8 = vcltq_s16(data_for_min_i16x8, min_i16x8);
1120
+ uint16x8_t greater_u16x8 = vcgtq_s16(data_for_max_i16x8, max_i16x8);
1121
+ min_i16x8 = vbslq_s16(less_u16x8, data_for_min_i16x8, min_i16x8);
1122
+ max_i16x8 = vbslq_s16(greater_u16x8, data_for_max_i16x8, max_i16x8);
1123
+ min_iter_u16x8 = vbslq_u16(less_u16x8, iter_u16x8, min_iter_u16x8);
1124
+ max_iter_u16x8 = vbslq_u16(greater_u16x8, iter_u16x8, max_iter_u16x8);
1125
+ }
1126
+ nk_i16_t min_value = vminvq_s16(min_i16x8), max_value = vmaxvq_s16(max_i16x8);
1127
+ uint16x8_t min_value_match_u16x8 = vceqq_s16(min_i16x8, vdupq_n_s16(min_value));
1128
+ uint16x8_t masked_min_iter_u16x8 = vbslq_u16(min_value_match_u16x8, min_iter_u16x8, vdupq_n_u16(0xFFFF));
1129
+ nk_u16_t earliest_min_cycle = vminvq_u16(masked_min_iter_u16x8);
1130
+ uint16x8_t max_value_match_u16x8 = vceqq_s16(max_i16x8, vdupq_n_s16(max_value));
1131
+ uint16x8_t masked_max_iter_u16x8 = vbslq_u16(max_value_match_u16x8, max_iter_u16x8, vdupq_n_u16(0xFFFF));
1132
+ nk_u16_t earliest_max_cycle = vminvq_u16(masked_max_iter_u16x8);
1133
+ uint16x8_t lane_indices_u16x8 = vcombine_u16(vreinterpret_u16_u64(vcreate_u64(0x0003000200010000ULL)),
1134
+ vreinterpret_u16_u64(vcreate_u64(0x0007000600050004ULL)));
1135
+ uint16x8_t min_cycle_match_u16x8 = vceqq_u16(min_iter_u16x8, vdupq_n_u16(earliest_min_cycle));
1136
+ uint16x8_t min_both_match_u16x8 = vandq_u16(min_value_match_u16x8, min_cycle_match_u16x8);
1137
+ uint16x8_t min_masked_lanes_u16x8 = vbslq_u16(min_both_match_u16x8, lane_indices_u16x8, vdupq_n_u16(0xFFFF));
1138
+ nk_u16_t min_lane_offset = vminvq_u16(min_masked_lanes_u16x8);
1139
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 8 + (nk_size_t)min_lane_offset;
1140
+ uint16x8_t max_cycle_match_u16x8 = vceqq_u16(max_iter_u16x8, vdupq_n_u16(earliest_max_cycle));
1141
+ uint16x8_t max_both_match_u16x8 = vandq_u16(max_value_match_u16x8, max_cycle_match_u16x8);
1142
+ uint16x8_t max_masked_lanes_u16x8 = vbslq_u16(max_both_match_u16x8, lane_indices_u16x8, vdupq_n_u16(0xFFFF));
1143
+ nk_u16_t max_lane_offset = vminvq_u16(max_masked_lanes_u16x8);
1144
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 8 + (nk_size_t)max_lane_offset;
1145
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
1146
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
1147
+ }
1148
+
1149
+ NK_INTERNAL void nk_reduce_minmax_i16_neon_strided_( //
1150
+ nk_i16_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
1151
+ nk_i16_t *min_value_ptr, nk_size_t *min_index_ptr, //
1152
+ nk_i16_t *max_value_ptr, nk_size_t *max_index_ptr) {
1153
+ int16x8_t min_i16x8 = vdupq_n_s16(NK_I16_MAX), max_i16x8 = vdupq_n_s16(NK_I16_MIN);
1154
+ uint16x8_t min_iter_u16x8 = vdupq_n_u16(0), max_iter_u16x8 = vdupq_n_u16(0);
1155
+ uint16x8_t iter_u16x8 = vdupq_n_u16(0), one_u16x8 = vdupq_n_u16(1);
1156
+ uint16x8_t lane_indices_u16x8 = vcombine_u16(vreinterpret_u16_u64(vcreate_u64(0x0003000200010000ULL)),
1157
+ vreinterpret_u16_u64(vcreate_u64(0x0007000600050004ULL)));
1158
+ nk_size_t idx = 0;
1159
+ int16x8_t data_for_min_i16x8, data_for_max_i16x8;
1160
+
1161
+ nk_reduce_minmax_i16_neon_cycle:
1162
+ if (stride_elements == 2 && idx + 8 < count) {
1163
+ int16x8x2_t loaded = vld2q_s16(data_ptr + idx * 2);
1164
+ data_for_min_i16x8 = loaded.val[0];
1165
+ data_for_max_i16x8 = loaded.val[0];
1166
+ idx += 8;
1167
+ }
1168
+ else if (stride_elements == 3 && idx + 8 < count) {
1169
+ int16x8x3_t loaded = vld3q_s16(data_ptr + idx * 3);
1170
+ data_for_min_i16x8 = loaded.val[0];
1171
+ data_for_max_i16x8 = loaded.val[0];
1172
+ idx += 8;
1173
+ }
1174
+ else if (stride_elements == 4 && idx + 8 < count) {
1175
+ int16x8x4_t loaded = vld4q_s16(data_ptr + idx * 4);
1176
+ data_for_min_i16x8 = loaded.val[0];
1177
+ data_for_max_i16x8 = loaded.val[0];
1178
+ idx += 8;
1179
+ }
1180
+ else if (idx < count) {
1181
+ nk_b128_vec_t tail_vec;
1182
+ nk_strided_load_b16x8_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
1183
+ uint16x8_t valid_u16x8 = vcltq_u16(lane_indices_u16x8, vdupq_n_u16((nk_u16_t)(count - idx)));
1184
+ data_for_min_i16x8 = vbslq_s16(valid_u16x8, tail_vec.i16x8, min_i16x8);
1185
+ data_for_max_i16x8 = vbslq_s16(valid_u16x8, tail_vec.i16x8, max_i16x8);
1186
+ idx = count;
1187
+ }
1188
+ else {
1189
+ nk_i16_t min_value = vminvq_s16(min_i16x8), max_value = vmaxvq_s16(max_i16x8);
1190
+ uint16x8_t min_value_match_u16x8 = vceqq_s16(min_i16x8, vdupq_n_s16(min_value));
1191
+ uint16x8_t masked_min_iter_u16x8 = vbslq_u16(min_value_match_u16x8, min_iter_u16x8, vdupq_n_u16(0xFFFF));
1192
+ nk_u16_t earliest_min_cycle = vminvq_u16(masked_min_iter_u16x8);
1193
+ uint16x8_t max_value_match_u16x8 = vceqq_s16(max_i16x8, vdupq_n_s16(max_value));
1194
+ uint16x8_t masked_max_iter_u16x8 = vbslq_u16(max_value_match_u16x8, max_iter_u16x8, vdupq_n_u16(0xFFFF));
1195
+ nk_u16_t earliest_max_cycle = vminvq_u16(masked_max_iter_u16x8);
1196
+ uint16x8_t min_cycle_match_u16x8 = vceqq_u16(min_iter_u16x8, vdupq_n_u16(earliest_min_cycle));
1197
+ uint16x8_t min_both_match_u16x8 = vandq_u16(min_value_match_u16x8, min_cycle_match_u16x8);
1198
+ uint16x8_t min_masked_lanes_u16x8 = vbslq_u16(min_both_match_u16x8, lane_indices_u16x8, vdupq_n_u16(0xFFFF));
1199
+ nk_u16_t min_lane_offset = vminvq_u16(min_masked_lanes_u16x8);
1200
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 8 + (nk_size_t)min_lane_offset;
1201
+ uint16x8_t max_cycle_match_u16x8 = vceqq_u16(max_iter_u16x8, vdupq_n_u16(earliest_max_cycle));
1202
+ uint16x8_t max_both_match_u16x8 = vandq_u16(max_value_match_u16x8, max_cycle_match_u16x8);
1203
+ uint16x8_t max_masked_lanes_u16x8 = vbslq_u16(max_both_match_u16x8, lane_indices_u16x8, vdupq_n_u16(0xFFFF));
1204
+ nk_u16_t max_lane_offset = vminvq_u16(max_masked_lanes_u16x8);
1205
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 8 + (nk_size_t)max_lane_offset;
1206
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
1207
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
1208
+ return;
1209
+ }
1210
+
1211
+ // Shared update body
1212
+ uint16x8_t less_u16x8 = vcltq_s16(data_for_min_i16x8, min_i16x8);
1213
+ uint16x8_t greater_u16x8 = vcgtq_s16(data_for_max_i16x8, max_i16x8);
1214
+ min_i16x8 = vbslq_s16(less_u16x8, data_for_min_i16x8, min_i16x8);
1215
+ max_i16x8 = vbslq_s16(greater_u16x8, data_for_max_i16x8, max_i16x8);
1216
+ min_iter_u16x8 = vbslq_u16(less_u16x8, iter_u16x8, min_iter_u16x8);
1217
+ max_iter_u16x8 = vbslq_u16(greater_u16x8, iter_u16x8, max_iter_u16x8);
1218
+ iter_u16x8 = vaddq_u16(iter_u16x8, one_u16x8);
1219
+ goto nk_reduce_minmax_i16_neon_cycle;
1220
+ }
1221
+
1222
+ NK_PUBLIC void nk_reduce_minmax_i16_neon( //
1223
+ nk_i16_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
1224
+ nk_i16_t *min_value_ptr, nk_size_t *min_index_ptr, //
1225
+ nk_i16_t *max_value_ptr, nk_size_t *max_index_ptr) {
1226
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_i16_t);
1227
+ int aligned = (stride_bytes % sizeof(nk_i16_t) == 0);
1228
+ if (count == 0)
1229
+ *min_value_ptr = NK_I16_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_I16_MIN,
1230
+ *max_index_ptr = NK_SIZE_MAX;
1231
+ else if (!aligned)
1232
+ nk_reduce_minmax_i16_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
1233
+ max_index_ptr);
1234
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 8) {
1235
+ nk_size_t left_count = count / 2;
1236
+ nk_i16_t left_min_value, right_min_value, left_max_value, right_max_value;
1237
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
1238
+ nk_reduce_minmax_i16_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index, &left_max_value,
1239
+ &left_max_index);
1240
+ nk_reduce_minmax_i16_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
1241
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
1242
+ if (right_min_value < left_min_value)
1243
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
1244
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
1245
+ if (right_max_value > left_max_value)
1246
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
1247
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
1248
+ }
1249
+ else if (stride_elements == 1)
1250
+ nk_reduce_minmax_i16_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
1251
+ max_index_ptr);
1252
+ else if (stride_elements <= 4)
1253
+ nk_reduce_minmax_i16_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
1254
+ max_value_ptr, max_index_ptr);
1255
+ else
1256
+ nk_reduce_minmax_i16_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
1257
+ max_index_ptr);
1258
+ }
1259
+
1260
+ NK_INTERNAL void nk_reduce_moments_u16_neon_contiguous_( //
1261
+ nk_u16_t const *data_ptr, nk_size_t count, //
1262
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1263
+ uint64x2_t sum_u64x2 = vdupq_n_u64(0);
1264
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
1265
+ nk_size_t idx = 0;
1266
+ for (; idx + 8 <= count; idx += 8) {
1267
+ uint16x8_t data_u16x8 = vld1q_u16(data_ptr + idx);
1268
+ uint32x4_t sum32_u32x4 = vpaddlq_u16(data_u16x8);
1269
+ sum_u64x2 = vaddq_u64(sum_u64x2, vpaddlq_u32(sum32_u32x4));
1270
+ uint32x4_t sq_low_u32x4 = vmull_u16(vget_low_u16(data_u16x8), vget_low_u16(data_u16x8));
1271
+ uint32x4_t sq_high_u32x4 = vmull_high_u16(data_u16x8, data_u16x8);
1272
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(sq_low_u32x4));
1273
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(sq_high_u32x4));
1274
+ }
1275
+ nk_u64_t sum = vgetq_lane_u64(sum_u64x2, 0) + vgetq_lane_u64(sum_u64x2, 1);
1276
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
1277
+ for (; idx < count; ++idx) {
1278
+ nk_u64_t value = (nk_u64_t)data_ptr[idx];
1279
+ sum += value, sumsq += value * value;
1280
+ }
1281
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
1282
+ }
1283
+
1284
+ NK_INTERNAL void nk_reduce_moments_u16_neon_strided_( //
1285
+ nk_u16_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
1286
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1287
+ uint64x2_t sum_u64x2 = vdupq_n_u64(0);
1288
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
1289
+ nk_size_t idx = 0;
1290
+
1291
+ if (stride_elements == 2) {
1292
+ for (; idx + 8 < count; idx += 8) {
1293
+ uint16x8x2_t loaded_u16x8x2 = vld2q_u16(data_ptr + idx * 2);
1294
+ uint16x8_t data_u16x8 = loaded_u16x8x2.val[0];
1295
+ uint32x4_t widened_sum_u32x4 = vpaddlq_u16(data_u16x8);
1296
+ sum_u64x2 = vaddq_u64(sum_u64x2, vpaddlq_u32(widened_sum_u32x4));
1297
+ uint32x4_t sq_low_u32x4 = vmull_u16(vget_low_u16(data_u16x8), vget_low_u16(data_u16x8));
1298
+ uint32x4_t sq_high_u32x4 = vmull_high_u16(data_u16x8, data_u16x8);
1299
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(sq_low_u32x4));
1300
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(sq_high_u32x4));
1301
+ }
1302
+ }
1303
+ else if (stride_elements == 3) {
1304
+ for (; idx + 8 < count; idx += 8) {
1305
+ uint16x8x3_t loaded_u16x8x3 = vld3q_u16(data_ptr + idx * 3);
1306
+ uint16x8_t data_u16x8 = loaded_u16x8x3.val[0];
1307
+ uint32x4_t widened_sum_u32x4 = vpaddlq_u16(data_u16x8);
1308
+ sum_u64x2 = vaddq_u64(sum_u64x2, vpaddlq_u32(widened_sum_u32x4));
1309
+ uint32x4_t sq_low_u32x4 = vmull_u16(vget_low_u16(data_u16x8), vget_low_u16(data_u16x8));
1310
+ uint32x4_t sq_high_u32x4 = vmull_high_u16(data_u16x8, data_u16x8);
1311
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(sq_low_u32x4));
1312
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(sq_high_u32x4));
1313
+ }
1314
+ }
1315
+ else {
1316
+ for (; idx + 8 < count; idx += 8) {
1317
+ uint16x8x4_t loaded_u16x8x4 = vld4q_u16(data_ptr + idx * 4);
1318
+ uint16x8_t data_u16x8 = loaded_u16x8x4.val[0];
1319
+ uint32x4_t widened_sum_u32x4 = vpaddlq_u16(data_u16x8);
1320
+ sum_u64x2 = vaddq_u64(sum_u64x2, vpaddlq_u32(widened_sum_u32x4));
1321
+ uint32x4_t sq_low_u32x4 = vmull_u16(vget_low_u16(data_u16x8), vget_low_u16(data_u16x8));
1322
+ uint32x4_t sq_high_u32x4 = vmull_high_u16(data_u16x8, data_u16x8);
1323
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(sq_low_u32x4));
1324
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(sq_high_u32x4));
1325
+ }
1326
+ }
1327
+
1328
+ nk_u64_t sum = vgetq_lane_u64(sum_u64x2, 0) + vgetq_lane_u64(sum_u64x2, 1);
1329
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
1330
+ for (; idx < count; ++idx) {
1331
+ nk_u64_t value = (nk_u64_t)data_ptr[idx * stride_elements];
1332
+ sum += value, sumsq += value * value;
1333
+ }
1334
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
1335
+ }
1336
+
1337
+ NK_PUBLIC void nk_reduce_moments_u16_neon( //
1338
+ nk_u16_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
1339
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1340
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_u16_t);
1341
+ int aligned = (stride_bytes % sizeof(nk_u16_t) == 0);
1342
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
1343
+ else if (!aligned) nk_reduce_moments_u16_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
1344
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 8) {
1345
+ nk_size_t left_count = count / 2;
1346
+ nk_u64_t left_sum, left_sumsq, right_sum, right_sumsq;
1347
+ nk_reduce_moments_u16_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
1348
+ nk_reduce_moments_u16_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
1349
+ &right_sum, &right_sumsq);
1350
+ *sum_ptr = nk_u64_saturating_add_serial(left_sum, right_sum);
1351
+ *sumsq_ptr = nk_u64_saturating_add_serial(left_sumsq, right_sumsq);
1352
+ }
1353
+ else if (stride_elements == 1) nk_reduce_moments_u16_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
1354
+ else if (stride_elements <= 4)
1355
+ nk_reduce_moments_u16_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
1356
+ else nk_reduce_moments_u16_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
1357
+ }
1358
+
1359
+ NK_INTERNAL void nk_reduce_minmax_u16_neon_contiguous_( //
1360
+ nk_u16_t const *data_ptr, nk_size_t count, //
1361
+ nk_u16_t *min_value_ptr, nk_size_t *min_index_ptr, //
1362
+ nk_u16_t *max_value_ptr, nk_size_t *max_index_ptr) {
1363
+ uint16x8_t min_u16x8 = vdupq_n_u16(NK_U16_MAX), max_u16x8 = vdupq_n_u16(0);
1364
+ uint16x8_t min_iter_u16x8 = vdupq_n_u16(0), max_iter_u16x8 = vdupq_n_u16(0);
1365
+ uint16x8_t iter_u16x8 = vdupq_n_u16(0), one_u16x8 = vdupq_n_u16(1);
1366
+ nk_size_t idx = 0;
1367
+ for (; idx + 8 <= count; idx += 8) {
1368
+ uint16x8_t data_u16x8 = vld1q_u16(data_ptr + idx);
1369
+ uint16x8_t less_u16x8 = vcltq_u16(data_u16x8, min_u16x8);
1370
+ uint16x8_t greater_u16x8 = vcgtq_u16(data_u16x8, max_u16x8);
1371
+ min_u16x8 = vbslq_u16(less_u16x8, data_u16x8, min_u16x8);
1372
+ max_u16x8 = vbslq_u16(greater_u16x8, data_u16x8, max_u16x8);
1373
+ min_iter_u16x8 = vbslq_u16(less_u16x8, iter_u16x8, min_iter_u16x8);
1374
+ max_iter_u16x8 = vbslq_u16(greater_u16x8, iter_u16x8, max_iter_u16x8);
1375
+ iter_u16x8 = vaddq_u16(iter_u16x8, one_u16x8);
1376
+ }
1377
+ nk_size_t remaining = count - idx;
1378
+ if (remaining > 0) {
1379
+ nk_b128_vec_t tail_vec;
1380
+ nk_partial_load_b16x8_serial_(data_ptr + idx, &tail_vec, remaining);
1381
+ uint16x8_t lane_indices_u16x8 = vcombine_u16(vreinterpret_u16_u64(vcreate_u64(0x0003000200010000ULL)),
1382
+ vreinterpret_u16_u64(vcreate_u64(0x0007000600050004ULL)));
1383
+ uint16x8_t valid_u16x8 = vcltq_u16(lane_indices_u16x8, vdupq_n_u16((nk_u16_t)remaining));
1384
+ uint16x8_t data_for_min_u16x8 = vbslq_u16(valid_u16x8, tail_vec.u16x8, vdupq_n_u16(NK_U16_MAX));
1385
+ uint16x8_t data_for_max_u16x8 = vbslq_u16(valid_u16x8, tail_vec.u16x8, vdupq_n_u16(0));
1386
+ uint16x8_t less_u16x8 = vcltq_u16(data_for_min_u16x8, min_u16x8);
1387
+ uint16x8_t greater_u16x8 = vcgtq_u16(data_for_max_u16x8, max_u16x8);
1388
+ min_u16x8 = vbslq_u16(less_u16x8, data_for_min_u16x8, min_u16x8);
1389
+ max_u16x8 = vbslq_u16(greater_u16x8, data_for_max_u16x8, max_u16x8);
1390
+ min_iter_u16x8 = vbslq_u16(less_u16x8, iter_u16x8, min_iter_u16x8);
1391
+ max_iter_u16x8 = vbslq_u16(greater_u16x8, iter_u16x8, max_iter_u16x8);
1392
+ }
1393
+ nk_u16_t min_value = vminvq_u16(min_u16x8), max_value = vmaxvq_u16(max_u16x8);
1394
+ uint16x8_t min_value_match_u16x8 = vceqq_u16(min_u16x8, vdupq_n_u16(min_value));
1395
+ uint16x8_t masked_min_iter_u16x8 = vbslq_u16(min_value_match_u16x8, min_iter_u16x8, vdupq_n_u16(0xFFFF));
1396
+ nk_u16_t earliest_min_cycle = vminvq_u16(masked_min_iter_u16x8);
1397
+ uint16x8_t max_value_match_u16x8 = vceqq_u16(max_u16x8, vdupq_n_u16(max_value));
1398
+ uint16x8_t masked_max_iter_u16x8 = vbslq_u16(max_value_match_u16x8, max_iter_u16x8, vdupq_n_u16(0xFFFF));
1399
+ nk_u16_t earliest_max_cycle = vminvq_u16(masked_max_iter_u16x8);
1400
+ uint16x8_t lane_indices_u16x8 = vcombine_u16(vreinterpret_u16_u64(vcreate_u64(0x0003000200010000ULL)),
1401
+ vreinterpret_u16_u64(vcreate_u64(0x0007000600050004ULL)));
1402
+ uint16x8_t min_cycle_match_u16x8 = vceqq_u16(min_iter_u16x8, vdupq_n_u16(earliest_min_cycle));
1403
+ uint16x8_t min_both_match_u16x8 = vandq_u16(min_value_match_u16x8, min_cycle_match_u16x8);
1404
+ uint16x8_t min_masked_lanes_u16x8 = vbslq_u16(min_both_match_u16x8, lane_indices_u16x8, vdupq_n_u16(0xFFFF));
1405
+ nk_u16_t min_lane_offset = vminvq_u16(min_masked_lanes_u16x8);
1406
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 8 + (nk_size_t)min_lane_offset;
1407
+ uint16x8_t max_cycle_match_u16x8 = vceqq_u16(max_iter_u16x8, vdupq_n_u16(earliest_max_cycle));
1408
+ uint16x8_t max_both_match_u16x8 = vandq_u16(max_value_match_u16x8, max_cycle_match_u16x8);
1409
+ uint16x8_t max_masked_lanes_u16x8 = vbslq_u16(max_both_match_u16x8, lane_indices_u16x8, vdupq_n_u16(0xFFFF));
1410
+ nk_u16_t max_lane_offset = vminvq_u16(max_masked_lanes_u16x8);
1411
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 8 + (nk_size_t)max_lane_offset;
1412
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
1413
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
1414
+ }
1415
+
1416
+ NK_INTERNAL void nk_reduce_minmax_u16_neon_strided_( //
1417
+ nk_u16_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
1418
+ nk_u16_t *min_value_ptr, nk_size_t *min_index_ptr, //
1419
+ nk_u16_t *max_value_ptr, nk_size_t *max_index_ptr) {
1420
+ uint16x8_t min_u16x8 = vdupq_n_u16(NK_U16_MAX), max_u16x8 = vdupq_n_u16(0);
1421
+ uint16x8_t min_iter_u16x8 = vdupq_n_u16(0), max_iter_u16x8 = vdupq_n_u16(0);
1422
+ uint16x8_t iter_u16x8 = vdupq_n_u16(0), one_u16x8 = vdupq_n_u16(1);
1423
+ uint16x8_t lane_indices_u16x8 = vcombine_u16(vreinterpret_u16_u64(vcreate_u64(0x0003000200010000ULL)),
1424
+ vreinterpret_u16_u64(vcreate_u64(0x0007000600050004ULL)));
1425
+ nk_size_t idx = 0;
1426
+ uint16x8_t data_for_min_u16x8, data_for_max_u16x8;
1427
+
1428
+ nk_reduce_minmax_u16_neon_cycle:
1429
+ if (stride_elements == 2 && idx + 8 < count) {
1430
+ uint16x8x2_t loaded = vld2q_u16((nk_u16_t const *)data_ptr + idx * 2);
1431
+ data_for_min_u16x8 = loaded.val[0];
1432
+ data_for_max_u16x8 = loaded.val[0];
1433
+ idx += 8;
1434
+ }
1435
+ else if (stride_elements == 3 && idx + 8 < count) {
1436
+ uint16x8x3_t loaded = vld3q_u16((nk_u16_t const *)data_ptr + idx * 3);
1437
+ data_for_min_u16x8 = loaded.val[0];
1438
+ data_for_max_u16x8 = loaded.val[0];
1439
+ idx += 8;
1440
+ }
1441
+ else if (stride_elements == 4 && idx + 8 < count) {
1442
+ uint16x8x4_t loaded = vld4q_u16((nk_u16_t const *)data_ptr + idx * 4);
1443
+ data_for_min_u16x8 = loaded.val[0];
1444
+ data_for_max_u16x8 = loaded.val[0];
1445
+ idx += 8;
1446
+ }
1447
+ else if (idx < count) {
1448
+ nk_b128_vec_t tail_vec;
1449
+ nk_strided_load_b16x8_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
1450
+ uint16x8_t valid_u16x8 = vcltq_u16(lane_indices_u16x8, vdupq_n_u16((nk_u16_t)(count - idx)));
1451
+ data_for_min_u16x8 = vbslq_u16(valid_u16x8, tail_vec.u16x8, min_u16x8);
1452
+ data_for_max_u16x8 = vbslq_u16(valid_u16x8, tail_vec.u16x8, max_u16x8);
1453
+ idx = count;
1454
+ }
1455
+ else {
1456
+ nk_u16_t min_value = vminvq_u16(min_u16x8), max_value = vmaxvq_u16(max_u16x8);
1457
+ uint16x8_t min_value_match_u16x8 = vceqq_u16(min_u16x8, vdupq_n_u16(min_value));
1458
+ uint16x8_t masked_min_iter_u16x8 = vbslq_u16(min_value_match_u16x8, min_iter_u16x8, vdupq_n_u16(0xFFFF));
1459
+ nk_u16_t earliest_min_cycle = vminvq_u16(masked_min_iter_u16x8);
1460
+ uint16x8_t max_value_match_u16x8 = vceqq_u16(max_u16x8, vdupq_n_u16(max_value));
1461
+ uint16x8_t masked_max_iter_u16x8 = vbslq_u16(max_value_match_u16x8, max_iter_u16x8, vdupq_n_u16(0xFFFF));
1462
+ nk_u16_t earliest_max_cycle = vminvq_u16(masked_max_iter_u16x8);
1463
+ uint16x8_t min_cycle_match_u16x8 = vceqq_u16(min_iter_u16x8, vdupq_n_u16(earliest_min_cycle));
1464
+ uint16x8_t min_both_match_u16x8 = vandq_u16(min_value_match_u16x8, min_cycle_match_u16x8);
1465
+ uint16x8_t min_masked_lanes_u16x8 = vbslq_u16(min_both_match_u16x8, lane_indices_u16x8, vdupq_n_u16(0xFFFF));
1466
+ nk_u16_t min_lane_offset = vminvq_u16(min_masked_lanes_u16x8);
1467
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 8 + (nk_size_t)min_lane_offset;
1468
+ uint16x8_t max_cycle_match_u16x8 = vceqq_u16(max_iter_u16x8, vdupq_n_u16(earliest_max_cycle));
1469
+ uint16x8_t max_both_match_u16x8 = vandq_u16(max_value_match_u16x8, max_cycle_match_u16x8);
1470
+ uint16x8_t max_masked_lanes_u16x8 = vbslq_u16(max_both_match_u16x8, lane_indices_u16x8, vdupq_n_u16(0xFFFF));
1471
+ nk_u16_t max_lane_offset = vminvq_u16(max_masked_lanes_u16x8);
1472
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 8 + (nk_size_t)max_lane_offset;
1473
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
1474
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
1475
+ return;
1476
+ }
1477
+
1478
+ // Shared update body
1479
+ uint16x8_t less_u16x8 = vcltq_u16(data_for_min_u16x8, min_u16x8);
1480
+ uint16x8_t greater_u16x8 = vcgtq_u16(data_for_max_u16x8, max_u16x8);
1481
+ min_u16x8 = vbslq_u16(less_u16x8, data_for_min_u16x8, min_u16x8);
1482
+ max_u16x8 = vbslq_u16(greater_u16x8, data_for_max_u16x8, max_u16x8);
1483
+ min_iter_u16x8 = vbslq_u16(less_u16x8, iter_u16x8, min_iter_u16x8);
1484
+ max_iter_u16x8 = vbslq_u16(greater_u16x8, iter_u16x8, max_iter_u16x8);
1485
+ iter_u16x8 = vaddq_u16(iter_u16x8, one_u16x8);
1486
+ goto nk_reduce_minmax_u16_neon_cycle;
1487
+ }
1488
+
1489
+ NK_PUBLIC void nk_reduce_minmax_u16_neon( //
1490
+ nk_u16_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
1491
+ nk_u16_t *min_value_ptr, nk_size_t *min_index_ptr, //
1492
+ nk_u16_t *max_value_ptr, nk_size_t *max_index_ptr) {
1493
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_u16_t);
1494
+ int aligned = (stride_bytes % sizeof(nk_u16_t) == 0);
1495
+ if (count == 0)
1496
+ *min_value_ptr = NK_U16_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = 0, *max_index_ptr = NK_SIZE_MAX;
1497
+ else if (!aligned)
1498
+ nk_reduce_minmax_u16_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
1499
+ max_index_ptr);
1500
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 8) {
1501
+ nk_size_t left_count = count / 2;
1502
+ nk_u16_t left_min_value, right_min_value, left_max_value, right_max_value;
1503
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
1504
+ nk_reduce_minmax_u16_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index, &left_max_value,
1505
+ &left_max_index);
1506
+ nk_reduce_minmax_u16_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
1507
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
1508
+ if (right_min_value < left_min_value)
1509
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
1510
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
1511
+ if (right_max_value > left_max_value)
1512
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
1513
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
1514
+ }
1515
+ else if (stride_elements == 1)
1516
+ nk_reduce_minmax_u16_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
1517
+ max_index_ptr);
1518
+ else if (stride_elements <= 4)
1519
+ nk_reduce_minmax_u16_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
1520
+ max_value_ptr, max_index_ptr);
1521
+ else
1522
+ nk_reduce_minmax_u16_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
1523
+ max_index_ptr);
1524
+ }
1525
+
1526
+ NK_INTERNAL void nk_reduce_moments_i32_neon_contiguous_( //
1527
+ nk_i32_t const *data_ptr, nk_size_t count, //
1528
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1529
+ // 128-bit accumulation: lower (u64) + upper (i64) per lane
1530
+ uint64x2_t sum_low_u64x2 = vdupq_n_u64(0);
1531
+ int64x2_t sum_high_i64x2 = vdupq_n_s64(0);
1532
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
1533
+ int sumsq_overflow = 0;
1534
+ // XOR sign-bit trick for unsigned u64 compare on NEON
1535
+ int64x2_t sign_bit_i64x2 = vdupq_n_s64((nk_i64_t)0x8000000000000000ULL);
1536
+ nk_size_t idx = 0;
1537
+ for (; idx + 4 <= count; idx += 4) {
1538
+ int32x4_t data_i32x4 = vld1q_s32(data_ptr + idx);
1539
+ // Sum: widen i32->i64 and accumulate with carry detection
1540
+ int64x2_t data_low_i64x2 = vmovl_s32(vget_low_s32(data_i32x4));
1541
+ uint64x2_t before_u64x2 = sum_low_u64x2;
1542
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(data_low_i64x2));
1543
+ int64x2_t result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
1544
+ int64x2_t before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(before_u64x2), sign_bit_i64x2);
1545
+ uint64x2_t carry_u64x2 = vcgtq_s64(before_biased_i64x2, result_biased_i64x2);
1546
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
1547
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, vshrq_n_s64(data_low_i64x2, 63));
1548
+
1549
+ int64x2_t data_high_i64x2 = vmovl_high_s32(data_i32x4);
1550
+ before_u64x2 = sum_low_u64x2;
1551
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(data_high_i64x2));
1552
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
1553
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(before_u64x2), sign_bit_i64x2);
1554
+ carry_u64x2 = vcgtq_s64(before_biased_i64x2, result_biased_i64x2);
1555
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
1556
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, vshrq_n_s64(data_high_i64x2, 63));
1557
+
1558
+ // Sumsq: widening multiply i32*i32 -> i64 (always non-negative for squares)
1559
+ int64x2_t squares_low_i64x2 = vmull_s32(vget_low_s32(data_i32x4), vget_low_s32(data_i32x4));
1560
+ int64x2_t squares_high_i64x2 = vmull_high_s32(data_i32x4, data_i32x4);
1561
+ uint64x2_t sq_before_u64x2 = sumsq_u64x2;
1562
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vreinterpretq_u64_s64(squares_low_i64x2));
1563
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
1564
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
1565
+ sumsq_overflow |=
1566
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
1567
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
1568
+ sq_before_u64x2 = sumsq_u64x2;
1569
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vreinterpretq_u64_s64(squares_high_i64x2));
1570
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
1571
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
1572
+ sumsq_overflow |=
1573
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
1574
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
1575
+ }
1576
+ // Sumsq horizontal saturating reduction
1577
+ nk_u64_t sumsq;
1578
+ if (sumsq_overflow) sumsq = NK_U64_MAX;
1579
+ else sumsq = nk_reduce_sadd_u64x2_neon_(sumsq_u64x2);
1580
+ // Sum: horizontal 128-bit reduction (2 lanes -> scalar)
1581
+ nk_b128_vec_t lower_vec, upper_vec;
1582
+ lower_vec.u64x2 = sum_low_u64x2;
1583
+ upper_vec.i64x2 = sum_high_i64x2;
1584
+ nk_u64_t sum_low = 0;
1585
+ nk_i64_t sum_high = 0;
1586
+ nk_u64_t sum_before = sum_low;
1587
+ sum_low += lower_vec.u64s[0], sum_high += (sum_low < sum_before) + upper_vec.i64s[0];
1588
+ sum_before = sum_low;
1589
+ sum_low += lower_vec.u64s[1], sum_high += (sum_low < sum_before) + upper_vec.i64s[1];
1590
+ // Scalar tail
1591
+ for (; idx < count; ++idx) {
1592
+ nk_i64_t value_i64 = (nk_i64_t)data_ptr[idx];
1593
+ sum_before = sum_low;
1594
+ sum_low += (nk_u64_t)value_i64;
1595
+ if (sum_low < sum_before) sum_high++;
1596
+ sum_high += (value_i64 >> 63);
1597
+ nk_i64_t product = nk_i64_saturating_mul_serial(value_i64, value_i64);
1598
+ nk_u64_t unsigned_product = (nk_u64_t)product;
1599
+ sumsq = nk_u64_saturating_add_serial(sumsq, unsigned_product);
1600
+ }
1601
+ // Clamp 128-bit sum to i64 range
1602
+ nk_i64_t sum_low_signed = (nk_i64_t)sum_low;
1603
+ if (sum_high == (sum_low_signed >> 63)) *sum_ptr = sum_low_signed;
1604
+ else if (sum_high >= 0) *sum_ptr = NK_I64_MAX;
1605
+ else *sum_ptr = NK_I64_MIN;
1606
+ *sumsq_ptr = sumsq;
1607
+ }
1608
+
1609
+ NK_INTERNAL void nk_reduce_moments_i32_neon_strided_( //
1610
+ nk_i32_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
1611
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1612
+ uint64x2_t sum_low_u64x2 = vdupq_n_u64(0);
1613
+ int64x2_t sum_high_i64x2 = vdupq_n_s64(0);
1614
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
1615
+ int sumsq_overflow = 0;
1616
+ int64x2_t sign_bit_i64x2 = vdupq_n_s64((nk_i64_t)0x8000000000000000ULL);
1617
+ nk_size_t idx = 0;
1618
+ if (stride_elements == 2) {
1619
+ for (; idx + 4 < count; idx += 4) {
1620
+ int32x4x2_t loaded_i32x4x2 = vld2q_s32(data_ptr + idx * 2);
1621
+ int32x4_t data_i32x4 = loaded_i32x4x2.val[0];
1622
+ int64x2_t low_i64x2 = vmovl_s32(vget_low_s32(data_i32x4));
1623
+ uint64x2_t before_u64x2 = sum_low_u64x2;
1624
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(low_i64x2));
1625
+ int64x2_t result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
1626
+ int64x2_t before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(before_u64x2), sign_bit_i64x2);
1627
+ uint64x2_t carry_u64x2 = vcgtq_s64(before_biased_i64x2, result_biased_i64x2);
1628
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
1629
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, vshrq_n_s64(low_i64x2, 63));
1630
+ int64x2_t high_i64x2 = vmovl_high_s32(data_i32x4);
1631
+ before_u64x2 = sum_low_u64x2;
1632
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(high_i64x2));
1633
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
1634
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(before_u64x2), sign_bit_i64x2);
1635
+ carry_u64x2 = vcgtq_s64(before_biased_i64x2, result_biased_i64x2);
1636
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
1637
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, vshrq_n_s64(high_i64x2, 63));
1638
+ int64x2_t squares_low_i64x2 = vmull_s32(vget_low_s32(data_i32x4), vget_low_s32(data_i32x4));
1639
+ int64x2_t squares_high_i64x2 = vmull_high_s32(data_i32x4, data_i32x4);
1640
+ uint64x2_t sq_before_u64x2 = sumsq_u64x2;
1641
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vreinterpretq_u64_s64(squares_low_i64x2));
1642
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
1643
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
1644
+ sumsq_overflow |=
1645
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
1646
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
1647
+ sq_before_u64x2 = sumsq_u64x2;
1648
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vreinterpretq_u64_s64(squares_high_i64x2));
1649
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
1650
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
1651
+ sumsq_overflow |=
1652
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
1653
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
1654
+ }
1655
+ }
1656
+ else if (stride_elements == 3) {
1657
+ for (; idx + 4 < count; idx += 4) {
1658
+ int32x4x3_t loaded_i32x4x3 = vld3q_s32(data_ptr + idx * 3);
1659
+ int32x4_t data_i32x4 = loaded_i32x4x3.val[0];
1660
+ int64x2_t low_i64x2 = vmovl_s32(vget_low_s32(data_i32x4));
1661
+ uint64x2_t before_u64x2 = sum_low_u64x2;
1662
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(low_i64x2));
1663
+ int64x2_t result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
1664
+ int64x2_t before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(before_u64x2), sign_bit_i64x2);
1665
+ uint64x2_t carry_u64x2 = vcgtq_s64(before_biased_i64x2, result_biased_i64x2);
1666
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
1667
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, vshrq_n_s64(low_i64x2, 63));
1668
+ int64x2_t high_i64x2 = vmovl_high_s32(data_i32x4);
1669
+ before_u64x2 = sum_low_u64x2;
1670
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(high_i64x2));
1671
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
1672
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(before_u64x2), sign_bit_i64x2);
1673
+ carry_u64x2 = vcgtq_s64(before_biased_i64x2, result_biased_i64x2);
1674
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
1675
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, vshrq_n_s64(high_i64x2, 63));
1676
+ int64x2_t squares_low_i64x2 = vmull_s32(vget_low_s32(data_i32x4), vget_low_s32(data_i32x4));
1677
+ int64x2_t squares_high_i64x2 = vmull_high_s32(data_i32x4, data_i32x4);
1678
+ uint64x2_t sq_before_u64x2 = sumsq_u64x2;
1679
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vreinterpretq_u64_s64(squares_low_i64x2));
1680
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
1681
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
1682
+ sumsq_overflow |=
1683
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
1684
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
1685
+ sq_before_u64x2 = sumsq_u64x2;
1686
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vreinterpretq_u64_s64(squares_high_i64x2));
1687
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
1688
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
1689
+ sumsq_overflow |=
1690
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
1691
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
1692
+ }
1693
+ }
1694
+ else {
1695
+ for (; idx + 4 < count; idx += 4) {
1696
+ int32x4x4_t loaded_i32x4x4 = vld4q_s32(data_ptr + idx * 4);
1697
+ int32x4_t data_i32x4 = loaded_i32x4x4.val[0];
1698
+ int64x2_t low_i64x2 = vmovl_s32(vget_low_s32(data_i32x4));
1699
+ uint64x2_t before_u64x2 = sum_low_u64x2;
1700
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(low_i64x2));
1701
+ int64x2_t result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
1702
+ int64x2_t before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(before_u64x2), sign_bit_i64x2);
1703
+ uint64x2_t carry_u64x2 = vcgtq_s64(before_biased_i64x2, result_biased_i64x2);
1704
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
1705
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, vshrq_n_s64(low_i64x2, 63));
1706
+ int64x2_t high_i64x2 = vmovl_high_s32(data_i32x4);
1707
+ before_u64x2 = sum_low_u64x2;
1708
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(high_i64x2));
1709
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
1710
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(before_u64x2), sign_bit_i64x2);
1711
+ carry_u64x2 = vcgtq_s64(before_biased_i64x2, result_biased_i64x2);
1712
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
1713
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, vshrq_n_s64(high_i64x2, 63));
1714
+ int64x2_t squares_low_i64x2 = vmull_s32(vget_low_s32(data_i32x4), vget_low_s32(data_i32x4));
1715
+ int64x2_t squares_high_i64x2 = vmull_high_s32(data_i32x4, data_i32x4);
1716
+ uint64x2_t sq_before_u64x2 = sumsq_u64x2;
1717
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vreinterpretq_u64_s64(squares_low_i64x2));
1718
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
1719
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
1720
+ sumsq_overflow |=
1721
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
1722
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
1723
+ sq_before_u64x2 = sumsq_u64x2;
1724
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vreinterpretq_u64_s64(squares_high_i64x2));
1725
+ result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
1726
+ before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
1727
+ sumsq_overflow |=
1728
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
1729
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
1730
+ }
1731
+ }
1732
+ nk_u64_t sumsq;
1733
+ if (sumsq_overflow) sumsq = NK_U64_MAX;
1734
+ else sumsq = nk_reduce_sadd_u64x2_neon_(sumsq_u64x2);
1735
+ nk_b128_vec_t lower_vec, upper_vec;
1736
+ lower_vec.u64x2 = sum_low_u64x2;
1737
+ upper_vec.i64x2 = sum_high_i64x2;
1738
+ nk_u64_t sum_low = 0;
1739
+ nk_i64_t sum_high = 0;
1740
+ nk_u64_t sum_before = sum_low;
1741
+ sum_low += lower_vec.u64s[0], sum_high += (sum_low < sum_before) + upper_vec.i64s[0];
1742
+ sum_before = sum_low;
1743
+ sum_low += lower_vec.u64s[1], sum_high += (sum_low < sum_before) + upper_vec.i64s[1];
1744
+ for (; idx < count; ++idx) {
1745
+ nk_i64_t val = (nk_i64_t) * (data_ptr + idx * stride_elements);
1746
+ sum_before = sum_low;
1747
+ sum_low += (nk_u64_t)val;
1748
+ if (sum_low < sum_before) sum_high++;
1749
+ sum_high += (val >> 63);
1750
+ nk_i64_t product = nk_i64_saturating_mul_serial(val, val);
1751
+ nk_u64_t unsigned_product = (nk_u64_t)product;
1752
+ sumsq = nk_u64_saturating_add_serial(sumsq, unsigned_product);
1753
+ }
1754
+ nk_i64_t sum_low_signed = (nk_i64_t)sum_low;
1755
+ if (sum_high == (sum_low_signed >> 63)) *sum_ptr = sum_low_signed;
1756
+ else if (sum_high >= 0) *sum_ptr = NK_I64_MAX;
1757
+ else *sum_ptr = NK_I64_MIN;
1758
+ *sumsq_ptr = sumsq;
1759
+ }
1760
+
1761
+ NK_PUBLIC void nk_reduce_moments_i32_neon( //
1762
+ nk_i32_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
1763
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1764
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_i32_t);
1765
+ int aligned = (stride_bytes % sizeof(nk_i32_t) == 0);
1766
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
1767
+ else if (!aligned) nk_reduce_moments_i32_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
1768
+ else if (stride_elements == 1) nk_reduce_moments_i32_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
1769
+ else if (stride_elements <= 4)
1770
+ nk_reduce_moments_i32_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
1771
+ else nk_reduce_moments_i32_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
1772
+ }
1773
+
1774
+ NK_INTERNAL void nk_reduce_minmax_i32_neon_contiguous_( //
1775
+ nk_i32_t const *data_ptr, nk_size_t count, //
1776
+ nk_i32_t *min_value_ptr, nk_size_t *min_index_ptr, //
1777
+ nk_i32_t *max_value_ptr, nk_size_t *max_index_ptr) {
1778
+ int32x4_t min_i32x4 = vdupq_n_s32(NK_I32_MAX), max_i32x4 = vdupq_n_s32(NK_I32_MIN);
1779
+ uint32x4_t min_iter_u32x4 = vdupq_n_u32(0), max_iter_u32x4 = vdupq_n_u32(0);
1780
+ uint32x4_t iter_u32x4 = vdupq_n_u32(0), one_u32x4 = vdupq_n_u32(1);
1781
+ nk_size_t idx = 0;
1782
+ for (; idx + 4 <= count; idx += 4) {
1783
+ int32x4_t data_i32x4 = vld1q_s32(data_ptr + idx);
1784
+ uint32x4_t less_u32x4 = vcltq_s32(data_i32x4, min_i32x4);
1785
+ uint32x4_t greater_u32x4 = vcgtq_s32(data_i32x4, max_i32x4);
1786
+ min_i32x4 = vbslq_s32(less_u32x4, data_i32x4, min_i32x4);
1787
+ max_i32x4 = vbslq_s32(greater_u32x4, data_i32x4, max_i32x4);
1788
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
1789
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
1790
+ iter_u32x4 = vaddq_u32(iter_u32x4, one_u32x4);
1791
+ }
1792
+ nk_size_t remaining = count - idx;
1793
+ if (remaining > 0) {
1794
+ nk_b128_vec_t tail_vec;
1795
+ nk_partial_load_b32x4_serial_(data_ptr + idx, &tail_vec, remaining);
1796
+ uint32x4_t lane_indices_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
1797
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
1798
+ uint32x4_t valid_u32x4 = vcltq_u32(lane_indices_u32x4, vdupq_n_u32((nk_u32_t)remaining));
1799
+ int32x4_t data_min_i32x4 = vbslq_s32(valid_u32x4, tail_vec.i32x4, vdupq_n_s32(NK_I32_MAX));
1800
+ int32x4_t data_max_i32x4 = vbslq_s32(valid_u32x4, tail_vec.i32x4, vdupq_n_s32(NK_I32_MIN));
1801
+ uint32x4_t less_u32x4 = vcltq_s32(data_min_i32x4, min_i32x4);
1802
+ uint32x4_t greater_u32x4 = vcgtq_s32(data_max_i32x4, max_i32x4);
1803
+ min_i32x4 = vbslq_s32(less_u32x4, data_min_i32x4, min_i32x4);
1804
+ max_i32x4 = vbslq_s32(greater_u32x4, data_max_i32x4, max_i32x4);
1805
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
1806
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
1807
+ }
1808
+ nk_i32_t min_value = vminvq_s32(min_i32x4), max_value = vmaxvq_s32(max_i32x4);
1809
+ uint32x4_t min_value_match_u32x4 = vceqq_s32(min_i32x4, vdupq_n_s32(min_value));
1810
+ uint32x4_t masked_min_iter_u32x4 = vbslq_u32(min_value_match_u32x4, min_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
1811
+ nk_u32_t earliest_min_cycle = vminvq_u32(masked_min_iter_u32x4);
1812
+ uint32x4_t max_value_match_u32x4 = vceqq_s32(max_i32x4, vdupq_n_s32(max_value));
1813
+ uint32x4_t masked_max_iter_u32x4 = vbslq_u32(max_value_match_u32x4, max_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
1814
+ nk_u32_t earliest_max_cycle = vminvq_u32(masked_max_iter_u32x4);
1815
+ uint32x4_t lane_indices_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
1816
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
1817
+ uint32x4_t min_cycle_match_u32x4 = vceqq_u32(min_iter_u32x4, vdupq_n_u32(earliest_min_cycle));
1818
+ uint32x4_t min_both_match_u32x4 = vandq_u32(min_value_match_u32x4, min_cycle_match_u32x4);
1819
+ uint32x4_t min_masked_lanes_u32x4 = vbslq_u32(min_both_match_u32x4, lane_indices_u32x4, vdupq_n_u32(NK_U32_MAX));
1820
+ nk_u32_t min_lane_offset = vminvq_u32(min_masked_lanes_u32x4);
1821
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 4 + (nk_size_t)min_lane_offset;
1822
+ uint32x4_t max_cycle_match_u32x4 = vceqq_u32(max_iter_u32x4, vdupq_n_u32(earliest_max_cycle));
1823
+ uint32x4_t max_both_match_u32x4 = vandq_u32(max_value_match_u32x4, max_cycle_match_u32x4);
1824
+ uint32x4_t max_masked_lanes_u32x4 = vbslq_u32(max_both_match_u32x4, lane_indices_u32x4, vdupq_n_u32(NK_U32_MAX));
1825
+ nk_u32_t max_lane_offset = vminvq_u32(max_masked_lanes_u32x4);
1826
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 4 + (nk_size_t)max_lane_offset;
1827
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
1828
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
1829
+ }
1830
+
1831
+ NK_INTERNAL void nk_reduce_minmax_i32_neon_strided_( //
1832
+ nk_i32_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
1833
+ nk_i32_t *min_value_ptr, nk_size_t *min_index_ptr, //
1834
+ nk_i32_t *max_value_ptr, nk_size_t *max_index_ptr) {
1835
+ int32x4_t min_i32x4 = vdupq_n_s32(NK_I32_MAX), max_i32x4 = vdupq_n_s32(NK_I32_MIN);
1836
+ uint32x4_t min_iter_u32x4 = vdupq_n_u32(0), max_iter_u32x4 = vdupq_n_u32(0);
1837
+ uint32x4_t iter_u32x4 = vdupq_n_u32(0), one_u32x4 = vdupq_n_u32(1);
1838
+ uint32x4_t lane_indices_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
1839
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
1840
+ nk_size_t idx = 0;
1841
+ int32x4_t data_for_min_i32x4, data_for_max_i32x4;
1842
+
1843
+ nk_reduce_minmax_i32_neon_cycle:
1844
+ if (stride_elements == 2 && idx + 4 < count) {
1845
+ int32x4x2_t loaded = vld2q_s32(data_ptr + idx * 2);
1846
+ data_for_min_i32x4 = loaded.val[0];
1847
+ data_for_max_i32x4 = loaded.val[0];
1848
+ idx += 4;
1849
+ }
1850
+ else if (stride_elements == 3 && idx + 4 < count) {
1851
+ int32x4x3_t loaded = vld3q_s32(data_ptr + idx * 3);
1852
+ data_for_min_i32x4 = loaded.val[0];
1853
+ data_for_max_i32x4 = loaded.val[0];
1854
+ idx += 4;
1855
+ }
1856
+ else if (stride_elements == 4 && idx + 4 < count) {
1857
+ int32x4x4_t loaded = vld4q_s32(data_ptr + idx * 4);
1858
+ data_for_min_i32x4 = loaded.val[0];
1859
+ data_for_max_i32x4 = loaded.val[0];
1860
+ idx += 4;
1861
+ }
1862
+ else if (idx < count) {
1863
+ nk_b128_vec_t tail_vec;
1864
+ nk_strided_load_b32x4_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
1865
+ uint32x4_t valid_u32x4 = vcltq_u32(lane_indices_u32x4, vdupq_n_u32((nk_u32_t)(count - idx)));
1866
+ data_for_min_i32x4 = vbslq_s32(valid_u32x4, tail_vec.i32x4, min_i32x4);
1867
+ data_for_max_i32x4 = vbslq_s32(valid_u32x4, tail_vec.i32x4, max_i32x4);
1868
+ idx = count;
1869
+ }
1870
+ else {
1871
+ nk_i32_t min_value = vminvq_s32(min_i32x4), max_value = vmaxvq_s32(max_i32x4);
1872
+ uint32x4_t min_value_match_u32x4 = vceqq_s32(min_i32x4, vdupq_n_s32(min_value));
1873
+ uint32x4_t masked_min_iter_u32x4 = vbslq_u32(min_value_match_u32x4, min_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
1874
+ nk_u32_t earliest_min_cycle = vminvq_u32(masked_min_iter_u32x4);
1875
+ uint32x4_t max_value_match_u32x4 = vceqq_s32(max_i32x4, vdupq_n_s32(max_value));
1876
+ uint32x4_t masked_max_iter_u32x4 = vbslq_u32(max_value_match_u32x4, max_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
1877
+ nk_u32_t earliest_max_cycle = vminvq_u32(masked_max_iter_u32x4);
1878
+ uint32x4_t min_cycle_match_u32x4 = vceqq_u32(min_iter_u32x4, vdupq_n_u32(earliest_min_cycle));
1879
+ uint32x4_t min_both_match_u32x4 = vandq_u32(min_value_match_u32x4, min_cycle_match_u32x4);
1880
+ uint32x4_t min_masked_lanes_u32x4 = vbslq_u32(min_both_match_u32x4, lane_indices_u32x4,
1881
+ vdupq_n_u32(NK_U32_MAX));
1882
+ nk_u32_t min_lane_offset = vminvq_u32(min_masked_lanes_u32x4);
1883
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 4 + (nk_size_t)min_lane_offset;
1884
+ uint32x4_t max_cycle_match_u32x4 = vceqq_u32(max_iter_u32x4, vdupq_n_u32(earliest_max_cycle));
1885
+ uint32x4_t max_both_match_u32x4 = vandq_u32(max_value_match_u32x4, max_cycle_match_u32x4);
1886
+ uint32x4_t max_masked_lanes_u32x4 = vbslq_u32(max_both_match_u32x4, lane_indices_u32x4,
1887
+ vdupq_n_u32(NK_U32_MAX));
1888
+ nk_u32_t max_lane_offset = vminvq_u32(max_masked_lanes_u32x4);
1889
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 4 + (nk_size_t)max_lane_offset;
1890
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
1891
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
1892
+ return;
1893
+ }
1894
+
1895
+ // Shared update body
1896
+ uint32x4_t less_u32x4 = vcltq_s32(data_for_min_i32x4, min_i32x4);
1897
+ uint32x4_t greater_u32x4 = vcgtq_s32(data_for_max_i32x4, max_i32x4);
1898
+ min_i32x4 = vbslq_s32(less_u32x4, data_for_min_i32x4, min_i32x4);
1899
+ max_i32x4 = vbslq_s32(greater_u32x4, data_for_max_i32x4, max_i32x4);
1900
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
1901
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
1902
+ iter_u32x4 = vaddq_u32(iter_u32x4, one_u32x4);
1903
+ goto nk_reduce_minmax_i32_neon_cycle;
1904
+ }
1905
+
1906
+ NK_PUBLIC void nk_reduce_minmax_i32_neon( //
1907
+ nk_i32_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
1908
+ nk_i32_t *min_value_ptr, nk_size_t *min_index_ptr, //
1909
+ nk_i32_t *max_value_ptr, nk_size_t *max_index_ptr) {
1910
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_i32_t);
1911
+ int aligned = (stride_bytes % sizeof(nk_i32_t) == 0);
1912
+ if (count == 0)
1913
+ *min_value_ptr = NK_I32_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_I32_MIN,
1914
+ *max_index_ptr = NK_SIZE_MAX;
1915
+ else if (!aligned)
1916
+ nk_reduce_minmax_i32_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
1917
+ max_index_ptr);
1918
+ else if (count > (nk_size_t)NK_U32_MAX * 4) {
1919
+ nk_size_t left_count = count / 2;
1920
+ nk_i32_t left_min_value, right_min_value, left_max_value, right_max_value;
1921
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
1922
+ nk_reduce_minmax_i32_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index, &left_max_value,
1923
+ &left_max_index);
1924
+ nk_reduce_minmax_i32_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
1925
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
1926
+ if (right_min_value < left_min_value)
1927
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
1928
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
1929
+ if (right_max_value > left_max_value)
1930
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
1931
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
1932
+ }
1933
+ else if (stride_elements == 1)
1934
+ nk_reduce_minmax_i32_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
1935
+ max_index_ptr);
1936
+ else if (stride_elements <= 4)
1937
+ nk_reduce_minmax_i32_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
1938
+ max_value_ptr, max_index_ptr);
1939
+ else
1940
+ nk_reduce_minmax_i32_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
1941
+ max_index_ptr);
1942
+ }
1943
+
1944
+ NK_INTERNAL void nk_reduce_moments_u32_neon_contiguous_( //
1945
+ nk_u32_t const *data_ptr, nk_size_t count, //
1946
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1947
+ uint64x2_t sum_u64x2 = vdupq_n_u64(0);
1948
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
1949
+ nk_size_t idx = 0;
1950
+ for (; idx + 4 <= count; idx += 4) {
1951
+ uint32x4_t data_u32x4 = vld1q_u32(data_ptr + idx);
1952
+ // Widen u32 -> u64 and accumulate sum
1953
+ sum_u64x2 = vaddq_u64(sum_u64x2, vmovl_u32(vget_low_u32(data_u32x4)));
1954
+ sum_u64x2 = vaddq_u64(sum_u64x2, vmovl_high_u32(data_u32x4));
1955
+ // Sumsq: widening multiply u32*u32 -> u64, saturating add
1956
+ uint64x2_t sq_low_u64x2 = vmull_u32(vget_low_u32(data_u32x4), vget_low_u32(data_u32x4));
1957
+ uint64x2_t sq_high_u64x2 = vmull_high_u32(data_u32x4, data_u32x4);
1958
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, sq_low_u64x2);
1959
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, sq_high_u64x2);
1960
+ }
1961
+ nk_u64_t sum = vgetq_lane_u64(sum_u64x2, 0) + vgetq_lane_u64(sum_u64x2, 1);
1962
+ nk_u64_t sumsq = nk_reduce_sadd_u64x2_neon_(sumsq_u64x2);
1963
+ for (; idx < count; ++idx) {
1964
+ nk_u64_t value = (nk_u64_t)data_ptr[idx];
1965
+ sum += value;
1966
+ nk_u64_t product = value * value;
1967
+ sumsq = nk_u64_saturating_add_serial(sumsq, product);
1968
+ }
1969
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
1970
+ }
1971
+
1972
+ NK_INTERNAL void nk_reduce_moments_u32_neon_strided_( //
1973
+ nk_u32_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
1974
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
1975
+ uint64x2_t sum_u64x2 = vdupq_n_u64(0);
1976
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
1977
+ nk_size_t idx = 0;
1978
+ if (stride_elements == 2) {
1979
+ for (; idx + 4 < count; idx += 4) {
1980
+ uint32x4x2_t loaded_u32x4x2 = vld2q_u32(data_ptr + idx * 2);
1981
+ uint32x4_t data_u32x4 = loaded_u32x4x2.val[0];
1982
+ sum_u64x2 = vaddq_u64(sum_u64x2, vmovl_u32(vget_low_u32(data_u32x4)));
1983
+ sum_u64x2 = vaddq_u64(sum_u64x2, vmovl_high_u32(data_u32x4));
1984
+ uint64x2_t squares_low_u64x2 = vmull_u32(vget_low_u32(data_u32x4), vget_low_u32(data_u32x4));
1985
+ uint64x2_t squares_high_u64x2 = vmull_high_u32(data_u32x4, data_u32x4);
1986
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, squares_low_u64x2);
1987
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, squares_high_u64x2);
1988
+ }
1989
+ }
1990
+ else if (stride_elements == 3) {
1991
+ for (; idx + 4 < count; idx += 4) {
1992
+ uint32x4x3_t loaded_u32x4x3 = vld3q_u32(data_ptr + idx * 3);
1993
+ uint32x4_t data_u32x4 = loaded_u32x4x3.val[0];
1994
+ sum_u64x2 = vaddq_u64(sum_u64x2, vmovl_u32(vget_low_u32(data_u32x4)));
1995
+ sum_u64x2 = vaddq_u64(sum_u64x2, vmovl_high_u32(data_u32x4));
1996
+ uint64x2_t squares_low_u64x2 = vmull_u32(vget_low_u32(data_u32x4), vget_low_u32(data_u32x4));
1997
+ uint64x2_t squares_high_u64x2 = vmull_high_u32(data_u32x4, data_u32x4);
1998
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, squares_low_u64x2);
1999
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, squares_high_u64x2);
2000
+ }
2001
+ }
2002
+ else {
2003
+ for (; idx + 4 < count; idx += 4) {
2004
+ uint32x4x4_t loaded_u32x4x4 = vld4q_u32(data_ptr + idx * 4);
2005
+ uint32x4_t data_u32x4 = loaded_u32x4x4.val[0];
2006
+ sum_u64x2 = vaddq_u64(sum_u64x2, vmovl_u32(vget_low_u32(data_u32x4)));
2007
+ sum_u64x2 = vaddq_u64(sum_u64x2, vmovl_high_u32(data_u32x4));
2008
+ uint64x2_t squares_low_u64x2 = vmull_u32(vget_low_u32(data_u32x4), vget_low_u32(data_u32x4));
2009
+ uint64x2_t squares_high_u64x2 = vmull_high_u32(data_u32x4, data_u32x4);
2010
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, squares_low_u64x2);
2011
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, squares_high_u64x2);
2012
+ }
2013
+ }
2014
+ nk_u64_t sum = vgetq_lane_u64(sum_u64x2, 0) + vgetq_lane_u64(sum_u64x2, 1);
2015
+ nk_u64_t sumsq = nk_reduce_sadd_u64x2_neon_(sumsq_u64x2);
2016
+ for (; idx < count; ++idx) {
2017
+ nk_u64_t val = (nk_u64_t) * (data_ptr + idx * stride_elements);
2018
+ sum += val;
2019
+ nk_u64_t product = val * val;
2020
+ sumsq = nk_u64_saturating_add_serial(sumsq, product);
2021
+ }
2022
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
2023
+ }
2024
+
2025
+ NK_PUBLIC void nk_reduce_moments_u32_neon( //
2026
+ nk_u32_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
2027
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
2028
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_u32_t);
2029
+ int aligned = (stride_bytes % sizeof(nk_u32_t) == 0);
2030
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
2031
+ else if (!aligned) nk_reduce_moments_u32_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
2032
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 4) {
2033
+ nk_size_t left_count = count / 2;
2034
+ nk_u64_t left_sum, left_sumsq, right_sum, right_sumsq;
2035
+ nk_reduce_moments_u32_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
2036
+ nk_reduce_moments_u32_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
2037
+ &right_sum, &right_sumsq);
2038
+ *sum_ptr = nk_u64_saturating_add_serial(left_sum, right_sum);
2039
+ *sumsq_ptr = nk_u64_saturating_add_serial(left_sumsq, right_sumsq);
2040
+ }
2041
+ else if (stride_elements == 1) nk_reduce_moments_u32_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
2042
+ else if (stride_elements <= 4)
2043
+ nk_reduce_moments_u32_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
2044
+ else nk_reduce_moments_u32_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
2045
+ }
2046
+
2047
+ NK_INTERNAL void nk_reduce_minmax_u32_neon_contiguous_( //
2048
+ nk_u32_t const *data_ptr, nk_size_t count, //
2049
+ nk_u32_t *min_value_ptr, nk_size_t *min_index_ptr, //
2050
+ nk_u32_t *max_value_ptr, nk_size_t *max_index_ptr) {
2051
+ uint32x4_t min_u32x4 = vdupq_n_u32(NK_U32_MAX), max_u32x4 = vdupq_n_u32(0);
2052
+ uint32x4_t min_iter_u32x4 = vdupq_n_u32(0), max_iter_u32x4 = vdupq_n_u32(0);
2053
+ uint32x4_t iter_u32x4 = vdupq_n_u32(0), one_u32x4 = vdupq_n_u32(1);
2054
+ nk_size_t idx = 0;
2055
+ for (; idx + 4 <= count; idx += 4) {
2056
+ uint32x4_t data_u32x4 = vld1q_u32(data_ptr + idx);
2057
+ uint32x4_t less_u32x4 = vcltq_u32(data_u32x4, min_u32x4);
2058
+ uint32x4_t greater_u32x4 = vcgtq_u32(data_u32x4, max_u32x4);
2059
+ min_u32x4 = vbslq_u32(less_u32x4, data_u32x4, min_u32x4);
2060
+ max_u32x4 = vbslq_u32(greater_u32x4, data_u32x4, max_u32x4);
2061
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
2062
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
2063
+ iter_u32x4 = vaddq_u32(iter_u32x4, one_u32x4);
2064
+ }
2065
+ nk_size_t remaining = count - idx;
2066
+ if (remaining > 0) {
2067
+ nk_b128_vec_t tail_vec;
2068
+ nk_partial_load_b32x4_serial_(data_ptr + idx, &tail_vec, remaining);
2069
+ uint32x4_t lane_indices_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
2070
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
2071
+ uint32x4_t valid_u32x4 = vcltq_u32(lane_indices_u32x4, vdupq_n_u32((nk_u32_t)remaining));
2072
+ uint32x4_t data_min_u32x4 = vbslq_u32(valid_u32x4, tail_vec.u32x4, vdupq_n_u32(NK_U32_MAX));
2073
+ uint32x4_t data_max_u32x4 = vbslq_u32(valid_u32x4, tail_vec.u32x4, vdupq_n_u32(0));
2074
+ uint32x4_t less_u32x4 = vcltq_u32(data_min_u32x4, min_u32x4);
2075
+ uint32x4_t greater_u32x4 = vcgtq_u32(data_max_u32x4, max_u32x4);
2076
+ min_u32x4 = vbslq_u32(less_u32x4, data_min_u32x4, min_u32x4);
2077
+ max_u32x4 = vbslq_u32(greater_u32x4, data_max_u32x4, max_u32x4);
2078
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
2079
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
2080
+ }
2081
+ nk_u32_t min_value = vminvq_u32(min_u32x4), max_value = vmaxvq_u32(max_u32x4);
2082
+ uint32x4_t min_value_match_u32x4 = vceqq_u32(min_u32x4, vdupq_n_u32(min_value));
2083
+ uint32x4_t masked_min_iter_u32x4 = vbslq_u32(min_value_match_u32x4, min_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
2084
+ nk_u32_t earliest_min_cycle = vminvq_u32(masked_min_iter_u32x4);
2085
+ uint32x4_t max_value_match_u32x4 = vceqq_u32(max_u32x4, vdupq_n_u32(max_value));
2086
+ uint32x4_t masked_max_iter_u32x4 = vbslq_u32(max_value_match_u32x4, max_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
2087
+ nk_u32_t earliest_max_cycle = vminvq_u32(masked_max_iter_u32x4);
2088
+ uint32x4_t lane_indices_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
2089
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
2090
+ uint32x4_t min_cycle_match_u32x4 = vceqq_u32(min_iter_u32x4, vdupq_n_u32(earliest_min_cycle));
2091
+ uint32x4_t min_both_match_u32x4 = vandq_u32(min_value_match_u32x4, min_cycle_match_u32x4);
2092
+ uint32x4_t min_masked_lanes_u32x4 = vbslq_u32(min_both_match_u32x4, lane_indices_u32x4, vdupq_n_u32(NK_U32_MAX));
2093
+ nk_u32_t min_lane_offset = vminvq_u32(min_masked_lanes_u32x4);
2094
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 4 + (nk_size_t)min_lane_offset;
2095
+ uint32x4_t max_cycle_match_u32x4 = vceqq_u32(max_iter_u32x4, vdupq_n_u32(earliest_max_cycle));
2096
+ uint32x4_t max_both_match_u32x4 = vandq_u32(max_value_match_u32x4, max_cycle_match_u32x4);
2097
+ uint32x4_t max_masked_lanes_u32x4 = vbslq_u32(max_both_match_u32x4, lane_indices_u32x4, vdupq_n_u32(NK_U32_MAX));
2098
+ nk_u32_t max_lane_offset = vminvq_u32(max_masked_lanes_u32x4);
2099
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 4 + (nk_size_t)max_lane_offset;
2100
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
2101
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
2102
+ }
2103
+
2104
+ NK_INTERNAL void nk_reduce_minmax_u32_neon_strided_( //
2105
+ nk_u32_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
2106
+ nk_u32_t *min_value_ptr, nk_size_t *min_index_ptr, //
2107
+ nk_u32_t *max_value_ptr, nk_size_t *max_index_ptr) {
2108
+ uint32x4_t min_u32x4 = vdupq_n_u32(NK_U32_MAX), max_u32x4 = vdupq_n_u32(0);
2109
+ uint32x4_t min_iter_u32x4 = vdupq_n_u32(0), max_iter_u32x4 = vdupq_n_u32(0);
2110
+ uint32x4_t iter_u32x4 = vdupq_n_u32(0), one_u32x4 = vdupq_n_u32(1);
2111
+ uint32x4_t lane_indices_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
2112
+ vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
2113
+ nk_size_t idx = 0;
2114
+ uint32x4_t data_for_min_u32x4, data_for_max_u32x4;
2115
+
2116
+ nk_reduce_minmax_u32_neon_cycle:
2117
+ if (stride_elements == 2 && idx + 4 < count) {
2118
+ uint32x4x2_t loaded = vld2q_u32(data_ptr + idx * 2);
2119
+ data_for_min_u32x4 = loaded.val[0];
2120
+ data_for_max_u32x4 = loaded.val[0];
2121
+ idx += 4;
2122
+ }
2123
+ else if (stride_elements == 3 && idx + 4 < count) {
2124
+ uint32x4x3_t loaded = vld3q_u32(data_ptr + idx * 3);
2125
+ data_for_min_u32x4 = loaded.val[0];
2126
+ data_for_max_u32x4 = loaded.val[0];
2127
+ idx += 4;
2128
+ }
2129
+ else if (stride_elements == 4 && idx + 4 < count) {
2130
+ uint32x4x4_t loaded = vld4q_u32(data_ptr + idx * 4);
2131
+ data_for_min_u32x4 = loaded.val[0];
2132
+ data_for_max_u32x4 = loaded.val[0];
2133
+ idx += 4;
2134
+ }
2135
+ else if (idx < count) {
2136
+ nk_b128_vec_t tail_vec;
2137
+ nk_strided_load_b32x4_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
2138
+ uint32x4_t valid_u32x4 = vcltq_u32(lane_indices_u32x4, vdupq_n_u32((nk_u32_t)(count - idx)));
2139
+ data_for_min_u32x4 = vbslq_u32(valid_u32x4, tail_vec.u32x4, min_u32x4);
2140
+ data_for_max_u32x4 = vbslq_u32(valid_u32x4, tail_vec.u32x4, max_u32x4);
2141
+ idx = count;
2142
+ }
2143
+ else {
2144
+ nk_u32_t min_value = vminvq_u32(min_u32x4), max_value = vmaxvq_u32(max_u32x4);
2145
+ uint32x4_t min_value_match_u32x4 = vceqq_u32(min_u32x4, vdupq_n_u32(min_value));
2146
+ uint32x4_t masked_min_iter_u32x4 = vbslq_u32(min_value_match_u32x4, min_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
2147
+ nk_u32_t earliest_min_cycle = vminvq_u32(masked_min_iter_u32x4);
2148
+ uint32x4_t max_value_match_u32x4 = vceqq_u32(max_u32x4, vdupq_n_u32(max_value));
2149
+ uint32x4_t masked_max_iter_u32x4 = vbslq_u32(max_value_match_u32x4, max_iter_u32x4, vdupq_n_u32(NK_U32_MAX));
2150
+ nk_u32_t earliest_max_cycle = vminvq_u32(masked_max_iter_u32x4);
2151
+ uint32x4_t min_cycle_match_u32x4 = vceqq_u32(min_iter_u32x4, vdupq_n_u32(earliest_min_cycle));
2152
+ uint32x4_t min_both_match_u32x4 = vandq_u32(min_value_match_u32x4, min_cycle_match_u32x4);
2153
+ uint32x4_t min_masked_lanes_u32x4 = vbslq_u32(min_both_match_u32x4, lane_indices_u32x4,
2154
+ vdupq_n_u32(NK_U32_MAX));
2155
+ nk_u32_t min_lane_offset = vminvq_u32(min_masked_lanes_u32x4);
2156
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 4 + (nk_size_t)min_lane_offset;
2157
+ uint32x4_t max_cycle_match_u32x4 = vceqq_u32(max_iter_u32x4, vdupq_n_u32(earliest_max_cycle));
2158
+ uint32x4_t max_both_match_u32x4 = vandq_u32(max_value_match_u32x4, max_cycle_match_u32x4);
2159
+ uint32x4_t max_masked_lanes_u32x4 = vbslq_u32(max_both_match_u32x4, lane_indices_u32x4,
2160
+ vdupq_n_u32(NK_U32_MAX));
2161
+ nk_u32_t max_lane_offset = vminvq_u32(max_masked_lanes_u32x4);
2162
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 4 + (nk_size_t)max_lane_offset;
2163
+ *min_value_ptr = min_value, *min_index_ptr = min_idx;
2164
+ *max_value_ptr = max_value, *max_index_ptr = max_idx;
2165
+ return;
2166
+ }
2167
+
2168
+ // Shared update body
2169
+ uint32x4_t less_u32x4 = vcltq_u32(data_for_min_u32x4, min_u32x4);
2170
+ uint32x4_t greater_u32x4 = vcgtq_u32(data_for_max_u32x4, max_u32x4);
2171
+ min_u32x4 = vbslq_u32(less_u32x4, data_for_min_u32x4, min_u32x4);
2172
+ max_u32x4 = vbslq_u32(greater_u32x4, data_for_max_u32x4, max_u32x4);
2173
+ min_iter_u32x4 = vbslq_u32(less_u32x4, iter_u32x4, min_iter_u32x4);
2174
+ max_iter_u32x4 = vbslq_u32(greater_u32x4, iter_u32x4, max_iter_u32x4);
2175
+ iter_u32x4 = vaddq_u32(iter_u32x4, one_u32x4);
2176
+ goto nk_reduce_minmax_u32_neon_cycle;
2177
+ }
2178
+
2179
+ NK_PUBLIC void nk_reduce_minmax_u32_neon( //
2180
+ nk_u32_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
2181
+ nk_u32_t *min_value_ptr, nk_size_t *min_index_ptr, //
2182
+ nk_u32_t *max_value_ptr, nk_size_t *max_index_ptr) {
2183
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_u32_t);
2184
+ int aligned = (stride_bytes % sizeof(nk_u32_t) == 0);
2185
+ if (count == 0)
2186
+ *min_value_ptr = NK_U32_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = 0, *max_index_ptr = NK_SIZE_MAX;
2187
+ else if (!aligned)
2188
+ nk_reduce_minmax_u32_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
2189
+ max_index_ptr);
2190
+ else if (count > (nk_size_t)NK_U32_MAX * 4) {
2191
+ nk_size_t left_count = count / 2;
2192
+ nk_u32_t left_min_value, right_min_value, left_max_value, right_max_value;
2193
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
2194
+ nk_reduce_minmax_u32_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index, &left_max_value,
2195
+ &left_max_index);
2196
+ nk_reduce_minmax_u32_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
2197
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
2198
+ if (right_min_value < left_min_value)
2199
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
2200
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
2201
+ if (right_max_value > left_max_value)
2202
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
2203
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
2204
+ }
2205
+ else if (stride_elements == 1)
2206
+ nk_reduce_minmax_u32_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
2207
+ max_index_ptr);
2208
+ else if (stride_elements <= 4)
2209
+ nk_reduce_minmax_u32_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
2210
+ max_value_ptr, max_index_ptr);
2211
+ else
2212
+ nk_reduce_minmax_u32_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
2213
+ max_index_ptr);
2214
+ }
2215
+
2216
+ NK_INTERNAL void nk_reduce_moments_i64_neon_contiguous_( //
2217
+ nk_i64_t const *data_ptr, nk_size_t count, //
2218
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
2219
+ uint64x2_t sum_low_u64x2 = vdupq_n_u64(0);
2220
+ int64x2_t sum_high_i64x2 = vdupq_n_s64(0);
2221
+ // NEON can still load/extract i64 vectors for sumsq via scalar nk_i64_smul_
2222
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
2223
+ int sumsq_overflow = 0;
2224
+ int64x2_t sign_bit_i64x2 = vdupq_n_s64((nk_i64_t)0x8000000000000000ULL);
2225
+ nk_size_t idx = 0;
2226
+ for (; idx + 2 <= count; idx += 2) {
2227
+ int64x2_t data_i64x2 = vld1q_s64(data_ptr + idx);
2228
+ // Sumsq via helper (scalar per-lane multiply)
2229
+ uint64x2_t sq_u64x2 = nk_i64_smul_sq_i64x2_neon_(data_i64x2);
2230
+ uint64x2_t sq_before_u64x2 = sumsq_u64x2;
2231
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, sq_u64x2);
2232
+ int64x2_t result_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sumsq_u64x2), sign_bit_i64x2);
2233
+ int64x2_t before_biased_i64x2 = veorq_s64(vreinterpretq_s64_u64(sq_before_u64x2), sign_bit_i64x2);
2234
+ sumsq_overflow |=
2235
+ (vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 0) |
2236
+ vgetq_lane_s64(vreinterpretq_s64_u64(vcgtq_s64(before_biased_i64x2, result_biased_i64x2)), 1));
2237
+ // Vectorized 128-bit carry-propagating sum
2238
+ uint64x2_t sum_before_u64x2 = sum_low_u64x2;
2239
+ sum_low_u64x2 = vaddq_u64(sum_low_u64x2, vreinterpretq_u64_s64(data_i64x2));
2240
+ int64x2_t sb_biased = veorq_s64(vreinterpretq_s64_u64(sum_before_u64x2), sign_bit_i64x2);
2241
+ int64x2_t sr_biased = veorq_s64(vreinterpretq_s64_u64(sum_low_u64x2), sign_bit_i64x2);
2242
+ uint64x2_t carry_u64x2 = vcgtq_s64(sb_biased, sr_biased);
2243
+ sum_high_i64x2 = vsubq_s64(sum_high_i64x2, vreinterpretq_s64_u64(carry_u64x2));
2244
+ int64x2_t sign_ext_i64x2 = vshrq_n_s64(data_i64x2, 63);
2245
+ sum_high_i64x2 = vaddq_s64(sum_high_i64x2, sign_ext_i64x2);
2246
+ }
2247
+ // Horizontal reduction of 2 lanes to scalar (sum_low, sum_high)
2248
+ nk_u64_t sum_low = vgetq_lane_u64(sum_low_u64x2, 0);
2249
+ nk_i64_t sum_high = vgetq_lane_s64(sum_high_i64x2, 0);
2250
+ {
2251
+ nk_u64_t before = sum_low;
2252
+ sum_low += vgetq_lane_u64(sum_low_u64x2, 1);
2253
+ if (sum_low < before) sum_high++;
2254
+ sum_high += vgetq_lane_s64(sum_high_i64x2, 1);
2255
+ }
2256
+ nk_u64_t sumsq;
2257
+ if (sumsq_overflow) sumsq = NK_U64_MAX;
2258
+ else sumsq = nk_reduce_sadd_u64x2_neon_(sumsq_u64x2);
2259
+ for (; idx < count; ++idx) {
2260
+ nk_i64_t val = data_ptr[idx];
2261
+ nk_i64_t product = nk_i64_saturating_mul_serial(val, val);
2262
+ nk_u64_t unsigned_product = (nk_u64_t)product;
2263
+ sumsq = nk_u64_saturating_add_serial(sumsq, unsigned_product);
2264
+ nk_u64_t before = sum_low;
2265
+ sum_low += (nk_u64_t)val;
2266
+ if (sum_low < before) sum_high++;
2267
+ sum_high += (val >> 63);
2268
+ }
2269
+ nk_i64_t sum_low_signed = (nk_i64_t)sum_low;
2270
+ if (sum_high == (sum_low_signed >> 63)) *sum_ptr = sum_low_signed;
2271
+ else if (sum_high >= 0) *sum_ptr = NK_I64_MAX;
2272
+ else *sum_ptr = NK_I64_MIN;
2273
+ *sumsq_ptr = sumsq;
2274
+ }
2275
+
2276
+ NK_PUBLIC void nk_reduce_moments_i64_neon( //
2277
+ nk_i64_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
2278
+ nk_i64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
2279
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_i64_t);
2280
+ int aligned = (stride_bytes % sizeof(nk_i64_t) == 0);
2281
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
2282
+ else if (!aligned) nk_reduce_moments_i64_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
2283
+ else if (stride_elements == 1) nk_reduce_moments_i64_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
2284
+ else nk_reduce_moments_i64_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
2285
+ }
2286
+
2287
+ NK_INTERNAL void nk_reduce_minmax_i64_neon_contiguous_( //
2288
+ nk_i64_t const *data_ptr, nk_size_t count, //
2289
+ nk_i64_t *min_value_ptr, nk_size_t *min_index_ptr, //
2290
+ nk_i64_t *max_value_ptr, nk_size_t *max_index_ptr) {
2291
+ int64x2_t min_i64x2 = vdupq_n_s64(NK_I64_MAX), max_i64x2 = vdupq_n_s64(NK_I64_MIN);
2292
+ uint64x2_t min_iter_u64x2 = vdupq_n_u64(0), max_iter_u64x2 = vdupq_n_u64(0);
2293
+ uint64x2_t iter_u64x2 = vdupq_n_u64(0), one_u64x2 = vdupq_n_u64(1);
2294
+ nk_size_t idx = 0;
2295
+ for (; idx + 2 <= count; idx += 2) {
2296
+ int64x2_t data_i64x2 = vld1q_s64(data_ptr + idx);
2297
+ uint64x2_t less_u64x2 = vcltq_s64(data_i64x2, min_i64x2);
2298
+ uint64x2_t greater_u64x2 = vcgtq_s64(data_i64x2, max_i64x2);
2299
+ min_i64x2 = vbslq_s64(less_u64x2, data_i64x2, min_i64x2);
2300
+ max_i64x2 = vbslq_s64(greater_u64x2, data_i64x2, max_i64x2);
2301
+ min_iter_u64x2 = vbslq_u64(less_u64x2, iter_u64x2, min_iter_u64x2);
2302
+ max_iter_u64x2 = vbslq_u64(greater_u64x2, iter_u64x2, max_iter_u64x2);
2303
+ iter_u64x2 = vaddq_u64(iter_u64x2, one_u64x2);
2304
+ }
2305
+ nk_b128_vec_t min_values_vec, max_values_vec, min_indices_vec, max_indices_vec;
2306
+ min_values_vec.i64x2 = min_i64x2;
2307
+ min_indices_vec.u64x2 = min_iter_u64x2;
2308
+ max_values_vec.i64x2 = max_i64x2;
2309
+ max_indices_vec.u64x2 = max_iter_u64x2;
2310
+ nk_i64_t min_value, max_value;
2311
+ nk_size_t min_index, max_index;
2312
+ if (min_values_vec.i64s[0] <= min_values_vec.i64s[1])
2313
+ min_value = min_values_vec.i64s[0], min_index = (nk_size_t)min_indices_vec.u64s[0] * 2;
2314
+ else min_value = min_values_vec.i64s[1], min_index = (nk_size_t)min_indices_vec.u64s[1] * 2 + 1;
2315
+ if (max_values_vec.i64s[0] >= max_values_vec.i64s[1])
2316
+ max_value = max_values_vec.i64s[0], max_index = (nk_size_t)max_indices_vec.u64s[0] * 2;
2317
+ else max_value = max_values_vec.i64s[1], max_index = (nk_size_t)max_indices_vec.u64s[1] * 2 + 1;
2318
+ for (; idx < count; ++idx) {
2319
+ nk_i64_t val = data_ptr[idx];
2320
+ if (val < min_value) min_value = val, min_index = idx;
2321
+ if (val > max_value) max_value = val, max_index = idx;
2322
+ }
2323
+ *min_value_ptr = min_value, *min_index_ptr = min_index;
2324
+ *max_value_ptr = max_value, *max_index_ptr = max_index;
2325
+ }
2326
+
2327
+ NK_PUBLIC void nk_reduce_minmax_i64_neon( //
2328
+ nk_i64_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
2329
+ nk_i64_t *min_value_ptr, nk_size_t *min_index_ptr, //
2330
+ nk_i64_t *max_value_ptr, nk_size_t *max_index_ptr) {
2331
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_i64_t);
2332
+ int aligned = (stride_bytes % sizeof(nk_i64_t) == 0);
2333
+ if (count == 0)
2334
+ *min_value_ptr = NK_I64_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_I64_MIN,
2335
+ *max_index_ptr = NK_SIZE_MAX;
2336
+ else if (!aligned)
2337
+ nk_reduce_minmax_i64_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
2338
+ max_index_ptr);
2339
+ else if (stride_elements == 1)
2340
+ nk_reduce_minmax_i64_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
2341
+ max_index_ptr);
2342
+ else
2343
+ nk_reduce_minmax_i64_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
2344
+ max_index_ptr);
2345
+ }
2346
+
2347
+ NK_INTERNAL void nk_reduce_moments_u64_neon_contiguous_( //
2348
+ nk_u64_t const *data_ptr, nk_size_t count, //
2349
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
2350
+ uint64x2_t sum_u64x2 = vdupq_n_u64(0);
2351
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
2352
+ nk_size_t idx = 0;
2353
+ for (; idx + 2 <= count; idx += 2) {
2354
+ uint64x2_t data_u64x2 = vld1q_u64(data_ptr + idx);
2355
+ sum_u64x2 = vqaddq_u64(sum_u64x2, data_u64x2);
2356
+ uint64x2_t sq_u64x2 = nk_u64_smul_sq_u64x2_neon_(data_u64x2);
2357
+ sumsq_u64x2 = vqaddq_u64(sumsq_u64x2, sq_u64x2);
2358
+ }
2359
+ nk_u64_t sum = nk_reduce_sadd_u64x2_neon_(sum_u64x2);
2360
+ nk_u64_t sumsq = nk_reduce_sadd_u64x2_neon_(sumsq_u64x2);
2361
+ for (; idx < count; ++idx) {
2362
+ nk_u64_t val = data_ptr[idx];
2363
+ sum = nk_u64_saturating_add_serial(sum, val);
2364
+ nk_u64_t product = nk_u64_saturating_mul_serial(val, val);
2365
+ sumsq = nk_u64_saturating_add_serial(sumsq, product);
2366
+ }
2367
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
2368
+ }
2369
+
2370
+ NK_PUBLIC void nk_reduce_moments_u64_neon( //
2371
+ nk_u64_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
2372
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
2373
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_u64_t);
2374
+ int aligned = (stride_bytes % sizeof(nk_u64_t) == 0);
2375
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
2376
+ else if (!aligned) nk_reduce_moments_u64_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
2377
+ else if (stride_elements == 1) nk_reduce_moments_u64_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
2378
+ else nk_reduce_moments_u64_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
2379
+ }
2380
+
2381
+ NK_INTERNAL void nk_reduce_minmax_u64_neon_contiguous_( //
2382
+ nk_u64_t const *data_ptr, nk_size_t count, //
2383
+ nk_u64_t *min_value_ptr, nk_size_t *min_index_ptr, //
2384
+ nk_u64_t *max_value_ptr, nk_size_t *max_index_ptr) {
2385
+ uint64x2_t min_u64x2 = vdupq_n_u64(NK_U64_MAX), max_u64x2 = vdupq_n_u64(0);
2386
+ uint64x2_t min_iter_u64x2 = vdupq_n_u64(0), max_iter_u64x2 = vdupq_n_u64(0);
2387
+ uint64x2_t iter_u64x2 = vdupq_n_u64(0), one_u64x2 = vdupq_n_u64(1);
2388
+ nk_size_t idx = 0;
2389
+ for (; idx + 2 <= count; idx += 2) {
2390
+ uint64x2_t data_u64x2 = vld1q_u64(data_ptr + idx);
2391
+ uint64x2_t less_u64x2 = vcltq_u64(data_u64x2, min_u64x2);
2392
+ uint64x2_t greater_u64x2 = vcgtq_u64(data_u64x2, max_u64x2);
2393
+ min_u64x2 = vbslq_u64(less_u64x2, data_u64x2, min_u64x2);
2394
+ max_u64x2 = vbslq_u64(greater_u64x2, data_u64x2, max_u64x2);
2395
+ min_iter_u64x2 = vbslq_u64(less_u64x2, iter_u64x2, min_iter_u64x2);
2396
+ max_iter_u64x2 = vbslq_u64(greater_u64x2, iter_u64x2, max_iter_u64x2);
2397
+ iter_u64x2 = vaddq_u64(iter_u64x2, one_u64x2);
2398
+ }
2399
+ nk_b128_vec_t min_values_vec, max_values_vec, min_indices_vec, max_indices_vec;
2400
+ min_values_vec.u64x2 = min_u64x2;
2401
+ min_indices_vec.u64x2 = min_iter_u64x2;
2402
+ max_values_vec.u64x2 = max_u64x2;
2403
+ max_indices_vec.u64x2 = max_iter_u64x2;
2404
+ nk_u64_t min_value, max_value;
2405
+ nk_size_t min_index, max_index;
2406
+ if (min_values_vec.u64s[0] <= min_values_vec.u64s[1])
2407
+ min_value = min_values_vec.u64s[0], min_index = (nk_size_t)min_indices_vec.u64s[0] * 2;
2408
+ else min_value = min_values_vec.u64s[1], min_index = (nk_size_t)min_indices_vec.u64s[1] * 2 + 1;
2409
+ if (max_values_vec.u64s[0] >= max_values_vec.u64s[1])
2410
+ max_value = max_values_vec.u64s[0], max_index = (nk_size_t)max_indices_vec.u64s[0] * 2;
2411
+ else max_value = max_values_vec.u64s[1], max_index = (nk_size_t)max_indices_vec.u64s[1] * 2 + 1;
2412
+ for (; idx < count; ++idx) {
2413
+ nk_u64_t val = data_ptr[idx];
2414
+ if (val < min_value) min_value = val, min_index = idx;
2415
+ if (val > max_value) max_value = val, max_index = idx;
2416
+ }
2417
+ *min_value_ptr = min_value, *min_index_ptr = min_index;
2418
+ *max_value_ptr = max_value, *max_index_ptr = max_index;
2419
+ }
2420
+
2421
+ NK_PUBLIC void nk_reduce_minmax_u64_neon( //
2422
+ nk_u64_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
2423
+ nk_u64_t *min_value_ptr, nk_size_t *min_index_ptr, //
2424
+ nk_u64_t *max_value_ptr, nk_size_t *max_index_ptr) {
2425
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_u64_t);
2426
+ int aligned = (stride_bytes % sizeof(nk_u64_t) == 0);
2427
+ if (count == 0)
2428
+ *min_value_ptr = NK_U64_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = 0, *max_index_ptr = NK_SIZE_MAX;
2429
+ else if (!aligned)
2430
+ nk_reduce_minmax_u64_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
2431
+ max_index_ptr);
2432
+ else if (stride_elements == 1)
2433
+ nk_reduce_minmax_u64_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
2434
+ max_index_ptr);
2435
+ else
2436
+ nk_reduce_minmax_u64_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
2437
+ max_index_ptr);
2438
+ }
2439
+
2440
+ /** @brief Convert 16 raw FP6 (e2m3/e3m2) sign-magnitude bytes to unsigned-comparable bytes.
2441
+ * FP6: sign bit 5, 5-bit magnitude. Positive maps to [0x20..0x3F], negative to [0x00..0x1F]. */
2442
+ NK_INTERNAL uint8x16_t nk_fp6x16_to_comparable_neon_(uint8x16_t raw_u8x16) {
2443
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2444
+ uint8x16_t sign_mask_u8x16 = vdupq_n_u8(0x20);
2445
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, sign_mask_u8x16);
2446
+ uint8x16_t positive_u8x16 = vorrq_u8(magnitude_u8x16, sign_mask_u8x16);
2447
+ uint8x16_t negative_u8x16 = vsubq_u8(vdupq_n_u8(0x1F), magnitude_u8x16);
2448
+ return vbslq_u8(is_negative_u8x16, negative_u8x16, positive_u8x16);
2449
+ }
2450
+
2451
+ /** @brief Convert a single comparable byte back to raw FP6 sign-magnitude byte. */
2452
+ NK_INTERNAL nk_u8_t nk_comparable_to_fp6_(nk_u8_t comparable) {
2453
+ if (comparable >= 0x20) return comparable ^ 0x20; // was positive
2454
+ else return (0x1F - comparable) | 0x20; // was negative
2455
+ }
2456
+
2457
+ /** @brief Convert 16 raw FP8 (e4m3/e5m2) sign-magnitude bytes to unsigned-comparable bytes. */
2458
+ NK_INTERNAL uint8x16_t nk_fp8x16_to_comparable_neon_(uint8x16_t raw_u8x16) {
2459
+ uint8x16_t sign_mask_u8x16 = vdupq_n_u8(0x80);
2460
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, sign_mask_u8x16);
2461
+ uint8x16_t flip_positive_u8x16 = veorq_u8(raw_u8x16, sign_mask_u8x16);
2462
+ uint8x16_t flip_negative_u8x16 = vmvnq_u8(raw_u8x16);
2463
+ return vbslq_u8(is_negative_u8x16, flip_negative_u8x16, flip_positive_u8x16);
2464
+ }
2465
+
2466
+ /** @brief Convert a single comparable byte back to raw FP8 sign-magnitude byte. */
2467
+ NK_INTERNAL nk_u8_t nk_comparable_to_fp8_(nk_u8_t comparable) {
2468
+ if (comparable >= 0x80) return comparable ^ 0x80; // was positive
2469
+ else return ~comparable; // was negative
2470
+ }
2471
+
2472
+ NK_INTERNAL void nk_reduce_moments_e2m3_neon_contiguous_( //
2473
+ nk_e2m3_t const *data_ptr, nk_size_t count, //
2474
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
2475
+ // VTBL LUT: maps 6-bit magnitude (0..31) to value×16 (unsigned), fits in u8
2476
+ uint8x16x2_t lut_e2m3_x16;
2477
+ // table[0]: values for magnitudes 0..15
2478
+ // 0x0E0C0A0806040200 → bytes [0..7] = 0,2,4,6,8,10,12,14
2479
+ // 0x1E1C1A1816141210 → bytes [8..15] = 16,18,20,22,24,26,28,30
2480
+ lut_e2m3_x16.val[0] = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0E0C0A0806040200ULL)),
2481
+ vreinterpret_u8_u64(vcreate_u64(0x1E1C1A1816141210ULL)));
2482
+ // table[1]: values for magnitudes 16..31
2483
+ // 0x3C3834302C282420 → bytes [0..7] = 32,36,40,44,48,52,56,60
2484
+ // 0x7870686058504840 → bytes [8..15] = 64,72,80,88,96,104,112,120
2485
+ lut_e2m3_x16.val[1] = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x3C3834302C282420ULL)),
2486
+ vreinterpret_u8_u64(vcreate_u64(0x7870686058504840ULL)));
2487
+ int32x4_t sum_i32x4 = vdupq_n_s32(0);
2488
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
2489
+ nk_size_t idx = 0;
2490
+ for (; idx + 16 <= count; idx += 16) {
2491
+ uint8x16_t raw_u8x16 = vld1q_u8((nk_u8_t const *)(data_ptr + idx));
2492
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2493
+ uint8x16_t unsigned_u8x16 = vqtbl2q_u8(lut_e2m3_x16, magnitude_u8x16);
2494
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, vdupq_n_u8(0x20));
2495
+ int8x16_t positive_i8x16 = vreinterpretq_s8_u8(unsigned_u8x16);
2496
+ int8x16_t negative_i8x16 = vnegq_s8(positive_i8x16);
2497
+ int8x16_t scaled_i8x16 = vbslq_s8(is_negative_u8x16, negative_i8x16, positive_i8x16);
2498
+ int16x8_t pairwise_i16x8 = vpaddlq_s8(scaled_i8x16);
2499
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(pairwise_i16x8));
2500
+ int16x8_t squares_low_i16x8 = vmull_s8(vget_low_s8(scaled_i8x16), vget_low_s8(scaled_i8x16));
2501
+ int16x8_t squares_high_i16x8 = vmull_high_s8(scaled_i8x16, scaled_i8x16);
2502
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_low_i16x8))));
2503
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_high_i16x8))));
2504
+ }
2505
+ nk_i64_t sum = vaddlvq_s32(sum_i32x4);
2506
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
2507
+ for (; idx < count; ++idx) {
2508
+ nk_f32_t value_f32;
2509
+ nk_e2m3_to_f32_serial(&data_ptr[idx], &value_f32);
2510
+ sum += (nk_i64_t)(value_f32 * 16.0f), sumsq += (nk_u64_t)(nk_i64_t)(value_f32 * value_f32 * 256.0f);
2511
+ }
2512
+ *sum_ptr = (nk_f32_t)sum / 16.0f, *sumsq_ptr = (nk_f32_t)sumsq / 256.0f;
2513
+ }
2514
+
2515
+ NK_INTERNAL void nk_reduce_moments_e2m3_neon_strided_( //
2516
+ nk_e2m3_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
2517
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
2518
+ uint8x16x2_t lut_e2m3_x16;
2519
+ // table[0]: values for magnitudes 0..15
2520
+ // 0x0E0C0A0806040200 → bytes [0..7] = 0,2,4,6,8,10,12,14
2521
+ // 0x1E1C1A1816141210 → bytes [8..15] = 16,18,20,22,24,26,28,30
2522
+ lut_e2m3_x16.val[0] = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0E0C0A0806040200ULL)),
2523
+ vreinterpret_u8_u64(vcreate_u64(0x1E1C1A1816141210ULL)));
2524
+ // table[1]: values for magnitudes 16..31
2525
+ // 0x3C3834302C282420 → bytes [0..7] = 32,36,40,44,48,52,56,60
2526
+ // 0x7870686058504840 → bytes [8..15] = 64,72,80,88,96,104,112,120
2527
+ lut_e2m3_x16.val[1] = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x3C3834302C282420ULL)),
2528
+ vreinterpret_u8_u64(vcreate_u64(0x7870686058504840ULL)));
2529
+ int32x4_t sum_i32x4 = vdupq_n_s32(0);
2530
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
2531
+ nk_size_t idx = 0;
2532
+ if (stride_elements == 2) {
2533
+ for (; idx + 16 < count; idx += 16) {
2534
+ uint8x16x2_t loaded_u8x16x2 = vld2q_u8((nk_u8_t const *)(data_ptr + idx * 2));
2535
+ uint8x16_t raw_u8x16 = loaded_u8x16x2.val[0];
2536
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2537
+ uint8x16_t unsigned_u8x16 = vqtbl2q_u8(lut_e2m3_x16, magnitude_u8x16);
2538
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, vdupq_n_u8(0x20));
2539
+ int8x16_t positive_i8x16 = vreinterpretq_s8_u8(unsigned_u8x16);
2540
+ int8x16_t negative_i8x16 = vnegq_s8(positive_i8x16);
2541
+ int8x16_t scaled_i8x16 = vbslq_s8(is_negative_u8x16, negative_i8x16, positive_i8x16);
2542
+ int16x8_t pairwise_i16x8 = vpaddlq_s8(scaled_i8x16);
2543
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(pairwise_i16x8));
2544
+ int16x8_t squares_low_i16x8 = vmull_s8(vget_low_s8(scaled_i8x16), vget_low_s8(scaled_i8x16));
2545
+ int16x8_t squares_high_i16x8 = vmull_high_s8(scaled_i8x16, scaled_i8x16);
2546
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_low_i16x8))));
2547
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_high_i16x8))));
2548
+ }
2549
+ }
2550
+ else if (stride_elements == 3) {
2551
+ for (; idx + 16 < count; idx += 16) {
2552
+ uint8x16x3_t loaded_u8x16x3 = vld3q_u8((nk_u8_t const *)(data_ptr + idx * 3));
2553
+ uint8x16_t raw_u8x16 = loaded_u8x16x3.val[0];
2554
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2555
+ uint8x16_t unsigned_u8x16 = vqtbl2q_u8(lut_e2m3_x16, magnitude_u8x16);
2556
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, vdupq_n_u8(0x20));
2557
+ int8x16_t positive_i8x16 = vreinterpretq_s8_u8(unsigned_u8x16);
2558
+ int8x16_t negative_i8x16 = vnegq_s8(positive_i8x16);
2559
+ int8x16_t scaled_i8x16 = vbslq_s8(is_negative_u8x16, negative_i8x16, positive_i8x16);
2560
+ int16x8_t pairwise_i16x8 = vpaddlq_s8(scaled_i8x16);
2561
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(pairwise_i16x8));
2562
+ int16x8_t squares_low_i16x8 = vmull_s8(vget_low_s8(scaled_i8x16), vget_low_s8(scaled_i8x16));
2563
+ int16x8_t squares_high_i16x8 = vmull_high_s8(scaled_i8x16, scaled_i8x16);
2564
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_low_i16x8))));
2565
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_high_i16x8))));
2566
+ }
2567
+ }
2568
+ else {
2569
+ for (; idx + 16 < count; idx += 16) {
2570
+ uint8x16x4_t loaded_u8x16x4 = vld4q_u8((nk_u8_t const *)(data_ptr + idx * 4));
2571
+ uint8x16_t raw_u8x16 = loaded_u8x16x4.val[0];
2572
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2573
+ uint8x16_t unsigned_u8x16 = vqtbl2q_u8(lut_e2m3_x16, magnitude_u8x16);
2574
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, vdupq_n_u8(0x20));
2575
+ int8x16_t positive_i8x16 = vreinterpretq_s8_u8(unsigned_u8x16);
2576
+ int8x16_t negative_i8x16 = vnegq_s8(positive_i8x16);
2577
+ int8x16_t scaled_i8x16 = vbslq_s8(is_negative_u8x16, negative_i8x16, positive_i8x16);
2578
+ int16x8_t pairwise_i16x8 = vpaddlq_s8(scaled_i8x16);
2579
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(pairwise_i16x8));
2580
+ int16x8_t squares_low_i16x8 = vmull_s8(vget_low_s8(scaled_i8x16), vget_low_s8(scaled_i8x16));
2581
+ int16x8_t squares_high_i16x8 = vmull_high_s8(scaled_i8x16, scaled_i8x16);
2582
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_low_i16x8))));
2583
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vpaddlq_u16(vreinterpretq_u16_s16(squares_high_i16x8))));
2584
+ }
2585
+ }
2586
+ nk_i64_t sum = vaddlvq_s32(sum_i32x4);
2587
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
2588
+ for (; idx < count; ++idx) {
2589
+ nk_f32_t val;
2590
+ nk_e2m3_to_f32_serial((nk_e2m3_t const *)(data_ptr + idx * stride_elements), &val);
2591
+ sum += (nk_i64_t)(val * 16.0f), sumsq += (nk_u64_t)(nk_i64_t)(val * val * 256.0f);
2592
+ }
2593
+ *sum_ptr = (nk_f32_t)sum / 16.0f, *sumsq_ptr = (nk_f32_t)sumsq / 256.0f;
2594
+ }
2595
+
2596
+ NK_PUBLIC void nk_reduce_moments_e2m3_neon( //
2597
+ nk_e2m3_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
2598
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
2599
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_e2m3_t);
2600
+ int aligned = (stride_bytes % sizeof(nk_e2m3_t) == 0);
2601
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
2602
+ else if (!aligned) nk_reduce_moments_e2m3_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
2603
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 16) {
2604
+ nk_size_t left_count = count / 2;
2605
+ nk_f32_t left_sum, left_sumsq, right_sum, right_sumsq;
2606
+ nk_reduce_moments_e2m3_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
2607
+ nk_reduce_moments_e2m3_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
2608
+ &right_sum, &right_sumsq);
2609
+ *sum_ptr = left_sum + right_sum, *sumsq_ptr = left_sumsq + right_sumsq;
2610
+ }
2611
+ else if (stride_elements == 1) nk_reduce_moments_e2m3_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
2612
+ else if (stride_elements <= 4)
2613
+ nk_reduce_moments_e2m3_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
2614
+ else nk_reduce_moments_e2m3_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
2615
+ }
2616
+
2617
+ NK_INTERNAL void nk_reduce_minmax_e2m3_neon_contiguous_( //
2618
+ nk_e2m3_t const *data_ptr, nk_size_t count, //
2619
+ nk_e2m3_t *min_value_ptr, nk_size_t *min_index_ptr, //
2620
+ nk_e2m3_t *max_value_ptr, nk_size_t *max_index_ptr) {
2621
+ // Handle initial chunk — partial or full
2622
+ uint8x16_t first_comparable_u8x16;
2623
+ nk_size_t first_count = count < 16 ? count : 16;
2624
+ if (count < 16) {
2625
+ nk_b128_vec_t first_vec;
2626
+ nk_partial_load_b8x16_serial_(data_ptr, &first_vec, count);
2627
+ first_comparable_u8x16 = nk_fp6x16_to_comparable_neon_(first_vec.u8x16);
2628
+ // Mask invalid lanes: min gets 0xFF (won't be selected), max gets 0x00
2629
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
2630
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
2631
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)count));
2632
+ first_comparable_u8x16 = vbslq_u8(valid_u8x16, first_comparable_u8x16, vdupq_n_u8(0));
2633
+ }
2634
+ else {
2635
+ uint8x16_t first_raw_u8x16 = vld1q_u8((nk_u8_t const *)data_ptr);
2636
+ first_comparable_u8x16 = nk_fp6x16_to_comparable_neon_(first_raw_u8x16);
2637
+ }
2638
+ // For min: invalid lanes (0x00) should not win, so initialize min from masked data where invalid = 0xFF
2639
+ // For max: invalid lanes (0x00) should not win, which is already correct since 0x00 won't beat real data
2640
+ uint8x16_t lane_indices_init_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
2641
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
2642
+ uint8x16_t valid_init_u8x16 = vcltq_u8(lane_indices_init_u8x16, vdupq_n_u8((nk_u8_t)first_count));
2643
+ uint8x16_t min_u8x16 = vbslq_u8(valid_init_u8x16, first_comparable_u8x16, vdupq_n_u8(0xFF));
2644
+ uint8x16_t max_u8x16 = first_comparable_u8x16; // invalid lanes are 0x00, safe for max
2645
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
2646
+ uint8x16_t iter_u8x16 = vdupq_n_u8(1), one_u8x16 = vdupq_n_u8(1);
2647
+ nk_size_t idx = first_count;
2648
+ for (; idx + 16 <= count; idx += 16) {
2649
+ uint8x16_t raw_u8x16 = vld1q_u8((nk_u8_t const *)(data_ptr + idx));
2650
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(raw_u8x16);
2651
+ uint8x16_t less_u8x16 = vcltq_u8(comparable_u8x16, min_u8x16);
2652
+ uint8x16_t greater_u8x16 = vcgtq_u8(comparable_u8x16, max_u8x16);
2653
+ min_u8x16 = vbslq_u8(less_u8x16, comparable_u8x16, min_u8x16);
2654
+ max_u8x16 = vbslq_u8(greater_u8x16, comparable_u8x16, max_u8x16);
2655
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
2656
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
2657
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
2658
+ }
2659
+ nk_size_t remaining = count - idx;
2660
+ if (remaining > 0) {
2661
+ nk_b128_vec_t tail_vec;
2662
+ nk_partial_load_b8x16_serial_(data_ptr + idx, &tail_vec, remaining);
2663
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(tail_vec.u8x16);
2664
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
2665
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
2666
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)remaining));
2667
+ uint8x16_t data_for_min_u8x16 = vbslq_u8(valid_u8x16, comparable_u8x16, vdupq_n_u8(0xFF));
2668
+ uint8x16_t data_for_max_u8x16 = vbslq_u8(valid_u8x16, comparable_u8x16, vdupq_n_u8(0));
2669
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
2670
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
2671
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
2672
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
2673
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
2674
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
2675
+ }
2676
+ nk_u8_t min_comparable = vminvq_u8(min_u8x16), max_comparable = vmaxvq_u8(max_u8x16);
2677
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_comparable));
2678
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
2679
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
2680
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_comparable));
2681
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
2682
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
2683
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
2684
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
2685
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
2686
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
2687
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
2688
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
2689
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
2690
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
2691
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
2692
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
2693
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
2694
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
2695
+ *min_value_ptr = nk_comparable_to_fp6_(min_comparable), *min_index_ptr = min_idx;
2696
+ *max_value_ptr = nk_comparable_to_fp6_(max_comparable), *max_index_ptr = max_idx;
2697
+ }
2698
+
2699
+ NK_INTERNAL void nk_reduce_minmax_e2m3_neon_strided_( //
2700
+ nk_e2m3_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
2701
+ nk_e2m3_t *min_value_ptr, nk_size_t *min_index_ptr, //
2702
+ nk_e2m3_t *max_value_ptr, nk_size_t *max_index_ptr) {
2703
+ uint8x16_t min_u8x16 = vdupq_n_u8(0xFF), max_u8x16 = vdupq_n_u8(0);
2704
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
2705
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
2706
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
2707
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
2708
+ nk_size_t idx = 0;
2709
+ uint8x16_t data_for_min_u8x16, data_for_max_u8x16;
2710
+
2711
+ nk_reduce_minmax_e2m3_neon_cycle:
2712
+ if (stride_elements == 2 && idx + 16 < count) {
2713
+ uint8x16x2_t loaded = vld2q_u8((nk_u8_t const *)(data_ptr + idx * 2));
2714
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(loaded.val[0]);
2715
+ data_for_min_u8x16 = comparable_u8x16;
2716
+ data_for_max_u8x16 = comparable_u8x16;
2717
+ idx += 16;
2718
+ }
2719
+ else if (stride_elements == 3 && idx + 16 < count) {
2720
+ uint8x16x3_t loaded = vld3q_u8((nk_u8_t const *)(data_ptr + idx * 3));
2721
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(loaded.val[0]);
2722
+ data_for_min_u8x16 = comparable_u8x16;
2723
+ data_for_max_u8x16 = comparable_u8x16;
2724
+ idx += 16;
2725
+ }
2726
+ else if (stride_elements == 4 && idx + 16 < count) {
2727
+ uint8x16x4_t loaded = vld4q_u8((nk_u8_t const *)(data_ptr + idx * 4));
2728
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(loaded.val[0]);
2729
+ data_for_min_u8x16 = comparable_u8x16;
2730
+ data_for_max_u8x16 = comparable_u8x16;
2731
+ idx += 16;
2732
+ }
2733
+ else if (idx < count) {
2734
+ nk_b128_vec_t tail_vec;
2735
+ nk_strided_load_b8x16_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
2736
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(tail_vec.u8x16);
2737
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)(count - idx)));
2738
+ data_for_min_u8x16 = vbslq_u8(valid_u8x16, comparable_u8x16, vdupq_n_u8(0xFF));
2739
+ data_for_max_u8x16 = vbslq_u8(valid_u8x16, comparable_u8x16, vdupq_n_u8(0x00));
2740
+ idx = count;
2741
+ }
2742
+ else {
2743
+ nk_u8_t min_comparable = vminvq_u8(min_u8x16), max_comparable = vmaxvq_u8(max_u8x16);
2744
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_comparable));
2745
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
2746
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
2747
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_comparable));
2748
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
2749
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
2750
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
2751
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
2752
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
2753
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
2754
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
2755
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
2756
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
2757
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
2758
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
2759
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
2760
+ *min_value_ptr = nk_comparable_to_fp6_(min_comparable), *min_index_ptr = min_idx;
2761
+ *max_value_ptr = nk_comparable_to_fp6_(max_comparable), *max_index_ptr = max_idx;
2762
+ return;
2763
+ }
2764
+
2765
+ // Shared update body
2766
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
2767
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
2768
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
2769
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
2770
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
2771
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
2772
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
2773
+ goto nk_reduce_minmax_e2m3_neon_cycle;
2774
+ }
2775
+
2776
+ NK_PUBLIC void nk_reduce_minmax_e2m3_neon( //
2777
+ nk_e2m3_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
2778
+ nk_e2m3_t *min_value_ptr, nk_size_t *min_index_ptr, //
2779
+ nk_e2m3_t *max_value_ptr, nk_size_t *max_index_ptr) {
2780
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_e2m3_t);
2781
+ int aligned = (stride_bytes % sizeof(nk_e2m3_t) == 0);
2782
+ if (count == 0)
2783
+ *min_value_ptr = NK_E2M3_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_E2M3_MIN,
2784
+ *max_index_ptr = NK_SIZE_MAX;
2785
+ else if (!aligned)
2786
+ nk_reduce_minmax_e2m3_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
2787
+ max_index_ptr);
2788
+ else if (count > (nk_size_t)256 * 16) {
2789
+ nk_size_t left_count = count / 2;
2790
+ nk_e2m3_t left_min_value, right_min_value, left_max_value, right_max_value;
2791
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
2792
+ nk_reduce_minmax_e2m3_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index,
2793
+ &left_max_value, &left_max_index);
2794
+ nk_reduce_minmax_e2m3_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
2795
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
2796
+ if (nk_e2m3_order_serial(right_min_value, left_min_value) < 0)
2797
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
2798
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
2799
+ if (nk_e2m3_order_serial(right_max_value, left_max_value) > 0)
2800
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
2801
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
2802
+ }
2803
+ else if (stride_elements == 1)
2804
+ nk_reduce_minmax_e2m3_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
2805
+ max_index_ptr);
2806
+ else if (stride_elements <= 4)
2807
+ nk_reduce_minmax_e2m3_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
2808
+ max_value_ptr, max_index_ptr);
2809
+ else
2810
+ nk_reduce_minmax_e2m3_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
2811
+ max_index_ptr);
2812
+ }
2813
+
2814
+ NK_INTERNAL void nk_reduce_moments_e3m2_neon_contiguous_( //
2815
+ nk_e3m2_t const *data_ptr, nk_size_t count, //
2816
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
2817
+ // VTBL LUT: maps 6-bit magnitude (0..31) to (value×16) low byte; max value×16 = 448 needs i16
2818
+ uint8x16x2_t lut_e3m2_low;
2819
+ // table[0]: low bytes for magnitudes 0..15
2820
+ // 0x0706050403020100 → bytes [0..7] = 0,1,2,3,4,5,6,7
2821
+ // 0x1C1814100E0C0A08 → bytes [8..15] = 8,10,12,14,16,20,24,28
2822
+ lut_e3m2_low.val[0] = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
2823
+ vreinterpret_u8_u64(vcreate_u64(0x1C1814100E0C0A08ULL)));
2824
+ // table[1]: low bytes for magnitudes 16..31
2825
+ // 0x7060504038302820 → bytes [0..7] = 32,40,48,56,64,80,96,112
2826
+ // 0xC0804000E0C0A080 → bytes [8..15] = 128,160,192,224,0,64,128,192
2827
+ lut_e3m2_low.val[1] = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x7060504038302820ULL)),
2828
+ vreinterpret_u8_u64(vcreate_u64(0xC0804000E0C0A080ULL)));
2829
+ int32x4_t sum_i32x4 = vdupq_n_s32(0);
2830
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
2831
+ nk_size_t idx = 0;
2832
+ for (; idx + 16 <= count; idx += 16) {
2833
+ uint8x16_t raw_u8x16 = vld1q_u8((nk_u8_t const *)(data_ptr + idx));
2834
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2835
+ uint8x16_t low_byte_u8x16 = vqtbl2q_u8(lut_e3m2_low, magnitude_u8x16);
2836
+ uint8x16_t high_byte_u8x16 = vandq_u8(vcgtq_u8(magnitude_u8x16, vdupq_n_u8(27)), vdupq_n_u8(1));
2837
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, vdupq_n_u8(0x20));
2838
+ // Interleave low+high bytes into i16 values (two halves of 8 each)
2839
+ uint16x8_t unsigned_low_u16x8 = vreinterpretq_u16_u8(vzip1q_u8(low_byte_u8x16, high_byte_u8x16));
2840
+ uint16x8_t unsigned_high_u16x8 = vreinterpretq_u16_u8(vzip2q_u8(low_byte_u8x16, high_byte_u8x16));
2841
+ // Sign-extend the per-byte negative mask to per-i16 lanes
2842
+ int8x8_t is_negative_low_i8x8 = vreinterpret_s8_u8(vget_low_u8(is_negative_u8x16));
2843
+ int8x8_t is_negative_high_i8x8 = vreinterpret_s8_u8(vget_high_u8(is_negative_u8x16));
2844
+ uint16x8_t is_negative_low_u16x8 = vreinterpretq_u16_s16(vmovl_s8(is_negative_low_i8x8));
2845
+ uint16x8_t is_negative_high_u16x8 = vreinterpretq_u16_s16(vmovl_s8(is_negative_high_i8x8));
2846
+ // Apply sign via conditional negate
2847
+ int16x8_t positive_low_i16x8 = vreinterpretq_s16_u16(unsigned_low_u16x8);
2848
+ int16x8_t scaled_low_i16x8 = vbslq_s16(is_negative_low_u16x8, vnegq_s16(positive_low_i16x8),
2849
+ positive_low_i16x8);
2850
+ int16x8_t positive_high_i16x8 = vreinterpretq_s16_u16(unsigned_high_u16x8);
2851
+ int16x8_t scaled_high_i16x8 = vbslq_s16(is_negative_high_u16x8, vnegq_s16(positive_high_i16x8),
2852
+ positive_high_i16x8);
2853
+ // Sum: i16→i32 widening, accumulate in i32x4
2854
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(scaled_low_i16x8));
2855
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(scaled_high_i16x8));
2856
+ // Sumsq: vmull_s16→i32 (always positive as squares), widen to u64
2857
+ int32x4_t squares_low_a_i32x4 = vmull_s16(vget_low_s16(scaled_low_i16x8), vget_low_s16(scaled_low_i16x8));
2858
+ int32x4_t squares_low_b_i32x4 = vmull_high_s16(scaled_low_i16x8, scaled_low_i16x8);
2859
+ int32x4_t squares_high_a_i32x4 = vmull_s16(vget_low_s16(scaled_high_i16x8), vget_low_s16(scaled_high_i16x8));
2860
+ int32x4_t squares_high_b_i32x4 = vmull_high_s16(scaled_high_i16x8, scaled_high_i16x8);
2861
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_a_i32x4)));
2862
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_b_i32x4)));
2863
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_a_i32x4)));
2864
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_b_i32x4)));
2865
+ }
2866
+ nk_i64_t sum = vaddlvq_s32(sum_i32x4);
2867
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
2868
+ for (; idx < count; ++idx) {
2869
+ nk_f32_t value_f32;
2870
+ nk_e3m2_to_f32_serial(&data_ptr[idx], &value_f32);
2871
+ sum += (nk_i64_t)(value_f32 * 16.0f), sumsq += (nk_u64_t)(nk_i64_t)(value_f32 * value_f32 * 256.0f);
2872
+ }
2873
+ *sum_ptr = (nk_f32_t)sum / 16.0f, *sumsq_ptr = (nk_f32_t)sumsq / 256.0f;
2874
+ }
2875
+
2876
+ NK_INTERNAL void nk_reduce_moments_e3m2_neon_strided_( //
2877
+ nk_e3m2_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
2878
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
2879
+ uint8x16x2_t lut_e3m2_low;
2880
+ // table[0]: low bytes for magnitudes 0..15
2881
+ // 0x0706050403020100 → bytes [0..7] = 0,1,2,3,4,5,6,7
2882
+ // 0x1C1814100E0C0A08 → bytes [8..15] = 8,10,12,14,16,20,24,28
2883
+ lut_e3m2_low.val[0] = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
2884
+ vreinterpret_u8_u64(vcreate_u64(0x1C1814100E0C0A08ULL)));
2885
+ // table[1]: low bytes for magnitudes 16..31
2886
+ // 0x7060504038302820 → bytes [0..7] = 32,40,48,56,64,80,96,112
2887
+ // 0xC0804000E0C0A080 → bytes [8..15] = 128,160,192,224,0,64,128,192
2888
+ lut_e3m2_low.val[1] = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x7060504038302820ULL)),
2889
+ vreinterpret_u8_u64(vcreate_u64(0xC0804000E0C0A080ULL)));
2890
+ int32x4_t sum_i32x4 = vdupq_n_s32(0);
2891
+ uint64x2_t sumsq_u64x2 = vdupq_n_u64(0);
2892
+ nk_size_t idx = 0;
2893
+ if (stride_elements == 2) {
2894
+ for (; idx + 16 < count; idx += 16) {
2895
+ uint8x16x2_t loaded_u8x16x2 = vld2q_u8((nk_u8_t const *)(data_ptr + idx * 2));
2896
+ uint8x16_t raw_u8x16 = loaded_u8x16x2.val[0];
2897
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2898
+ uint8x16_t low_byte_u8x16 = vqtbl2q_u8(lut_e3m2_low, magnitude_u8x16);
2899
+ uint8x16_t high_byte_u8x16 = vandq_u8(vcgtq_u8(magnitude_u8x16, vdupq_n_u8(27)), vdupq_n_u8(1));
2900
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, vdupq_n_u8(0x20));
2901
+ uint16x8_t unsigned_low_u16x8 = vreinterpretq_u16_u8(vzip1q_u8(low_byte_u8x16, high_byte_u8x16));
2902
+ uint16x8_t unsigned_high_u16x8 = vreinterpretq_u16_u8(vzip2q_u8(low_byte_u8x16, high_byte_u8x16));
2903
+ int8x8_t is_negative_low_i8x8 = vreinterpret_s8_u8(vget_low_u8(is_negative_u8x16));
2904
+ int8x8_t is_negative_high_i8x8 = vreinterpret_s8_u8(vget_high_u8(is_negative_u8x16));
2905
+ uint16x8_t is_negative_low_u16x8 = vreinterpretq_u16_s16(vmovl_s8(is_negative_low_i8x8));
2906
+ uint16x8_t is_negative_high_u16x8 = vreinterpretq_u16_s16(vmovl_s8(is_negative_high_i8x8));
2907
+ int16x8_t positive_low_i16x8 = vreinterpretq_s16_u16(unsigned_low_u16x8);
2908
+ int16x8_t scaled_low_i16x8 = vbslq_s16(is_negative_low_u16x8, vnegq_s16(positive_low_i16x8),
2909
+ positive_low_i16x8);
2910
+ int16x8_t positive_high_i16x8 = vreinterpretq_s16_u16(unsigned_high_u16x8);
2911
+ int16x8_t scaled_high_i16x8 = vbslq_s16(is_negative_high_u16x8, vnegq_s16(positive_high_i16x8),
2912
+ positive_high_i16x8);
2913
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(scaled_low_i16x8));
2914
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(scaled_high_i16x8));
2915
+ int32x4_t squares_low_a_i32x4 = vmull_s16(vget_low_s16(scaled_low_i16x8), vget_low_s16(scaled_low_i16x8));
2916
+ int32x4_t squares_low_b_i32x4 = vmull_high_s16(scaled_low_i16x8, scaled_low_i16x8);
2917
+ int32x4_t squares_high_a_i32x4 = vmull_s16(vget_low_s16(scaled_high_i16x8),
2918
+ vget_low_s16(scaled_high_i16x8));
2919
+ int32x4_t squares_high_b_i32x4 = vmull_high_s16(scaled_high_i16x8, scaled_high_i16x8);
2920
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_a_i32x4)));
2921
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_b_i32x4)));
2922
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_a_i32x4)));
2923
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_b_i32x4)));
2924
+ }
2925
+ }
2926
+ else if (stride_elements == 3) {
2927
+ for (; idx + 16 < count; idx += 16) {
2928
+ uint8x16x3_t loaded_u8x16x3 = vld3q_u8((nk_u8_t const *)(data_ptr + idx * 3));
2929
+ uint8x16_t raw_u8x16 = loaded_u8x16x3.val[0];
2930
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2931
+ uint8x16_t low_byte_u8x16 = vqtbl2q_u8(lut_e3m2_low, magnitude_u8x16);
2932
+ uint8x16_t high_byte_u8x16 = vandq_u8(vcgtq_u8(magnitude_u8x16, vdupq_n_u8(27)), vdupq_n_u8(1));
2933
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, vdupq_n_u8(0x20));
2934
+ uint16x8_t unsigned_low_u16x8 = vreinterpretq_u16_u8(vzip1q_u8(low_byte_u8x16, high_byte_u8x16));
2935
+ uint16x8_t unsigned_high_u16x8 = vreinterpretq_u16_u8(vzip2q_u8(low_byte_u8x16, high_byte_u8x16));
2936
+ int8x8_t is_negative_low_i8x8 = vreinterpret_s8_u8(vget_low_u8(is_negative_u8x16));
2937
+ int8x8_t is_negative_high_i8x8 = vreinterpret_s8_u8(vget_high_u8(is_negative_u8x16));
2938
+ uint16x8_t is_negative_low_u16x8 = vreinterpretq_u16_s16(vmovl_s8(is_negative_low_i8x8));
2939
+ uint16x8_t is_negative_high_u16x8 = vreinterpretq_u16_s16(vmovl_s8(is_negative_high_i8x8));
2940
+ int16x8_t positive_low_i16x8 = vreinterpretq_s16_u16(unsigned_low_u16x8);
2941
+ int16x8_t scaled_low_i16x8 = vbslq_s16(is_negative_low_u16x8, vnegq_s16(positive_low_i16x8),
2942
+ positive_low_i16x8);
2943
+ int16x8_t positive_high_i16x8 = vreinterpretq_s16_u16(unsigned_high_u16x8);
2944
+ int16x8_t scaled_high_i16x8 = vbslq_s16(is_negative_high_u16x8, vnegq_s16(positive_high_i16x8),
2945
+ positive_high_i16x8);
2946
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(scaled_low_i16x8));
2947
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(scaled_high_i16x8));
2948
+ int32x4_t squares_low_a_i32x4 = vmull_s16(vget_low_s16(scaled_low_i16x8), vget_low_s16(scaled_low_i16x8));
2949
+ int32x4_t squares_low_b_i32x4 = vmull_high_s16(scaled_low_i16x8, scaled_low_i16x8);
2950
+ int32x4_t squares_high_a_i32x4 = vmull_s16(vget_low_s16(scaled_high_i16x8),
2951
+ vget_low_s16(scaled_high_i16x8));
2952
+ int32x4_t squares_high_b_i32x4 = vmull_high_s16(scaled_high_i16x8, scaled_high_i16x8);
2953
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_a_i32x4)));
2954
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_b_i32x4)));
2955
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_a_i32x4)));
2956
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_b_i32x4)));
2957
+ }
2958
+ }
2959
+ else {
2960
+ for (; idx + 16 < count; idx += 16) {
2961
+ uint8x16x4_t loaded_u8x16x4 = vld4q_u8((nk_u8_t const *)(data_ptr + idx * 4));
2962
+ uint8x16_t raw_u8x16 = loaded_u8x16x4.val[0];
2963
+ uint8x16_t magnitude_u8x16 = vandq_u8(raw_u8x16, vdupq_n_u8(0x1F));
2964
+ uint8x16_t low_byte_u8x16 = vqtbl2q_u8(lut_e3m2_low, magnitude_u8x16);
2965
+ uint8x16_t high_byte_u8x16 = vandq_u8(vcgtq_u8(magnitude_u8x16, vdupq_n_u8(27)), vdupq_n_u8(1));
2966
+ uint8x16_t is_negative_u8x16 = vtstq_u8(raw_u8x16, vdupq_n_u8(0x20));
2967
+ uint16x8_t unsigned_low_u16x8 = vreinterpretq_u16_u8(vzip1q_u8(low_byte_u8x16, high_byte_u8x16));
2968
+ uint16x8_t unsigned_high_u16x8 = vreinterpretq_u16_u8(vzip2q_u8(low_byte_u8x16, high_byte_u8x16));
2969
+ int8x8_t is_negative_low_i8x8 = vreinterpret_s8_u8(vget_low_u8(is_negative_u8x16));
2970
+ int8x8_t is_negative_high_i8x8 = vreinterpret_s8_u8(vget_high_u8(is_negative_u8x16));
2971
+ uint16x8_t is_negative_low_u16x8 = vreinterpretq_u16_s16(vmovl_s8(is_negative_low_i8x8));
2972
+ uint16x8_t is_negative_high_u16x8 = vreinterpretq_u16_s16(vmovl_s8(is_negative_high_i8x8));
2973
+ int16x8_t positive_low_i16x8 = vreinterpretq_s16_u16(unsigned_low_u16x8);
2974
+ int16x8_t scaled_low_i16x8 = vbslq_s16(is_negative_low_u16x8, vnegq_s16(positive_low_i16x8),
2975
+ positive_low_i16x8);
2976
+ int16x8_t positive_high_i16x8 = vreinterpretq_s16_u16(unsigned_high_u16x8);
2977
+ int16x8_t scaled_high_i16x8 = vbslq_s16(is_negative_high_u16x8, vnegq_s16(positive_high_i16x8),
2978
+ positive_high_i16x8);
2979
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(scaled_low_i16x8));
2980
+ sum_i32x4 = vaddq_s32(sum_i32x4, vpaddlq_s16(scaled_high_i16x8));
2981
+ int32x4_t squares_low_a_i32x4 = vmull_s16(vget_low_s16(scaled_low_i16x8), vget_low_s16(scaled_low_i16x8));
2982
+ int32x4_t squares_low_b_i32x4 = vmull_high_s16(scaled_low_i16x8, scaled_low_i16x8);
2983
+ int32x4_t squares_high_a_i32x4 = vmull_s16(vget_low_s16(scaled_high_i16x8),
2984
+ vget_low_s16(scaled_high_i16x8));
2985
+ int32x4_t squares_high_b_i32x4 = vmull_high_s16(scaled_high_i16x8, scaled_high_i16x8);
2986
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_a_i32x4)));
2987
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_low_b_i32x4)));
2988
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_a_i32x4)));
2989
+ sumsq_u64x2 = vaddq_u64(sumsq_u64x2, vpaddlq_u32(vreinterpretq_u32_s32(squares_high_b_i32x4)));
2990
+ }
2991
+ }
2992
+ nk_i64_t sum = vaddlvq_s32(sum_i32x4);
2993
+ nk_u64_t sumsq = vgetq_lane_u64(sumsq_u64x2, 0) + vgetq_lane_u64(sumsq_u64x2, 1);
2994
+ for (; idx < count; ++idx) {
2995
+ nk_f32_t val;
2996
+ nk_e3m2_to_f32_serial((nk_e3m2_t const *)(data_ptr + idx * stride_elements), &val);
2997
+ sum += (nk_i64_t)(val * 16.0f), sumsq += (nk_u64_t)(nk_i64_t)(val * val * 256.0f);
2998
+ }
2999
+ *sum_ptr = (nk_f32_t)sum / 16.0f, *sumsq_ptr = (nk_f32_t)sumsq / 256.0f;
3000
+ }
3001
+
3002
+ NK_PUBLIC void nk_reduce_moments_e3m2_neon( //
3003
+ nk_e3m2_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
3004
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3005
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_e3m2_t);
3006
+ int aligned = (stride_bytes % sizeof(nk_e3m2_t) == 0);
3007
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
3008
+ else if (!aligned) nk_reduce_moments_e3m2_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3009
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 16) {
3010
+ nk_size_t left_count = count / 2;
3011
+ nk_f32_t left_sum, left_sumsq, right_sum, right_sumsq;
3012
+ nk_reduce_moments_e3m2_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
3013
+ nk_reduce_moments_e3m2_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
3014
+ &right_sum, &right_sumsq);
3015
+ *sum_ptr = left_sum + right_sum, *sumsq_ptr = left_sumsq + right_sumsq;
3016
+ }
3017
+ else if (stride_elements == 1) nk_reduce_moments_e3m2_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
3018
+ else if (stride_elements <= 4)
3019
+ nk_reduce_moments_e3m2_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
3020
+ else nk_reduce_moments_e3m2_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3021
+ }
3022
+
3023
+ NK_INTERNAL void nk_reduce_minmax_e3m2_neon_contiguous_( //
3024
+ nk_e3m2_t const *data_ptr, nk_size_t count, //
3025
+ nk_e3m2_t *min_value_ptr, nk_size_t *min_index_ptr, //
3026
+ nk_e3m2_t *max_value_ptr, nk_size_t *max_index_ptr) {
3027
+ // Handle initial chunk — partial or full
3028
+ uint8x16_t first_comparable_u8x16;
3029
+ nk_size_t first_count = count < 16 ? count : 16;
3030
+ if (count < 16) {
3031
+ nk_b128_vec_t first_vec;
3032
+ nk_partial_load_b8x16_serial_(data_ptr, &first_vec, count);
3033
+ first_comparable_u8x16 = nk_fp6x16_to_comparable_neon_(first_vec.u8x16);
3034
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3035
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3036
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)count));
3037
+ first_comparable_u8x16 = vbslq_u8(valid_u8x16, first_comparable_u8x16, vdupq_n_u8(0));
3038
+ }
3039
+ else {
3040
+ uint8x16_t first_raw_u8x16 = vld1q_u8((nk_u8_t const *)data_ptr);
3041
+ first_comparable_u8x16 = nk_fp6x16_to_comparable_neon_(first_raw_u8x16);
3042
+ }
3043
+ uint8x16_t lane_indices_init_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3044
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3045
+ uint8x16_t valid_init_u8x16 = vcltq_u8(lane_indices_init_u8x16, vdupq_n_u8((nk_u8_t)first_count));
3046
+ uint8x16_t min_u8x16 = vbslq_u8(valid_init_u8x16, first_comparable_u8x16, vdupq_n_u8(0xFF));
3047
+ uint8x16_t max_u8x16 = first_comparable_u8x16;
3048
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
3049
+ uint8x16_t iter_u8x16 = vdupq_n_u8(1), one_u8x16 = vdupq_n_u8(1);
3050
+ nk_size_t idx = first_count;
3051
+ for (; idx + 16 <= count; idx += 16) {
3052
+ uint8x16_t raw_u8x16 = vld1q_u8((nk_u8_t const *)(data_ptr + idx));
3053
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(raw_u8x16);
3054
+ uint8x16_t less_u8x16 = vcltq_u8(comparable_u8x16, min_u8x16);
3055
+ uint8x16_t greater_u8x16 = vcgtq_u8(comparable_u8x16, max_u8x16);
3056
+ min_u8x16 = vbslq_u8(less_u8x16, comparable_u8x16, min_u8x16);
3057
+ max_u8x16 = vbslq_u8(greater_u8x16, comparable_u8x16, max_u8x16);
3058
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3059
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3060
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
3061
+ }
3062
+ nk_size_t remaining = count - idx;
3063
+ if (remaining > 0) {
3064
+ nk_b128_vec_t tail_vec;
3065
+ nk_partial_load_b8x16_serial_(data_ptr + idx, &tail_vec, remaining);
3066
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(tail_vec.u8x16);
3067
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3068
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3069
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)remaining));
3070
+ uint8x16_t data_for_min_u8x16 = vbslq_u8(valid_u8x16, comparable_u8x16, vdupq_n_u8(0xFF));
3071
+ uint8x16_t data_for_max_u8x16 = vbslq_u8(valid_u8x16, comparable_u8x16, vdupq_n_u8(0));
3072
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
3073
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
3074
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
3075
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
3076
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3077
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3078
+ }
3079
+ nk_u8_t min_comparable = vminvq_u8(min_u8x16), max_comparable = vmaxvq_u8(max_u8x16);
3080
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_comparable));
3081
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
3082
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
3083
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_comparable));
3084
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
3085
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
3086
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3087
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3088
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
3089
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
3090
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3091
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
3092
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
3093
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
3094
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
3095
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3096
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
3097
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
3098
+ *min_value_ptr = nk_comparable_to_fp6_(min_comparable), *min_index_ptr = min_idx;
3099
+ *max_value_ptr = nk_comparable_to_fp6_(max_comparable), *max_index_ptr = max_idx;
3100
+ }
3101
+
3102
+ NK_INTERNAL void nk_reduce_minmax_e3m2_neon_strided_( //
3103
+ nk_e3m2_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
3104
+ nk_e3m2_t *min_value_ptr, nk_size_t *min_index_ptr, //
3105
+ nk_e3m2_t *max_value_ptr, nk_size_t *max_index_ptr) {
3106
+ uint8x16_t min_u8x16 = vdupq_n_u8(0xFF), max_u8x16 = vdupq_n_u8(0);
3107
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
3108
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
3109
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3110
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3111
+ nk_size_t idx = 0;
3112
+ uint8x16_t data_for_min_u8x16, data_for_max_u8x16;
3113
+
3114
+ nk_reduce_minmax_e3m2_neon_cycle:
3115
+ if (stride_elements == 2 && idx + 16 < count) {
3116
+ uint8x16x2_t loaded = vld2q_u8((nk_u8_t const *)(data_ptr + idx * 2));
3117
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(loaded.val[0]);
3118
+ data_for_min_u8x16 = comparable_u8x16;
3119
+ data_for_max_u8x16 = comparable_u8x16;
3120
+ idx += 16;
3121
+ }
3122
+ else if (stride_elements == 3 && idx + 16 < count) {
3123
+ uint8x16x3_t loaded = vld3q_u8((nk_u8_t const *)(data_ptr + idx * 3));
3124
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(loaded.val[0]);
3125
+ data_for_min_u8x16 = comparable_u8x16;
3126
+ data_for_max_u8x16 = comparable_u8x16;
3127
+ idx += 16;
3128
+ }
3129
+ else if (stride_elements == 4 && idx + 16 < count) {
3130
+ uint8x16x4_t loaded = vld4q_u8((nk_u8_t const *)(data_ptr + idx * 4));
3131
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(loaded.val[0]);
3132
+ data_for_min_u8x16 = comparable_u8x16;
3133
+ data_for_max_u8x16 = comparable_u8x16;
3134
+ idx += 16;
3135
+ }
3136
+ else if (idx < count) {
3137
+ nk_b128_vec_t tail_vec;
3138
+ nk_strided_load_b8x16_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
3139
+ uint8x16_t comparable_u8x16 = nk_fp6x16_to_comparable_neon_(tail_vec.u8x16);
3140
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)(count - idx)));
3141
+ data_for_min_u8x16 = vbslq_u8(valid_u8x16, comparable_u8x16, vdupq_n_u8(0xFF));
3142
+ data_for_max_u8x16 = vbslq_u8(valid_u8x16, comparable_u8x16, vdupq_n_u8(0x00));
3143
+ idx = count;
3144
+ }
3145
+ else {
3146
+ nk_u8_t min_comparable = vminvq_u8(min_u8x16), max_comparable = vmaxvq_u8(max_u8x16);
3147
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_comparable));
3148
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
3149
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
3150
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_comparable));
3151
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
3152
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
3153
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
3154
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
3155
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3156
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
3157
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
3158
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
3159
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
3160
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3161
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
3162
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
3163
+ *min_value_ptr = nk_comparable_to_fp6_(min_comparable), *min_index_ptr = min_idx;
3164
+ *max_value_ptr = nk_comparable_to_fp6_(max_comparable), *max_index_ptr = max_idx;
3165
+ return;
3166
+ }
3167
+
3168
+ // Shared update body
3169
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
3170
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
3171
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
3172
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
3173
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3174
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3175
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
3176
+ goto nk_reduce_minmax_e3m2_neon_cycle;
3177
+ }
3178
+
3179
+ NK_PUBLIC void nk_reduce_minmax_e3m2_neon( //
3180
+ nk_e3m2_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
3181
+ nk_e3m2_t *min_value_ptr, nk_size_t *min_index_ptr, //
3182
+ nk_e3m2_t *max_value_ptr, nk_size_t *max_index_ptr) {
3183
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_e3m2_t);
3184
+ int aligned = (stride_bytes % sizeof(nk_e3m2_t) == 0);
3185
+ if (count == 0)
3186
+ *min_value_ptr = NK_E3M2_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_E3M2_MIN,
3187
+ *max_index_ptr = NK_SIZE_MAX;
3188
+ else if (!aligned)
3189
+ nk_reduce_minmax_e3m2_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
3190
+ max_index_ptr);
3191
+ else if (count > (nk_size_t)256 * 16) {
3192
+ nk_size_t left_count = count / 2;
3193
+ nk_e3m2_t left_min_value, right_min_value, left_max_value, right_max_value;
3194
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
3195
+ nk_reduce_minmax_e3m2_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index,
3196
+ &left_max_value, &left_max_index);
3197
+ nk_reduce_minmax_e3m2_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
3198
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
3199
+ if (nk_e3m2_order_serial(right_min_value, left_min_value) < 0)
3200
+ *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
3201
+ else *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
3202
+ if (nk_e3m2_order_serial(right_max_value, left_max_value) > 0)
3203
+ *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
3204
+ else *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
3205
+ }
3206
+ else if (stride_elements == 1)
3207
+ nk_reduce_minmax_e3m2_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
3208
+ max_index_ptr);
3209
+ else if (stride_elements <= 4)
3210
+ nk_reduce_minmax_e3m2_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
3211
+ max_value_ptr, max_index_ptr);
3212
+ else
3213
+ nk_reduce_minmax_e3m2_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
3214
+ max_index_ptr);
3215
+ }
3216
+
3217
+ NK_INTERNAL void nk_reduce_moments_e4m3_neon_contiguous_( //
3218
+ nk_e4m3_t const *data_ptr, nk_size_t count, //
3219
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3220
+ float32x4_t sum_f32x4 = vdupq_n_f32(0), sumsq_f32x4 = vdupq_n_f32(0);
3221
+ nk_size_t idx = 0;
3222
+ for (; idx + 16 <= count; idx += 16) {
3223
+ uint8x16_t raw_u8x16 = vld1q_u8((nk_u8_t const *)(data_ptr + idx));
3224
+ float16x8_t half_low_f16x8, half_high_f16x8;
3225
+ nk_e4m3x16_to_f16x8x2_neon_(raw_u8x16, &half_low_f16x8, &half_high_f16x8);
3226
+ float32x4_t a_f32x4 = vcvt_f32_f16(vget_low_f16(half_low_f16x8));
3227
+ float32x4_t b_f32x4 = vcvt_high_f32_f16(half_low_f16x8);
3228
+ float32x4_t c_f32x4 = vcvt_f32_f16(vget_low_f16(half_high_f16x8));
3229
+ float32x4_t d_f32x4 = vcvt_high_f32_f16(half_high_f16x8);
3230
+ sum_f32x4 = vaddq_f32(vaddq_f32(sum_f32x4, vaddq_f32(a_f32x4, b_f32x4)), vaddq_f32(c_f32x4, d_f32x4));
3231
+ sumsq_f32x4 = vfmaq_f32(vfmaq_f32(vfmaq_f32(vfmaq_f32( //
3232
+ sumsq_f32x4, a_f32x4, a_f32x4),
3233
+ b_f32x4, b_f32x4),
3234
+ c_f32x4, c_f32x4),
3235
+ d_f32x4, d_f32x4);
3236
+ }
3237
+ nk_f32_t sum = vaddvq_f32(sum_f32x4), sumsq = vaddvq_f32(sumsq_f32x4);
3238
+ for (; idx < count; ++idx) {
3239
+ nk_f32_t value_f32;
3240
+ nk_e4m3_to_f32_serial(&data_ptr[idx], &value_f32);
3241
+ sum += value_f32, sumsq += value_f32 * value_f32;
3242
+ }
3243
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
3244
+ }
3245
+
3246
+ NK_INTERNAL void nk_reduce_moments_e4m3_neon_strided_( //
3247
+ nk_e4m3_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
3248
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3249
+ float32x4_t sum_f32x4 = vdupq_n_f32(0), sumsq_f32x4 = vdupq_n_f32(0);
3250
+ nk_size_t idx = 0;
3251
+ if (stride_elements == 2) {
3252
+ for (; idx + 16 < count; idx += 16) {
3253
+ uint8x16x2_t loaded_u8x16x2 = vld2q_u8((nk_u8_t const *)(data_ptr + idx * 2));
3254
+ float16x8_t half_low_f16x8, half_high_f16x8;
3255
+ nk_e4m3x16_to_f16x8x2_neon_(loaded_u8x16x2.val[0], &half_low_f16x8, &half_high_f16x8);
3256
+ float32x4_t a_f32x4 = vcvt_f32_f16(vget_low_f16(half_low_f16x8));
3257
+ float32x4_t b_f32x4 = vcvt_high_f32_f16(half_low_f16x8);
3258
+ float32x4_t c_f32x4 = vcvt_f32_f16(vget_low_f16(half_high_f16x8));
3259
+ float32x4_t d_f32x4 = vcvt_high_f32_f16(half_high_f16x8);
3260
+ sum_f32x4 = vaddq_f32(vaddq_f32(sum_f32x4, vaddq_f32(a_f32x4, b_f32x4)), vaddq_f32(c_f32x4, d_f32x4));
3261
+ sumsq_f32x4 = vfmaq_f32(vfmaq_f32(vfmaq_f32(vfmaq_f32( //
3262
+ sumsq_f32x4, a_f32x4, a_f32x4),
3263
+ b_f32x4, b_f32x4),
3264
+ c_f32x4, c_f32x4),
3265
+ d_f32x4, d_f32x4);
3266
+ }
3267
+ }
3268
+ else if (stride_elements == 3) {
3269
+ for (; idx + 16 < count; idx += 16) {
3270
+ uint8x16x3_t loaded_u8x16x3 = vld3q_u8((nk_u8_t const *)(data_ptr + idx * 3));
3271
+ float16x8_t half_low_f16x8, half_high_f16x8;
3272
+ nk_e4m3x16_to_f16x8x2_neon_(loaded_u8x16x3.val[0], &half_low_f16x8, &half_high_f16x8);
3273
+ float32x4_t a_f32x4 = vcvt_f32_f16(vget_low_f16(half_low_f16x8));
3274
+ float32x4_t b_f32x4 = vcvt_high_f32_f16(half_low_f16x8);
3275
+ float32x4_t c_f32x4 = vcvt_f32_f16(vget_low_f16(half_high_f16x8));
3276
+ float32x4_t d_f32x4 = vcvt_high_f32_f16(half_high_f16x8);
3277
+ sum_f32x4 = vaddq_f32(vaddq_f32(sum_f32x4, vaddq_f32(a_f32x4, b_f32x4)), vaddq_f32(c_f32x4, d_f32x4));
3278
+ sumsq_f32x4 = vfmaq_f32(vfmaq_f32(vfmaq_f32(vfmaq_f32( //
3279
+ sumsq_f32x4, a_f32x4, a_f32x4),
3280
+ b_f32x4, b_f32x4),
3281
+ c_f32x4, c_f32x4),
3282
+ d_f32x4, d_f32x4);
3283
+ }
3284
+ }
3285
+ else {
3286
+ for (; idx + 16 < count; idx += 16) {
3287
+ uint8x16x4_t loaded_u8x16x4 = vld4q_u8((nk_u8_t const *)(data_ptr + idx * 4));
3288
+ float16x8_t half_low_f16x8, half_high_f16x8;
3289
+ nk_e4m3x16_to_f16x8x2_neon_(loaded_u8x16x4.val[0], &half_low_f16x8, &half_high_f16x8);
3290
+ float32x4_t a_f32x4 = vcvt_f32_f16(vget_low_f16(half_low_f16x8));
3291
+ float32x4_t b_f32x4 = vcvt_high_f32_f16(half_low_f16x8);
3292
+ float32x4_t c_f32x4 = vcvt_f32_f16(vget_low_f16(half_high_f16x8));
3293
+ float32x4_t d_f32x4 = vcvt_high_f32_f16(half_high_f16x8);
3294
+ sum_f32x4 = vaddq_f32(vaddq_f32(sum_f32x4, vaddq_f32(a_f32x4, b_f32x4)), vaddq_f32(c_f32x4, d_f32x4));
3295
+ sumsq_f32x4 = vfmaq_f32(vfmaq_f32(vfmaq_f32(vfmaq_f32( //
3296
+ sumsq_f32x4, a_f32x4, a_f32x4),
3297
+ b_f32x4, b_f32x4),
3298
+ c_f32x4, c_f32x4),
3299
+ d_f32x4, d_f32x4);
3300
+ }
3301
+ }
3302
+ nk_f32_t sum = vaddvq_f32(sum_f32x4), sumsq = vaddvq_f32(sumsq_f32x4);
3303
+ for (; idx < count; ++idx) {
3304
+ nk_f32_t val;
3305
+ nk_e4m3_to_f32_serial((nk_e4m3_t const *)(data_ptr + idx * stride_elements), &val);
3306
+ sum += val, sumsq += val * val;
3307
+ }
3308
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
3309
+ }
3310
+
3311
+ NK_PUBLIC void nk_reduce_moments_e4m3_neon( //
3312
+ nk_e4m3_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
3313
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3314
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_e4m3_t);
3315
+ int aligned = (stride_bytes % sizeof(nk_e4m3_t) == 0);
3316
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
3317
+ else if (!aligned) nk_reduce_moments_e4m3_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3318
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 16) {
3319
+ nk_size_t left_count = count / 2;
3320
+ nk_f32_t left_sum, left_sumsq, right_sum, right_sumsq;
3321
+ nk_reduce_moments_e4m3_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
3322
+ nk_reduce_moments_e4m3_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
3323
+ &right_sum, &right_sumsq);
3324
+ *sum_ptr = left_sum + right_sum, *sumsq_ptr = left_sumsq + right_sumsq;
3325
+ }
3326
+ else if (stride_elements == 1) nk_reduce_moments_e4m3_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
3327
+ else if (stride_elements <= 4)
3328
+ nk_reduce_moments_e4m3_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
3329
+ else nk_reduce_moments_e4m3_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3330
+ }
3331
+
3332
+ NK_INTERNAL void nk_reduce_minmax_e4m3_neon_contiguous_( //
3333
+ nk_e4m3_t const *data_ptr, nk_size_t count, //
3334
+ nk_e4m3_t *min_value_ptr, nk_size_t *min_index_ptr, //
3335
+ nk_e4m3_t *max_value_ptr, nk_size_t *max_index_ptr) {
3336
+ uint8x16_t min_u8x16 = vdupq_n_u8(0xFF), max_u8x16 = vdupq_n_u8(0);
3337
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
3338
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
3339
+ nk_size_t idx = 0;
3340
+ for (; idx + 16 <= count; idx += 16) {
3341
+ uint8x16_t raw_u8x16 = vld1q_u8((nk_u8_t const *)(data_ptr + idx));
3342
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(raw_u8x16);
3343
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vceqq_u8(comparable_u8x16, vdupq_n_u8(0x00)),
3344
+ vceqq_u8(comparable_u8x16, vdupq_n_u8(0xFF)));
3345
+ uint8x16_t data_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3346
+ uint8x16_t data_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3347
+ uint8x16_t less_u8x16 = vcltq_u8(data_min_u8x16, min_u8x16);
3348
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_max_u8x16, max_u8x16);
3349
+ min_u8x16 = vbslq_u8(less_u8x16, data_min_u8x16, min_u8x16);
3350
+ max_u8x16 = vbslq_u8(greater_u8x16, data_max_u8x16, max_u8x16);
3351
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3352
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3353
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
3354
+ }
3355
+ nk_size_t remaining = count - idx;
3356
+ if (remaining > 0) {
3357
+ nk_b128_vec_t tail_vec;
3358
+ nk_partial_load_b8x16_serial_(data_ptr + idx, &tail_vec, remaining);
3359
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(tail_vec.u8x16);
3360
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vceqq_u8(comparable_u8x16, vdupq_n_u8(0x00)),
3361
+ vceqq_u8(comparable_u8x16, vdupq_n_u8(0xFF)));
3362
+ uint8x16_t nan_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3363
+ uint8x16_t nan_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3364
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3365
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3366
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)remaining));
3367
+ uint8x16_t data_for_min_u8x16 = vbslq_u8(valid_u8x16, nan_min_u8x16, vdupq_n_u8(0xFF));
3368
+ uint8x16_t data_for_max_u8x16 = vbslq_u8(valid_u8x16, nan_max_u8x16, vdupq_n_u8(0));
3369
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
3370
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
3371
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
3372
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
3373
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3374
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3375
+ }
3376
+ nk_u8_t min_comparable = vminvq_u8(min_u8x16), max_comparable = vmaxvq_u8(max_u8x16);
3377
+ // If min stayed at 0xFF, all values were NaN
3378
+ if (min_comparable == 0xFF) {
3379
+ *min_value_ptr = (nk_e4m3_t)NK_E4M3_MAX, *min_index_ptr = NK_SIZE_MAX;
3380
+ *max_value_ptr = (nk_e4m3_t)NK_E4M3_MIN, *max_index_ptr = NK_SIZE_MAX;
3381
+ return;
3382
+ }
3383
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_comparable));
3384
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
3385
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
3386
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_comparable));
3387
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
3388
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
3389
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3390
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3391
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
3392
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
3393
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3394
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
3395
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
3396
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
3397
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
3398
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3399
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
3400
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
3401
+ *min_value_ptr = nk_comparable_to_fp8_(min_comparable), *min_index_ptr = min_idx;
3402
+ *max_value_ptr = nk_comparable_to_fp8_(max_comparable), *max_index_ptr = max_idx;
3403
+ }
3404
+
3405
+ NK_INTERNAL void nk_reduce_minmax_e4m3_neon_strided_( //
3406
+ nk_e4m3_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
3407
+ nk_e4m3_t *min_value_ptr, nk_size_t *min_index_ptr, //
3408
+ nk_e4m3_t *max_value_ptr, nk_size_t *max_index_ptr) {
3409
+ uint8x16_t min_u8x16 = vdupq_n_u8(0xFF), max_u8x16 = vdupq_n_u8(0);
3410
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
3411
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
3412
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3413
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3414
+ nk_size_t idx = 0;
3415
+ uint8x16_t data_for_min_u8x16, data_for_max_u8x16;
3416
+
3417
+ nk_reduce_minmax_e4m3_neon_cycle:
3418
+ if (stride_elements == 2 && idx + 16 < count) {
3419
+ uint8x16x2_t loaded = vld2q_u8((nk_u8_t const *)(data_ptr + idx * 2));
3420
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(loaded.val[0]);
3421
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vceqq_u8(comparable_u8x16, vdupq_n_u8(0x00)),
3422
+ vceqq_u8(comparable_u8x16, vdupq_n_u8(0xFF)));
3423
+ data_for_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3424
+ data_for_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3425
+ idx += 16;
3426
+ }
3427
+ else if (stride_elements == 3 && idx + 16 < count) {
3428
+ uint8x16x3_t loaded = vld3q_u8((nk_u8_t const *)(data_ptr + idx * 3));
3429
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(loaded.val[0]);
3430
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vceqq_u8(comparable_u8x16, vdupq_n_u8(0x00)),
3431
+ vceqq_u8(comparable_u8x16, vdupq_n_u8(0xFF)));
3432
+ data_for_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3433
+ data_for_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3434
+ idx += 16;
3435
+ }
3436
+ else if (stride_elements == 4 && idx + 16 < count) {
3437
+ uint8x16x4_t loaded = vld4q_u8((nk_u8_t const *)(data_ptr + idx * 4));
3438
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(loaded.val[0]);
3439
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vceqq_u8(comparable_u8x16, vdupq_n_u8(0x00)),
3440
+ vceqq_u8(comparable_u8x16, vdupq_n_u8(0xFF)));
3441
+ data_for_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3442
+ data_for_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3443
+ idx += 16;
3444
+ }
3445
+ else if (idx < count) {
3446
+ nk_b128_vec_t tail_vec;
3447
+ nk_strided_load_b8x16_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
3448
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(tail_vec.u8x16);
3449
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vceqq_u8(comparable_u8x16, vdupq_n_u8(0x00)),
3450
+ vceqq_u8(comparable_u8x16, vdupq_n_u8(0xFF)));
3451
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)(count - idx)));
3452
+ uint8x16_t invalid_or_nan_u8x16 = vornq_u8(is_nan_u8x16, valid_u8x16);
3453
+ data_for_min_u8x16 = vbslq_u8(invalid_or_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3454
+ data_for_max_u8x16 = vbslq_u8(invalid_or_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3455
+ idx = count;
3456
+ }
3457
+ else {
3458
+ nk_u8_t min_comparable = vminvq_u8(min_u8x16), max_comparable = vmaxvq_u8(max_u8x16);
3459
+ if (min_comparable == 0xFF) {
3460
+ *min_value_ptr = (nk_e4m3_t)NK_E4M3_MAX, *min_index_ptr = NK_SIZE_MAX;
3461
+ *max_value_ptr = (nk_e4m3_t)NK_E4M3_MIN, *max_index_ptr = NK_SIZE_MAX;
3462
+ return;
3463
+ }
3464
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_comparable));
3465
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
3466
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
3467
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_comparable));
3468
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
3469
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
3470
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
3471
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
3472
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3473
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
3474
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
3475
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
3476
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
3477
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3478
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
3479
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
3480
+ *min_value_ptr = nk_comparable_to_fp8_(min_comparable), *min_index_ptr = min_idx;
3481
+ *max_value_ptr = nk_comparable_to_fp8_(max_comparable), *max_index_ptr = max_idx;
3482
+ return;
3483
+ }
3484
+
3485
+ // Shared update body
3486
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
3487
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
3488
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
3489
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
3490
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3491
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3492
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
3493
+ goto nk_reduce_minmax_e4m3_neon_cycle;
3494
+ }
3495
+
3496
+ NK_PUBLIC void nk_reduce_minmax_e4m3_neon( //
3497
+ nk_e4m3_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
3498
+ nk_e4m3_t *min_value_ptr, nk_size_t *min_index_ptr, //
3499
+ nk_e4m3_t *max_value_ptr, nk_size_t *max_index_ptr) {
3500
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_e4m3_t);
3501
+ int aligned = (stride_bytes % sizeof(nk_e4m3_t) == 0);
3502
+ if (count == 0)
3503
+ *min_value_ptr = NK_E4M3_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_E4M3_MIN,
3504
+ *max_index_ptr = NK_SIZE_MAX;
3505
+ else if (!aligned)
3506
+ nk_reduce_minmax_e4m3_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
3507
+ max_index_ptr);
3508
+ else if (count > (nk_size_t)256 * 16) {
3509
+ nk_size_t left_count = count / 2;
3510
+ nk_e4m3_t left_min_value, right_min_value, left_max_value, right_max_value;
3511
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
3512
+ nk_reduce_minmax_e4m3_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index,
3513
+ &left_max_value, &left_max_index);
3514
+ nk_reduce_minmax_e4m3_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
3515
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
3516
+ if (left_min_index == NK_SIZE_MAX)
3517
+ *min_value_ptr = right_min_value,
3518
+ *min_index_ptr = right_min_index == NK_SIZE_MAX ? NK_SIZE_MAX : left_count + right_min_index;
3519
+ else if (right_min_index == NK_SIZE_MAX || nk_e4m3_order_serial(left_min_value, right_min_value) <= 0)
3520
+ *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
3521
+ else *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
3522
+ if (left_max_index == NK_SIZE_MAX)
3523
+ *max_value_ptr = right_max_value,
3524
+ *max_index_ptr = right_max_index == NK_SIZE_MAX ? NK_SIZE_MAX : left_count + right_max_index;
3525
+ else if (right_max_index == NK_SIZE_MAX || nk_e4m3_order_serial(right_max_value, left_max_value) <= 0)
3526
+ *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
3527
+ else *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
3528
+ }
3529
+ else if (stride_elements == 1)
3530
+ nk_reduce_minmax_e4m3_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
3531
+ max_index_ptr);
3532
+ else if (stride_elements <= 4)
3533
+ nk_reduce_minmax_e4m3_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
3534
+ max_value_ptr, max_index_ptr);
3535
+ else
3536
+ nk_reduce_minmax_e4m3_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
3537
+ max_index_ptr);
3538
+ }
3539
+
3540
+ NK_INTERNAL void nk_reduce_moments_e5m2_neon_contiguous_( //
3541
+ nk_e5m2_t const *data_ptr, nk_size_t count, //
3542
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3543
+ float32x4_t sum_f32x4 = vdupq_n_f32(0), sumsq_f32x4 = vdupq_n_f32(0);
3544
+ nk_size_t idx = 0;
3545
+ for (; idx + 8 <= count; idx += 8) {
3546
+ uint8x8_t raw_u8x8 = vld1_u8((nk_u8_t const *)(data_ptr + idx));
3547
+ float16x8_t half_f16x8 = nk_e5m2x8_to_f16x8_neon_(raw_u8x8);
3548
+ float32x4_t low_f32x4 = vcvt_f32_f16(vget_low_f16(half_f16x8));
3549
+ float32x4_t high_f32x4 = vcvt_high_f32_f16(half_f16x8);
3550
+ sum_f32x4 = vaddq_f32(sum_f32x4, vaddq_f32(low_f32x4, high_f32x4));
3551
+ sumsq_f32x4 = vfmaq_f32(vfmaq_f32(sumsq_f32x4, low_f32x4, low_f32x4), high_f32x4, high_f32x4);
3552
+ }
3553
+ nk_f32_t sum = vaddvq_f32(sum_f32x4), sumsq = vaddvq_f32(sumsq_f32x4);
3554
+ for (; idx < count; ++idx) {
3555
+ nk_f32_t value_f32;
3556
+ nk_e5m2_to_f32_serial(&data_ptr[idx], &value_f32);
3557
+ sum += value_f32, sumsq += value_f32 * value_f32;
3558
+ }
3559
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
3560
+ }
3561
+
3562
+ NK_INTERNAL void nk_reduce_moments_e5m2_neon_strided_( //
3563
+ nk_e5m2_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
3564
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3565
+ float32x4_t sum_f32x4 = vdupq_n_f32(0), sumsq_f32x4 = vdupq_n_f32(0);
3566
+ nk_size_t idx = 0;
3567
+ if (stride_elements == 2) {
3568
+ for (; idx + 8 < count; idx += 8) {
3569
+ uint8x8x2_t loaded_u8x8x2 = vld2_u8((nk_u8_t const *)(data_ptr + idx * 2));
3570
+ float16x8_t half_f16x8 = nk_e5m2x8_to_f16x8_neon_(loaded_u8x8x2.val[0]);
3571
+ float32x4_t low_f32x4 = vcvt_f32_f16(vget_low_f16(half_f16x8));
3572
+ float32x4_t high_f32x4 = vcvt_high_f32_f16(half_f16x8);
3573
+ sum_f32x4 = vaddq_f32(sum_f32x4, vaddq_f32(low_f32x4, high_f32x4));
3574
+ sumsq_f32x4 = vfmaq_f32(vfmaq_f32(sumsq_f32x4, low_f32x4, low_f32x4), high_f32x4, high_f32x4);
3575
+ }
3576
+ }
3577
+ else if (stride_elements == 3) {
3578
+ for (; idx + 8 < count; idx += 8) {
3579
+ uint8x8x3_t loaded_u8x8x3 = vld3_u8((nk_u8_t const *)(data_ptr + idx * 3));
3580
+ float16x8_t half_f16x8 = nk_e5m2x8_to_f16x8_neon_(loaded_u8x8x3.val[0]);
3581
+ float32x4_t low_f32x4 = vcvt_f32_f16(vget_low_f16(half_f16x8));
3582
+ float32x4_t high_f32x4 = vcvt_high_f32_f16(half_f16x8);
3583
+ sum_f32x4 = vaddq_f32(sum_f32x4, vaddq_f32(low_f32x4, high_f32x4));
3584
+ sumsq_f32x4 = vfmaq_f32(vfmaq_f32(sumsq_f32x4, low_f32x4, low_f32x4), high_f32x4, high_f32x4);
3585
+ }
3586
+ }
3587
+ else {
3588
+ for (; idx + 8 < count; idx += 8) {
3589
+ uint8x8x4_t loaded_u8x8x4 = vld4_u8((nk_u8_t const *)(data_ptr + idx * 4));
3590
+ float16x8_t half_f16x8 = nk_e5m2x8_to_f16x8_neon_(loaded_u8x8x4.val[0]);
3591
+ float32x4_t low_f32x4 = vcvt_f32_f16(vget_low_f16(half_f16x8));
3592
+ float32x4_t high_f32x4 = vcvt_high_f32_f16(half_f16x8);
3593
+ sum_f32x4 = vaddq_f32(sum_f32x4, vaddq_f32(low_f32x4, high_f32x4));
3594
+ sumsq_f32x4 = vfmaq_f32(vfmaq_f32(sumsq_f32x4, low_f32x4, low_f32x4), high_f32x4, high_f32x4);
3595
+ }
3596
+ }
3597
+ nk_f32_t sum = vaddvq_f32(sum_f32x4), sumsq = vaddvq_f32(sumsq_f32x4);
3598
+ for (; idx < count; ++idx) {
3599
+ nk_f32_t val;
3600
+ nk_e5m2_to_f32_serial((nk_e5m2_t const *)(data_ptr + idx * stride_elements), &val);
3601
+ sum += val, sumsq += val * val;
3602
+ }
3603
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
3604
+ }
3605
+
3606
+ NK_PUBLIC void nk_reduce_moments_e5m2_neon( //
3607
+ nk_e5m2_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
3608
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3609
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_e5m2_t);
3610
+ int aligned = (stride_bytes % sizeof(nk_e5m2_t) == 0);
3611
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
3612
+ else if (!aligned) nk_reduce_moments_e5m2_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3613
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 8) {
3614
+ nk_size_t left_count = count / 2;
3615
+ nk_f32_t left_sum, left_sumsq, right_sum, right_sumsq;
3616
+ nk_reduce_moments_e5m2_neon(data_ptr, left_count, stride_bytes, &left_sum, &left_sumsq);
3617
+ nk_reduce_moments_e5m2_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
3618
+ &right_sum, &right_sumsq);
3619
+ *sum_ptr = left_sum + right_sum, *sumsq_ptr = left_sumsq + right_sumsq;
3620
+ }
3621
+ else if (stride_elements == 1) nk_reduce_moments_e5m2_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
3622
+ else if (stride_elements <= 4)
3623
+ nk_reduce_moments_e5m2_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
3624
+ else nk_reduce_moments_e5m2_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3625
+ }
3626
+
3627
+ NK_INTERNAL void nk_reduce_minmax_e5m2_neon_contiguous_( //
3628
+ nk_e5m2_t const *data_ptr, nk_size_t count, //
3629
+ nk_e5m2_t *min_value_ptr, nk_size_t *min_index_ptr, //
3630
+ nk_e5m2_t *max_value_ptr, nk_size_t *max_index_ptr) {
3631
+ uint8x16_t min_u8x16 = vdupq_n_u8(0xFF), max_u8x16 = vdupq_n_u8(0);
3632
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
3633
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
3634
+ nk_size_t idx = 0;
3635
+ for (; idx + 16 <= count; idx += 16) {
3636
+ uint8x16_t raw_u8x16 = vld1q_u8((nk_u8_t const *)(data_ptr + idx));
3637
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(raw_u8x16);
3638
+ uint8x16_t is_nan_low_u8x16 = vcleq_u8(comparable_u8x16, vdupq_n_u8(0x02));
3639
+ uint8x16_t is_nan_high_u8x16 = vcgeq_u8(comparable_u8x16, vdupq_n_u8(0xFD));
3640
+ uint8x16_t is_nan_u8x16 = vorrq_u8(is_nan_low_u8x16, is_nan_high_u8x16);
3641
+ uint8x16_t data_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3642
+ uint8x16_t data_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3643
+ uint8x16_t less_u8x16 = vcltq_u8(data_min_u8x16, min_u8x16);
3644
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_max_u8x16, max_u8x16);
3645
+ min_u8x16 = vbslq_u8(less_u8x16, data_min_u8x16, min_u8x16);
3646
+ max_u8x16 = vbslq_u8(greater_u8x16, data_max_u8x16, max_u8x16);
3647
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3648
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3649
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
3650
+ }
3651
+ nk_size_t remaining = count - idx;
3652
+ if (remaining > 0) {
3653
+ nk_b128_vec_t tail_vec;
3654
+ nk_partial_load_b8x16_serial_(data_ptr + idx, &tail_vec, remaining);
3655
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(tail_vec.u8x16);
3656
+ uint8x16_t is_nan_low_u8x16 = vcleq_u8(comparable_u8x16, vdupq_n_u8(0x02));
3657
+ uint8x16_t is_nan_high_u8x16 = vcgeq_u8(comparable_u8x16, vdupq_n_u8(0xFD));
3658
+ uint8x16_t is_nan_u8x16 = vorrq_u8(is_nan_low_u8x16, is_nan_high_u8x16);
3659
+ uint8x16_t nan_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3660
+ uint8x16_t nan_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3661
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3662
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3663
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)remaining));
3664
+ uint8x16_t data_for_min_u8x16 = vbslq_u8(valid_u8x16, nan_min_u8x16, vdupq_n_u8(0xFF));
3665
+ uint8x16_t data_for_max_u8x16 = vbslq_u8(valid_u8x16, nan_max_u8x16, vdupq_n_u8(0));
3666
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
3667
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
3668
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
3669
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
3670
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3671
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3672
+ }
3673
+ nk_u8_t min_comparable = vminvq_u8(min_u8x16), max_comparable = vmaxvq_u8(max_u8x16);
3674
+ // If min stayed at 0xFF, all values were NaN
3675
+ if (min_comparable == 0xFF) {
3676
+ *min_value_ptr = (nk_e5m2_t)NK_E5M2_MAX, *min_index_ptr = NK_SIZE_MAX;
3677
+ *max_value_ptr = (nk_e5m2_t)NK_E5M2_MIN, *max_index_ptr = NK_SIZE_MAX;
3678
+ return;
3679
+ }
3680
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_comparable));
3681
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
3682
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
3683
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_comparable));
3684
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
3685
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
3686
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3687
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3688
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
3689
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
3690
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3691
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
3692
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
3693
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
3694
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
3695
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3696
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
3697
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
3698
+ *min_value_ptr = nk_comparable_to_fp8_(min_comparable), *min_index_ptr = min_idx;
3699
+ *max_value_ptr = nk_comparable_to_fp8_(max_comparable), *max_index_ptr = max_idx;
3700
+ }
3701
+
3702
+ NK_INTERNAL void nk_reduce_minmax_e5m2_neon_strided_( //
3703
+ nk_e5m2_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
3704
+ nk_e5m2_t *min_value_ptr, nk_size_t *min_index_ptr, //
3705
+ nk_e5m2_t *max_value_ptr, nk_size_t *max_index_ptr) {
3706
+ uint8x16_t min_u8x16 = vdupq_n_u8(0xFF), max_u8x16 = vdupq_n_u8(0);
3707
+ uint8x16_t min_iter_u8x16 = vdupq_n_u8(0), max_iter_u8x16 = vdupq_n_u8(0);
3708
+ uint8x16_t iter_u8x16 = vdupq_n_u8(0), one_u8x16 = vdupq_n_u8(1);
3709
+ uint8x16_t lane_indices_u8x16 = vcombine_u8(vreinterpret_u8_u64(vcreate_u64(0x0706050403020100ULL)),
3710
+ vreinterpret_u8_u64(vcreate_u64(0x0F0E0D0C0B0A0908ULL)));
3711
+ nk_size_t idx = 0;
3712
+ uint8x16_t data_for_min_u8x16, data_for_max_u8x16;
3713
+
3714
+ nk_reduce_minmax_e5m2_neon_cycle:
3715
+ if (stride_elements == 2 && idx + 16 < count) {
3716
+ uint8x16x2_t loaded = vld2q_u8((nk_u8_t const *)(data_ptr + idx * 2));
3717
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(loaded.val[0]);
3718
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vcleq_u8(comparable_u8x16, vdupq_n_u8(0x02)),
3719
+ vcgeq_u8(comparable_u8x16, vdupq_n_u8(0xFD)));
3720
+ data_for_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3721
+ data_for_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3722
+ idx += 16;
3723
+ }
3724
+ else if (stride_elements == 3 && idx + 16 < count) {
3725
+ uint8x16x3_t loaded = vld3q_u8((nk_u8_t const *)(data_ptr + idx * 3));
3726
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(loaded.val[0]);
3727
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vcleq_u8(comparable_u8x16, vdupq_n_u8(0x02)),
3728
+ vcgeq_u8(comparable_u8x16, vdupq_n_u8(0xFD)));
3729
+ data_for_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3730
+ data_for_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3731
+ idx += 16;
3732
+ }
3733
+ else if (stride_elements == 4 && idx + 16 < count) {
3734
+ uint8x16x4_t loaded = vld4q_u8((nk_u8_t const *)(data_ptr + idx * 4));
3735
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(loaded.val[0]);
3736
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vcleq_u8(comparable_u8x16, vdupq_n_u8(0x02)),
3737
+ vcgeq_u8(comparable_u8x16, vdupq_n_u8(0xFD)));
3738
+ data_for_min_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3739
+ data_for_max_u8x16 = vbslq_u8(is_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3740
+ idx += 16;
3741
+ }
3742
+ else if (idx < count) {
3743
+ nk_b128_vec_t tail_vec;
3744
+ nk_strided_load_b8x16_serial_(data_ptr + idx * stride_elements, stride_elements, &tail_vec, count - idx);
3745
+ uint8x16_t comparable_u8x16 = nk_fp8x16_to_comparable_neon_(tail_vec.u8x16);
3746
+ uint8x16_t is_nan_u8x16 = vorrq_u8(vcleq_u8(comparable_u8x16, vdupq_n_u8(0x02)),
3747
+ vcgeq_u8(comparable_u8x16, vdupq_n_u8(0xFD)));
3748
+ uint8x16_t valid_u8x16 = vcltq_u8(lane_indices_u8x16, vdupq_n_u8((nk_u8_t)(count - idx)));
3749
+ uint8x16_t invalid_or_nan_u8x16 = vornq_u8(is_nan_u8x16, valid_u8x16);
3750
+ data_for_min_u8x16 = vbslq_u8(invalid_or_nan_u8x16, vdupq_n_u8(0xFF), comparable_u8x16);
3751
+ data_for_max_u8x16 = vbslq_u8(invalid_or_nan_u8x16, vdupq_n_u8(0x00), comparable_u8x16);
3752
+ idx = count;
3753
+ }
3754
+ else {
3755
+ nk_u8_t min_comparable = vminvq_u8(min_u8x16), max_comparable = vmaxvq_u8(max_u8x16);
3756
+ if (min_comparable == 0xFF) {
3757
+ *min_value_ptr = (nk_e5m2_t)NK_E5M2_MAX, *min_index_ptr = NK_SIZE_MAX;
3758
+ *max_value_ptr = (nk_e5m2_t)NK_E5M2_MIN, *max_index_ptr = NK_SIZE_MAX;
3759
+ return;
3760
+ }
3761
+ uint8x16_t min_value_match_u8x16 = vceqq_u8(min_u8x16, vdupq_n_u8(min_comparable));
3762
+ uint8x16_t masked_min_iter_u8x16 = vbslq_u8(min_value_match_u8x16, min_iter_u8x16, vdupq_n_u8(0xFF));
3763
+ nk_u8_t earliest_min_cycle = vminvq_u8(masked_min_iter_u8x16);
3764
+ uint8x16_t max_value_match_u8x16 = vceqq_u8(max_u8x16, vdupq_n_u8(max_comparable));
3765
+ uint8x16_t masked_max_iter_u8x16 = vbslq_u8(max_value_match_u8x16, max_iter_u8x16, vdupq_n_u8(0xFF));
3766
+ nk_u8_t earliest_max_cycle = vminvq_u8(masked_max_iter_u8x16);
3767
+ uint8x16_t min_cycle_match_u8x16 = vceqq_u8(min_iter_u8x16, vdupq_n_u8(earliest_min_cycle));
3768
+ uint8x16_t min_both_match_u8x16 = vandq_u8(min_value_match_u8x16, min_cycle_match_u8x16);
3769
+ uint8x16_t min_masked_lanes_u8x16 = vbslq_u8(min_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3770
+ nk_u8_t min_lane_offset = vminvq_u8(min_masked_lanes_u8x16);
3771
+ nk_size_t min_idx = (nk_size_t)earliest_min_cycle * 16 + (nk_size_t)min_lane_offset;
3772
+ uint8x16_t max_cycle_match_u8x16 = vceqq_u8(max_iter_u8x16, vdupq_n_u8(earliest_max_cycle));
3773
+ uint8x16_t max_both_match_u8x16 = vandq_u8(max_value_match_u8x16, max_cycle_match_u8x16);
3774
+ uint8x16_t max_masked_lanes_u8x16 = vbslq_u8(max_both_match_u8x16, lane_indices_u8x16, vdupq_n_u8(0xFF));
3775
+ nk_u8_t max_lane_offset = vminvq_u8(max_masked_lanes_u8x16);
3776
+ nk_size_t max_idx = (nk_size_t)earliest_max_cycle * 16 + (nk_size_t)max_lane_offset;
3777
+ *min_value_ptr = nk_comparable_to_fp8_(min_comparable), *min_index_ptr = min_idx;
3778
+ *max_value_ptr = nk_comparable_to_fp8_(max_comparable), *max_index_ptr = max_idx;
3779
+ return;
3780
+ }
3781
+
3782
+ // Shared update body
3783
+ uint8x16_t less_u8x16 = vcltq_u8(data_for_min_u8x16, min_u8x16);
3784
+ uint8x16_t greater_u8x16 = vcgtq_u8(data_for_max_u8x16, max_u8x16);
3785
+ min_u8x16 = vbslq_u8(less_u8x16, data_for_min_u8x16, min_u8x16);
3786
+ max_u8x16 = vbslq_u8(greater_u8x16, data_for_max_u8x16, max_u8x16);
3787
+ min_iter_u8x16 = vbslq_u8(less_u8x16, iter_u8x16, min_iter_u8x16);
3788
+ max_iter_u8x16 = vbslq_u8(greater_u8x16, iter_u8x16, max_iter_u8x16);
3789
+ iter_u8x16 = vaddq_u8(iter_u8x16, one_u8x16);
3790
+ goto nk_reduce_minmax_e5m2_neon_cycle;
3791
+ }
3792
+
3793
+ NK_PUBLIC void nk_reduce_minmax_e5m2_neon( //
3794
+ nk_e5m2_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
3795
+ nk_e5m2_t *min_value_ptr, nk_size_t *min_index_ptr, //
3796
+ nk_e5m2_t *max_value_ptr, nk_size_t *max_index_ptr) {
3797
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_e5m2_t);
3798
+ int aligned = (stride_bytes % sizeof(nk_e5m2_t) == 0);
3799
+ if (count == 0)
3800
+ *min_value_ptr = NK_E5M2_MAX, *min_index_ptr = NK_SIZE_MAX, *max_value_ptr = NK_E5M2_MIN,
3801
+ *max_index_ptr = NK_SIZE_MAX;
3802
+ else if (!aligned)
3803
+ nk_reduce_minmax_e5m2_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
3804
+ max_index_ptr);
3805
+ else if (count > (nk_size_t)256 * 16) {
3806
+ nk_size_t left_count = count / 2;
3807
+ nk_e5m2_t left_min_value, right_min_value, left_max_value, right_max_value;
3808
+ nk_size_t left_min_index, right_min_index, left_max_index, right_max_index;
3809
+ nk_reduce_minmax_e5m2_neon(data_ptr, left_count, stride_bytes, &left_min_value, &left_min_index,
3810
+ &left_max_value, &left_max_index);
3811
+ nk_reduce_minmax_e5m2_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
3812
+ &right_min_value, &right_min_index, &right_max_value, &right_max_index);
3813
+ if (left_min_index == NK_SIZE_MAX)
3814
+ *min_value_ptr = right_min_value,
3815
+ *min_index_ptr = right_min_index == NK_SIZE_MAX ? NK_SIZE_MAX : left_count + right_min_index;
3816
+ else if (right_min_index == NK_SIZE_MAX || nk_e5m2_order_serial(left_min_value, right_min_value) <= 0)
3817
+ *min_value_ptr = left_min_value, *min_index_ptr = left_min_index;
3818
+ else *min_value_ptr = right_min_value, *min_index_ptr = left_count + right_min_index;
3819
+ if (left_max_index == NK_SIZE_MAX)
3820
+ *max_value_ptr = right_max_value,
3821
+ *max_index_ptr = right_max_index == NK_SIZE_MAX ? NK_SIZE_MAX : left_count + right_max_index;
3822
+ else if (right_max_index == NK_SIZE_MAX || nk_e5m2_order_serial(right_max_value, left_max_value) <= 0)
3823
+ *max_value_ptr = left_max_value, *max_index_ptr = left_max_index;
3824
+ else *max_value_ptr = right_max_value, *max_index_ptr = left_count + right_max_index;
3825
+ }
3826
+ else if (stride_elements == 1)
3827
+ nk_reduce_minmax_e5m2_neon_contiguous_(data_ptr, count, min_value_ptr, min_index_ptr, max_value_ptr,
3828
+ max_index_ptr);
3829
+ else if (stride_elements <= 4)
3830
+ nk_reduce_minmax_e5m2_neon_strided_(data_ptr, count, stride_elements, min_value_ptr, min_index_ptr,
3831
+ max_value_ptr, max_index_ptr);
3832
+ else
3833
+ nk_reduce_minmax_e5m2_serial(data_ptr, count, stride_bytes, min_value_ptr, min_index_ptr, max_value_ptr,
3834
+ max_index_ptr);
3835
+ }
3836
+
3837
+ NK_INTERNAL void nk_reduce_moments_f16_neon_contiguous_( //
3838
+ nk_f16_t const *data_ptr, nk_size_t count, //
3839
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3840
+ float32x4_t sum_f32x4 = vdupq_n_f32(0);
3841
+ float32x4_t sumsq_f32x4 = vdupq_n_f32(0);
3842
+ nk_size_t idx = 0;
3843
+
3844
+ for (; idx + 8 <= count; idx += 8) {
3845
+ float16x8_t data_f16x8 = vreinterpretq_f16_u16(vld1q_u16((nk_u16_t const *)(data_ptr + idx)));
3846
+ float32x4_t low_f32x4 = vcvt_f32_f16(vget_low_f16(data_f16x8));
3847
+ float32x4_t high_f32x4 = vcvt_high_f32_f16(data_f16x8);
3848
+ sum_f32x4 = vaddq_f32(sum_f32x4, low_f32x4);
3849
+ sum_f32x4 = vaddq_f32(sum_f32x4, high_f32x4);
3850
+ sumsq_f32x4 = vfmaq_f32(sumsq_f32x4, low_f32x4, low_f32x4);
3851
+ sumsq_f32x4 = vfmaq_f32(sumsq_f32x4, high_f32x4, high_f32x4);
3852
+ }
3853
+
3854
+ nk_f32_t sum = vaddvq_f32(sum_f32x4);
3855
+ nk_f32_t sumsq = vaddvq_f32(sumsq_f32x4);
3856
+ for (; idx < count; ++idx) {
3857
+ nk_f32_t value_f32;
3858
+ nk_f16_to_f32_serial(data_ptr + idx, &value_f32);
3859
+ sum += value_f32, sumsq += value_f32 * value_f32;
3860
+ }
3861
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
3862
+ }
3863
+
3864
+ NK_INTERNAL void nk_reduce_moments_f16_neon_strided_( //
3865
+ nk_f16_t const *data_ptr, nk_size_t count, nk_size_t stride_elements, //
3866
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3867
+ float32x4_t sum_f32x4 = vdupq_n_f32(0);
3868
+ float32x4_t sumsq_f32x4 = vdupq_n_f32(0);
3869
+ nk_size_t idx = 0;
3870
+
3871
+ if (stride_elements == 2) {
3872
+ for (; idx + 8 < count; idx += 8) {
3873
+ uint16x8x2_t loaded_u16x8x2 = vld2q_u16((nk_u16_t const *)(data_ptr + idx * 2));
3874
+ float16x8_t data_f16x8 = vreinterpretq_f16_u16(loaded_u16x8x2.val[0]);
3875
+ float32x4_t low_f32x4 = vcvt_f32_f16(vget_low_f16(data_f16x8));
3876
+ float32x4_t high_f32x4 = vcvt_high_f32_f16(data_f16x8);
3877
+ sum_f32x4 = vaddq_f32(sum_f32x4, low_f32x4);
3878
+ sum_f32x4 = vaddq_f32(sum_f32x4, high_f32x4);
3879
+ sumsq_f32x4 = vfmaq_f32(sumsq_f32x4, low_f32x4, low_f32x4);
3880
+ sumsq_f32x4 = vfmaq_f32(sumsq_f32x4, high_f32x4, high_f32x4);
3881
+ }
3882
+ }
3883
+ else if (stride_elements == 3) {
3884
+ for (; idx + 8 < count; idx += 8) {
3885
+ uint16x8x3_t loaded_u16x8x3 = vld3q_u16((nk_u16_t const *)(data_ptr + idx * 3));
3886
+ float16x8_t data_f16x8 = vreinterpretq_f16_u16(loaded_u16x8x3.val[0]);
3887
+ float32x4_t low_f32x4 = vcvt_f32_f16(vget_low_f16(data_f16x8));
3888
+ float32x4_t high_f32x4 = vcvt_high_f32_f16(data_f16x8);
3889
+ sum_f32x4 = vaddq_f32(sum_f32x4, low_f32x4);
3890
+ sum_f32x4 = vaddq_f32(sum_f32x4, high_f32x4);
3891
+ sumsq_f32x4 = vfmaq_f32(sumsq_f32x4, low_f32x4, low_f32x4);
3892
+ sumsq_f32x4 = vfmaq_f32(sumsq_f32x4, high_f32x4, high_f32x4);
3893
+ }
3894
+ }
3895
+ else if (stride_elements == 4) {
3896
+ for (; idx + 8 < count; idx += 8) {
3897
+ uint16x8x4_t loaded_u16x8x4 = vld4q_u16((nk_u16_t const *)(data_ptr + idx * 4));
3898
+ float16x8_t data_f16x8 = vreinterpretq_f16_u16(loaded_u16x8x4.val[0]);
3899
+ float32x4_t low_f32x4 = vcvt_f32_f16(vget_low_f16(data_f16x8));
3900
+ float32x4_t high_f32x4 = vcvt_high_f32_f16(data_f16x8);
3901
+ sum_f32x4 = vaddq_f32(sum_f32x4, low_f32x4);
3902
+ sum_f32x4 = vaddq_f32(sum_f32x4, high_f32x4);
3903
+ sumsq_f32x4 = vfmaq_f32(sumsq_f32x4, low_f32x4, low_f32x4);
3904
+ sumsq_f32x4 = vfmaq_f32(sumsq_f32x4, high_f32x4, high_f32x4);
3905
+ }
3906
+ }
3907
+
3908
+ nk_f32_t sum = vaddvq_f32(sum_f32x4);
3909
+ nk_f32_t sumsq = vaddvq_f32(sumsq_f32x4);
3910
+ for (; idx < count; ++idx) {
3911
+ nk_f32_t value_f32;
3912
+ nk_f16_to_f32_serial((nk_f16_t const *)(data_ptr + idx * stride_elements), &value_f32);
3913
+ sum += value_f32, sumsq += value_f32 * value_f32;
3914
+ }
3915
+ *sum_ptr = sum, *sumsq_ptr = sumsq;
3916
+ }
3917
+
3918
+ NK_PUBLIC void nk_reduce_moments_f16_neon( //
3919
+ nk_f16_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
3920
+ nk_f32_t *sum_ptr, nk_f32_t *sumsq_ptr) {
3921
+ nk_size_t stride_elements = stride_bytes / sizeof(nk_f16_t);
3922
+ int aligned = (stride_bytes % sizeof(nk_f16_t) == 0);
3923
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
3924
+ else if (!aligned) nk_reduce_moments_f16_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3925
+ else if (count > (nk_size_t)(NK_U16_MAX + 1) * 8) {
3926
+ nk_size_t left_count = count / 2;
3927
+ nk_f32_t left_sum_value, left_sumsq_value, right_sum_value, right_sumsq_value;
3928
+ nk_reduce_moments_f16_neon(data_ptr, left_count, stride_bytes, &left_sum_value, &left_sumsq_value);
3929
+ nk_reduce_moments_f16_neon(data_ptr + left_count * stride_elements, count - left_count, stride_bytes,
3930
+ &right_sum_value, &right_sumsq_value);
3931
+ *sum_ptr = left_sum_value + right_sum_value, *sumsq_ptr = left_sumsq_value + right_sumsq_value;
3932
+ }
3933
+ else if (stride_elements == 1) nk_reduce_moments_f16_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
3934
+ else if (stride_elements <= 4)
3935
+ nk_reduce_moments_f16_neon_strided_(data_ptr, count, stride_elements, sum_ptr, sumsq_ptr);
3936
+ else nk_reduce_moments_f16_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3937
+ }
3938
+
3939
+ NK_INTERNAL void nk_reduce_moments_u1_neon_contiguous_( //
3940
+ nk_u1x8_t const *data_ptr, nk_size_t count, //
3941
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
3942
+ nk_size_t byte_count = nk_size_divide_round_up_(count, NK_BITS_PER_BYTE);
3943
+ nk_u64_t sum = 0;
3944
+ nk_size_t idx = 0;
3945
+ // Each vcntq_u8 produces values 0-8 per lane; accumulate at u8 level
3946
+ // for up to 31 iterations (31 × 8 = 248, fits in u8) before widening.
3947
+ while (idx + 16 <= byte_count) {
3948
+ uint8x16_t popcount_u8x16 = vdupq_n_u8(0);
3949
+ for (nk_size_t cycle = 0; cycle < 31 && idx + 16 <= byte_count; ++cycle, idx += 16) {
3950
+ uint8x16_t data_u8x16 = vld1q_u8((nk_u8_t const *)data_ptr + idx);
3951
+ popcount_u8x16 = vaddq_u8(popcount_u8x16, vcntq_u8(data_u8x16));
3952
+ }
3953
+ sum += (nk_u64_t)vaddlvq_u8(popcount_u8x16);
3954
+ }
3955
+ for (; idx < byte_count; ++idx) sum += nk_u1x8_popcount_(((nk_u8_t const *)data_ptr)[idx]);
3956
+ *sum_ptr = sum, *sumsq_ptr = sum;
3957
+ }
3958
+
3959
+ NK_PUBLIC void nk_reduce_moments_u1_neon( //
3960
+ nk_u1x8_t const *data_ptr, nk_size_t count, nk_size_t stride_bytes, //
3961
+ nk_u64_t *sum_ptr, nk_u64_t *sumsq_ptr) {
3962
+ count = nk_size_round_up_to_multiple_(count, 8);
3963
+ if (count == 0) *sum_ptr = 0, *sumsq_ptr = 0;
3964
+ else if (stride_bytes == 1) nk_reduce_moments_u1_neon_contiguous_(data_ptr, count, sum_ptr, sumsq_ptr);
3965
+ else nk_reduce_moments_u1_serial(data_ptr, count, stride_bytes, sum_ptr, sumsq_ptr);
3966
+ }
3967
+
3968
+ #if defined(__clang__)
3969
+ #pragma clang attribute pop
3970
+ #elif defined(__GNUC__)
3971
+ #pragma GCC pop_options
3972
+ #endif
3973
+
3974
+ #if defined(__cplusplus)
3975
+ } // extern "C"
3976
+ #endif
3977
+
3978
+ #endif // NK_TARGET_NEON
3979
+ #endif // NK_TARGET_ARM64_
3980
+ #endif // NK_REDUCE_NEON_H