numkong 7.4.4 → 7.5.0

package/README.md

@@ -341,6 +341,7 @@ NumKong provides two dispatch mechanisms.
 __Compile-time dispatch__ selects the fastest kernel supported by the target platform at build time — thinner binaries, no indirection overhead, but requires knowing your deployment hardware.
 __Run-time dispatch__ compiles every supported kernel into the binary and picks the best one on the target machine via `nk_capabilities()` — one pointer indirection per call, but a single binary runs everywhere.
 The run-time path is common in DBMS products (ClickHouse), web browsers (Chromium), and other upstream projects that ship to heterogeneous fleets.
+Distributed artifacts (Rust crate, Python wheels, JS native modules, shared libs from the default CMake build) pin the translation-unit baseline to each architecture's ABI floor so the library runs on any CPU matching the ABI, not just the build host — see [CONTRIBUTING.md](CONTRIBUTING.md#target-baseline-policy) for the per-arch table and the `NK_MARCH_NATIVE` override used for host-tuned local builds.
 
 All kernel names follow the pattern `nk_{operation}_{type}_{backend}`.
 If you need to resolve the best kernel manually, use `nk_find_kernel_punned` with a `nk_kernel_kind_t`, `nk_dtype_t`, and a viable capabilities mask:

package/binding.gyp

@@ -39,11 +39,13 @@
             "defines": [
                 "NK_NATIVE_F16=0",
                 "NK_NATIVE_BF16=0",
-                "NK_DYNAMIC_DISPATCH=1"
+                "NK_DYNAMIC_DISPATCH=1",
+                "NK_USE_OPENMP=1"
             ],
             "cflags": [
                 "-std=c11",
                 "-O3",
+                "-fopenmp",
                 "-Wno-unknown-pragmas",
                 "-Wno-maybe-uninitialized",
                 "-Wno-cast-function-type",
@@ -52,31 +54,98 @@
                 "-include",
                 "<(module_root_dir)/nk_probes.h",
             ],
+            "ldflags": [
+                "-fopenmp"
+            ],
             "msvs_settings": {
                 "VCCLCompilerTool": {
                     "ForcedIncludeFiles": [
                         "<(module_root_dir)/nk_probes.h"
                     ],
                     "AdditionalOptions": [
-                        "/Zc:preprocessor"
+                        "/Zc:preprocessor",
+                        "/openmp:llvm"
                     ],
                 },
             },
             "conditions": [
+                # Pin TU baseline to each arch's ABI floor; SIMD kernels use per-function pragmas.
+                [
+                    "OS!='win' and target_arch=='arm64'",
+                    {
+                        "cflags": [
+                            "-march=armv8-a"
+                        ]
+                    }
+                ],
+                [
+                    "OS!='win' and target_arch=='x64'",
+                    {
+                        "cflags": [
+                            "-march=x86-64"
+                        ]
+                    }
+                ],
+                [
+                    "OS!='win' and target_arch=='riscv64'",
+                    {
+                        "cflags": [
+                            "-march=rv64gc"
+                        ]
+                    }
+                ],
+                # Forbid auto-vectorization so serial fallbacks don't get silently
+                # promoted to NEON/SSE2/VSX. SIMD kernels use explicit intrinsics
+                # and per-function `target` pragmas; unaffected. MSVC has no
+                # command-line vectorizer toggle.
+                [
+                    "OS!='win'",
+                    {
+                        "cflags": [
+                            "-fno-tree-vectorize",
+                            "-fno-tree-slp-vectorize"
+                        ]
+                    }
+                ],
                 [
                     "OS=='mac'",
                     {
                         "xcode_settings": {
-                            "MACOSX_DEPLOYMENT_TARGET": "11.0"
+                            "MACOSX_DEPLOYMENT_TARGET": "11.0",
+                            # Apple Clang ships no `omp.h`; the CI step
+                            # `brew install libomp` makes it keg-only under
+                            # `/opt/homebrew/opt/libomp` (arm64) or
+                            # `/usr/local/opt/libomp` (x86_64). Clang silently
+                            # ignores `-I` / `-L` dirs that don't exist, so
+                            # listing both keeps the file arch-agnostic.
+                            "OTHER_CFLAGS": [
+                                "-Xpreprocessor",
+                                "-fopenmp",
+                                "-I/opt/homebrew/opt/libomp/include",
+                                "-I/usr/local/opt/libomp/include"
+                            ],
+                            "OTHER_LDFLAGS": [
+                                "-lomp",
+                                "-L/opt/homebrew/opt/libomp/lib",
+                                "-L/usr/local/opt/libomp/lib"
+                            ]
                         }
                     }
                 ],
+                # MSVC: no per-function target pragma; these match defaults.
                 [
                     "OS=='win' and target_arch=='arm64'",
                     {
                         "defines": [
                             "_ARM64_"
-                        ]
+                        ],
+                        "msvs_settings": {
+                            "VCCLCompilerTool": {
+                                "AdditionalOptions": [
+                                    "/arch:armv8.0"
+                                ]
+                            }
+                        }
                     }
                 ],
                 [
@@ -84,7 +153,14 @@
                     {
                         "defines": [
                             "_AMD64_"
-                        ]
+                        ],
+                        "msvs_settings": {
+                            "VCCLCompilerTool": {
+                                "AdditionalOptions": [
+                                    "/arch:SSE2"
+                                ]
+                            }
+                        }
                     }
                 ],
             ],

