usearch 2.26.1 → 2.26.2
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +1 -1
- package/include/usearch/index.hpp +5 -1
- package/numkong/CITATION.cff +13 -9
- package/numkong/CONTRIBUTING.md +29 -24
- package/numkong/Cargo.toml +1 -1
- package/numkong/README.md +21 -7
- package/numkong/binding.gyp +47 -31
- package/numkong/build.rs +15 -5
- package/numkong/c/dispatch.h +0 -1
- package/numkong/c/parallel.c +139 -0
- package/numkong/c/parallel.h +47 -0
- package/numkong/cmake/toolchain-wasi.cmake +39 -5
- package/numkong/include/README.md +24 -17
- package/numkong/include/numkong/attention/sapphireamx.h +6 -6
- package/numkong/include/numkong/attention/sme.h +2 -2
- package/numkong/include/numkong/capabilities.h +2 -2
- package/numkong/include/numkong/cast/README.md +9 -7
- package/numkong/include/numkong/cast/loongsonasx.h +48 -40
- package/numkong/include/numkong/cast/serial.h +89 -11
- package/numkong/include/numkong/cast.h +3 -3
- package/numkong/include/numkong/curved/README.md +10 -10
- package/numkong/include/numkong/curved/neon.h +1 -0
- package/numkong/include/numkong/curved/neonbfdot.h +1 -0
- package/numkong/include/numkong/curved/smef64.h +1 -0
- package/numkong/include/numkong/dot/README.md +18 -18
- package/numkong/include/numkong/dot/alder.h +8 -8
- package/numkong/include/numkong/dot/haswell.h +8 -8
- package/numkong/include/numkong/dot/icelake.h +12 -12
- package/numkong/include/numkong/dot/loongsonasx.h +6 -6
- package/numkong/include/numkong/dot/neon.h +2 -1
- package/numkong/include/numkong/dot/neonbfdot.h +1 -0
- package/numkong/include/numkong/dot/neonfhm.h +1 -0
- package/numkong/include/numkong/dot/neonsdot.h +1 -0
- package/numkong/include/numkong/dot/powervsx.h +3 -3
- package/numkong/include/numkong/dot/serial.h +19 -6
- package/numkong/include/numkong/dot/skylake.h +4 -4
- package/numkong/include/numkong/dot/svehalf.h +1 -1
- package/numkong/include/numkong/dot/v128relaxed.h +24 -23
- package/numkong/include/numkong/dots/README.md +26 -26
- package/numkong/include/numkong/dots/genoa.h +3 -1
- package/numkong/include/numkong/dots/icelake.h +1 -0
- package/numkong/include/numkong/dots/powervsx.h +1 -1
- package/numkong/include/numkong/dots/serial.h +32 -33
- package/numkong/include/numkong/dots/sierra.h +4 -3
- package/numkong/include/numkong/each/README.md +5 -3
- package/numkong/include/numkong/each/sapphire.h +2 -0
- package/numkong/include/numkong/each/skylake.h +1 -0
- package/numkong/include/numkong/geospatial/README.md +6 -6
- package/numkong/include/numkong/geospatial/v128relaxed.h +1 -0
- package/numkong/include/numkong/maxsim/README.md +5 -5
- package/numkong/include/numkong/maxsim/genoa.h +4 -0
- package/numkong/include/numkong/maxsim/icelake.h +1 -0
- package/numkong/include/numkong/maxsim/serial.h +22 -9
- package/numkong/include/numkong/maxsim/sme.h +3 -2
- package/numkong/include/numkong/mesh/README.md +7 -4
- package/numkong/include/numkong/mesh/genoa.h +2 -2
- package/numkong/include/numkong/mesh/haswell.h +8 -8
- package/numkong/include/numkong/mesh/neon.h +6 -6
- package/numkong/include/numkong/mesh/neonbfdot.h +3 -2
- package/numkong/include/numkong/mesh/neonfhm.h +3 -2
- package/numkong/include/numkong/mesh/rvv.h +4 -4
- package/numkong/include/numkong/mesh/serial.h +17 -17
- package/numkong/include/numkong/mesh/skylake.h +4 -4
- package/numkong/include/numkong/mesh/v128relaxed.h +4 -4
- package/numkong/include/numkong/mesh.hpp +2 -2
- package/numkong/include/numkong/numkong.h +1 -1
- package/numkong/include/numkong/numkong.hpp +1 -1
- package/numkong/include/numkong/probability/README.md +17 -15
- package/numkong/include/numkong/probability/skylake.h +1 -1
- package/numkong/include/numkong/reduce/README.md +27 -21
- package/numkong/include/numkong/reduce/genoa.h +2 -1
- package/numkong/include/numkong/reduce/icelake.h +1 -0
- package/numkong/include/numkong/reduce/neon.h +1 -1
- package/numkong/include/numkong/reduce/serial.h +9 -9
- package/numkong/include/numkong/reduce/skylake.h +4 -4
- package/numkong/include/numkong/reduce.hpp +38 -0
- package/numkong/include/numkong/scalar/README.md +29 -28
- package/numkong/include/numkong/set/README.md +17 -14
- package/numkong/include/numkong/set/haswell.h +15 -16
- package/numkong/include/numkong/set/icelake.h +15 -16
- package/numkong/include/numkong/set/loongsonasx.h +12 -10
- package/numkong/include/numkong/set/neon.h +15 -19
- package/numkong/include/numkong/set/powervsx.h +15 -21
- package/numkong/include/numkong/set/serial.h +19 -20
- package/numkong/include/numkong/set/v128relaxed.h +10 -8
- package/numkong/include/numkong/sets/icelake.h +1 -0
- package/numkong/include/numkong/sparse/README.md +4 -4
- package/numkong/include/numkong/sparse/icelake.h +1 -0
- package/numkong/include/numkong/sparse/neon.h +1 -0
- package/numkong/include/numkong/sparse/serial.h +10 -3
- package/numkong/include/numkong/sparse/turin.h +1 -0
- package/numkong/include/numkong/spatial/README.md +25 -24
- package/numkong/include/numkong/spatial/alder.h +2 -1
- package/numkong/include/numkong/spatial/haswell.h +48 -40
- package/numkong/include/numkong/spatial/icelake.h +2 -1
- package/numkong/include/numkong/spatial/neon.h +54 -46
- package/numkong/include/numkong/spatial/neonbfdot.h +1 -1
- package/numkong/include/numkong/spatial/powervsx.h +54 -46
- package/numkong/include/numkong/spatial/serial.h +56 -46
- package/numkong/include/numkong/spatial/sierra.h +4 -3
- package/numkong/include/numkong/spatial/skylake.h +27 -22
- package/numkong/include/numkong/spatial/sve.h +6 -6
- package/numkong/include/numkong/spatial/svebfdot.h +1 -1
- package/numkong/include/numkong/spatial/svehalf.h +1 -1
- package/numkong/include/numkong/spatial/v128relaxed.h +36 -30
- package/numkong/include/numkong/spatials/README.md +15 -15
- package/numkong/include/numkong/spatials/genoa.h +3 -0
- package/numkong/include/numkong/spatials/graniteamx.h +2 -0
- package/numkong/include/numkong/spatials/icelake.h +1 -0
- package/numkong/include/numkong/spatials/sierra.h +1 -0
- package/numkong/include/numkong/tensor.hpp +236 -123
- package/numkong/include/numkong/trigonometry/README.md +4 -4
- package/numkong/include/numkong/trigonometry/serial.h +6 -6
- package/numkong/include/numkong/trigonometry/v128relaxed.h +1 -0
- package/numkong/include/numkong/types.hpp +46 -16
- package/numkong/include/numkong/vector.hpp +116 -23
- package/numkong/javascript/@numkong-darwin-arm64/package.json +1 -1
- package/numkong/javascript/@numkong-darwin-x64/package.json +1 -1
- package/numkong/javascript/@numkong-linux-arm64/package.json +1 -1
- package/numkong/javascript/@numkong-linux-x64/package.json +1 -1
- package/numkong/javascript/@numkong-win32-arm64/package.json +1 -1
- package/numkong/javascript/@numkong-win32-x64/package.json +1 -1
- package/numkong/javascript/README.md +18 -16
- package/numkong/javascript/numkong.c +70 -40
- package/numkong/package.json +8 -8
- package/numkong/rust/README.md +19 -9
- package/numkong/rust/cast.rs +52 -1
- package/numkong/rust/each.rs +480 -7
- package/numkong/rust/matrix.rs +877 -216
- package/numkong/rust/numkong.rs +12 -8
- package/numkong/rust/reduce.rs +433 -5
- package/numkong/rust/tensor.rs +450 -842
- package/numkong/rust/vector.rs +520 -101
- package/numkong/swift/README.md +5 -3
- package/numkong/test/README.md +4 -4
- package/numkong/test/test-wasi.mjs +17 -0
- package/numkong/test/test_each.cpp +30 -1
- package/numkong/test/test_each.py +74 -1
- package/numkong/test/test_reduce.cpp +51 -0
- package/numkong/test/test_tensor.cpp +2 -63
- package/numkong/test/test_tensor.py +93 -0
- package/package.json +1 -1
- package/prebuilds/darwin-arm64+x64/usearch.node +0 -0
- package/prebuilds/linux-arm64/usearch.node +0 -0
- package/prebuilds/linux-x64/usearch.node +0 -0
- package/prebuilds/win32-x64/usearch.node +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
# Batched Dot Products in NumKong
|
|
2
2
|
|
|
3
|
-
NumKong implements batched GEMM computing C = A × Bᵀ (packed) and C = A × Aᵀ (symmetric). B is pre-packed once and reused across queries. This is the foundation for the spatials
|
|
3
|
+
NumKong implements batched GEMM computing C = A × Bᵀ (packed) and C = A × Aᵀ (symmetric). B is pre-packed once and reused across queries. This is the foundation for the spatials and sets modules.
|
|
4
4
|
|
|
5
5
|
Packed dot product computes the full cross-product matrix:
|
|
6
6
|
|
|
@@ -28,21 +28,21 @@ def dots_symmetric(a: np.ndarray) -> np.ndarray:
|
|
|
28
28
|
|
|
29
29
|
## Input & Output Types
|
|
30
30
|
|
|
31
|
-
| Input Type | Output Type | Description
|
|
32
|
-
| ---------- | ----------- |
|
|
33
|
-
| `f64` | `f64` | 64-bit IEEE 754 double precision
|
|
34
|
-
| `f32` | `
|
|
35
|
-
| `f16` | `f32` | 16-bit IEEE 754 half precision, widened output
|
|
36
|
-
| `bf16` | `f32` | 16-bit brain float, widened output
|
|
37
|
-
| `e4m3` | `f32` | 8-bit Float8: 4 exponent, 3 mantissa bits
|
|
38
|
-
| `e5m2` | `f32` | 8-bit Float8: 5 exponent, 2 mantissa bits
|
|
39
|
-
| `e2m3` | `f32` | 8-bit MX format: 2 exponent, 3 mantissa bits
|
|
40
|
-
| `e3m2` | `f32` | 8-bit MX format: 3 exponent, 2 mantissa bits
|
|
41
|
-
| `i8` | `i32` | 8-bit signed integers
|
|
42
|
-
| `u8` | `u32` | 8-bit unsigned integers
|
|
43
|
-
| `i4` | `i32` | 4-bit signed integers, packed nibble pairs
|
|
44
|
-
| `u4` | `u32` | 4-bit unsigned integers, packed nibble pairs
|
|
45
|
-
| `u1` | `u32` | 1-bit binary packed octets, popcount of AND
|
|
31
|
+
| Input Type | Output Type | Description |
|
|
32
|
+
| ---------- | ----------- | ------------------------------------------------ |
|
|
33
|
+
| `f64` | `f64` | 64-bit IEEE 754 double precision |
|
|
34
|
+
| `f32` | `f64` | 32-bit IEEE 754 single precision, widened output |
|
|
35
|
+
| `f16` | `f32` | 16-bit IEEE 754 half precision, widened output |
|
|
36
|
+
| `bf16` | `f32` | 16-bit brain float, widened output |
|
|
37
|
+
| `e4m3` | `f32` | 8-bit Float8: 4 exponent, 3 mantissa bits |
|
|
38
|
+
| `e5m2` | `f32` | 8-bit Float8: 5 exponent, 2 mantissa bits |
|
|
39
|
+
| `e2m3` | `f32` | 8-bit MX format: 2 exponent, 3 mantissa bits |
|
|
40
|
+
| `e3m2` | `f32` | 8-bit MX format: 3 exponent, 2 mantissa bits |
|
|
41
|
+
| `i8` | `i32` | 8-bit signed integers |
|
|
42
|
+
| `u8` | `u32` | 8-bit unsigned integers |
|
|
43
|
+
| `i4` | `i32` | 4-bit signed integers, packed nibble pairs |
|
|
44
|
+
| `u4` | `u32` | 4-bit unsigned integers, packed nibble pairs |
|
|
45
|
+
| `u1` | `u32` | 1-bit binary packed octets, popcount of AND |
|
|
46
46
|
|
|
47
47
|
## Optimizations
|
|
48
48
|
|
|
@@ -73,7 +73,7 @@ Tile configuration via `LDTILECFG` sets row counts and column byte-widths per ti
|
|
|
73
73
|
Morton Z-curve ordering for tile traversal improves cache reuse when both A and B exceed L2.
|
|
74
74
|
This eliminates the explicit M×N×K loop nesting and register file pressure of vector ISAs — the entire dot-product reduction happens inside the tile instruction.
|
|
75
75
|
FP8 inputs on Sapphire AMX go through an on-the-fly E4M3/E5M2 → BF16 pack via the Ice Lake `VPERMI2W` LUT helpers — port-5-bound but the simplest correct route to feed `TDPBF16PS` tiles.
|
|
76
|
-
Granite Rapids adds `TDPFP16PS` (same tile shape, FP16 operands); the E5M2 variant widens inputs with a
|
|
76
|
+
Granite Rapids adds `TDPFP16PS` (same tile shape, FP16 operands); the E5M2 variant widens inputs with `VPMOVZXBW` and a left shift by 8 into FP16 tiles at pack time and then reuses the native FP16 compute loop — keeps the intermediate at FP16 precision instead of truncating to BF16 like the Sapphire path.
|
|
77
77
|
|
|
78
78
|
### SME Outer-Product Streaming
|
|
79
79
|
|
|
@@ -105,17 +105,17 @@ The 4×4 tile accumulates 16 dot products in registers, then stores results 4-wi
|
|
|
105
105
|
A finalizer function pointer processes 4 results simultaneously, amortizing horizontal reductions and type conversions:
|
|
106
106
|
|
|
107
107
|
```
|
|
108
|
-
// 4-wide finalizer signature
|
|
109
|
-
void finalizer(nk_b128_vec_t
|
|
110
|
-
nk_f32_t query_norm,
|
|
111
|
-
nk_b128_vec_t
|
|
112
|
-
nk_b128_vec_t *
|
|
108
|
+
// 4-wide finalizer signature — per-lane arrays always pass as `nk_bXXX_vec_t const *`
|
|
109
|
+
void finalizer(nk_b128_vec_t const *dots_vec, // 4 dot products
|
|
110
|
+
nk_f32_t query_norm, // precomputed query squared-norm (scalar)
|
|
111
|
+
nk_b128_vec_t const *target_norms_vec, // 4 target squared-norms
|
|
112
|
+
nk_b128_vec_t *result_vec) // 4 output distances
|
|
113
113
|
|
|
114
114
|
// Angular: 4 divisions + 4 subtractions in one call
|
|
115
|
-
|
|
115
|
+
result_vec->f32s[i] = 1 - dots_vec->f32s[i] / sqrt(query_norm * target_norms_vec->f32s[i])
|
|
116
116
|
|
|
117
117
|
// Euclidean: 4 sqrt(a² + b² - 2ab) in one call
|
|
118
|
-
|
|
118
|
+
result_vec->f32s[i] = sqrt(query_norm + target_norms_vec->f32s[i] - 2 * dots_vec->f32s[i])
|
|
119
119
|
```
|
|
120
120
|
|
|
121
121
|
The 4×4 tile emits 4 rows of 4 results each — the finalizer is called 4 times per tile, once per query row.
|
|
@@ -359,8 +359,8 @@ Measured with Wasmtime v42 (Cranelift backend).
|
|
|
359
359
|
| __f16__ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ |
|
|
360
360
|
| `nk_dots_packed_f16_serial` | 14.8 gso/s, 204 ulp | 14.2 gso/s, 36 ulp | 14.8 gso/s, 326 ulp |
|
|
361
361
|
| `nk_dots_symmetric_f16_serial` | 24.3 gso/s, 13 ulp | 24.9 gso/s, 24.6 ulp | 26.7 gso/s, 506 ulp |
|
|
362
|
-
| `
|
|
363
|
-
| `
|
|
362
|
+
| `nk_dots_packed_f16_neon` | 77.0 gso/s, 16.8 ulp | 79.1 gso/s, 25.5 ulp | 84.2 gso/s, 618 ulp |
|
|
363
|
+
| `nk_dots_symmetric_f16_neon` | 20.5 gso/s, 12.1 ulp | 20.4 gso/s, 25.0 ulp | 22.5 gso/s, 506 ulp |
|
|
364
364
|
| `nk_dots_packed_f16_neonfhm` | 104 gso/s, 16.7 ulp | 110 gso/s, 25.5 ulp | 118 gso/s, 618 ulp |
|
|
365
365
|
| `nk_dots_symmetric_f16_neonfhm` | 34.5 gso/s, 12.1 ulp | 40.4 gso/s, 25.0 ulp | 41.5 gso/s, 506 ulp |
|
|
366
366
|
| `nk_dots_packed_f16_sme` | 1,106 gso/s, 14.8 ulp | 1,213 gso/s, 28.2 ulp | 1,190 gso/s, 28.2 ulp |
|
|
@@ -13,7 +13,9 @@
|
|
|
13
13
|
#if NK_TARGET_GENOA
|
|
14
14
|
|
|
15
15
|
#include "numkong/dot/genoa.h"
|
|
16
|
-
#include "numkong/cast/icelake.h"
|
|
16
|
+
#include "numkong/cast/icelake.h" // `nk_load_e4m3x32_to_bf16x32_icelake_`
|
|
17
|
+
#include "numkong/dot/skylake.h" // `nk_dot_through_f32_finalize_skylake_`
|
|
18
|
+
#include "numkong/reduce/skylake.h" // `nk_reduce_add_f32x16_skylake_`
|
|
17
19
|
|
|
18
20
|
#if defined(__cplusplus)
|
|
19
21
|
extern "C" {
|
|
@@ -19,7 +19,6 @@
|
|
|
19
19
|
* GEMM kernels use tiled dot products with 4-way parallel accumulation to hide FMA latency.
|
|
20
20
|
* Type-specific tile sizes: f32 uses depth_simd_dimensions=4, f64 uses depth_simd_dimensions=2,
|
|
21
21
|
* bf16/f16 use depth_simd_dimensions=8, u1 uses depth_simd_dimensions=128.
|
|
22
|
-
* Load/store helpers are defined in cast/powervsx.h and included transitively via dot/powervsx.h.
|
|
23
22
|
*/
|
|
24
23
|
#ifndef NK_DOTS_POWERVSX_H
|
|
25
24
|
#define NK_DOTS_POWERVSX_H
|
|
@@ -27,6 +26,7 @@
|
|
|
27
26
|
#if NK_TARGET_POWER64_
|
|
28
27
|
#if NK_TARGET_POWERVSX
|
|
29
28
|
|
|
29
|
+
#include "numkong/cast/powervsx.h" // `nk_load_b128_powervsx_`, `nk_partial_load_b32x4_powervsx_`
|
|
30
30
|
#include "numkong/dot/powervsx.h"
|
|
31
31
|
|
|
32
32
|
#if defined(__cplusplus)
|
|
@@ -8,8 +8,8 @@
|
|
|
8
8
|
*
|
|
9
9
|
* This file provides two macro families for generating GEMM kernels:
|
|
10
10
|
*
|
|
11
|
-
* -
|
|
12
|
-
* -
|
|
11
|
+
* - nk_define_cross_packed_: vectorized inner-products between rows of A and Bᵀ
|
|
12
|
+
* - nk_define_cross_symmetric_: vectorized inner-products between rows and columns of A
|
|
13
13
|
*
|
|
14
14
|
* Both use the same B packing format (see below), enabling pack-once-use-anywhere.
|
|
15
15
|
*
|
|
@@ -229,15 +229,15 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
229
229
|
* @brief Generates function to calculate packed B matrix buffer size for GEMM micro-kernels.
|
|
230
230
|
*
|
|
231
231
|
* Memory layout: B_packed[column_count, depth_padded] with header storing metadata.
|
|
232
|
-
* Buffer size: sizeof(header) + column_count × depth_padded × sizeof(
|
|
232
|
+
* Buffer size: sizeof(header) + column_count × depth_padded × sizeof(packed_value_type) + column_count × sizeof(norm)
|
|
233
233
|
* Depth padding logic: Round up to `depth_simd_dimensions` multiple, then add `depth_simd_dimensions`
|
|
234
234
|
* if stride is power-of-2.
|
|
235
235
|
*
|
|
236
236
|
* @param api_name Operation name (hammings, dots)
|
|
237
237
|
* @param input_type_name Original type's name of B matrix values (i4, f16, bf16, e4m3, e5m2, f32, etc.)
|
|
238
238
|
* @param isa_suffix Platform Instruct Set Architecture suffix (serial, haswell, icelake, etc.)
|
|
239
|
-
* @param
|
|
240
|
-
* @param
|
|
239
|
+
* @param input_value_type Original type of B matrix values (i4x2, f16, bf16, e4m3, e5m2, f32, etc.)
|
|
240
|
+
* @param packed_value_type Internal storage type in packed buffer (often bf16 or f32 for mixed precision)
|
|
241
241
|
* @param norm_value_type Type of per-column norm values (f32, f64, u32) appended after packed data
|
|
242
242
|
* @param depth_simd_dimensions SIMD vector width in values for this platform/type combination
|
|
243
243
|
* @param dimensions_per_value Number of logical dimensions in a single value of input_type_name.
|
|
@@ -492,21 +492,20 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
492
492
|
* @param api_name Operation family (dots, hammings, jaccards) for codegen namespace
|
|
493
493
|
* @param input_type_name Type identifier for codegen (f32, bf16, i8, u1, etc.)
|
|
494
494
|
* @param isa_suffix ISA backend identifier (serial, haswell, neon, sve, icelake, etc.)
|
|
495
|
-
* @param
|
|
496
|
-
* @param
|
|
497
|
-
* @param
|
|
498
|
-
* @param vec_type SIMD vector type for depth dimension (e.g., __m256,
|
|
495
|
+
* @param input_value_type C type of input matrix values (f32, bf16, i8, u1x8, etc.)
|
|
496
|
+
* @param packed_value_type Storage type in packed B buffer (often bf16 or f32 for mixed precision)
|
|
497
|
+
* @param result_value_type C type of output matrix C values (f32, u32, f64, etc.)
|
|
498
|
+
* @param vec_type SIMD vector type for depth dimension (e.g., __m256, nk_b256_vec_t)
|
|
499
499
|
* @param state_type Accumulator state type (often vec_type or wider, e.g., __m256 or __m512)
|
|
500
500
|
* @param result_vec_type SIMD vector type for reduction results (e.g., __m128 for 4 f32 results)
|
|
501
501
|
* @param init_accumulator_fn Initialize accumulator: void fn(state_type*)
|
|
502
|
-
* @param load_a_vec_fn Full A vector load: vec_type fn(
|
|
502
|
+
* @param load_a_vec_fn Full A vector load: vec_type fn(input_value_type const*, nk_size_t offset)
|
|
503
503
|
* @param partial_load_a_vec_fn Partial A load for remainder
|
|
504
|
-
* @param load_b_vec_fn Full B vector load: vec_type fn(
|
|
504
|
+
* @param load_b_vec_fn Full B vector load: vec_type fn(packed_value_type const*, nk_size_t offset)
|
|
505
505
|
* @param partial_load_b_vec_fn Partial B load for remainder
|
|
506
506
|
* @param inner_product_fn Inner product accumulate
|
|
507
507
|
* @param reduce_accumulators_fn Reduce 4 accumulators
|
|
508
508
|
* @param store_fn Full-width store for results
|
|
509
|
-
* @param store_fn Full-width store for results
|
|
510
509
|
* @param partial_store_fn Partial store for results
|
|
511
510
|
* @param depth_simd_dimensions SIMD vector width in logical dimensions (e.g., 8 for f32 on AVX2, 128 for u1 on serial)
|
|
512
511
|
* @param dimensions_per_value Packing ratio: dimensions per storage value (1 for f32, 2 for i4x2, 8 for u1x8)
|
|
@@ -1190,7 +1189,7 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
1190
1189
|
result_vec_type result_vector; \
|
|
1191
1190
|
for (nk_size_t r = 0; r < register_row_count; ++r) { \
|
|
1192
1191
|
compensated_finalize_fn(&acc[r][0], &acc[r][1], &acc[r][2], &acc[r][3], depth, a_sums[r], \
|
|
1193
|
-
b_sum_vec, &result_vector);
|
|
1192
|
+
&b_sum_vec, &result_vector); \
|
|
1194
1193
|
nk_##result_value_type##_t *c_row = \
|
|
1195
1194
|
(nk_##result_value_type##_t *)((char *)c_matrix + (tr + r) * c_stride_in_bytes); \
|
|
1196
1195
|
store_fn(&result_vector, c_row + tc); \
|
|
@@ -1270,9 +1269,9 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
1270
1269
|
result_vec_type rv; \
|
|
1271
1270
|
nk_##result_value_type##_t *c_row = (nk_##result_value_type##_t *)((char *)c_matrix + \
|
|
1272
1271
|
ri * c_stride_in_bytes); \
|
|
1273
|
-
compensated_finalize_fn(&s0, &s1, &s2, &s3, depth, a_sum_val, b_sum_low, &rv);
|
|
1272
|
+
compensated_finalize_fn(&s0, &s1, &s2, &s3, depth, a_sum_val, &b_sum_low, &rv); \
|
|
1274
1273
|
store_fn(&rv, c_row + tc); \
|
|
1275
|
-
compensated_finalize_fn(&s4, &s5, &s6, &s7, depth, a_sum_val, b_sum_high, &rv);
|
|
1274
|
+
compensated_finalize_fn(&s4, &s5, &s6, &s7, depth, a_sum_val, &b_sum_high, &rv); \
|
|
1276
1275
|
store_fn(&rv, c_row + tc + 4); \
|
|
1277
1276
|
} \
|
|
1278
1277
|
} \
|
|
@@ -1445,7 +1444,7 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
1445
1444
|
for (nk_size_t rr = 0; rr < tile_row_count; ++rr) { \
|
|
1446
1445
|
result_vec_type rv; \
|
|
1447
1446
|
compensated_finalize_fn(&acc[rr][0], &acc[rr][1], &acc[rr][2], &acc[rr][3], depth, \
|
|
1448
|
-
a_sums[rr], b_sum_vec, &rv);
|
|
1447
|
+
a_sums[rr], &b_sum_vec, &rv); \
|
|
1449
1448
|
nk_##result_value_type##_t *c_row = \
|
|
1450
1449
|
(nk_##result_value_type##_t *)((char *)c_matrix + (tr + rr) * c_stride_in_bytes); \
|
|
1451
1450
|
partial_store_fn(&rv, c_row + tc, tile_col_count); \
|
|
@@ -1715,7 +1714,7 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
1715
1714
|
compensated_finalize_fn( \
|
|
1716
1715
|
&accumulators[row][is_diag ? row : 0], &accumulators[row][(is_diag ? row : 0) + 1], \
|
|
1717
1716
|
&accumulators[row][(is_diag ? row : 0) + 2], &accumulators[row][(is_diag ? row : 0) + 3], \
|
|
1718
|
-
depth, row_sums[row], col_sum_vec, &rv);
|
|
1717
|
+
depth, row_sums[row], &col_sum_vec, &rv); \
|
|
1719
1718
|
nk_size_t global_row = i_macro + tile_row_start + row; \
|
|
1720
1719
|
nk_size_t global_col_start = i_macro + tile_col_start + (is_diag ? row : 0); \
|
|
1721
1720
|
nk_size_t store_count = is_diag ? (tile_cols - row) : tile_cols; \
|
|
@@ -1875,7 +1874,7 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
1875
1874
|
for (nk_size_t row = 0; row < tile_rows; row++) { \
|
|
1876
1875
|
result_vec_type rv; \
|
|
1877
1876
|
compensated_finalize_fn(&accumulators[row][0], &accumulators[row][1], &accumulators[row][2], \
|
|
1878
|
-
&accumulators[row][3], depth, row_sums[tile_row_start + row], cs_vec,
|
|
1877
|
+
&accumulators[row][3], depth, row_sums[tile_row_start + row], &cs_vec, \
|
|
1879
1878
|
&rv); \
|
|
1880
1879
|
nk_##result_value_type##_t *dest = result + \
|
|
1881
1880
|
(i_macro + tile_row_start + row) * result_stride_values + \
|
|
@@ -2016,13 +2015,13 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
2016
2015
|
* @param api_name Operation family (dots, hammings, jaccards) for codegen namespace
|
|
2017
2016
|
* @param input_type_name Type identifier for codegen (f32, bf16, i8, u1, etc.)
|
|
2018
2017
|
* @param isa_suffix ISA backend identifier (serial, haswell, neon, sve, icelake, etc.)
|
|
2019
|
-
* @param
|
|
2020
|
-
* @param
|
|
2021
|
-
* @param vec_type SIMD vector type for input vectors (e.g., __m256,
|
|
2018
|
+
* @param input_value_type C type of input matrix values (f32, bf16, i8, u1x8, etc.)
|
|
2019
|
+
* @param result_value_type C type of output matrix values (f32, u32, f64, etc.)
|
|
2020
|
+
* @param vec_type SIMD vector type for input vectors (e.g., __m256, nk_b256_vec_t)
|
|
2022
2021
|
* @param state_type Accumulator state type (often vec_type or wider, e.g., __m256 or __m512)
|
|
2023
2022
|
* @param result_vec_type SIMD vector type for reduction results (e.g., __m128 for 4 f32 results)
|
|
2024
2023
|
* @param init_accumulator_fn Initialize accumulator: void fn(state_type*)
|
|
2025
|
-
* @param load_vec_fn Full vector load: vec_type fn(
|
|
2024
|
+
* @param load_vec_fn Full vector load: vec_type fn(input_value_type const*, nk_size_t offset)
|
|
2026
2025
|
* @param partial_load_vec_fn Partial vector load for remainder
|
|
2027
2026
|
* @param inner_product_fn Inner product accumulate
|
|
2028
2027
|
* @param reduce_accumulators_fn Reduce 4 accumulators
|
|
@@ -2435,10 +2434,12 @@ NK_INTERNAL nk_i32_t nk_dots_reduce_sum_i4_(nk_i4x2_t const *data, nk_size_t cou
|
|
|
2435
2434
|
* Without this, -O3 + LTO can vectorize or clone the serial kernels under AVX-512
|
|
2436
2435
|
* callers in dispatch_*.c, which wastes ~1 MB of binary and — more importantly —
|
|
2437
2436
|
* breaks the nk_*_serial-as-scalar-oracle contract that tests and the numerical-
|
|
2438
|
-
* stability docs in this header rely on.
|
|
2439
|
-
|
|
2440
|
-
|
|
2441
|
-
|
|
2437
|
+
* stability docs in this header rely on.
|
|
2438
|
+
*
|
|
2439
|
+
* Clang gets no blanket region here: one expansion of `nk_define_cross_packed_` /
|
|
2440
|
+
* `nk_define_cross_symmetric_` emits `always_inline` `_aligned_` fast paths next to the kernel, and
|
|
2441
|
+
* clang rejects `noinline` on those. `no-ipa-cp-clone` is the knob that stops the cloning anyway. */
|
|
2442
|
+
#if defined(__GNUC__) && !defined(__clang__)
|
|
2442
2443
|
#pragma GCC push_options
|
|
2443
2444
|
#pragma GCC optimize("no-tree-vectorize", "no-tree-slp-vectorize", "no-ipa-cp-clone", "no-inline")
|
|
2444
2445
|
#endif
|
|
@@ -2698,9 +2699,7 @@ nk_define_cross_packed_(dots, u1, serial, u1x8, u1x8, u32, nk_b128_vec_t, nk_dot
|
|
|
2698
2699
|
#endif
|
|
2699
2700
|
#endif
|
|
2700
2701
|
|
|
2701
|
-
#if defined(__clang__)
|
|
2702
|
-
#pragma clang attribute pop
|
|
2703
|
-
#elif defined(__GNUC__)
|
|
2702
|
+
#if defined(__GNUC__) && !defined(__clang__)
|
|
2704
2703
|
#pragma GCC pop_options
|
|
2705
2704
|
#endif
|
|
2706
2705
|
|
|
@@ -2786,14 +2785,14 @@ NK_PUBLIC void nk_dots_compact_i8_serial(void *c, nk_size_t row_count, nk_size_t
|
|
|
2786
2785
|
vec_type dots_vec, norms_vec, results_vec; \
|
|
2787
2786
|
load_fn(r_row_dots + column_index, &dots_vec); \
|
|
2788
2787
|
load_fn(b_norms + column_index, &norms_vec); \
|
|
2789
|
-
from_dot_fn(dots_vec, query_norm, norms_vec, &results_vec);
|
|
2788
|
+
from_dot_fn(&dots_vec, query_norm, &norms_vec, &results_vec); \
|
|
2790
2789
|
store_fn(&results_vec, r_row_out + column_index); \
|
|
2791
2790
|
} \
|
|
2792
2791
|
if (column_index < column_count) { \
|
|
2793
2792
|
vec_type dots_vec = {0}, norms_vec = {0}, results_vec; \
|
|
2794
2793
|
partial_load_fn(r_row_dots + column_index, &dots_vec, column_count - column_index); \
|
|
2795
2794
|
partial_load_fn(b_norms + column_index, &norms_vec, column_count - column_index); \
|
|
2796
|
-
from_dot_fn(dots_vec, query_norm, norms_vec, &results_vec);
|
|
2795
|
+
from_dot_fn(&dots_vec, query_norm, &norms_vec, &results_vec); \
|
|
2797
2796
|
partial_store_fn(&results_vec, r_row_out + column_index, column_count - column_index); \
|
|
2798
2797
|
} \
|
|
2799
2798
|
} \
|
|
@@ -2849,7 +2848,7 @@ NK_PUBLIC void nk_dots_compact_i8_serial(void *c, nk_size_t row_count, nk_size_t
|
|
|
2849
2848
|
load_fn(&column_norms[j - column_chunk_start], &target_norms_vec); \
|
|
2850
2849
|
vec_type dots_vec, results_vec; \
|
|
2851
2850
|
load_fn(r_dots + j, &dots_vec); \
|
|
2852
|
-
from_dot_fn(dots_vec, sumsq_i, target_norms_vec, &results_vec);
|
|
2851
|
+
from_dot_fn(&dots_vec, sumsq_i, &target_norms_vec, &results_vec); \
|
|
2853
2852
|
store_fn(&results_vec, r_out + j); \
|
|
2854
2853
|
} \
|
|
2855
2854
|
/* Remainder */ \
|
|
@@ -2857,7 +2856,7 @@ NK_PUBLIC void nk_dots_compact_i8_serial(void *c, nk_size_t row_count, nk_size_t
|
|
|
2857
2856
|
vec_type dots_vec = {0}, norms_vec = {0}, results_vec; \
|
|
2858
2857
|
partial_load_fn(r_dots + j, &dots_vec, column_chunk_end - j); \
|
|
2859
2858
|
partial_load_fn(&column_norms[j - column_chunk_start], &norms_vec, column_chunk_end - j); \
|
|
2860
|
-
from_dot_fn(dots_vec, sumsq_i, norms_vec, &results_vec);
|
|
2859
|
+
from_dot_fn(&dots_vec, sumsq_i, &norms_vec, &results_vec); \
|
|
2861
2860
|
partial_store_fn(&results_vec, r_out + j, column_chunk_end - j); \
|
|
2862
2861
|
} \
|
|
2863
2862
|
} \
|
|
@@ -16,9 +16,10 @@
|
|
|
16
16
|
#if NK_TARGET_X8664_
|
|
17
17
|
#if NK_TARGET_SIERRA
|
|
18
18
|
|
|
19
|
-
#include "numkong/dot/sierra.h"
|
|
20
|
-
#include "numkong/dot/haswell.h"
|
|
21
|
-
#include "numkong/dots/serial.h"
|
|
19
|
+
#include "numkong/dot/sierra.h" // Sierra-specific dot product helpers
|
|
20
|
+
#include "numkong/dot/haswell.h" // Haswell partial load functions
|
|
21
|
+
#include "numkong/dots/serial.h" // GEMM macro definitions
|
|
22
|
+
#include "numkong/reduce/haswell.h" // `nk_reduce_add_i32x8_haswell_`
|
|
22
23
|
|
|
23
24
|
#if defined(__cplusplus)
|
|
24
25
|
extern "C" {
|
|
@@ -60,6 +60,8 @@ Real and integer element-wise operations:
|
|
|
60
60
|
| `bf16` | `bf16` | 16-bit brain float |
|
|
61
61
|
| `e4m3` | `e4m3` | 8-bit Float8: 4 exponent, 3 mantissa bits |
|
|
62
62
|
| `e5m2` | `e5m2` | 8-bit Float8: 5 exponent, 2 mantissa bits |
|
|
63
|
+
| `e3m2` | `e3m2` | 6-bit Float6: 3 exponent, 2 mantissa bits |
|
|
64
|
+
| `e2m3` | `e2m3` | 6-bit Float6: 2 exponent, 3 mantissa bits |
|
|
63
65
|
| `i8` | `i8` | 8-bit signed integers, saturating |
|
|
64
66
|
| `u8` | `u8` | 8-bit unsigned integers, saturating |
|
|
65
67
|
| `i16` | `i16` | 16-bit signed integers |
|
|
@@ -86,7 +88,7 @@ Haswell processes 8 Float16 elements per cycle: `VCVTPH2PS` (widen) -> `VFMADD23
|
|
|
86
88
|
|
|
87
89
|
### Saturating Integer Arithmetic
|
|
88
90
|
|
|
89
|
-
`nk_each_sum_i8_haswell`, `
|
|
91
|
+
`nk_each_sum_i8_haswell`, `nk_each_sum_u8_neon` use saturating addition — clamping to type bounds instead of wrapping on overflow.
|
|
90
92
|
Haswell uses `VPADDSB` / `VPADDUSB` for signed/unsigned 8-bit saturation in a single instruction (32 elements per cycle at YMM width).
|
|
91
93
|
Serial fallback implements saturation via branch-free min/max: `result = min(max(a + b, TYPE_MIN), TYPE_MAX)` with overflow detection through sign-bit comparison.
|
|
92
94
|
|
|
@@ -387,7 +389,7 @@ Workloads that significantly degrade CPU frequencies (Intel AMX, Apple SME) run
|
|
|
387
389
|
| `nk_each_fma_e3m2_serial` | 0.506 gb/s, ? ulp | 0.433 gb/s, ? ulp | 0.436 gb/s, ? ulp |
|
|
388
390
|
| __i8__ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ |
|
|
389
391
|
| `nk_each_sum_i8_serial` | 94.6 gb/s | 55.5 gb/s | 42.7 gb/s |
|
|
390
|
-
| `
|
|
392
|
+
| `nk_each_sum_i8_neon` | 85.5 gb/s | 60.5 gb/s | 39.6 gb/s |
|
|
391
393
|
| `nk_each_scale_i8_serial` | 0.171 gb/s | 0.145 gb/s | 0.148 gb/s |
|
|
392
394
|
| `nk_each_scale_i8_neonhalf` | 22.6 gb/s | 22.4 gb/s | 20.8 gb/s |
|
|
393
395
|
| `nk_each_blend_i8_serial` | 0.302 gb/s | 0.254 gb/s | 0.272 gb/s |
|
|
@@ -395,7 +397,7 @@ Workloads that significantly degrade CPU frequencies (Intel AMX, Apple SME) run
|
|
|
395
397
|
| `nk_each_fma_i8_serial` | 0.436 gb/s | 0.359 gb/s | 0.376 gb/s |
|
|
396
398
|
| __u8__ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ |
|
|
397
399
|
| `nk_each_sum_u8_serial` | 13.2 gb/s | 12.3 gb/s | 12.3 gb/s |
|
|
398
|
-
| `
|
|
400
|
+
| `nk_each_sum_u8_neon` | 85.3 gb/s | 84.8 gb/s | 67.3 gb/s |
|
|
399
401
|
| `nk_each_scale_u8_serial` | 0.135 gb/s | 0.120 gb/s | 0.119 gb/s |
|
|
400
402
|
| `nk_each_scale_u8_neonhalf` | 23.0 gb/s | 21.9 gb/s | 21.2 gb/s |
|
|
401
403
|
| `nk_each_blend_u8_serial` | 0.256 gb/s | 0.218 gb/s | 0.220 gb/s |
|
|
@@ -29,6 +29,8 @@
|
|
|
29
29
|
|
|
30
30
|
#include "numkong/types.h"
|
|
31
31
|
#include "numkong/cast/sapphire.h" // `nk_f32_to_f16_sapphire`
|
|
32
|
+
#include "numkong/cast/icelake.h" // `nk_cast_icelake`
|
|
33
|
+
#include "numkong/each/icelake.h" // `nk_each_sum_i8_icelake`, `nk_each_sum_u8_icelake`
|
|
32
34
|
|
|
33
35
|
#if defined(__cplusplus)
|
|
34
36
|
extern "C" {
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
# Geospatial Distances in NumKong
|
|
2
2
|
|
|
3
|
-
NumKong implements geodesic distance functions for points on Earth's surface: Haversine computes great-circle distance on a perfect sphere, while Vincenty solves the inverse geodesic problem on the
|
|
3
|
+
NumKong implements geodesic distance functions for points on Earth's surface: Haversine computes great-circle distance on a perfect sphere, while Vincenty solves the inverse geodesic problem on an oblate spheroid, using the IERS-2003 ellipsoid constants by default.
|
|
4
4
|
Both operate on arrays of latitude/longitude pairs in radians and produce distances in meters.
|
|
5
5
|
|
|
6
6
|
The Haversine formula computes the great-circle distance between two points:
|
|
@@ -9,7 +9,7 @@ $$
|
|
|
9
9
|
\text{haversine}(\phi_1, \lambda_1, \phi_2, \lambda_2) = 2R \arcsin\sqrt{\sin^2\frac{\phi_2 - \phi_1}{2} + \cos\phi_1 \cos\phi_2 \sin^2\frac{\lambda_2 - \lambda_1}{2}}
|
|
10
10
|
$$
|
|
11
11
|
|
|
12
|
-
where $R$ is Earth
|
|
12
|
+
where $R$ is the Earth radius from `NK_EARTH_MEDIATORIAL_RADIUS`, 6,335,439 m by default, and $(\phi, \lambda)$ are latitude and longitude in radians.
|
|
13
13
|
|
|
14
14
|
Vincenty's formula solves the inverse geodesic problem on an oblate spheroid, iteratively refining the reduced latitude difference until convergence:
|
|
15
15
|
|
|
@@ -17,14 +17,14 @@ $$
|
|
|
17
17
|
\text{vincenty}(\phi_1, \lambda_1, \phi_2, \lambda_2) = b \cdot A \cdot (\sigma - \Delta\sigma)
|
|
18
18
|
$$
|
|
19
19
|
|
|
20
|
-
where $a$ and $b$ are the equatorial and polar semi-axes of the
|
|
20
|
+
where $a$ and $b$ are the equatorial and polar semi-axes of the configured Earth ellipsoid, $\sigma$ is the angular separation, and $\Delta\sigma$ is the correction term computed through iterative convergence.
|
|
21
21
|
|
|
22
22
|
Reformulating as Python pseudocode:
|
|
23
23
|
|
|
24
24
|
```python
|
|
25
25
|
import numpy as np
|
|
26
26
|
|
|
27
|
-
def haversine(lat1, lon1, lat2, lon2, R=
|
|
27
|
+
def haversine(lat1, lon1, lat2, lon2, R=6335439):
|
|
28
28
|
dlat = lat2 - lat1
|
|
29
29
|
dlon = lon2 - lon1
|
|
30
30
|
a = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
|
|
@@ -45,8 +45,8 @@ Input coordinates are in radians, output distances are in meters.
|
|
|
45
45
|
### Trigonometric Polynomial Approximations
|
|
46
46
|
|
|
47
47
|
`nk_haversine_f32_haswell`, `nk_haversine_f32_skylake`, `nk_haversine_f32_neon` replace `libm` sin, cos, atan2, and asin with SIMD polynomial approximations achieving approximately 3.5 ULP accuracy.
|
|
48
|
-
Range reduction maps the input angle to $[-\pi/
|
|
49
|
-
The `f32` kernels use 5-term polynomials while `f64` kernels use
|
|
48
|
+
Range reduction maps the input angle to $[-\pi/2, \pi/2]$ using Cody-Waite extended precision constants, then an odd-degree minimax polynomial evaluates sin, and cos reuses the same polynomial after shifting the argument by $\pi/2$.
|
|
49
|
+
The `f32` kernels use 5-term polynomials while `f64` kernels use 10-term polynomials for the extra precision required by double-precision inputs.
|
|
50
50
|
This avoids the latency of scalar `libm` calls — each trigonometric evaluation would otherwise serialize through a single execution port, while the polynomial chains pipeline across multiple FMA units.
|
|
51
51
|
|
|
52
52
|
### Vincenty Iterative Convergence with Masked Lanes
|
|
@@ -33,6 +33,7 @@
|
|
|
33
33
|
|
|
34
34
|
#include "numkong/types.h"
|
|
35
35
|
#include "numkong/trigonometry/v128relaxed.h"
|
|
36
|
+
#include "numkong/cast/serial.h" // `nk_partial_load_b32x4_serial_`, `nk_partial_load_b64x2_serial_`, `nk_partial_store_b32x4_serial_`
|
|
36
37
|
|
|
37
38
|
#if defined(__cplusplus)
|
|
38
39
|
extern "C" {
|
|
@@ -38,11 +38,11 @@ def maxsim(queries: np.ndarray, documents: np.ndarray) -> float:
|
|
|
38
38
|
|
|
39
39
|
## Input & Output Types
|
|
40
40
|
|
|
41
|
-
| Input Type | Output Type | Description
|
|
42
|
-
| ---------- | ----------- |
|
|
43
|
-
| `bf16` | `f32` | 16-bit brain float, widened output
|
|
44
|
-
| `f32` | `
|
|
45
|
-
| `f16` | `f32` | 16-bit IEEE 754 half precision
|
|
41
|
+
| Input Type | Output Type | Description |
|
|
42
|
+
| ---------- | ----------- | ------------------------------------------------ |
|
|
43
|
+
| `bf16` | `f32` | 16-bit brain float, widened output |
|
|
44
|
+
| `f32` | `f64` | 32-bit IEEE 754 single precision, widened output |
|
|
45
|
+
| `f16` | `f32` | 16-bit IEEE 754 half precision |
|
|
46
46
|
|
|
47
47
|
## Optimizations
|
|
48
48
|
|
|
@@ -21,6 +21,10 @@
|
|
|
21
21
|
#include "numkong/types.h"
|
|
22
22
|
#include "numkong/maxsim/icelake.h" // `nk_maxsim_coarse_argmax_icelake_`
|
|
23
23
|
#include "numkong/dot.h" // `nk_dot_bf16`
|
|
24
|
+
#include "numkong/cast/icelake.h" // `nk_e5m2x32_to_bf16x32_icelake_`
|
|
25
|
+
#include "numkong/dot/skylake.h" // `nk_dot_through_f32_finalize_skylake_`
|
|
26
|
+
#include "numkong/maxsim/serial.h" // `nk_maxsim_packed_header_setup_`, `nk_maxsim_packed_size_`
|
|
27
|
+
#include "numkong/reduce/skylake.h" // `nk_reduce_add_f32x16_skylake_`
|
|
24
28
|
|
|
25
29
|
#if defined(__cplusplus)
|
|
26
30
|
extern "C" {
|
|
@@ -30,6 +30,7 @@
|
|
|
30
30
|
#include "numkong/dot.h" // `nk_dot_f32`, `nk_dot_f16`
|
|
31
31
|
#include "numkong/cast/haswell.h" // `nk_f16_to_f32_haswell`
|
|
32
32
|
#include "numkong/spatial/haswell.h" // `nk_f32_sqrt_haswell`
|
|
33
|
+
#include "numkong/reduce/skylake.h" // `nk_reduce_add_f32x16_skylake_`
|
|
33
34
|
|
|
34
35
|
#if defined(__cplusplus)
|
|
35
36
|
extern "C" {
|
|
@@ -71,15 +71,6 @@ NK_STATIC_ASSERT(sizeof(nk_maxsim_vector_metadata_t) == 12, nk_maxsim_vector_met
|
|
|
71
71
|
*/
|
|
72
72
|
typedef void (*nk_maxsim_to_f32_t)(void const *source, nk_f32_t *destination);
|
|
73
73
|
|
|
74
|
-
/* Keep the serial instantiations below actually scalar, regardless of build type.
|
|
75
|
-
* See dots/serial.h for rationale. */
|
|
76
|
-
#if defined(__clang__)
|
|
77
|
-
#pragma clang attribute push(__attribute__((noinline)), apply_to = function)
|
|
78
|
-
#elif defined(__GNUC__)
|
|
79
|
-
#pragma GCC push_options
|
|
80
|
-
#pragma GCC optimize("no-tree-vectorize", "no-tree-slp-vectorize", "no-ipa-cp-clone", "no-inline")
|
|
81
|
-
#endif
|
|
82
|
-
|
|
83
74
|
/** @brief Identity conversion for f32 sources — just a typed memcpy. */
|
|
84
75
|
NK_INTERNAL void nk_f32_to_f32_(void const *source, nk_f32_t *destination) { *destination = *(nk_f32_t const *)source; }
|
|
85
76
|
|
|
@@ -234,6 +225,15 @@ NK_INTERNAL nk_size_t nk_maxsim_packed_size_( //
|
|
|
234
225
|
return header_size + i8_region_size + metadata_region_size + originals_region_size;
|
|
235
226
|
}
|
|
236
227
|
|
|
228
|
+
/* Keep the serial instantiations below actually scalar, regardless of build type.
|
|
229
|
+
* See dots/serial.h for rationale. */
|
|
230
|
+
#if defined(__clang__)
|
|
231
|
+
#pragma clang attribute push(__attribute__((noinline)), apply_to = function)
|
|
232
|
+
#elif defined(__GNUC__)
|
|
233
|
+
#pragma GCC push_options
|
|
234
|
+
#pragma GCC optimize("no-tree-vectorize", "no-tree-slp-vectorize", "no-ipa-cp-clone", "no-inline")
|
|
235
|
+
#endif
|
|
236
|
+
|
|
237
237
|
NK_PUBLIC nk_size_t nk_maxsim_packed_size_bf16_serial(nk_size_t vector_count, nk_size_t depth) {
|
|
238
238
|
return nk_maxsim_packed_size_(vector_count, depth, sizeof(nk_bf16_t), 1);
|
|
239
239
|
}
|
|
@@ -323,6 +323,12 @@ NK_PUBLIC void nk_maxsim_pack_f16_serial( //
|
|
|
323
323
|
}
|
|
324
324
|
}
|
|
325
325
|
|
|
326
|
+
#if defined(__clang__)
|
|
327
|
+
#pragma clang attribute pop
|
|
328
|
+
#elif defined(__GNUC__)
|
|
329
|
+
#pragma GCC pop_options
|
|
330
|
+
#endif
|
|
331
|
+
|
|
326
332
|
/**
|
|
327
333
|
* @brief Dtype-agnostic coarse i8 argmax kernel for the serial backend.
|
|
328
334
|
* Produces per-query best document indices using signed i8×i8 dot products.
|
|
@@ -389,6 +395,13 @@ NK_INTERNAL void nk_maxsim_coarse_argmax_serial_( //
|
|
|
389
395
|
}
|
|
390
396
|
}
|
|
391
397
|
|
|
398
|
+
#if defined(__clang__)
|
|
399
|
+
#pragma clang attribute push(__attribute__((noinline)), apply_to = function)
|
|
400
|
+
#elif defined(__GNUC__)
|
|
401
|
+
#pragma GCC push_options
|
|
402
|
+
#pragma GCC optimize("no-tree-vectorize", "no-tree-slp-vectorize", "no-ipa-cp-clone", "no-inline")
|
|
403
|
+
#endif
|
|
404
|
+
|
|
392
405
|
NK_PUBLIC void nk_maxsim_packed_bf16_serial( //
|
|
393
406
|
void const *query_packed, void const *document_packed, nk_size_t query_count, nk_size_t document_count,
|
|
394
407
|
nk_size_t depth, nk_f32_t *result) {
|
|
@@ -46,8 +46,9 @@
|
|
|
46
46
|
#if NK_TARGET_ARM64_
|
|
47
47
|
#if NK_TARGET_SME
|
|
48
48
|
|
|
49
|
-
#include "numkong/dots/sme.h"
|
|
50
|
-
#include "numkong/reduce/sve.h"
|
|
49
|
+
#include "numkong/dots/sme.h" // `nk_dots_sme_packed_header_t`
|
|
50
|
+
#include "numkong/reduce/sve.h" // `nk_svaddv_f64_`
|
|
51
|
+
#include "numkong/scalar/neon.h" // `nk_f64_rsqrt_neon`
|
|
51
52
|
|
|
52
53
|
#if defined(__cplusplus)
|
|
53
54
|
extern "C" {
|
|
@@ -51,10 +51,13 @@ def rmsd(a: np.ndarray, b: np.ndarray) -> float:
|
|
|
51
51
|
| Input Type | Output Type | Description |
|
|
52
52
|
| ---------- | ----------- | ---------------------------------------------- |
|
|
53
53
|
| `f64` | `f64` | 64-bit IEEE 754 double precision |
|
|
54
|
-
| `f32` | `
|
|
54
|
+
| `f32` | `f64` | 32-bit IEEE 754 single precision, widened RMSD |
|
|
55
55
|
| `f16` | `f32` | 16-bit IEEE 754 half precision, widened output |
|
|
56
56
|
| `bf16` | `f32` | 16-bit brain float, widened output |
|
|
57
57
|
|
|
58
|
+
The output type above is the RMSD result.
|
|
59
|
+
The centroid, rotation, and scale outputs use the transform type instead: `f64` for `f64` inputs, `f32` for `f32`, `f16`, and `bf16` inputs.
|
|
60
|
+
|
|
58
61
|
## Optimizations
|
|
59
62
|
|
|
60
63
|
### McAdams Branching-Free 3×3 SVD
|
|
@@ -199,9 +202,9 @@ Measured with Wasmtime v43 (Cranelift backend), WASI-SDK 24, `-msimd128 -mrelaxe
|
|
|
199
202
|
| `nk_kabsch_bf16_neonbfdot` | 180 mp/s, 0.7 ulp | 448 mp/s, 0.9 ulp | 726 mp/s, 1.3 ulp |
|
|
200
203
|
| `nk_umeyama_bf16_neonbfdot` | 176 mp/s, 0.2 ulp | 433 mp/s, 0.4 ulp | 705 mp/s, 0.8 ulp |
|
|
201
204
|
| __f16__ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ | ░░░░░░░░░░░░░░░░░░░░░░░░ |
|
|
202
|
-
| `
|
|
203
|
-
| `
|
|
204
|
-
| `
|
|
205
|
+
| `nk_rmsd_f16_neon` | 2,998 mp/s, 0.4 ulp | 3,215 mp/s, 1.7 ulp | 3,216 mp/s, 4.6 ulp |
|
|
206
|
+
| `nk_kabsch_f16_neon` | 178 mp/s, 0.9 ulp | 443 mp/s, 1.3 ulp | 711 mp/s, 2.4 ulp |
|
|
207
|
+
| `nk_umeyama_f16_neon` | 175 mp/s, 0.4 ulp | 408 mp/s, 0.8 ulp | 620 mp/s, 1.5 ulp |
|
|
205
208
|
|
|
206
209
|
#### WASM
|
|
207
210
|
|
|
@@ -410,7 +410,7 @@ NK_PUBLIC void nk_umeyama_bf16_genoa(nk_bf16_t const *a, nk_bf16_t const *b, nk_
|
|
|
410
410
|
optimal_rotation[3] = 0.0f, optimal_rotation[4] = 1.0f, optimal_rotation[5] = 0.0f;
|
|
411
411
|
optimal_rotation[6] = 0.0f, optimal_rotation[7] = 0.0f, optimal_rotation[8] = 1.0f;
|
|
412
412
|
trace_rotation_covariance = cross_covariance[0] + cross_covariance[4] + cross_covariance[8];
|
|
413
|
-
c =
|
|
413
|
+
c = trace_rotation_covariance / centered_norm_squared_a;
|
|
414
414
|
}
|
|
415
415
|
else {
|
|
416
416
|
nk_f32_t svd_left[9], svd_diagonal[9], svd_right[9];
|
|
@@ -422,7 +422,7 @@ NK_PUBLIC void nk_umeyama_bf16_genoa(nk_bf16_t const *a, nk_bf16_t const *b, nk_
|
|
|
422
422
|
nk_f32_t d3 = det < 0.0f ? -1.0f : 1.0f;
|
|
423
423
|
nk_f32_t trace_ds = nk_sum_three_products_f32_(svd_diagonal[0], 1.0f, svd_diagonal[4], 1.0f, svd_diagonal[8],
|
|
424
424
|
d3);
|
|
425
|
-
c =
|
|
425
|
+
c = trace_ds / centered_norm_squared_a;
|
|
426
426
|
|
|
427
427
|
if (det < 0.0f) {
|
|
428
428
|
svd_right[2] = -svd_right[2], svd_right[5] = -svd_right[5], svd_right[8] = -svd_right[8];
|