perf(backend-cpu): dense FP32 products through the platform BLAS, on by default on macOS

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Changed
+
+- **Dense FP32 products on the JVM go to the platform BLAS where one is on, and on macOS it is on by default** (`skainet-backend-cpu`): `JvmBlas` never reached a BLAS before — it looked a framework up by its bare name, which `dlopen` cannot resolve, and passed heap arrays to a downcall, which JDK 21 rejects — so the `skainet.cpu.blas.enabled` switch was a no-op. It now resolves `cblas_sgemm` from `skainet.cpu.blas.library` / `SKAINET_CPU_BLAS_LIBRARY`, the Accelerate framework by path on macOS, or the common Linux sonames, and has an off-heap face. Three products use it when `blasEnabled` is true and the product is above a small-work floor: `A · B` on heap arrays (staged off-heap, as the native heap kernel does), `A · B` on two segment-backed operands (read in place), and the direct `x · Wᵀ` on a resident dense weight (`sgemm` with the weight transposed, no copy of it). On macOS the default is now on — Accelerate's SGEMM runs on the matrix units of Apple silicon and threads itself, several times the NEON kernel's throughput — and `-Dskainet.cpu.blas.enabled=false` restores the bundled kernels; elsewhere the default stays off. Results move within rounding against the bundled kernels, as they already do between lanes.
+
 ## [0.62.3] - 2026-10-25
 
 Headline: **opening a tokenizer GGUF and building its tokenizer is several times cheaper, and the dense FP32 `A · B`

docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc

@@ -36,9 +36,11 @@ You can toggle acceleration paths at runtime using system properties or environm
 * Vector acceleration:
 ** `-Dskainet.cpu.vector.enabled=true++|++false`
 ** or `SKAINET++_++CPU++_++VECTOR++_++ENABLED=true++|++false`
-* BLAS via Panama (matmul heuristic for larger sizes):
+* Platform BLAS via Panama for dense FP32 products (`A · B` and `x · Wᵀ` above a small-work floor). On by default on macOS, where it is the Accelerate framework; off elsewhere:
 ** `-Dskainet.cpu.blas.enabled=true++|++false`
 ** or `SKAINET++_++CPU++_++BLAS++_++ENABLED=true++|++false`
+** `-Dskainet.cpu.blas.library=<path or soname>` / `SKAINET++_++CPU++_++BLAS++_++LIBRARY` names the library to load (default: Accelerate on macOS, the OpenBLAS / reference BLAS sonames on Linux)
+** To see which library resolved, run `./gradlew :skainet-backends:skainet-backend-cpu:jvmTest --tests '*JvmBlasTest*'` and look for the `[blas] <library>` line in `build/test-results/jvmTest/TEST-sk.ainet.exec.tensor.ops.JvmBlasTest.xml`; `[skip]` there means no `cblas_sgemm` was found and the bundled kernels served every product.
 
 Each benchmark also exposes `@Param` to toggle these flags without modifying Gradle args.
 

skainet-backends/skainet-backend-cpu/build.gradle.kts

@@ -76,8 +76,15 @@ ksp {
     arg("skainet.module", "skainet-backend-cpu")
 }
 
