numkong 7.0.0 → 7.4.2

package/include/numkong/mesh.h

@@ -82,17 +82,17 @@
  *
  *  The SIMD kernels are dominated by FMA, permutes, and gathers:
  *
- *      Intrinsic                     Instruction        Notes
- *      _mm256_fmadd_ps/pd            VFMADD*            FMA on FP ports (Haswell/Skylake: ports 0/1)
- *      _mm256_i32gather_ps           VGATHERDPS         High-latency; memory-bound
- *      _mm512_permutex2var_ps/pd     VPERMT2*           Shuffle-heavy; can bottleneck on shuffle ports
- *      _mm512_reduce_add_ps/pd       (sequence)         Implemented via shuffles + adds
+ *      Intrinsic                  Instruction  Notes
+ *      _mm256_fmadd_ps/pd         VFMADD*      FMA on FP ports (Haswell/Skylake: ports 0/1)
+ *      _mm256_i32gather_ps        VGATHERDPS   High-latency; memory-bound
+ *      _mm512_permutex2var_ps/pd  VPERMT2*     Shuffle-heavy; can bottleneck on shuffle ports
+ *      _mm512_reduce_add_ps/pd    (sequence)   Implemented via shuffles + adds
  *
  *  Gather-heavy tails are intentionally isolated to keep the steady-state loop on contiguous loads.
  *
  *  @section references References
  *
- *  - x86 intrinsics: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/
+ *  - x86 intrinsics: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html
  *  - Arm intrinsics: https://developer.arm.com/architectures/instruction-sets/intrinsics/
  *
  */
