feat(native): vendored NeoGPU ternary f32 LUT kernel + FFM downcall (#1137)

skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt

@@ -18,6 +18,10 @@ set(SKAINET_KERNEL_SOURCES
     ${SKAINET_KERNELS_ROOT}/src/q5k_matmul.c
     ${SKAINET_KERNELS_ROOT}/src/q6k_matmul.c
     ${SKAINET_KERNELS_ROOT}/src/bitnet_gemv.c
+    ${SKAINET_KERNELS_ROOT}/src/skainet_ternary_f32.c
+    # Vendored NeoGPU ternary LUT kernel (MIT, verbatim). NDK toolchains are
+    # clang, and bionic bundles pthreads — no Threads link needed.
+    ${SKAINET_KERNELS_ROOT}/src/vendor/neogpu/hs_ml_ternary_neon.c
 )
 
 set(SKAINET_JNI_SHIM ${CMAKE_CURRENT_SOURCE_DIR}/skainet_jni.c)
@@ -34,6 +38,7 @@ function(skainet_add_jni_lib target)
     add_library(${target} SHARED ${SKAINET_JNI_SHIM} ${SKAINET_KERNEL_SOURCES})
     target_include_directories(${target} PRIVATE ${SKAINET_KERNELS_ROOT}/include)
     target_compile_options(${target} PRIVATE ${SKAINET_C_FLAGS})
+    target_compile_definitions(${target} PRIVATE SKAINET_HAVE_NEOGPU_TERNARY)
     target_link_options(${target} PRIVATE ${SKAINET_LINK_FLAGS})
 endfunction()
 

skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt

@@ -23,8 +23,17 @@ set(SKAINET_KERNEL_SOURCES
     src/q4_0_matmul.c
     src/q5_0_matmul.c
     src/q5_1_matmul.c
+    src/skainet_ternary_f32.c
 )
 
+# Vendored NeoGPU ternary LUT kernel (MIT, verbatim — src/vendor/neogpu/).
+# Needs pthreads + GNU/Clang builtins, so MSVC builds get the portable
+# scalar fallback inside skainet_ternary_f32.c instead.
+if(NOT MSVC)
+    list(APPEND SKAINET_KERNEL_SOURCES src/vendor/neogpu/hs_ml_ternary_neon.c)
+    find_package(Threads REQUIRED)
+endif()
+
 # SHARED: consumed by the JVM via java.lang.foreign (FFM), bundled as a JAR
 #         resource (libskainet_kernels.{so,dylib,dll}).
 # STATIC: consumed by Kotlin/Native via cinterop, linked into the K/N binary
@@ -49,6 +58,12 @@ endif()
 foreach(tgt IN LISTS SKAINET_KERNEL_TARGETS)
     target_include_directories(${tgt} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
 
+    if(NOT MSVC)
+        # The adapter forwards to the vendored NeoGPU kernel (pthreads).
+        target_compile_definitions(${tgt} PRIVATE SKAINET_HAVE_NEOGPU_TERNARY)
+        target_link_libraries(${tgt} PRIVATE Threads::Threads)
+    endif()
+
     # Strip the "lib" prefix on Windows so the shared artifact name matches
     # the resource-bundle path skainet_kernels.{dll,so,dylib}.
     if(WIN32)
@@ -83,6 +98,16 @@ foreach(tgt IN LISTS SKAINET_KERNEL_TARGETS)
             target_compile_options(${tgt} PRIVATE
                 -march=armv8.2-a+fp16+dotprod
             )
+            # The vendored NeoGPU ternary LUT kernel exists FOR dotprod-less
+            # cores (Cortex-A72/Pi-4): pin it and its adapter back to the
+            # armv8-a baseline so the library-wide v8.2 flags above cannot
+            # emit instructions those cores lack. Source-level COMPILE_OPTIONS
+            # append after target options, so the later -march wins.
+            set_source_files_properties(
+                src/vendor/neogpu/hs_ml_ternary_neon.c
+                src/skainet_ternary_f32.c
+                PROPERTIES COMPILE_OPTIONS "-march=armv8-a"
+            )
         endif()
         set_target_properties(${tgt} PROPERTIES C_VISIBILITY_PRESET hidden)
     elseif(CMAKE_C_COMPILER_ID MATCHES "MSVC")

skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h

@@ -291,6 +291,54 @@ SKAINET_API void skainet_bitnet_gemv_tq2_0(
     int32_t input_dim, int32_t output_dim,
     float* output, int32_t output_offset);
 
+/*
+ * Ternary f32 GEMV — exact FP32 activations against sequentially-packed
+ * ternary weights (the BitNet b1.58 / `BITNET_B1_58` payload: 4 codes per
+ * byte, low bit-pair first, code {0,1,2} → {-1,0,+1}; byte code 3 decodes
+ * to +2, loaders reject it at import).
+ *
+ *   output[output_offset + o] = sum_j input[input_offset + j] *
+ *                                decode(weight row o)
+ *
+ * NO scale is applied — the caller owns the per-tensor scale. Weights are
+ * row-major, input_dim/4 bytes per output row, at
+ *   weight + weight_byte_offset + o * (input_dim / 4)
+ *
+ * input_dim must be a multiple of 4. Unlike the int8 `bitnet_gemv` path
+ * there is no activation quantization: results are exact. Backed by the
+ * vendored NeoGPU LUT kernel (baseline NEON, no dotprod — the fast path for
+ * Cortex-A72/Pi-4 class cores); it threads internally with pthreads once
+ * output_dim >= 512. A portable scalar build stands in where the vendored
+ * file cannot compile (MSVC).
+ */
+SKAINET_API void skainet_ternary_f32_gemv(
+    const float* input, int32_t input_offset,
+    const uint8_t* weight, int32_t weight_byte_offset,
+    int32_t input_dim, int32_t output_dim,
+    float* output, int32_t output_offset);
+
+/*
+ * Fused 4-plane ternary lm_head, Stage 1 (NeoGPU multi-plane trit format).
+ *
+ *   output[output_offset + o] = f16(row_scale[row_scale_offset + o]) *
+ *       sum_{p=0}^{3} (1/3^p) * sum_j input[..+j] * decode(plane p, row o)
+ *
+ * Each plane is a full sequentially-packed ternary matrix (same payload rule
+ * as skainet_ternary_f32_gemv); plane p starts at
+ *   planes + planes_byte_offset + p * plane_stride_bytes
+ * so a single buffer of four concatenated planes uses
+ * plane_stride_bytes == output_dim * input_dim / 4. row_scale holds raw
+ * little-endian IEEE binary16 bit patterns, one per output row (sign
+ * ignored — encoders store max|row| >= 0). input_dim must be a multiple
+ * of 4. Threads internally with pthreads at any output_dim.
+ */
+SKAINET_API void skainet_ternary_lmhead_stage1(
+    const float* input, int32_t input_offset,
+    const uint8_t* planes, int32_t planes_byte_offset, int32_t plane_stride_bytes,
+    const uint16_t* row_scale, int32_t row_scale_offset,
+    int32_t input_dim, int32_t output_dim,
+    float* output, int32_t output_offset);
+
 #ifdef __cplusplus
 }
 #endif

skainet-backends/skainet-backend-native-cpu/native/src/skainet_ternary_f32.c

@@ -0,0 +1,152 @@
+/*
+ * Adapter over the vendored NeoGPU ternary LUT kernel
+ * (src/vendor/neogpu/hs_ml_ternary_neon.c, MIT, vendored verbatim — see the
+ * vendor README). Exposes the two useful symbols under the skainet_kernels
+ * ABI and closes the vendored file's non-atomic LUT-init guard with a
+ * pthread_once warm-up, so first use is race-free under any threading.
+ *
+ * The vendored file needs pthreads and GNU/Clang builtins, so it is only in
+ * the build where SKAINET_HAVE_NEOGPU_TERNARY is defined (non-MSVC). The
+ * fallback branch below is a portable scalar mirror with identical semantics,
+ * including the byte-code-3 → +2.0 decode (loaders reject code 3 at import;
+ * the kernel contract is "garbage in, garbage out", never a crash).
+ */
+
+#include "skainet_kernels.h"
+
+#include <string.h>
+
+#ifdef SKAINET_HAVE_NEOGPU_TERNARY
+
+#include <pthread.h>
+
+/* The vendored file ships no header; these mirror its public definitions. */
+extern void hs_ml_ternary_f32_proj(float *out, const float *in,
+                                   const uint8_t *W, uint32_t N, uint32_t K);
+extern void hs_ml_lmhead_stage1(float *out, const float *in,
+                                const uint8_t *P0, const uint8_t *P1,
+                                const uint8_t *P2, const uint8_t *P3,
+                                const uint16_t *row_scale,
+                                uint32_t N, uint32_t K);
+
+static pthread_once_t skainet_ternary_lut_once = PTHREAD_ONCE_INIT;
+
+static void skainet_ternary_lut_warmup(void) {
+    /* Any call builds the 256-entry LUT; N=1 stays on the calling thread
+     * (far below the vendored THREAD_THRESHOLD of 512). */
+    static const float in[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
+    static const uint8_t w[1] = { 0 };
+    float out;
+    hs_ml_ternary_f32_proj(&out, in, w, 1u, 4u);
+}
+
+void skainet_ternary_f32_gemv(
+    const float* input, int32_t input_offset,
+    const uint8_t* weight, int32_t weight_byte_offset,
+    int32_t input_dim, int32_t output_dim,
+    float* output, int32_t output_offset)
+{
+    if (output_dim <= 0) return;
+    pthread_once(&skainet_ternary_lut_once, skainet_ternary_lut_warmup);
+    /* input_dim == 0 falls through: the vendored kernel writes 0.0f per row. */
+    hs_ml_ternary_f32_proj(
+        output + output_offset,
+        input + input_offset,
+        weight + weight_byte_offset,
+        (uint32_t)output_dim,
+        (uint32_t)input_dim);
+}
+
+void skainet_ternary_lmhead_stage1(
+    const float* input, int32_t input_offset,
+    const uint8_t* planes, int32_t planes_byte_offset, int32_t plane_stride_bytes,
+    const uint16_t* row_scale, int32_t row_scale_offset,
+    int32_t input_dim, int32_t output_dim,
+    float* output, int32_t output_offset)
+{
+    if (output_dim <= 0) return;
+    pthread_once(&skainet_ternary_lut_once, skainet_ternary_lut_warmup);
+    const uint8_t* base = planes + planes_byte_offset;
+    hs_ml_lmhead_stage1(
+        output + output_offset,
+        input + input_offset,
+        base,
+        base + (size_t)plane_stride_bytes,
+        base + (size_t)plane_stride_bytes * 2,
+        base + (size_t)plane_stride_bytes * 3,
+        row_scale + row_scale_offset,
+        (uint32_t)output_dim,
+        (uint32_t)input_dim);
+}
+
+#else /* !SKAINET_HAVE_NEOGPU_TERNARY — portable scalar mirror (MSVC etc.) */
+
+/* decode(byte, lane) = ((byte >> (lane*2)) & 3) - 1, exactly the vendored LUT
+ * (so byte code 3 yields +2, matching TernaryCodec.decodeBitNet). */
+static float skainet_ternary_decode(uint8_t b, int lane) {
+    return (float)((int)((b >> (lane * 2)) & 3) - 1);
+}
+
+void skainet_ternary_f32_gemv(
+    const float* input, int32_t input_offset,
+    const uint8_t* weight, int32_t weight_byte_offset,
+    int32_t input_dim, int32_t output_dim,
+    float* output, int32_t output_offset)
+{
+    if (output_dim <= 0) return;
+    const float* in = input + input_offset;
+    const uint8_t* w = weight + weight_byte_offset;
+    const int32_t row_bytes = input_dim / 4;
+    for (int32_t n = 0; n < output_dim; n++) {
+        const uint8_t* wrow = w + (size_t)n * row_bytes;
+        float acc = 0.0f;
+        for (int32_t bi = 0; bi < row_bytes; bi++) {
+            const uint8_t b = wrow[bi];
+            acc += skainet_ternary_decode(b, 0) * in[bi * 4    ];
+            acc += skainet_ternary_decode(b, 1) * in[bi * 4 + 1];
+            acc += skainet_ternary_decode(b, 2) * in[bi * 4 + 2];
+            acc += skainet_ternary_decode(b, 3) * in[bi * 4 + 3];
+        }
+        output[output_offset + n] = acc;
+    }
+}
+
+void skainet_ternary_lmhead_stage1(
+    const float* input, int32_t input_offset,
+    const uint8_t* planes, int32_t planes_byte_offset, int32_t plane_stride_bytes,
+    const uint16_t* row_scale, int32_t row_scale_offset,
+    int32_t input_dim, int32_t output_dim,
+    float* output, int32_t output_offset)
+{
+    if (output_dim <= 0) return;
+    const float* in = input + input_offset;
+    const uint8_t* base = planes + planes_byte_offset;
+    const uint16_t* rs = row_scale + row_scale_offset;
+    const int32_t row_bytes = input_dim / 4;
+    const float pw[4] = { 1.0f, 1.0f / 3.0f, 1.0f / 9.0f, 1.0f / 27.0f };
+    for (int32_t n = 0; n < output_dim; n++) {
+        float acc = 0.0f;
+        for (int32_t bi = 0; bi < row_bytes; bi++) {
+            for (int lane = 0; lane < 4; lane++) {
+                float u = 0.0f;
+                for (int p = 0; p < 4; p++) {
+                    const uint8_t b =
+                        base[(size_t)plane_stride_bytes * p + (size_t)n * row_bytes + bi];
+                    u += skainet_ternary_decode(b, lane) * pw[p];
+                }
+                acc += u * in[bi * 4 + lane];
+            }
+        }
+        /* FP16 row scale decode, same bit arithmetic as the vendored file
+         * (sign ignored — encoders store max|row| >= 0). */
+        const uint32_t h16 = rs[n];
+        const uint32_t eu = (h16 >> 10) & 0x1Fu;
+        const uint32_t mu = h16 & 0x03FFu;
+        const uint32_t fu = ((eu + 112u) << 23) | (mu << 13);
+        float rsc;
+        memcpy(&rsc, &fu, 4);
+        output[output_offset + n] = acc * rsc;
+    }
+}
+
+#endif /* SKAINET_HAVE_NEOGPU_TERNARY */

skainet-backends/skainet-backend-native-cpu/native/src/vendor/neogpu/README.md

@@ -0,0 +1,53 @@
+# Vendored: NeoGPU ternary LUT kernel
+
+Byte-identical copy of one file from the NeoGPU project, vendored under its
+MIT license (see the REUSE `.license` sidecar; `LICENSES/MIT.txt` at the repo
+root carries the license text). Agreed with upstream in
+[anjaustin/neogpu#1](https://github.com/anjaustin/neogpu/issues/1); SKaiNET
+tracking issue: [#1136](https://github.com/SKaiNET-developers/SKaiNET/issues/1136).
+
+| | |
+|---|---|
+| Upstream | <https://github.com/anjaustin/neogpu> |
+| File | `src/hs_ml_ternary_neon.c` |
+| Vendored at commit | `0846b24ceb9f76a610b3efe2967ff4ead2ef10e6` |
+| SHA-256 | `a560ffcf4d5a2e2f600d715b8193da999a92982160c625e5da48076d2257d587` |
+| Local modifications | **none** |
+
+## What it is
+
+An f32-activation × ternary-weight ({-1, 0, +1}) matmul for AArch64 using a
+256-entry × 4-float decode LUT (4 KB, L1-resident) — `vld1q_f32(lut[w[b]])` +
+`vfmaq_f32` inner loop. Needs only baseline NEON (`armv8-a+simd`), no
+`FEAT_DotProd`: it is the fast path for Cortex-A72 / Raspberry Pi 4-class
+cores. Carries its own exact scalar fallback for non-ARM builds. The 2-bit
+payload rule (4 codes per byte, low bit-pair first, code = value + 1) is
+byte-identical to SKaiNET's `BITNET_B1_58` / `TernaryPacked` encoding.
+
+Exports three symbols; SKaiNET wraps the first two via
+`src/skainet_ternary_f32.c` and never calls the third
+(`hs_ml_lmhead_stage1_i8` uses a thread-unsafe static scratch buffer):
+
+- `hs_ml_ternary_f32_proj` — single-plane projection GEMV
+- `hs_ml_lmhead_stage1` — fused 4-plane lm_head with FP16 row scales
+- `hs_ml_lmhead_stage1_i8` — ABI-compat stub, **do not use**
+
+## Constraints the adapter enforces / works around
+
+- `build_lut()` uses a non-atomic init guard → the adapter warms it up under
+  `pthread_once` before any concurrent use.
+- pthreads: threading is NeoGPU's own (4 threads once `N >= 512` for the proj;
+  the lm_head always threads). No MSVC build — the adapter compiles a portable
+  scalar fallback there instead (`SKAINET_HAVE_NEOGPU_TERNARY` unset).
+- Byte code 3 decodes to +2.0 (matches `TernaryCodec.decodeBitNet`); loaders
+  must reject code 3 at import, the kernel never validates.
+- Must be compiled at `-march=armv8-a` on non-Apple AArch64 — the library
+  default `-march=armv8.2-a+fp16+dotprod` would defeat its purpose
+  (dotprod-less targets).
+
+## Re-vendoring
+
+Copy the file byte-identical from upstream, update the commit + SHA-256 here,
+keep "Local modifications: none" true (fixes belong in the adapter or
+upstream), and re-run the ternary goldens in
+`src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernelTest.kt`.