From 8e2ab347a1ee40a98b59edcdc6887b6ccabec1a4 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Sun, 25 Oct 2026 14:44:53 +0100 Subject: [PATCH 1/2] perf(backend-cpu): dense FP32 products through the platform BLAS, on by default on macOS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JvmBlas never reached a BLAS: it looked the Accelerate framework up by its bare name, which dlopen cannot resolve, and passed heap arrays to a downcall, which JDK 21 rejects, so skainet.cpu.blas.enabled was a no-op. It now resolves cblas_sgemm from skainet.cpu.blas.library (path or soname), the Accelerate framework by path on macOS, or the common Linux sonames, and has an off-heap face with both transposes. Three products use it when blasEnabled and above a small-work floor: A·B on heap arrays (staged, offsets allowed), A·B on two segment-backed operands (in place) and the direct x·Wᵀ on a resident dense weight (sgemm with the weight transposed, no copy of it), ahead of the schedule split since the BLAS threads itself. On macOS the default is now on — Accelerate's SGEMM runs on the matrix units of Apple silicon, several times the NEON kernel's throughput; -Dskainet.cpu.blas.enabled=false restores the bundled kernels. Elsewhere the default stays off. JvmBlasTest runs wherever a cblas_sgemm resolves and reports itself skipped otherwise; the test task forwards the BLAS properties so a suite can be run against a BLAS. Verified here against the reference BLAS on Linux (same cblas ABI); the Accelerate path itself needs a Mac. --- CHANGELOG.md | 4 + .../ROOT/pages/explanation/perf/jvm-cpu.adoc | 3 +- .../skainet-backend-cpu/build.gradle.kts | 7 + .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 69 +++++- .../sk/ainet/exec/tensor/ops/JvmBlas.kt | 206 ++++++++++++------ .../exec/tensor/ops/JvmCpuBackendConfig.kt | 14 +- .../sk/ainet/exec/tensor/ops/JvmBlasTest.kt | 176 +++++++++++++++ 7 files changed, 401 insertions(+), 78 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/JvmBlasTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ea3a60ba..d003bf4ea 100644 --- a/CHANGELOG.md +++ b/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.2] - 2026-10-24 Headline: **the dense FP32 `x · Wᵀ` product on the JVM native lane uses every core and, on x86-64, AVX2/FMA** diff --git a/docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc b/docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc index 71060e1b5..7a8dedbb3 100644 --- a/docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc +++ b/docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc @@ -36,9 +36,10 @@ 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=` / `SKAINET++_++CPU++_++BLAS++_++LIBRARY` names the library to load (default: Accelerate on macOS, the OpenBLAS / reference BLAS sonames on Linux) Each benchmark also exposes `@Param` to toggle these flags without modifying Gradle args. diff --git a/skainet-backends/skainet-backend-cpu/build.gradle.kts b/skainet-backends/skainet-backend-cpu/build.gradle.kts index 469f1ff12..fbe91d915 100644 --- a/skainet-backends/skainet-backend-cpu/build.gradle.kts +++ b/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=`). +val blasEnabledProperty = providers.systemProperty("skainet.cpu.blas.enabled") +val blasLibraryProperty = providers.systemProperty("skainet.cpu.blas.library") + tasks.withType().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().configureEach { diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index aefd68ab6..2570d1ef2 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/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, + resident: Pair, + 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(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, 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 diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmBlas.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmBlas.kt index d1b68cbe9..373e90311 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmBlas.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmBlas.kt @@ -5,105 +5,183 @@ import java.lang.foreign.FunctionDescriptor import java.lang.foreign.Linker import java.lang.foreign.MemorySegment import java.lang.foreign.SymbolLookup +import java.lang.foreign.ValueLayout import java.lang.invoke.MethodHandle +import java.nio.file.Files +import java.nio.file.Path /** - * Minimal Panama FFM wrapper for BLAS `cblas_sgemm`. - * This is optional and guarded by availability checks and config flags. + * The platform BLAS's `cblas_sgemm` through Panama FFM — on macOS the Accelerate framework, whose + * SGEMM runs on the matrix units of Apple silicon and threads itself; elsewhere OpenBLAS or the + * reference BLAS when one is installed. Optional: [isAvailable] is false when no library resolves, + * and the ops layer consults [JvmCpuBackendConfig.blasEnabled] before asking. + * + * Library resolution, first hit wins: the `skainet.cpu.blas.library` system property or the + * `SKAINET_CPU_BLAS_LIBRARY` environment variable (a path, or a name `dlopen` resolves), the + * Accelerate framework by its path on macOS, then the common Linux sonames. A bare `"Accelerate"` + * never resolves — `dlopen` searches library paths for a leaf name, and a framework lives under + * `/System/Library/Frameworks` — and a heap `MemorySegment.ofArray` cannot be passed to a downcall + * on JDK 21, which is why the previous version of this object never reached a BLAS; both faces below + * take off-heap segments and the heap face stages its arrays itself. + * + * Row-major throughout. Dimensions are the C-side ints; a product with more than `Int.MAX_VALUE` + * elements in an operand is the caller's problem long before it reaches here. */ internal object JvmBlas { - private const val CBLAS_LAYOUT_ROW_MAJOR = 101 - private const val CBLAS_LAYOUT_COL_MAJOR = 102 + private const val CBLAS_ROW_MAJOR = 101 private const val CBLAS_NO_TRANS = 111 private const val CBLAS_TRANS = 112 + private const val PROP_LIBRARY = "skainet.cpu.blas.library" + private const val ENV_LIBRARY = "SKAINET_CPU_BLAS_LIBRARY" + private const val ACCELERATE_PATH = "/System/Library/Frameworks/Accelerate.framework/Accelerate" + private val linker: Linker = Linker.nativeLinker() + /** The library name or path `cblas_sgemm` was resolved from, or `null` when none has (yet). */ + @Volatile + var libraryDescription: String? = null + private set + @Volatile private var lookedUp: Boolean = false @Volatile private var sgemmHandle: MethodHandle? = null - /** Try to find BLAS library and resolve cblas_sgemm symbol. */ + /** Whether a `cblas_sgemm` has been resolved; the first call does the lookup. */ fun isAvailable(): Boolean { if (lookedUp) return sgemmHandle != null synchronized(this) { if (lookedUp) return sgemmHandle != null - val libsToTry = listOf( - // macOS Accelerate provides cblas_sgemm - "Accelerate", - // Common Linux names for OpenBLAS/ATLAS - "blas", - "openblas", - "cblas" - ) - val arena = Arena.ofAuto() - var handle: MethodHandle? = null - for (lib in libsToTry) { - val lookup = runCatching { SymbolLookup.libraryLookup(lib, arena) }.getOrNull() ?: continue - val symbol = lookup.find("cblas_sgemm").orElse(null) ?: continue - handle = downcallSgemm(symbol) - if (handle != null) break - } - sgemmHandle = handle + resolve() lookedUp = true return sgemmHandle != null } } - private fun downcallSgemm(symbol: MemorySegment): MethodHandle? { - // cblas_sgemm layout: - // void cblas_sgemm(const enum CBLAS_LAYOUT Layout, const enum CBLAS_TRANSPOSE TransA, - // const enum CBLAS_TRANSPOSE TransB, const int M, const int N, const int K, - // const float alpha, const float *A, const int lda, const float *B, const int ldb, - // const float beta, float *C, const int ldc); - val desc = FunctionDescriptor.ofVoid( - java.lang.foreign.ValueLayout.JAVA_INT, // Layout - java.lang.foreign.ValueLayout.JAVA_INT, // TransA - java.lang.foreign.ValueLayout.JAVA_INT, // TransB - java.lang.foreign.ValueLayout.JAVA_INT, // M - java.lang.foreign.ValueLayout.JAVA_INT, // N - java.lang.foreign.ValueLayout.JAVA_INT, // K - java.lang.foreign.ValueLayout.JAVA_FLOAT, // alpha - java.lang.foreign.ValueLayout.ADDRESS, // A - java.lang.foreign.ValueLayout.JAVA_INT, // lda - java.lang.foreign.ValueLayout.ADDRESS, // B - java.lang.foreign.ValueLayout.JAVA_INT, // ldb - java.lang.foreign.ValueLayout.JAVA_FLOAT, // beta - java.lang.foreign.ValueLayout.ADDRESS, // C - java.lang.foreign.ValueLayout.JAVA_INT, // ldc - ) - return runCatching { linker.downcallHandle(symbol, desc) }.getOrNull() + /** Drops the resolved handle so the next [isAvailable] resolves again (tests change the library property). */ + fun reload() { + synchronized(this) { + sgemmHandle = null + libraryDescription = null + lookedUp = false + } + } + + private fun candidates(): List { + val explicit = runCatching { System.getProperty(PROP_LIBRARY) }.getOrNull()?.takeIf { it.isNotBlank() } + ?: runCatching { System.getenv(ENV_LIBRARY) }.getOrNull()?.takeIf { it.isNotBlank() } + val os = runCatching { System.getProperty("os.name") }.getOrNull()?.lowercase() ?: "" + val platform = when { + os.contains("mac") || os.contains("darwin") -> listOf(ACCELERATE_PATH) + os.contains("win") -> listOf("libopenblas.dll", "openblas.dll") + else -> listOf("libopenblas.so.0", "libopenblas.so", "libblas.so.3", "libcblas.so.3") + } + return listOfNotNull(explicit) + platform } + private fun resolve() { + val arena = Arena.ofAuto() + for (candidate in candidates()) { + val lookup = runCatching { + val asPath = Path.of(candidate) + if (candidate.contains('/') || candidate.contains('\\')) { + if (!Files.exists(asPath)) return@runCatching null + SymbolLookup.libraryLookup(asPath, arena) + } else { + SymbolLookup.libraryLookup(candidate, arena) + } + }.getOrNull() ?: continue + val symbol = lookup.find("cblas_sgemm").orElse(null) ?: continue + val handle = runCatching { linker.downcallHandle(symbol, SGEMM) }.getOrNull() ?: continue + sgemmHandle = handle + libraryDescription = candidate + return + } + } + + // void cblas_sgemm(Layout, TransA, TransB, M, N, K, alpha, const float* A, lda, + // const float* B, ldb, beta, float* C, ldc) + private val SGEMM: FunctionDescriptor = FunctionDescriptor.ofVoid( + ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, + ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, + ValueLayout.JAVA_FLOAT, + ValueLayout.ADDRESS, ValueLayout.JAVA_INT, + ValueLayout.ADDRESS, ValueLayout.JAVA_INT, + ValueLayout.JAVA_FLOAT, + ValueLayout.ADDRESS, ValueLayout.JAVA_INT, + ) + /** - * Row-major SGEMM (NoTrans x NoTrans) only. Other forms can be added as needed. - * C := alpha * A x B + beta * C, where A is MxK, B is KxN, C is MxN. + * `C = op(A) · op(B)` over off-heap segments, row-major, `beta = 0` (C is overwritten, never read). + * `op(A)` is `m × k` and `op(B)` is `k × n`: with [transA] false A is stored `m × k` with leading + * dimension [lda] ≥ k, with it true A is stored `k × m` and `lda` ≥ m; likewise B is `k × n` + * (`ldb` ≥ n) or, transposed, `n × k` (`ldb` ≥ k) — the shape a Linear layer stores its `[out, in]` + * weight in, so `x · Wᵀ` is `sgemm(false, true, …, ldb = k)` with no copy of the weight. Offsets + * are in bytes. Returns false when no BLAS is available or the call threw. */ - fun sgemmRowMajorNN(m: Int, n: Int, k: Int, alpha: Float, - a: FloatArray, b: FloatArray, c: FloatArray, - lda: Int = k, ldb: Int = n, ldc: Int = n, - beta: Float = 0f): Boolean { + fun sgemm( + transA: Boolean, transB: Boolean, + m: Int, n: Int, k: Int, + a: MemorySegment, aByteOffset: Long, lda: Int, + b: MemorySegment, bByteOffset: Long, ldb: Int, + c: MemorySegment, cByteOffset: Long, ldc: Int, + ): Boolean { val handle = sgemmHandle ?: return false - val arena = Arena.ofConfined() - val aSeg = MemorySegment.ofArray(a) - val bSeg = MemorySegment.ofArray(b) - val cSeg = MemorySegment.ofArray(c) + if (m <= 0 || n <= 0) return true + if (k <= 0) { + c.asSlice(cByteOffset).fill(0) + return true + } return runCatching { handle.invoke( - CBLAS_LAYOUT_ROW_MAJOR, - CBLAS_NO_TRANS, - CBLAS_NO_TRANS, + CBLAS_ROW_MAJOR, + if (transA) CBLAS_TRANS else CBLAS_NO_TRANS, + if (transB) CBLAS_TRANS else CBLAS_NO_TRANS, m, n, k, - alpha, - aSeg, lda, - bSeg, ldb, - beta, - cSeg, ldc, + 1.0f, + a.asSlice(aByteOffset), lda, + b.asSlice(bByteOffset), ldb, + 0.0f, + c.asSlice(cByteOffset), ldc, ) - // copy back into c in case the segment backed by a distinct array; MemorySegment.ofArray wraps the array true }.getOrElse { false } } + + /** + * The heap face of [sgemm] for `Fp32MatmulKernel`-shaped operands: `C(m, n) = A(m, k) · B(k, n)` + * with element offsets and row strides in floats. A, B and the `m × n` block of C are staged + * off-heap for the call — the weight-sized copy the native heap kernel pays too — so callers keep + * this for products large enough to amortise it. Only the `m × n` block of `out` is written. + */ + fun sgemmRowMajorNN( + m: Int, n: Int, k: Int, + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + ): Boolean { + if (sgemmHandle == null) return false + if (m <= 0 || n <= 0) return true + val aFloats = if (k == 0) 0 else (m - 1) * aStride + k + val bFloats = if (k == 0) 0 else (k - 1) * bStride + n + Arena.ofConfined().use { arena -> + val aSeg = arena.allocate(maxOf(aFloats, 1).toLong() * Float.SIZE_BYTES, 64) + val bSeg = arena.allocate(maxOf(bFloats, 1).toLong() * Float.SIZE_BYTES, 64) + val cSeg = arena.allocate(m.toLong() * n * Float.SIZE_BYTES, 64) + if (aFloats > 0) MemorySegment.copy(a, aOffset, aSeg, ValueLayout.JAVA_FLOAT, 0L, aFloats) + if (bFloats > 0) MemorySegment.copy(b, bOffset, bSeg, ValueLayout.JAVA_FLOAT, 0L, bFloats) + val ok = sgemm(false, false, m, n, k, aSeg, 0L, aStride, bSeg, 0L, bStride, cSeg, 0L, n) + if (!ok) return false + if (outStride == n) { + MemorySegment.copy(cSeg, ValueLayout.JAVA_FLOAT, 0L, out, outOffset, m * n) + } else { + for (i in 0 until m) { + MemorySegment.copy(cSeg, ValueLayout.JAVA_FLOAT, i.toLong() * n * Float.SIZE_BYTES, out, outOffset + i * outStride, n) + } + } + return true + } + } } diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmCpuBackendConfig.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmCpuBackendConfig.kt index d08c0e925..2c0e9d5a0 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmCpuBackendConfig.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmCpuBackendConfig.kt @@ -8,8 +8,20 @@ internal object JvmCpuBackendConfig { val vectorEnabled: Boolean get() = readFlag(ENV_VECTOR, PROP_VECTOR, default = true) + /** + * Whether dense FP32 products go to the platform BLAS ([JvmBlas]) when one resolves. On by + * default on macOS, where the BLAS is Accelerate and runs on the matrix units of Apple silicon; + * off elsewhere, where a BLAS is an optional install and the bundled kernels are the known + * quantity. `-Dskainet.cpu.blas.enabled` / `SKAINET_CPU_BLAS_ENABLED` override either way, and + * `-Dskainet.cpu.blas.library` / `SKAINET_CPU_BLAS_LIBRARY` name the library to load. + */ val blasEnabled: Boolean - get() = readFlag(ENV_BLAS, PROP_BLAS, default = false) + get() = readFlag(ENV_BLAS, PROP_BLAS, default = isMacOs) + + private val isMacOs: Boolean by lazy { + val os = runCatching { System.getProperty("os.name") }.getOrNull()?.lowercase() ?: "" + os.contains("mac") || os.contains("darwin") + } private fun readFlag(envKey: String, propKey: String, default: Boolean): Boolean { val sysProp = runCatching { System.getProperty(propKey) }.getOrNull() diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/JvmBlasTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/JvmBlasTest.kt new file mode 100644 index 000000000..875664932 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/JvmBlasTest.kt @@ -0,0 +1,176 @@ +package sk.ainet.exec.tensor.ops + +import java.lang.foreign.Arena +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.MemorySegmentTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.FP32 + +/** + * The platform BLAS binding and the three dense FP32 products the JVM ops send to it. Runs wherever a + * `cblas_sgemm` resolves — Accelerate on macOS, OpenBLAS or the reference BLAS on Linux, or the library + * named by `-Dskainet.cpu.blas.library` — and reports itself skipped otherwise, so a CI host without a + * BLAS neither fails nor pretends. The reference is a double-accumulated product; the bar is the + * rounding tolerance the kernel parity tests use, since a BLAS orders its sums as it likes. + */ +class JvmBlasTest { + + private val dataFactory = DenseTensorDataFactory() + private val ops = DefaultCpuOpsJvm(dataFactory) + + private fun available(): Boolean { + if (JvmBlas.isAvailable()) return true + println("[skip] no cblas_sgemm resolved (set -Dskainet.cpu.blas.library=); BLAS tests not run") + return false + } + + private fun reference( + a: FloatArray, aOff: Int, lda: Int, transA: Boolean, + b: FloatArray, bOff: Int, ldb: Int, transB: Boolean, + m: Int, n: Int, k: Int, + ): FloatArray { + val out = FloatArray(m * n) + for (i in 0 until m) for (j in 0 until n) { + var acc = 0.0 + for (p in 0 until k) { + val av = if (transA) a[aOff + p * lda + i] else a[aOff + i * lda + p] + val bv = if (transB) b[bOff + j * ldb + p] else b[bOff + p * ldb + j] + acc += av.toDouble() * bv + } + out[i * n + j] = acc.toFloat() + } + return out + } + + private fun assertClose(expected: FloatArray, actual: FloatArray, what: String) { + assertEquals(expected.size, actual.size, "$what size") + var maxAbs = 0f + for (v in expected) maxAbs = maxOf(maxAbs, abs(v)) + val tol = maxAbs * 1e-5f + 1e-6f + for (i in expected.indices) { + assertTrue(abs(expected[i] - actual[i]) <= tol, "$what[$i]: ${actual[i]} vs ${expected[i]} (tol $tol)") + } + } + + private fun segmentOf(arena: Arena, values: FloatArray): MemorySegment { + val seg = arena.allocate(values.size.toLong() * Float.SIZE_BYTES, 64) + MemorySegment.copy(values, 0, seg, ValueLayout.JAVA_FLOAT, 0L, values.size) + return seg + } + + @Test + fun segmentFaceMatchesTheReferenceForBothTransposes() { + if (!available()) return + println("[blas] ${JvmBlas.libraryDescription}") + val rng = Random(31) + Arena.ofConfined().use { arena -> + for ((m, n, k) in listOf(Triple(13, 37, 300), Triple(1, 1024, 256), Triple(6, 16, 1), Triple(63, 1021, 1023))) { + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + // B as [k, n] for the plain product and as [n, k] (a Linear weight) for x · Wᵀ. + val bKN = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val bNK = FloatArray(n * k) { rng.nextFloat() - 0.5f } + val aSeg = segmentOf(arena, a) + val cSeg = arena.allocate(m.toLong() * n * Float.SIZE_BYTES, 64) + val out = FloatArray(m * n) + + assertTrue(JvmBlas.sgemm(false, false, m, n, k, aSeg, 0L, k, segmentOf(arena, bKN), 0L, n, cSeg, 0L, n)) + MemorySegment.copy(cSeg, ValueLayout.JAVA_FLOAT, 0L, out, 0, m * n) + assertClose(reference(a, 0, k, false, bKN, 0, n, false, m, n, k), out, "NN m=$m n=$n k=$k") + + assertTrue(JvmBlas.sgemm(false, true, m, n, k, aSeg, 0L, k, segmentOf(arena, bNK), 0L, k, cSeg, 0L, n)) + MemorySegment.copy(cSeg, ValueLayout.JAVA_FLOAT, 0L, out, 0, m * n) + assertClose(reference(a, 0, k, false, bNK, 0, k, true, m, n, k), out, "NT m=$m n=$n k=$k") + } + // k = 0 zeroes the block rather than reading anything. + val cSeg = arena.allocate(3L * 4 * Float.SIZE_BYTES, 64).also { it.fill(1) } + assertTrue(JvmBlas.sgemm(false, false, 3, 4, 0, cSeg, 0L, 1, cSeg, 0L, 4, cSeg, 0L, 4)) + for (i in 0 until 12) assertEquals(0f, cSeg.getAtIndex(ValueLayout.JAVA_FLOAT, i.toLong())) + } + } + + @Test + fun heapFaceTakesOffsetsAndStridesAndLeavesThePaddingAlone() { + if (!available()) return + val rng = Random(32) + val m = 7; val n = 19; val k = 33 + val aStride = k + 5; val bStride = n + 3; val outStride = n + 7 + val aOff = 2; val bOff = 1; val outOff = 3 + val a = FloatArray(aOff + m * aStride) { rng.nextFloat() - 0.5f } + val b = FloatArray(bOff + k * bStride) { rng.nextFloat() - 0.5f } + val out = FloatArray(outOff + m * outStride) { 123f } + assertTrue(JvmBlas.sgemmRowMajorNN(m, n, k, a, aOff, aStride, b, bOff, bStride, out, outOff, outStride)) + val block = FloatArray(m * n) + for (i in 0 until m) System.arraycopy(out, outOff + i * outStride, block, i * n, n) + assertClose(reference(a, aOff, aStride, false, b, bOff, bStride, false, m, n, k), block, "heap NN") + for (idx in 0 until outOff) assertEquals(123f, out[idx], "prefix $idx") + for (i in 0 until m) for (j in n until outStride) { + if (outOff + i * outStride + j < out.size) assertEquals(123f, out[outOff + i * outStride + j], "pad [$i, $j]") + } + } + + private fun heapTensor(shape: Shape, values: FloatArray): Tensor = + VoidOpsTensor(dataFactory.fromFloatArray(shape, FP32::class, values), FP32::class) + + @Suppress("UNCHECKED_CAST") + private fun segmentTensor(shape: Shape, values: FloatArray, arena: Arena): Tensor { + val data = MemorySegmentTensorData(shape, arena) + data.copyFromFloatArray(values) + return VoidOpsTensor(data as TensorData, FP32::class) + } + + private fun floatsOf(t: Tensor<*, *>): FloatArray = when (val d = t.data) { + is FloatArrayTensorData<*> -> d.buffer.copyOf(t.shape.volume) + is MemorySegmentTensorData<*> -> d.copyToFloatArray() + else -> error("unexpected data ${d::class}") + } + + /** + * With BLAS switched on, the three routed products — heap `A · B`, segment `A · B`, and `x · Wᵀ` on + * a dense `[out, in]` weight — agree with the reference, and the same products with BLAS switched + * off (the bundled kernels) agree with them within rounding. Shapes sit above the small-work floor. + */ + @Test + fun opsProductsAgreeWithBlasOnAndOff() { + if (!available()) return + val rng = Random(33) + val m = 48; val n = 320; val k = 256 // 3.9 M multiply-adds: above the BLAS floor, below the small-shape loop's ceiling + val x = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val bKN = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val w = FloatArray(n * k) { rng.nextFloat() - 0.5f } + val expectedNN = reference(x, 0, k, false, bKN, 0, n, false, m, n, k) + val expectedNT = reference(x, 0, k, false, w, 0, k, true, m, n, k) + + val previous = System.getProperty("skainet.cpu.blas.enabled") + try { + val results = mutableMapOf() + for (blas in listOf(true, false)) { + System.setProperty("skainet.cpu.blas.enabled", blas.toString()) + assertEquals(blas, JvmCpuBackendConfig.blasEnabled) + Arena.ofConfined().use { arena -> + val heap = ops.matmul(heapTensor(Shape(m, k), x), heapTensor(Shape(k, n), bKN)) + val seg = ops.matmul(segmentTensor(Shape(m, k), x, arena), segmentTensor(Shape(k, n), bKN, arena)) + val wt = ops.matmulWeightTransposed(heapTensor(Shape(m, k), x), heapTensor(Shape(n, k), w)) + results["heap blas=$blas"] = floatsOf(heap).also { assertClose(expectedNN, it, "heap A·B blas=$blas") } + results["segment blas=$blas"] = floatsOf(seg).also { assertClose(expectedNN, it, "segment A·B blas=$blas") } + results["wt blas=$blas"] = floatsOf(wt).also { assertClose(expectedNT, it, "x·Wᵀ blas=$blas") } + } + } + for (face in listOf("heap", "segment", "wt")) { + assertClose(results["$face blas=false"]!!, results["$face blas=true"]!!, "$face on vs off") + } + } finally { + if (previous == null) System.clearProperty("skainet.cpu.blas.enabled") else System.setProperty("skainet.cpu.blas.enabled", previous) + } + } +} From 6dbb61e0dca53854f7f8d2cac82b2c051095f452 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Sun, 25 Oct 2026 15:09:03 +0100 Subject: [PATCH 2/2] docs(perf): how to see which BLAS the JVM resolved --- docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc b/docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc index 7a8dedbb3..24cc83fce 100644 --- a/docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc +++ b/docs/modules/ROOT/pages/explanation/perf/jvm-cpu.adoc @@ -40,6 +40,7 @@ You can toggle acceleration paths at runtime using system properties or environm ** `-Dskainet.cpu.blas.enabled=true++|++false` ** or `SKAINET++_++CPU++_++BLAS++_++ENABLED=true++|++false` ** `-Dskainet.cpu.blas.library=` / `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] ` 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.