@@ -245,6 +245,25 @@ NK_PUBLIC void nk_kabsch_f64_skylake(nk_f64_t const *a, nk_f64_t const *b, nk_si
 /** @copydoc nk_umeyama_f64 */
 NK_PUBLIC void nk_umeyama_f64_skylake(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *a_centroid,
                                       nk_f64_t *b_centroid, nk_f64_t *rotation, nk_f64_t *scale, nk_f64_t *result);
+
+/** @copydoc nk_rmsd_f16 */
+NK_PUBLIC void nk_rmsd_f16_skylake(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                   nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
+/** @copydoc nk_kabsch_f16 */
+NK_PUBLIC void nk_kabsch_f16_skylake(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                     nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
+/** @copydoc nk_umeyama_f16 */
+NK_PUBLIC void nk_umeyama_f16_skylake(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                      nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
+/** @copydoc nk_rmsd_bf16 */
+NK_PUBLIC void nk_rmsd_bf16_skylake(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                    nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
+/** @copydoc nk_kabsch_bf16 */
+NK_PUBLIC void nk_kabsch_bf16_skylake(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                      nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
+/** @copydoc nk_umeyama_bf16 */
+NK_PUBLIC void nk_umeyama_bf16_skylake(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                       nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
 #endif // NK_TARGET_SKYLAKE
 
 /*  SIMD-powered backends for AVX2 CPUs of Haswell generation and newer.
@@ -313,21 +332,16 @@ NK_PUBLIC void nk_kabsch_f64_neon(nk_f64_t const *a, nk_f64_t const *b, nk_size_
 /** @copydoc nk_umeyama_f64 */
 NK_PUBLIC void nk_umeyama_f64_neon(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *a_centroid,
                                    nk_f64_t *b_centroid, nk_f64_t *rotation, nk_f64_t *scale, nk_f64_t *result);
-#endif // NK_TARGET_NEON
-
-/*  SIMD-powered backends for Arm NEON FP16 CPUs.
- */
-#if NK_TARGET_NEONHALF
 /** @copydoc nk_rmsd_f16 */
-NK_PUBLIC void nk_rmsd_f16_neonhalf(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
-                                    nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
+NK_PUBLIC void nk_rmsd_f16_neon(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
 /** @copydoc nk_kabsch_f16 */
-NK_PUBLIC void nk_kabsch_f16_neonhalf(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
-                                      nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
+NK_PUBLIC void nk_kabsch_f16_neon(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                  nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
 /** @copydoc nk_umeyama_f16 */
-NK_PUBLIC void nk_umeyama_f16_neonhalf(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
-                                       nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
-#endif // NK_TARGET_NEONHALF
+NK_PUBLIC void nk_umeyama_f16_neon(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
+                                   nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result);
+#endif // NK_TARGET_NEON
 
 /*  SIMD-powered backends for Arm NEON BF16 CPUs.
  */
@@ -406,22 +420,10 @@ NK_PUBLIC void nk_umeyama_f64_v128relaxed(nk_f64_t const *a, nk_f64_t const *b,
 #endif // NK_TARGET_V128RELAXED
 
 /**
- *  @brief  Returns the output dtype for RMSD.
- */
-NK_INTERNAL nk_dtype_t nk_rmsd_output_dtype(nk_dtype_t dtype) {
-    switch (dtype) {
-    case nk_f64_k: return nk_f64_k;
-    case nk_f32_k: return nk_f64_k;
-    case nk_f16_k: return nk_f32_k;
-    case nk_bf16_k: return nk_f32_k;
-    default: return nk_dtype_unknown_k;
-    }
-}
-
-/**
- *  @brief  Returns the output dtype for Kabsch alignment.
+ *  @brief  Returns the metric output dtype for mesh alignment operations.
+ *  Matches the C++ `mesh_metric_t` alias in types.hpp.
  */
-NK_INTERNAL nk_dtype_t nk_kabsch_output_dtype(nk_dtype_t dtype) {
+NK_INTERNAL nk_dtype_t nk_mesh_metric_dtype(nk_dtype_t dtype) {
     switch (dtype) {
     case nk_f64_k: return nk_f64_k;
     case nk_f32_k: return nk_f64_k;
@@ -432,12 +434,13 @@ NK_INTERNAL nk_dtype_t nk_kabsch_output_dtype(nk_dtype_t dtype) {
 }
 
 /**
- *  @brief  Returns the output dtype for Umeyama alignment.
+ *  @brief  Returns the transform output dtype for mesh alignment operations.
+ *  Matches the C++ `mesh_transform_t` alias in types.hpp.
  */
-NK_INTERNAL nk_dtype_t nk_umeyama_output_dtype(nk_dtype_t dtype) {
+NK_INTERNAL nk_dtype_t nk_mesh_transform_dtype(nk_dtype_t dtype) {
     switch (dtype) {
     case nk_f64_k: return nk_f64_k;
-    case nk_f32_k: return nk_f64_k;
+    case nk_f32_k: return nk_f32_k;
     case nk_f16_k: return nk_f32_k;
     case nk_bf16_k: return nk_f32_k;
     default: return nk_dtype_unknown_k;
@@ -450,7 +453,6 @@ NK_INTERNAL nk_dtype_t nk_umeyama_output_dtype(nk_dtype_t dtype) {
 
 #include "numkong/mesh/serial.h"
 #include "numkong/mesh/neon.h"
-#include "numkong/mesh/neonhalf.h"
 #include "numkong/mesh/neonbfdot.h"
 #include "numkong/mesh/haswell.h"
 #include "numkong/mesh/skylake.h"
@@ -499,10 +501,12 @@ NK_PUBLIC void nk_rmsd_f32(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk
 
 NK_PUBLIC void nk_rmsd_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
                            nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result) {
-#if NK_TARGET_HASWELL
+#if NK_TARGET_SKYLAKE
+    nk_rmsd_f16_skylake(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_HASWELL
     nk_rmsd_f16_haswell(a, b, n, a_centroid, b_centroid, rotation, scale, result);
-#elif NK_TARGET_NEONHALF
-    nk_rmsd_f16_neonhalf(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_NEON
+    nk_rmsd_f16_neon(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #elif NK_TARGET_RVV
     nk_rmsd_f16_rvv(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #else
@@ -512,7 +516,9 @@ NK_PUBLIC void nk_rmsd_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk
 
 NK_PUBLIC void nk_rmsd_bf16(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
                             nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result) {
-#if NK_TARGET_HASWELL
+#if NK_TARGET_SKYLAKE
+    nk_rmsd_bf16_skylake(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_HASWELL
     nk_rmsd_bf16_haswell(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #elif NK_TARGET_NEONBFDOT
     nk_rmsd_bf16_neonbfdot(a, b, n, a_centroid, b_centroid, rotation, scale, result);
@@ -559,10 +565,12 @@ NK_PUBLIC void nk_kabsch_f32(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n,
 
 NK_PUBLIC void nk_kabsch_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
                              nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result) {
-#if NK_TARGET_HASWELL
+#if NK_TARGET_SKYLAKE
+    nk_kabsch_f16_skylake(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_HASWELL
     nk_kabsch_f16_haswell(a, b, n, a_centroid, b_centroid, rotation, scale, result);
-#elif NK_TARGET_NEONHALF
-    nk_kabsch_f16_neonhalf(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_NEON
+    nk_kabsch_f16_neon(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #elif NK_TARGET_RVV
     nk_kabsch_f16_rvv(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #else
@@ -572,7 +580,9 @@ NK_PUBLIC void nk_kabsch_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n,
 
 NK_PUBLIC void nk_kabsch_bf16(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
                               nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result) {
-#if NK_TARGET_HASWELL
+#if NK_TARGET_SKYLAKE
+    nk_kabsch_bf16_skylake(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_HASWELL
     nk_kabsch_bf16_haswell(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #elif NK_TARGET_NEONBFDOT
     nk_kabsch_bf16_neonbfdot(a, b, n, a_centroid, b_centroid, rotation, scale, result);
@@ -619,10 +629,12 @@ NK_PUBLIC void nk_umeyama_f32(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n,
 
 NK_PUBLIC void nk_umeyama_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
                               nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result) {
-#if NK_TARGET_HASWELL
+#if NK_TARGET_SKYLAKE
+    nk_umeyama_f16_skylake(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_HASWELL
     nk_umeyama_f16_haswell(a, b, n, a_centroid, b_centroid, rotation, scale, result);
-#elif NK_TARGET_NEONHALF
-    nk_umeyama_f16_neonhalf(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_NEON
+    nk_umeyama_f16_neon(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #elif NK_TARGET_RVV
     nk_umeyama_f16_rvv(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #else
@@ -632,7 +644,9 @@ NK_PUBLIC void nk_umeyama_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n,
 
 NK_PUBLIC void nk_umeyama_bf16(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *a_centroid,
                                nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result) {
-#if NK_TARGET_HASWELL
+#if NK_TARGET_SKYLAKE
+    nk_umeyama_bf16_skylake(a, b, n, a_centroid, b_centroid, rotation, scale, result);
+#elif NK_TARGET_HASWELL
     nk_umeyama_bf16_haswell(a, b, n, a_centroid, b_centroid, rotation, scale, result);
 #elif NK_TARGET_NEONBFDOT
     nk_umeyama_bf16_neonbfdot(a, b, n, a_centroid, b_centroid, rotation, scale, result);

package/include/numkong/mesh.hpp

@@ -17,7 +17,7 @@
 
 namespace ashvardanian::numkong {
 
-#pragma region - SVD Helpers for Scalar Fallbacks
+#pragma region SVD Helpers for Scalar Fallbacks
 
 /** @brief 3x3 matrix determinant. */
 template <typename scalar_type_>
@@ -313,9 +313,9 @@ void svd3x3_(scalar_type_ const *a, scalar_type_ *svd_u, scalar_type_ *svd_s, sc
     svd_s[8] = s3_sq.sqrt();
 }
 
-#pragma endregion - SVD Helpers for Scalar Fallbacks
+#pragma endregion SVD Helpers for Scalar Fallbacks
 
-#pragma region - Mesh Alignment Kernels
+#pragma region Mesh Alignment Kernels
 
 /**
  *  @brief Root Mean Square Deviation between two 3D point clouds (no alignment)
@@ -755,7 +755,7 @@ void umeyama(in_type_ const *a, in_type_ const *b, std::size_t n, transform_type
     }
 }
 
-#pragma endregion - Mesh Alignment Kernels
+#pragma endregion Mesh Alignment Kernels
 
 } // namespace ashvardanian::numkong
 

package/include/numkong/numkong.h

@@ -62,9 +62,9 @@ NK_PUBLIC nk_dtype_t nk_kernel_output_dtype(nk_kernel_kind_t kind, nk_dtype_t in
     case nk_kernel_vincenty_k: return nk_vincenty_output_dtype(input);
     case nk_kernel_kld_k:
     case nk_kernel_jsd_k: return nk_probability_output_dtype(input);
-    case nk_kernel_rmsd_k: return nk_rmsd_output_dtype(input);
-    case nk_kernel_kabsch_k: return nk_kabsch_output_dtype(input);
-    case nk_kernel_umeyama_k: return nk_umeyama_output_dtype(input);
+    case nk_kernel_rmsd_k:
+    case nk_kernel_kabsch_k:
+    case nk_kernel_umeyama_k: return nk_mesh_metric_dtype(input);
     case nk_kernel_sparse_dot_k: return nk_sparse_dot_output_dtype(input);
     case nk_kernel_maxsim_packed_k: return nk_maxsim_output_dtype(input);
     default: return nk_dtype_unknown_k;

package/include/numkong/numkong.hpp

@@ -37,6 +37,7 @@
 #define NK_NUMKONG_HPP
 
 #include "numkong/random.hpp"
+#include "numkong/cast.hpp"
 #include "numkong/dot.hpp"
 #include "numkong/spatial.hpp"
 #include "numkong/spatials.hpp"

package/include/numkong/probability/README.md

@@ -5,17 +5,21 @@ These are used in variational inference, topic modeling, and distribution compar
 
 Kullback-Leibler divergence from $P$ to $Q$:
 
-```math
+$$
 \text{KLD}(P \| Q) = \sum_{i=0}^{n-1} P(i) \log_2 \frac{P(i)}{Q(i)}
-```
+$$
 
 Jensen-Shannon distance is the square root of the symmetrized KLD through a mixture:
 
-$$\text{JSD}(P, Q) = \frac{1}{2} \text{KLD}(P \| M) + \frac{1}{2} \text{KLD}(Q \| M)$$
+$$
+\text{JSD}(P, Q) = \frac{1}{2} \text{KLD}(P \| M) + \frac{1}{2} \text{KLD}(Q \| M)
+$$
 
 where $M = \frac{P + Q}{2}$, yielding the distance:
 
-$$d_{JS}(P, Q) = \sqrt{\text{JSD}(P, Q)}$$
+$$
+d_{JS}(P, Q) = \sqrt{\text{JSD}(P, Q)}
+$$
 
 Unlike the raw divergence, $d_{JS}$ is a true metric satisfying the triangle inequality.
 
@@ -35,9 +39,9 @@ def jsd(p: np.ndarray, q: np.ndarray) -> float:
 
 ## Use Cases
 
-__Kullback-Leibler divergence__ is the workhorse of variational inference (ELBO objective), knowledge distillation between neural networks, information gain in decision trees, and measuring fit between a model and observed data.
+__Kullback-Leibler divergence__ is widely used in variational inference (ELBO objective), knowledge distillation between neural networks, information gain in decision trees, and measuring fit between a model and observed data.
 
-__Jensen-Shannon distance__ sees primary use in microbiome community comparison (enterotyping), where its metric property enables clustering with standard algorithms. It also appears in distribution drift detection, topic model evaluation, and as the theoretical foundation of the original GAN objective — though in practice GAN training uses proxy losses rather than computing JSD directly.
+__Jensen-Shannon distance__ is commonly used in microbiome community comparison (enterotyping), where its metric property enables clustering with standard algorithms. It also appears in distribution drift detection, topic model evaluation, and as the theoretical foundation of the original GAN objective — though in practice GAN training uses proxy losses rather than computing JSD directly.
 
 ## Input & Output Types
 
@@ -149,25 +153,25 @@ Measured with Wasmtime v42 (Cranelift backend).
 | `nk_kld_f16_serial`  |    0.118 gb/s, 1.04K ulp |    0.127 gb/s, 4.53K ulp |    0.111 gb/s, 18.3K ulp |
 | `nk_jsd_f16_serial`  |     0.0748 gb/s, 1.4 ulp |     0.0681 gb/s, 2.6 ulp |     0.0857 gb/s, 9.7 ulp |
 
-### Apple M4
+### Apple M5
 
 #### Native
 
 | Kernel                |                      256 |                     1024 |                     4096 |
 | :-------------------- | -----------------------: | -----------------------: | -----------------------: |
 | __f64__               | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ |
-| `nk_kld_f64_serial`   |      2.21 gb/s, 5.6K ulp |       2.22 gb/s, 25K ulp |       2.18 gb/s, 99K ulp |
-| `nk_jsd_f64_serial`   |       1.40 gb/s, 0.4 ulp |       1.45 gb/s, 0.4 ulp |       1.45 gb/s, 0.5 ulp |
+| `nk_kld_f64_serial`   |      3.22 gb/s, 5.6K ulp |       3.36 gb/s, 25K ulp |       3.32 gb/s, 99K ulp |
+| `nk_jsd_f64_serial`   |       2.06 gb/s, 0.4 ulp |       2.17 gb/s, 0.4 ulp |       2.17 gb/s, 0.5 ulp |
 | __f32__               | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ |
-| `nk_kld_f32_serial`   |      6.29 gb/s, 1.0K ulp |      6.35 gb/s, 4.5K ulp |       6.22 gb/s, 18K ulp |
-| `nk_jsd_f32_serial`   |       1.21 gb/s, 0.4 ulp |       1.20 gb/s, 0.4 ulp |       1.20 gb/s, 4.6 ulp |
-| `nk_kld_f32_neon`     |      14.5 gb/s, 1.0K ulp |      14.4 gb/s, 4.5K ulp |       12.8 gb/s, 18K ulp |
-| `nk_jsd_f32_neon`     |        6.81 gb/s, 15 ulp |        7.04 gb/s, 14 ulp |       6.78 gb/s, 9.9 ulp |
+| `nk_kld_f32_serial`   |      9.26 gb/s, 1.0K ulp |      8.73 gb/s, 4.5K ulp |       9.10 gb/s, 18K ulp |
+| `nk_jsd_f32_serial`   |       2.08 gb/s, 0.4 ulp |       2.16 gb/s, 0.4 ulp |       2.13 gb/s, 4.6 ulp |
+| `nk_kld_f32_neon`     |      19.0 gb/s, 1.0K ulp |      17.4 gb/s, 4.5K ulp |       18.1 gb/s, 18K ulp |
+| `nk_jsd_f32_neon`     |        9.75 gb/s, 15 ulp |        9.32 gb/s, 14 ulp |       9.62 gb/s, 9.9 ulp |
 | __bf16__              | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ |
-| `nk_kld_bf16_serial`  |      3.16 gb/s, 1.0K ulp |      2.96 gb/s, 4.5K ulp |       3.16 gb/s, 18K ulp |
-| `nk_jsd_bf16_serial`  |      0.611 gb/s, 1.4 ulp |      0.595 gb/s, 2.9 ulp |      0.613 gb/s, 9.7 ulp |
+| `nk_kld_bf16_serial`  |      4.58 gb/s, 1.0K ulp |      4.47 gb/s, 4.5K ulp |       4.65 gb/s, 18K ulp |
+| `nk_jsd_bf16_serial`  |       1.08 gb/s, 1.4 ulp |       1.07 gb/s, 2.9 ulp |       1.09 gb/s, 9.7 ulp |
 | __f16__               | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ |
-| `nk_kld_f16_serial`   |      3.15 gb/s, 1.0K ulp |      3.14 gb/s, 4.5K ulp |       2.81 gb/s, 18K ulp |
-| `nk_jsd_f16_serial`   |      0.610 gb/s, 1.4 ulp |      0.611 gb/s, 2.7 ulp |      0.602 gb/s, 8.7 ulp |
-| `nk_kld_f16_neonhalf` |      6.78 gb/s, 1.0K ulp |      6.72 gb/s, 4.5K ulp |       6.09 gb/s, 18K ulp |
-| `nk_jsd_f16_neonhalf` |        3.42 gb/s, 15 ulp |        3.40 gb/s, 14 ulp |       3.14 gb/s, 9.9 ulp |
+| `nk_kld_f16_serial`   |      4.63 gb/s, 1.0K ulp |      4.45 gb/s, 4.5K ulp |       4.55 gb/s, 18K ulp |
+| `nk_jsd_f16_serial`   |       1.03 gb/s, 1.4 ulp |      0.962 gb/s, 2.7 ulp |      0.976 gb/s, 8.7 ulp |
+| `nk_kld_f16_neonhalf` |      10.2 gb/s, 1.0K ulp |      9.67 gb/s, 4.5K ulp |       9.99 gb/s, 18K ulp |
+| `nk_jsd_f16_neonhalf` |        5.00 gb/s, 15 ulp |        4.79 gb/s, 14 ulp |       4.94 gb/s, 9.9 ulp |

package/include/numkong/probability/neon.h

@@ -57,8 +57,8 @@ NK_PUBLIC float32x4_t nk_log2_f32x4_neon_(float32x4_t x) {
 NK_PUBLIC void nk_kld_f32_neon(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result) {
     nk_f32_t epsilon = NK_F32_DIVISION_EPSILON;
     float32x4_t epsilon_f32x4 = vdupq_n_f32(epsilon);
-    float64x2_t sum_lower_f64x2 = vdupq_n_f64(0.0);
-    float64x2_t sum_upper_f64x2 = vdupq_n_f64(0.0);
+    float64x2_t sum_low_f64x2 = vdupq_n_f64(0.0);
+    float64x2_t sum_high_f64x2 = vdupq_n_f64(0.0);
     float32x4_t a_f32x4, b_f32x4;
 
 nk_kld_f32_neon_cycle:
@@ -79,20 +79,20 @@ nk_kld_f32_neon_cycle:
     float32x4_t ratio_f32x4 = vdivq_f32(vaddq_f32(a_f32x4, epsilon_f32x4), vaddq_f32(b_f32x4, epsilon_f32x4));
     float32x4_t log_ratio_f32x4 = nk_log2_f32x4_neon_(ratio_f32x4);
     float32x4_t contribution_f32x4 = vmulq_f32(a_f32x4, log_ratio_f32x4);
-    sum_lower_f64x2 = vaddq_f64(sum_lower_f64x2, vcvt_f64_f32(vget_low_f32(contribution_f32x4)));
-    sum_upper_f64x2 = vaddq_f64(sum_upper_f64x2, vcvt_f64_f32(vget_high_f32(contribution_f32x4)));
+    sum_low_f64x2 = vaddq_f64(sum_low_f64x2, vcvt_f64_f32(vget_low_f32(contribution_f32x4)));
+    sum_high_f64x2 = vaddq_f64(sum_high_f64x2, vcvt_high_f64_f32(contribution_f32x4));
     if (n != 0) goto nk_kld_f32_neon_cycle;
 
     nk_f64_t log2_normalizer = 0.6931471805599453;
-    nk_f64_t sum = vaddvq_f64(vaddq_f64(sum_lower_f64x2, sum_upper_f64x2)) * log2_normalizer;
+    nk_f64_t sum = vaddvq_f64(vaddq_f64(sum_low_f64x2, sum_high_f64x2)) * log2_normalizer;
     *result = sum;
 }
 
 NK_PUBLIC void nk_jsd_f32_neon(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result) {
     nk_f32_t epsilon = NK_F32_DIVISION_EPSILON;
     float32x4_t epsilon_f32x4 = vdupq_n_f32(epsilon);
-    float64x2_t sum_lower_f64x2 = vdupq_n_f64(0.0);
-    float64x2_t sum_upper_f64x2 = vdupq_n_f64(0.0);
+    float64x2_t sum_low_f64x2 = vdupq_n_f64(0.0);
+    float64x2_t sum_high_f64x2 = vdupq_n_f64(0.0);
     float32x4_t a_f32x4, b_f32x4;
 
 nk_jsd_f32_neon_cycle:
@@ -118,12 +118,12 @@ nk_jsd_f32_neon_cycle:
     float32x4_t contribution_a_f32x4 = vmulq_f32(a_f32x4, log_ratio_a_f32x4);
     float32x4_t contribution_b_f32x4 = vmulq_f32(b_f32x4, log_ratio_b_f32x4);
     float32x4_t contribution_f32x4 = vaddq_f32(contribution_a_f32x4, contribution_b_f32x4);
-    sum_lower_f64x2 = vaddq_f64(sum_lower_f64x2, vcvt_f64_f32(vget_low_f32(contribution_f32x4)));
-    sum_upper_f64x2 = vaddq_f64(sum_upper_f64x2, vcvt_f64_f32(vget_high_f32(contribution_f32x4)));
+    sum_low_f64x2 = vaddq_f64(sum_low_f64x2, vcvt_f64_f32(vget_low_f32(contribution_f32x4)));
+    sum_high_f64x2 = vaddq_f64(sum_high_f64x2, vcvt_high_f64_f32(contribution_f32x4));
     if (n != 0) goto nk_jsd_f32_neon_cycle;
 
     nk_f64_t log2_normalizer = 0.6931471805599453;
-    nk_f64_t sum = vaddvq_f64(vaddq_f64(sum_lower_f64x2, sum_upper_f64x2)) * log2_normalizer / 2.0;
+    nk_f64_t sum = vaddvq_f64(vaddq_f64(sum_low_f64x2, sum_high_f64x2)) * log2_normalizer / 2.0;
     *result = sum > 0 ? nk_f64_sqrt_neon(sum) : 0;
 }
 
@@ -134,76 +134,106 @@ nk_jsd_f32_neon_cycle:
 #endif
 #endif // NK_TARGET_NEON
 
-#if NK_TARGET_NEONHALF
+#if NK_TARGET_NEON
 #if defined(__clang__)
-#pragma clang attribute push(__attribute__((target("arch=armv8.2-a+simd+fp16"))), apply_to = function)
+#pragma clang attribute push(__attribute__((target("arch=armv8.2-a+simd"))), apply_to = function)
 #elif defined(__GNUC__)
 #pragma GCC push_options
-#pragma GCC target("arch=armv8.2-a+simd+fp16")
+#pragma GCC target("arch=armv8.2-a+simd")
 #endif
 
-NK_PUBLIC void nk_kld_f16_neonhalf(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
+NK_PUBLIC void nk_kld_f16_neon(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
     float32x4_t sum_f32x4 = vdupq_n_f32(0);
     nk_f32_t epsilon = NK_F32_DIVISION_EPSILON;
     float32x4_t epsilon_f32x4 = vdupq_n_f32(epsilon);
-    float32x4_t a_f32x4, b_f32x4;
+    float32x4_t a_low_f32x4, a_high_f32x4, b_low_f32x4, b_high_f32x4;
 
-nk_kld_f16_neonhalf_cycle:
-    if (n < 4) {
-        nk_b64_vec_t a_vec, b_vec;
-        nk_partial_load_b16x4_serial_(a, &a_vec, n);
-        nk_partial_load_b16x4_serial_(b, &b_vec, n);
-        a_f32x4 = vcvt_f32_f16(vreinterpret_f16_u16(a_vec.u16x4));
-        b_f32x4 = vcvt_f32_f16(vreinterpret_f16_u16(b_vec.u16x4));
+nk_kld_f16_neon_cycle:
+    if (n < 8) {
+        nk_b128_vec_t a_vec, b_vec;
+        nk_partial_load_b16x8_serial_(a, &a_vec, n);
+        nk_partial_load_b16x8_serial_(b, &b_vec, n);
+        float16x8_t a_f16x8 = vreinterpretq_f16_u16(a_vec.u16x8);
+        float16x8_t b_f16x8 = vreinterpretq_f16_u16(b_vec.u16x8);
+        a_low_f32x4 = vcvt_f32_f16(vget_low_f16(a_f16x8));
+        a_high_f32x4 = vcvt_high_f32_f16(a_f16x8);
+        b_low_f32x4 = vcvt_f32_f16(vget_low_f16(b_f16x8));
+        b_high_f32x4 = vcvt_high_f32_f16(b_f16x8);
         n = 0;
     }
     else {
-        a_f32x4 = vcvt_f32_f16(vld1_f16((nk_f16_for_arm_simd_t const *)a));
-        b_f32x4 = vcvt_f32_f16(vld1_f16((nk_f16_for_arm_simd_t const *)b));
-        n -= 4, a += 4, b += 4;
+        float16x8_t a_f16x8 = vreinterpretq_f16_u16(vld1q_u16((nk_u16_t const *)a));
+        float16x8_t b_f16x8 = vreinterpretq_f16_u16(vld1q_u16((nk_u16_t const *)b));
+        a_low_f32x4 = vcvt_f32_f16(vget_low_f16(a_f16x8));
+        a_high_f32x4 = vcvt_high_f32_f16(a_f16x8);
+        b_low_f32x4 = vcvt_f32_f16(vget_low_f16(b_f16x8));
+        b_high_f32x4 = vcvt_high_f32_f16(b_f16x8);
+        n -= 8, a += 8, b += 8;
     }
 
-    float32x4_t ratio_f32x4 = vdivq_f32(vaddq_f32(a_f32x4, epsilon_f32x4), vaddq_f32(b_f32x4, epsilon_f32x4));
-    float32x4_t log_ratio_f32x4 = nk_log2_f32x4_neon_(ratio_f32x4);
-    float32x4_t contribution_f32x4 = vmulq_f32(a_f32x4, log_ratio_f32x4);
-    sum_f32x4 = vaddq_f32(sum_f32x4, contribution_f32x4);
-    if (n) goto nk_kld_f16_neonhalf_cycle;
+    float32x4_t ratio_low_f32x4 = vdivq_f32(vaddq_f32(a_low_f32x4, epsilon_f32x4),
+                                            vaddq_f32(b_low_f32x4, epsilon_f32x4));
+    float32x4_t ratio_high_f32x4 = vdivq_f32(vaddq_f32(a_high_f32x4, epsilon_f32x4),
+                                             vaddq_f32(b_high_f32x4, epsilon_f32x4));
+    float32x4_t log_ratio_low_f32x4 = nk_log2_f32x4_neon_(ratio_low_f32x4);
+    float32x4_t log_ratio_high_f32x4 = nk_log2_f32x4_neon_(ratio_high_f32x4);
+    sum_f32x4 = vfmaq_f32(sum_f32x4, a_low_f32x4, log_ratio_low_f32x4);
+    sum_f32x4 = vfmaq_f32(sum_f32x4, a_high_f32x4, log_ratio_high_f32x4);
+    if (n) goto nk_kld_f16_neon_cycle;
 
     nk_f32_t log2_normalizer = 0.693147181f;
     nk_f32_t sum = vaddvq_f32(sum_f32x4) * log2_normalizer;
     *result = sum;
 }
 
-NK_PUBLIC void nk_jsd_f16_neonhalf(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
+NK_PUBLIC void nk_jsd_f16_neon(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
     float32x4_t sum_f32x4 = vdupq_n_f32(0);
     nk_f32_t epsilon = NK_F32_DIVISION_EPSILON;
     float32x4_t epsilon_f32x4 = vdupq_n_f32(epsilon);
-    float32x4_t a_f32x4, b_f32x4;
+    float32x4_t a_low_f32x4, a_high_f32x4, b_low_f32x4, b_high_f32x4;
 
-nk_jsd_f16_neonhalf_cycle:
-    if (n < 4) {
-        nk_b64_vec_t a_vec, b_vec;
-        nk_partial_load_b16x4_serial_(a, &a_vec, n);
-        nk_partial_load_b16x4_serial_(b, &b_vec, n);
-        a_f32x4 = vcvt_f32_f16(vreinterpret_f16_u16(a_vec.u16x4));
-        b_f32x4 = vcvt_f32_f16(vreinterpret_f16_u16(b_vec.u16x4));
+nk_jsd_f16_neon_cycle:
+    if (n < 8) {
+        nk_b128_vec_t a_vec, b_vec;
+        nk_partial_load_b16x8_serial_(a, &a_vec, n);
+        nk_partial_load_b16x8_serial_(b, &b_vec, n);
+        float16x8_t a_f16x8 = vreinterpretq_f16_u16(a_vec.u16x8);
+        float16x8_t b_f16x8 = vreinterpretq_f16_u16(b_vec.u16x8);
+        a_low_f32x4 = vcvt_f32_f16(vget_low_f16(a_f16x8));
+        a_high_f32x4 = vcvt_high_f32_f16(a_f16x8);
+        b_low_f32x4 = vcvt_f32_f16(vget_low_f16(b_f16x8));
+        b_high_f32x4 = vcvt_high_f32_f16(b_f16x8);
         n = 0;
     }
     else {
-        a_f32x4 = vcvt_f32_f16(vld1_f16((nk_f16_for_arm_simd_t const *)a));
-        b_f32x4 = vcvt_f32_f16(vld1_f16((nk_f16_for_arm_simd_t const *)b));
-        n -= 4, a += 4, b += 4;
+        float16x8_t a_f16x8 = vreinterpretq_f16_u16(vld1q_u16((nk_u16_t const *)a));
+        float16x8_t b_f16x8 = vreinterpretq_f16_u16(vld1q_u16((nk_u16_t const *)b));
+        a_low_f32x4 = vcvt_f32_f16(vget_low_f16(a_f16x8));
+        a_high_f32x4 = vcvt_high_f32_f16(a_f16x8);
+        b_low_f32x4 = vcvt_f32_f16(vget_low_f16(b_f16x8));
+        b_high_f32x4 = vcvt_high_f32_f16(b_f16x8);
+        n -= 8, a += 8, b += 8;
     }
 
-    float32x4_t mean_f32x4 = vmulq_n_f32(vaddq_f32(a_f32x4, b_f32x4), 0.5f);
-    float32x4_t ratio_a_f32x4 = vdivq_f32(vaddq_f32(a_f32x4, epsilon_f32x4), vaddq_f32(mean_f32x4, epsilon_f32x4));
-    float32x4_t ratio_b_f32x4 = vdivq_f32(vaddq_f32(b_f32x4, epsilon_f32x4), vaddq_f32(mean_f32x4, epsilon_f32x4));
-    float32x4_t log_ratio_a_f32x4 = nk_log2_f32x4_neon_(ratio_a_f32x4);
-    float32x4_t log_ratio_b_f32x4 = nk_log2_f32x4_neon_(ratio_b_f32x4);
-    float32x4_t contribution_a_f32x4 = vmulq_f32(a_f32x4, log_ratio_a_f32x4);
-    float32x4_t contribution_b_f32x4 = vmulq_f32(b_f32x4, log_ratio_b_f32x4);
-    sum_f32x4 = vaddq_f32(sum_f32x4, vaddq_f32(contribution_a_f32x4, contribution_b_f32x4));
-    if (n) goto nk_jsd_f16_neonhalf_cycle;
+    float32x4_t mean_low_f32x4 = vmulq_n_f32(vaddq_f32(a_low_f32x4, b_low_f32x4), 0.5f);
+    float32x4_t mean_high_f32x4 = vmulq_n_f32(vaddq_f32(a_high_f32x4, b_high_f32x4), 0.5f);
+    float32x4_t ratio_a_low_f32x4 = vdivq_f32(vaddq_f32(a_low_f32x4, epsilon_f32x4),
+                                              vaddq_f32(mean_low_f32x4, epsilon_f32x4));
+    float32x4_t ratio_a_high_f32x4 = vdivq_f32(vaddq_f32(a_high_f32x4, epsilon_f32x4),
+                                               vaddq_f32(mean_high_f32x4, epsilon_f32x4));
+    float32x4_t ratio_b_low_f32x4 = vdivq_f32(vaddq_f32(b_low_f32x4, epsilon_f32x4),
+                                              vaddq_f32(mean_low_f32x4, epsilon_f32x4));
+    float32x4_t ratio_b_high_f32x4 = vdivq_f32(vaddq_f32(b_high_f32x4, epsilon_f32x4),
+                                               vaddq_f32(mean_high_f32x4, epsilon_f32x4));
+    float32x4_t log_ratio_a_low_f32x4 = nk_log2_f32x4_neon_(ratio_a_low_f32x4);
+    float32x4_t log_ratio_a_high_f32x4 = nk_log2_f32x4_neon_(ratio_a_high_f32x4);
+    float32x4_t log_ratio_b_low_f32x4 = nk_log2_f32x4_neon_(ratio_b_low_f32x4);
+    float32x4_t log_ratio_b_high_f32x4 = nk_log2_f32x4_neon_(ratio_b_high_f32x4);
+    sum_f32x4 = vfmaq_f32(sum_f32x4, a_low_f32x4, log_ratio_a_low_f32x4);
+    sum_f32x4 = vfmaq_f32(sum_f32x4, a_high_f32x4, log_ratio_a_high_f32x4);
+    sum_f32x4 = vfmaq_f32(sum_f32x4, b_low_f32x4, log_ratio_b_low_f32x4);
+    sum_f32x4 = vfmaq_f32(sum_f32x4, b_high_f32x4, log_ratio_b_high_f32x4);
+    if (n) goto nk_jsd_f16_neon_cycle;
 
     nk_f32_t log2_normalizer = 0.693147181f;
     nk_f32_t sum = vaddvq_f32(sum_f32x4) * log2_normalizer / 2;
@@ -215,7 +245,7 @@ nk_jsd_f16_neonhalf_cycle:
 #elif defined(__GNUC__)
 #pragma GCC pop_options
 #endif
-#endif // NK_TARGET_NEONHALF
+#endif // NK_TARGET_NEON
 
 #if defined(__cplusplus)
 } // extern "C"