+// Forward the BLAS switches from the Gradle CLI to the forked test JVM, so a suite can be run against
+// the platform BLAS (`-Dskainet.cpu.blas.enabled=true`, optionally `-Dskainet.cpu.blas.library=<path>`).
+val blasEnabledProperty = providers.systemProperty("skainet.cpu.blas.enabled")
+val blasLibraryProperty = providers.systemProperty("skainet.cpu.blas.library")
+
 tasks.withType<Test>().configureEach {
     jvmArgs("--enable-preview", "--add-modules", "jdk.incubator.vector")
+    blasEnabledProperty.orNull?.let { systemProperty("skainet.cpu.blas.enabled", it) }
+    blasLibraryProperty.orNull?.let { systemProperty("skainet.cpu.blas.library", it) }
 }
 
 tasks.withType<JavaExec>().configureEach {

skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt

@@ -330,7 +330,7 @@ internal class DefaultCpuOpsJvm(
         val activation = heapActivation(x) ?: return null
         val m = if (k == 0) 0 else x.shape.volume / k
         val out = FloatArray(m * n)
-        if (m > 0 && n > 0) {
+        if (m > 0 && n > 0 && !blasWeightTransposed(activation, resident, out, m, n, k)) {
             val (xArr, xOff) = activation
             val (wSeg, wOff) = resident
             val tasks = if (m >= 2 && m.toLong() * n >= PARALLEL_MATMUL_MIN_OUTPUT && isSharedSegment(wSeg)) {
@@ -371,6 +371,38 @@ internal class DefaultCpuOpsJvm(
         else -> null
     }
 
+    /**
+     * `x · Wᵀ` through the platform BLAS when one is enabled and resolved ([JvmBlas]): `sgemm` with B
+     * transposed reads the resident `[out, in]` weight where it lies, the activation is staged
+     * off-heap (`m × k`, small) and the `m × n` result copied back. On macOS this is Accelerate, whose
+     * SGEMM runs on the matrix units of Apple silicon and threads itself — several times the NEON
+     * kernel's throughput — so it takes the product before the schedule split. False when BLAS is
+     * off, absent, the product is too small to pay the call, or the call failed; the caller then
+     * runs the native kernel.
+     */
+    private fun blasWeightTransposed(
+        activation: Pair<FloatArray, Int>,
+        resident: Pair<java.lang.foreign.MemorySegment, Long>,
+        out: FloatArray,
+        m: Int, n: Int, k: Int,
+    ): Boolean {
+        if (!blasWanted(m.toLong() * n * k)) return false
+        val (xArr, xOff) = activation
+        val (wSeg, wOff) = resident
+        Arena.ofConfined().use { arena ->
+            val xSeg = arena.allocate(m.toLong() * k * java.lang.Float.BYTES, 64L)
+            java.lang.foreign.MemorySegment.copy(xArr, xOff, xSeg, java.lang.foreign.ValueLayout.JAVA_FLOAT, 0L, m * k)
+            val cSeg = arena.allocate(m.toLong() * n * java.lang.Float.BYTES, 64L)
+            if (!JvmBlas.sgemm(false, true, m, n, k, xSeg, 0L, k, wSeg, wOff, k, cSeg, 0L, n)) return false
+            java.lang.foreign.MemorySegment.copy(cSeg, java.lang.foreign.ValueLayout.JAVA_FLOAT, 0L, out, 0, m * n)
+            return true
+        }
+    }
+
+    /** Whether a dense FP32 product of [work] multiply-adds goes to the platform BLAS: enabled, resolved, and large enough to pay the call and the staging copies. */
+    private fun blasWanted(work: Long): Boolean =
+        work >= BLAS_MIN_WORK && JvmCpuBackendConfig.blasEnabled && JvmBlas.isAvailable()
+
     /**
      * Whether [segment] may be read from a thread other than the current one — a shared, automatic
      * or global arena. A confined segment (another caller's `Arena.ofConfined()`) is not, and a
@@ -542,6 +574,12 @@ internal class DefaultCpuOpsJvm(
      */
     private val SMALL_FP32_MATMUL_WORK: Long = 4_000_000L
 
+    /**
+     * Below this many multiply-adds a dense FP32 product stays on the in-process kernels even when a
+     * BLAS is on: the call, and on the heap faces the staging copies, cost more than the arithmetic.
+     */
+    private val BLAS_MIN_WORK: Long = 1L shl 18
+
     /** The heap packed types whose JVM kernels read input-block-major bytes. */
     private fun isHeapPackedWeightForJvm(data: sk.ainet.lang.tensor.data.TensorData<*, *>): Boolean =
         data is sk.ainet.lang.tensor.data.Q4_KTensorData || data is sk.ainet.lang.tensor.data.Q5_KTensorData ||
@@ -1338,6 +1376,18 @@ internal class DefaultCpuOpsJvm(
             // forward pass exhausts the JVM direct-memory cap.
             val arena = Arena.ofAuto()
             val result = MemorySegmentTensorData<T>(Shape(m, n), arena)
+            // Both operands already off-heap: the platform BLAS reads them in place, no staging.
+            if (blasWanted(m.toLong() * n * k) &&
+                JvmBlas.sgemm(
+                    false, false, m, n, k,
+                    aMemSeg.segment, aMemSeg.segmentByteOffset, k,
+                    bMemSeg.segment, bMemSeg.segmentByteOffset, n,
+                    result.segment, result.segmentByteOffset, n,
+                )
+            ) {
+                @Suppress("UNCHECKED_CAST")
+                return CpuTensor(result as TensorData<T, V>, this, a.dtype)
+            }
             val blockedThresholdMS = 16 * 16
             if (m >= blockedThresholdMS || n >= blockedThresholdMS || k >= blockedThresholdMS) {
                 JvmVectorKernels.matmulFloatBlockedMemSeg(
@@ -1366,17 +1416,12 @@ internal class DefaultCpuOpsJvm(
         val work = m.toLong() * n.toLong() * k.toLong()
         val outBuffer = FloatArray(m * n)
 
-        // Try BLAS for large sizes if enabled and available. The JNI shim takes whole arrays, so
-        // it only serves offset-0 operands; slab windows go through the kernel SPI below, which
-        // takes offsets natively.
-        if (aWin.off == 0 && bWin.off == 0 && JvmCpuBackendConfig.blasEnabled && JvmBlas.isAvailable()) {
-            val blasThreshold = 512L * 512L * 256L // tuneable
-            if (work >= blasThreshold) {
-                val ok = JvmBlas.sgemmRowMajorNN(m, n, k, 1f, aWin.arr, bWin.arr, outBuffer)
-                if (ok) {
-                    return floatResult(Shape(m, n), a.dtype, outBuffer)
-                }
-            }
+        // The platform BLAS first when it is on and resolved (Accelerate on macOS by default): the
+        // heap face stages the operands off-heap, which the native heap kernel below does too.
+        if (blasWanted(work) &&
+            JvmBlas.sgemmRowMajorNN(m, n, k, aWin.arr, aWin.off, k, bWin.arr, bWin.off, n, outBuffer, 0, n)
+        ) {
+            return floatResult(Shape(m, n), a.dtype, outBuffer)
         }
 
         // Small shapes are overhead-bound, not throughput-bound. The tile-blocked SPI kernel below