numkong 7.0.0 → 7.4.2

package/include/numkong/mesh/neonbfdot.h

@@ -8,16 +8,15 @@
  *
  *  @section mesh_neonbfdot_instructions ARM NEON BF16 Instructions (ARMv8.6-BF16)
  *
- *      Intrinsic                   Instruction                     Latency     Throughput
- *                                                                              A76         M4+/V1+/Oryon
- *      vld3_u16                    LD3 (V.4H x 3)                  6cy         1/cy        2/cy
- *      vshll_n_u16                 USHLL (V.4S, V.4H, #16)         2cy         2/cy        4/cy
- *      vfmaq_f32                   FMLA (V.4S, V.4S, V.4S)         4cy         2/cy        4/cy
- *      vaddq_f32                   FADD (V.4S, V.4S, V.4S)         2cy         2/cy        4/cy
- *      vsubq_f32                   FSUB (V.4S, V.4S, V.4S)         2cy         2/cy        4/cy
- *      vmulq_f32                   FMUL (V.4S, V.4S, V.4S)         3cy         2/cy        4/cy
- *      vdupq_n_f32                 DUP (V.4S, scalar)              2cy         2/cy        4/cy
- *      vaddvq_f32                  FADDP+FADDP (V.4S)              4cy         1/cy        2/cy
+ *      Intrinsic    Instruction              A76       M5
+ *      vld3_u16     LD3 (V.4H x 3)           4cy @ 1p  4cy @ 1p
+ *      vshll_n_u16  USHLL (V.4S, V.4H, #16)  2cy @ 2p  2cy @ 4p
+ *      vfmaq_f32    FMLA (V.4S, V.4S, V.4S)  4cy @ 2p  3cy @ 4p
+ *      vaddq_f32    FADD (V.4S, V.4S, V.4S)  2cy @ 2p  2cy @ 4p
+ *      vsubq_f32    FSUB (V.4S, V.4S, V.4S)  2cy @ 2p  2cy @ 4p
+ *      vmulq_f32    FMUL (V.4S, V.4S, V.4S)  3cy @ 2p  3cy @ 4p
+ *      vdupq_n_f32  DUP (V.4S, scalar)       2cy @ 2p  2cy @ 4p
+ *      vaddvq_f32   FADDP+FADDP (V.4S)       5cy @ 1p  8cy @ 1p
  *
  *  The ARMv8.6-BF16 extension enables BF16 storage with F32 computation for 3D mesh alignment
  *  operations. BF16's wider exponent range (matching F32) prevents overflow in geometric calculations
@@ -57,14 +56,14 @@ extern "C" {
 NK_INTERNAL void nk_deinterleave_bf16x4_to_f32x4_neonbfdot_(nk_bf16_t const *ptr, float32x4_t *x_out,
                                                             float32x4_t *y_out, float32x4_t *z_out) {
     // Load 12 bf16 values and de-interleave into x, y, z components
-    uint16x4x3_t xyz = vld3_u16((uint16_t const *)ptr);
+    uint16x4x3_t xyz_u16x4x3 = vld3_u16((nk_u16_t const *)ptr);
     // Convert bf16 to f32 by zero-extending to lower 16 bits, then shifting left by 16
-    uint32x4_t x_u32 = vshll_n_u16(xyz.val[0], 16);
-    uint32x4_t y_u32 = vshll_n_u16(xyz.val[1], 16);
-    uint32x4_t z_u32 = vshll_n_u16(xyz.val[2], 16);
-    *x_out = vreinterpretq_f32_u32(x_u32);
-    *y_out = vreinterpretq_f32_u32(y_u32);
-    *z_out = vreinterpretq_f32_u32(z_u32);
+    uint32x4_t x_u32x4 = vshll_n_u16(xyz_u16x4x3.val[0], 16);
+    uint32x4_t y_u32x4 = vshll_n_u16(xyz_u16x4x3.val[1], 16);
+    uint32x4_t z_u32x4 = vshll_n_u16(xyz_u16x4x3.val[2], 16);
+    *x_out = vreinterpretq_f32_u32(x_u32x4);
+    *y_out = vreinterpretq_f32_u32(y_u32x4);
+    *z_out = vreinterpretq_f32_u32(z_u32x4);
 }
 
 NK_INTERNAL void nk_partial_deinterleave_bf16_to_f32x4_neonbfdot_(nk_bf16_t const *ptr, nk_size_t n_points,
@@ -216,8 +215,9 @@ NK_INTERNAL nk_f32_t nk_transformed_ssd_bf16_neonbfdot_(nk_bf16_t const *a, nk_b
         nk_partial_deinterleave_bf16_to_f32x4_neonbfdot_(b + j * 3, n - j, &b_x_f32x4, &b_y_f32x4, &b_z_f32x4);
 
         // Mask invalid lanes to zero BEFORE centering
-        uint32x4_t lane_u32x4 = {0, 1, 2, 3};
-        uint32x4_t valid_u32x4 = vcltq_u32(lane_u32x4, vdupq_n_u32((uint32_t)(n - j)));
+        uint32x4_t lane_u32x4 = vcombine_u32(vreinterpret_u32_u64(vcreate_u64(0x0000000100000000ULL)),
+                                             vreinterpret_u32_u64(vcreate_u64(0x0000000300000002ULL)));
+        uint32x4_t valid_u32x4 = vcltq_u32(lane_u32x4, vdupq_n_u32((nk_u32_t)(n - j)));
         float32x4_t zero_f32x4 = vdupq_n_f32(0);
         a_x_f32x4 = vbslq_f32(valid_u32x4, a_x_f32x4, zero_f32x4);
         a_y_f32x4 = vbslq_f32(valid_u32x4, a_y_f32x4, zero_f32x4);
@@ -262,12 +262,10 @@ NK_INTERNAL nk_f32_t nk_transformed_ssd_bf16_neonbfdot_(nk_bf16_t const *a, nk_b
 
 NK_PUBLIC void nk_rmsd_bf16_neonbfdot(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) {
-    /* RMSD uses identity rotation and scale=1.0 */
-    if (rotation) {
-        rotation[0] = 1, rotation[1] = 0, rotation[2] = 0;
-        rotation[3] = 0, rotation[4] = 1, rotation[5] = 0;
+    // RMSD uses identity rotation and scale=1.0
+    if (rotation)
+        rotation[0] = 1, rotation[1] = 0, rotation[2] = 0, rotation[3] = 0, rotation[4] = 1, rotation[5] = 0,
         rotation[6] = 0, rotation[7] = 0, rotation[8] = 1;
-    }
     if (scale) *scale = 1.0f;
 
     float32x4_t const zeros_f32x4 = vdupq_n_f32(0);
@@ -343,16 +341,8 @@ NK_PUBLIC void nk_rmsd_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b, nk
     nk_f32_t centroid_b_y = total_by * inv_n;
     nk_f32_t centroid_b_z = total_bz * inv_n;
 
-    if (a_centroid) {
-        a_centroid[0] = centroid_a_x;
-        a_centroid[1] = centroid_a_y;
-        a_centroid[2] = centroid_a_z;
-    }
-    if (b_centroid) {
-        b_centroid[0] = centroid_b_x;
-        b_centroid[1] = centroid_b_y;
-        b_centroid[2] = centroid_b_z;
-    }
+    if (a_centroid) a_centroid[0] = centroid_a_x, a_centroid[1] = centroid_a_y, a_centroid[2] = centroid_a_z;
+    if (b_centroid) b_centroid[0] = centroid_b_x, b_centroid[1] = centroid_b_y, b_centroid[2] = centroid_b_z;
 
     // Compute RMSD
     nk_f32_t mean_diff_x = centroid_a_x - centroid_b_x;
@@ -368,7 +358,7 @@ NK_PUBLIC void nk_kabsch_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b,
                                         nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result) {
     float32x4_t const zeros_f32x4 = vdupq_n_f32(0);
 
-    /*  2x unrolling with dual accumulators to hide FMA latency. */
+    // 2x unrolling with dual accumulators to hide FMA latency.
     float32x4_t sum_a_x_a_f32x4 = zeros_f32x4, sum_a_y_a_f32x4 = zeros_f32x4, sum_a_z_a_f32x4 = zeros_f32x4;
     float32x4_t sum_b_x_a_f32x4 = zeros_f32x4, sum_b_y_a_f32x4 = zeros_f32x4, sum_b_z_a_f32x4 = zeros_f32x4;
     float32x4_t sum_a_x_b_f32x4 = zeros_f32x4, sum_a_y_b_f32x4 = zeros_f32x4, sum_a_z_b_f32x4 = zeros_f32x4;
@@ -512,16 +502,8 @@ NK_PUBLIC void nk_kabsch_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b,
     nk_f32_t centroid_b_y = sum_b_y * inv_n;
     nk_f32_t centroid_b_z = sum_b_z * inv_n;
 
-    if (a_centroid) {
-        a_centroid[0] = centroid_a_x;
-        a_centroid[1] = centroid_a_y;
-        a_centroid[2] = centroid_a_z;
-    }
-    if (b_centroid) {
-        b_centroid[0] = centroid_b_x;
-        b_centroid[1] = centroid_b_y;
-        b_centroid[2] = centroid_b_z;
-    }
+    if (a_centroid) a_centroid[0] = centroid_a_x, a_centroid[1] = centroid_a_y, a_centroid[2] = centroid_a_z;
+    if (b_centroid) b_centroid[0] = centroid_b_x, b_centroid[1] = centroid_b_y, b_centroid[2] = centroid_b_z;
 
     // Apply centering correction: H_centered = H - n * centroid_a * centroid_bᵀ
     covariance_x_x -= n * centroid_a_x * centroid_b_x;
@@ -554,9 +536,7 @@ NK_PUBLIC void nk_kabsch_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b,
 
     // Handle reflection: if det(R) < 0, negate third column of V and recompute R
     if (nk_det3x3_f32_(r) < 0) {
-        svd_v[2] = -svd_v[2];
-        svd_v[5] = -svd_v[5];
-        svd_v[8] = -svd_v[8];
+        svd_v[2] = -svd_v[2], svd_v[5] = -svd_v[5], svd_v[8] = -svd_v[8];
         r[0] = svd_v[0] * svd_u[0] + svd_v[1] * svd_u[1] + svd_v[2] * svd_u[2];
         r[1] = svd_v[0] * svd_u[3] + svd_v[1] * svd_u[4] + svd_v[2] * svd_u[5];
         r[2] = svd_v[0] * svd_u[6] + svd_v[1] * svd_u[7] + svd_v[2] * svd_u[8];
@@ -568,10 +548,9 @@ NK_PUBLIC void nk_kabsch_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b,
         r[8] = svd_v[6] * svd_u[6] + svd_v[7] * svd_u[7] + svd_v[8] * svd_u[8];
     }
 
-    /* Output rotation matrix and scale=1.0 */
-    if (rotation) {
+    // Output rotation matrix and scale=1.0
+    if (rotation)
         for (int j = 0; j < 9; ++j) rotation[j] = r[j];
-    }
     if (scale) *scale = 1.0f;
 
     // Compute RMSD after optimal rotation
@@ -584,7 +563,7 @@ NK_PUBLIC void nk_umeyama_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b,
                                          nk_f32_t *b_centroid, nk_f32_t *rotation, nk_f32_t *scale, nk_f32_t *result) {
     float32x4_t const zeros_f32x4 = vdupq_n_f32(0);
 
-    /*  2x unrolling with dual accumulators to hide FMA latency. */
+    // 2x unrolling with dual accumulators to hide FMA latency.
     float32x4_t sum_a_x_a_f32x4 = zeros_f32x4, sum_a_y_a_f32x4 = zeros_f32x4, sum_a_z_a_f32x4 = zeros_f32x4;
     float32x4_t sum_b_x_a_f32x4 = zeros_f32x4, sum_b_y_a_f32x4 = zeros_f32x4, sum_b_z_a_f32x4 = zeros_f32x4;
     float32x4_t sum_a_x_b_f32x4 = zeros_f32x4, sum_a_y_b_f32x4 = zeros_f32x4, sum_a_z_b_f32x4 = zeros_f32x4;
@@ -749,16 +728,8 @@ NK_PUBLIC void nk_umeyama_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b,
     nk_f32_t centroid_b_y = sum_b_y * inv_n;
     nk_f32_t centroid_b_z = sum_b_z * inv_n;
 
-    if (a_centroid) {
-        a_centroid[0] = centroid_a_x;
-        a_centroid[1] = centroid_a_y;
-        a_centroid[2] = centroid_a_z;
-    }
-    if (b_centroid) {
-        b_centroid[0] = centroid_b_x;
-        b_centroid[1] = centroid_b_y;
-        b_centroid[2] = centroid_b_z;
-    }
+    if (a_centroid) a_centroid[0] = centroid_a_x, a_centroid[1] = centroid_a_y, a_centroid[2] = centroid_a_z;
+    if (b_centroid) b_centroid[0] = centroid_b_x, b_centroid[1] = centroid_b_y, b_centroid[2] = centroid_b_z;
 
     // Compute centered variance of A
     nk_f32_t variance_a = variance_a_sum * inv_n -
@@ -802,9 +773,7 @@ NK_PUBLIC void nk_umeyama_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b,
     if (scale) *scale = c;
 
     if (rotation_det < 0) {
-        svd_v[2] = -svd_v[2];
-        svd_v[5] = -svd_v[5];
-        svd_v[8] = -svd_v[8];
+        svd_v[2] = -svd_v[2], svd_v[5] = -svd_v[5], svd_v[8] = -svd_v[8];
         r[0] = svd_v[0] * svd_u[0] + svd_v[1] * svd_u[1] + svd_v[2] * svd_u[2];
         r[1] = svd_v[0] * svd_u[3] + svd_v[1] * svd_u[4] + svd_v[2] * svd_u[5];
         r[2] = svd_v[0] * svd_u[6] + svd_v[1] * svd_u[7] + svd_v[2] * svd_u[8];
@@ -816,10 +785,9 @@ NK_PUBLIC void nk_umeyama_bf16_neonbfdot(nk_bf16_t const *a, nk_bf16_t const *b,
         r[8] = svd_v[6] * svd_u[6] + svd_v[7] * svd_u[7] + svd_v[8] * svd_u[8];
     }
 
-    /* Output rotation matrix */
-    if (rotation) {
+    // Output rotation matrix
+    if (rotation)
         for (int j = 0; j < 9; ++j) rotation[j] = r[j];
-    }
 
     // Compute RMSD after similarity transform: ‖c × R × a - b‖
     nk_f32_t sum_squared = nk_transformed_ssd_bf16_neonbfdot_(a, b, n, r, c, centroid_a_x, centroid_a_y, centroid_a_z,