package/c/dispatch_f16.c

@@ -137,6 +137,29 @@ void nk_dispatch_f16_find_(nk_capability_t v, nk_kernel_kind_t k, nk_kernel_punn
         default: break;
         }
 #endif
+#if NK_TARGET_GRANITEAMX
+    if (v & nk_cap_graniteamx_k) switch (k) {
+        case nk_kernel_dots_packed_size_k:
+            *m = (m_t)&nk_dots_packed_size_f16_graniteamx, *c = nk_cap_graniteamx_k;
+            return;
+        case nk_kernel_dots_pack_k: *m = (m_t)&nk_dots_pack_f16_graniteamx, *c = nk_cap_graniteamx_k; return;
+        case nk_kernel_dots_packed_k: *m = (m_t)&nk_dots_packed_f16_graniteamx, *c = nk_cap_graniteamx_k; return;
+        case nk_kernel_dots_symmetric_k: *m = (m_t)&nk_dots_symmetric_f16_graniteamx, *c = nk_cap_graniteamx_k; return;
+        case nk_kernel_angulars_packed_k:
+            *m = (m_t)&nk_angulars_packed_f16_graniteamx, *c = nk_cap_graniteamx_k;
+            return;
+        case nk_kernel_angulars_symmetric_k:
+            *m = (m_t)&nk_angulars_symmetric_f16_graniteamx, *c = nk_cap_graniteamx_k;
+            return;
+        case nk_kernel_euclideans_packed_k:
+            *m = (m_t)&nk_euclideans_packed_f16_graniteamx, *c = nk_cap_graniteamx_k;
+            return;
+        case nk_kernel_euclideans_symmetric_k:
+            *m = (m_t)&nk_euclideans_symmetric_f16_graniteamx, *c = nk_cap_graniteamx_k;
+            return;
+        default: break;
+        }
+#endif
 #if NK_TARGET_SAPPHIREAMX
     if (v & nk_cap_sapphireamx_k) switch (k) {
         case nk_kernel_maxsim_packed_size_k:

package/c/numkong.c

@@ -935,19 +935,6 @@ BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpReserved) {
 #endif
 #endif
 
-// SME ABI runtime stubs — provide the lazy-ZA-save helpers that compiler-rt
-// may not ship (e.g., Apple's toolchain). Called by compiler-generated code
-// in __arm_new("za") prologues/epilogues (used by dots streaming functions).
-//
-// In NumKong, TPIDR2_EL0 is always null at entry because no NK_PUBLIC function
-// carries ZA state. So __arm_tpidr2_save is always a no-op and
-// __arm_tpidr2_restore has nothing to restore.
-// Weak linkage lets a real compiler-rt override these if available.
-#if NK_TARGET_ARM64_ && NK_TARGET_SME
-__attribute__((weak, visibility("default"))) void __arm_tpidr2_save(void) {}
-__attribute__((weak, visibility("default"))) void __arm_tpidr2_restore(void *blk) { nk_unused_(blk); }
-#endif
-
 #ifdef __cplusplus
 }
 #endif

package/include/numkong/attention/sme.h

@@ -249,10 +249,9 @@ NK_PUBLIC nk_size_t nk_attention_packed_kv_size_f16_sme(nk_size_t num_kv_heads,
     return nk_attention_packed_kv_size_bf16_sme(num_kv_heads, head_dim, max_seq_len);
 }
 
-__arm_locally_streaming static void nk_attention_pack_kv_bf16_sme_streaming_(nk_bf16_t const *k, nk_bf16_t const *v,
-                                                                             nk_size_t num_kv_heads, nk_size_t head_dim,
-                                                                             nk_size_t seq_len, nk_size_t k_stride,
-                                                                             nk_size_t v_stride, void *kv_packed) {
+static void nk_attention_pack_kv_bf16_sme_ssve_( //
+    nk_bf16_t const *k, nk_bf16_t const *v, nk_size_t num_kv_heads, nk_size_t head_dim, nk_size_t seq_len,
+    nk_size_t k_stride, nk_size_t v_stride, void *kv_packed) NK_STREAMING_ {
 
     nk_attention_sme_packed_header_t *header = (nk_attention_sme_packed_header_t *)kv_packed;
     nk_size_t head_dim_padded = (head_dim + 31) / 32 * 32;
@@ -315,16 +314,17 @@ __arm_locally_streaming static void nk_attention_pack_kv_bf16_sme_streaming_(nk_
     }
 }
 
-NK_PUBLIC void nk_attention_pack_kv_bf16_sme(nk_bf16_t const *k, nk_bf16_t const *v, nk_size_t num_kv_heads,
-                                             nk_size_t head_dim, nk_size_t seq_len, nk_size_t k_stride,
-                                             nk_size_t v_stride, void *kv_packed) {
-    nk_attention_pack_kv_bf16_sme_streaming_(k, v, num_kv_heads, head_dim, seq_len, k_stride, v_stride, kv_packed);
+NK_PUBLIC void nk_attention_pack_kv_bf16_sme( //
+    nk_bf16_t const *k, nk_bf16_t const *v, nk_size_t num_kv_heads, nk_size_t head_dim, nk_size_t seq_len,
+    nk_size_t k_stride, nk_size_t v_stride, void *kv_packed) {
+    nk_sme_start_streaming_();
+    nk_attention_pack_kv_bf16_sme_ssve_(k, v, num_kv_heads, head_dim, seq_len, k_stride, v_stride, kv_packed);
+    nk_sme_stop_streaming_();
 }
 
-__arm_locally_streaming static void nk_attention_pack_kv_f16_sme_streaming_(nk_f16_t const *k, nk_f16_t const *v,
-                                                                            nk_size_t num_kv_heads, nk_size_t head_dim,
-                                                                            nk_size_t seq_len, nk_size_t k_stride,
-                                                                            nk_size_t v_stride, void *kv_packed) {
+static void nk_attention_pack_kv_f16_sme_ssve_( //
+    nk_f16_t const *k, nk_f16_t const *v, nk_size_t num_kv_heads, nk_size_t head_dim, nk_size_t seq_len,
+    nk_size_t k_stride, nk_size_t v_stride, void *kv_packed) NK_STREAMING_ {
 
     nk_attention_sme_packed_header_t *header = (nk_attention_sme_packed_header_t *)kv_packed;
     nk_size_t head_dim_padded = (head_dim + 31) / 32 * 32;
@@ -385,10 +385,12 @@ __arm_locally_streaming static void nk_attention_pack_kv_f16_sme_streaming_(nk_f
     }
 }
 
-NK_PUBLIC void nk_attention_pack_kv_f16_sme(nk_f16_t const *k, nk_f16_t const *v, nk_size_t num_kv_heads,
-                                            nk_size_t head_dim, nk_size_t seq_len, nk_size_t k_stride,
-                                            nk_size_t v_stride, void *kv_packed) {
-    nk_attention_pack_kv_f16_sme_streaming_(k, v, num_kv_heads, head_dim, seq_len, k_stride, v_stride, kv_packed);
+NK_PUBLIC void nk_attention_pack_kv_f16_sme( //
+    nk_f16_t const *k, nk_f16_t const *v, nk_size_t num_kv_heads, nk_size_t head_dim, nk_size_t seq_len,
+    nk_size_t k_stride, nk_size_t v_stride, void *kv_packed) {
+    nk_sme_start_streaming_();
+    nk_attention_pack_kv_f16_sme_ssve_(k, v, num_kv_heads, head_dim, seq_len, k_stride, v_stride, kv_packed);
+    nk_sme_stop_streaming_();
 }
 
 /**
@@ -402,13 +404,13 @@ NK_PUBLIC void nk_attention_pack_kv_f16_sme(nk_f16_t const *k, nk_f16_t const *v
  *  - Correction skip when running max is unchanged
  *  - Decode path (valid_query_count==1) remains element-wise SVE (BFMOPA overhead too high)
  */
-__arm_locally_streaming __arm_new("za") static void nk_attention_bf16_sme_streaming_(
+__arm_new("za") static void nk_attention_bf16_sme_streaming_(
     nk_bf16_t const *q,        // [query_len, head_dim]
     nk_bf16_t const *k,        // [kv_len, head_dim_padded] BFMOPA-interleaved
     nk_bf16_t const *v_packed, // BFMOPA-interleaved V for this KV head
     nk_bf16_t *output,         // [query_len, head_dim]
     nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim, nk_size_t head_dim_padded, nk_size_t dim_tile_count,
-    nk_f32_t scale) {
+    nk_f32_t scale) NK_STREAMING_ {
 
     svbool_t const predicate_all_b32x = svptrue_b32();
     svbool_t const predicate_all_b16x = svptrue_b16();
@@ -1184,9 +1186,9 @@ __arm_locally_streaming __arm_new("za") static void nk_attention_bf16_sme_stream
     }
 }
 
-NK_PUBLIC void nk_attention_bf16_sme(nk_bf16_t const *q, void const *kv_packed, nk_bf16_t *output, nk_size_t num_heads,
-                                     nk_size_t num_kv_heads, nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim,
-                                     nk_f32_t scale) {
+NK_PUBLIC void nk_attention_bf16_sme( //
+    nk_bf16_t const *q, void const *kv_packed, nk_bf16_t *output, nk_size_t num_heads, nk_size_t num_kv_heads,
+    nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim, nk_f32_t scale) {
 
     nk_attention_sme_packed_header_t const *header = (nk_attention_sme_packed_header_t const *)kv_packed;
     nk_size_t head_dim_padded = header->head_dim_padded;
@@ -1199,6 +1201,7 @@ NK_PUBLIC void nk_attention_bf16_sme(nk_bf16_t const *q, void const *kv_packed,
 
     nk_size_t group_size = (num_kv_heads > 0) ? num_heads / num_kv_heads : 1;
 
+    nk_sme_start_streaming_();
     for (nk_size_t q_head = 0; q_head < num_heads; q_head++) {
         nk_size_t kv_head = q_head / group_size;
 
@@ -1214,15 +1217,13 @@ NK_PUBLIC void nk_attention_bf16_sme(nk_bf16_t const *q, void const *kv_packed,
                                              q_block_len, kv_len, head_dim, head_dim_padded, dim_tile_count, scale);
         }
     }
+    nk_sme_stop_streaming_();
 }
 
-__arm_locally_streaming __arm_new("za") static void nk_attention_f16_sme_streaming_(
-    nk_f16_t const *q,        // [query_len, head_dim]
-    nk_f16_t const *k,        // [kv_len, head_dim_padded] FMOPA-interleaved
-    nk_f16_t const *v_packed, // FMOPA-interleaved V for this KV head
-    nk_f16_t *output,         // [query_len, head_dim]
-    nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim, nk_size_t head_dim_padded, nk_size_t dim_tile_count,
-    nk_f32_t scale) {
+__arm_new("za") static void nk_attention_f16_sme_streaming_( //
+    nk_f16_t const *q, nk_f16_t const *k, nk_f16_t const *v_packed, nk_f16_t *output, nk_size_t query_len,
+    nk_size_t kv_len, nk_size_t head_dim, nk_size_t head_dim_padded, nk_size_t dim_tile_count,
+    nk_f32_t scale) NK_STREAMING_ {
 
     svbool_t const predicate_all_b32x = svptrue_b32();
     svbool_t const predicate_all_b16x = svptrue_b16();
@@ -2008,9 +2009,9 @@ __arm_locally_streaming __arm_new("za") static void nk_attention_f16_sme_streami
     }
 }
 
-NK_PUBLIC void nk_attention_f16_sme(nk_f16_t const *q, void const *kv_packed, nk_f16_t *output, nk_size_t num_heads,
-                                    nk_size_t num_kv_heads, nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim,
-                                    nk_f32_t scale) {
+NK_PUBLIC void nk_attention_f16_sme( //
+    nk_f16_t const *q, void const *kv_packed, nk_f16_t *output, nk_size_t num_heads, nk_size_t num_kv_heads,
+    nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim, nk_f32_t scale) {
 
     nk_attention_sme_packed_header_t const *header = (nk_attention_sme_packed_header_t const *)kv_packed;
     nk_size_t head_dim_padded = header->head_dim_padded;
@@ -2024,6 +2025,7 @@ NK_PUBLIC void nk_attention_f16_sme(nk_f16_t const *q, void const *kv_packed, nk
 
     nk_size_t group_size = (num_kv_heads > 0) ? num_heads / num_kv_heads : 1;
 
+    nk_sme_start_streaming_();
     for (nk_size_t q_head = 0; q_head < num_heads; q_head++) {
         nk_size_t kv_head = q_head / group_size;
 
@@ -2039,6 +2041,7 @@ NK_PUBLIC void nk_attention_f16_sme(nk_f16_t const *q, void const *kv_packed, nk
                                             q_block_len, kv_len, head_dim, head_dim_padded, dim_tile_count, scale);
         }
     }
+    nk_sme_stop_streaming_();
 }
 
 NK_PUBLIC void nk_attention_causal_bf16_sme(nk_bf16_t const *q, void const *kv_packed, nk_bf16_t *output,

package/include/numkong/capabilities.h

@@ -95,8 +95,8 @@
 #include "numkong/types.h" // `nk_u64_t`, `NK_DEFINED_LINUX_`
 
 #define NK_VERSION_MAJOR 7
-#define NK_VERSION_MINOR 4
-#define NK_VERSION_PATCH 4
+#define NK_VERSION_MINOR 5
+#define NK_VERSION_PATCH 0
 
 /**
  *  @brief  Removes compile-time dispatching, and replaces it with runtime dispatching.
@@ -500,13 +500,6 @@ NK_PUBLIC nk_capability_t nk_capabilities_x8664_(void) {
 
 #if NK_TARGET_ARM64_
 
-#if defined(__clang__)
-#pragma clang attribute push(__attribute__((target("arch=armv8.5-a+sve"))), apply_to = function)
-#elif defined(__GNUC__)
-#pragma GCC push_options
-#pragma GCC target("arch=armv8.5-a+sve")
-#endif
-
 #if NK_HAS_POSIX_EXTENSIONS_
 static sigjmp_buf nk_mrs_arm64_jump_buffer_;
 static void nk_mrs_arm64_sigill_handler_(int sig) {
@@ -716,12 +709,6 @@ NK_PUBLIC nk_capability_t nk_capabilities_arm64_(void) {
 #endif
 }
 
-#if defined(__clang__)
-#pragma clang attribute pop
-#elif defined(__GNUC__)
-#pragma GCC pop_options
-#endif
-
 #endif // NK_TARGET_ARM64_
 
 #if NK_TARGET_RISCV64_

package/include/numkong/cast/neon.h

@@ -104,6 +104,21 @@ NK_INTERNAL void nk_store_b256_neon_(nk_b256_vec_t const *src, void *dst) {
 /** @brief Type-agnostic 64-bit full load (NEON). */
 NK_INTERNAL void nk_load_b64_neon_(void const *src, nk_b64_vec_t *dst) { dst->u8x8 = vld1_u8((nk_u8_t const *)src); }
 
+/**
+ *  @brief 8-lane `uint16x8_t` splat that hides the source from the optimizer.
+ *
+ *  GCC 13 lowers `vdupq_n_u16(X)` to `fmov v.8h, #imm` (a FEAT_FP16 encoding) whenever X matches a
+ *  representable FP16 immediate, including bf16 bit patterns like 1.0 (`0x3F80`). That fails to
+ *  assemble under a `+bf16`-only pragma. The empty `__asm__` constraint forces `mov w; dup v.8h, w`
+ *  instead, valid on plain `armv8-a+simd`. No-op on Clang; skipped on MSVC (neither is affected).
+ */
+NK_INTERNAL uint16x8_t nk_u16x8_splat_(nk_u16_t bits) {
+#if defined(__GNUC__) || defined(__clang__)
+    __asm__("" : "+r"(bits));
+#endif
+    return vdupq_n_u16(bits);
+}
+
 #pragma endregion Type Punned Loads and Stores
 
 #pragma region Vectorized Conversions