diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt index 1f24ce0d..548f8f11 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt @@ -11,6 +11,7 @@ import sk.ainet.lang.memory.trace.TraceEvent import sk.ainet.lang.memory.trace.TraceSink import sk.ainet.lang.memory.trace.kernel as traceKernel import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.TensorEncoding /** * Kernel selection on declared descriptors instead of an `is`-ladder over Kotlin classes @@ -41,6 +42,22 @@ public object KernelDispatch { public fun clearForTesting() { kernels.clear() } + /** + * Encodings a [MappedCapableKernel] registered right now serves as a `BLOCKED_ROW_MAJOR` + * weight — derived from actual registrations (#1193), not a hand-kept list. Call after + * installing every kernel pack the platform has; a pack that isn't installed (or isn't + * available on this platform) simply contributes nothing, which is why this is a guard-test + * tool rather than the sole source of truth `StorageCapabilities.mappedServableEncodings` + * (`skainet-lang-core`) uses at runtime — that module cannot depend on this one. + */ + public fun mappedServableEncodings(): Set = kernels + .asSequence() + .filter { it is MappedCapableKernel } + .flatMap { it.key.operands } + .filter { it.layout == LayoutClass.BLOCKED_ROW_MAJOR } + .map { it.format.encoding } + .toSet() + /** * Normalize a matmul operand pair to rank 2 as **views** (rule 5, §5.1 "rank handling happens * once"): `[k]` becomes `[1, k]`, `[b, s, k]` becomes `[b*s, k]` when contiguous. Returns the diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/MappedCapableKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/MappedCapableKernel.kt new file mode 100644 index 00000000..b43c183f --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/MappedCapableKernel.kt @@ -0,0 +1,21 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * Marks a [ViewKernel] that serves its `BLOCKED_ROW_MAJOR` weight operand straight from off-heap + * or mapped storage (a mmap'd GGUF page, a direct buffer) as well as from heap bytes — not just + * the heap arm every packed kernel has. [FfmRowMajorMatmulKernel] and [JniRowMajorMatmulKernel] + * are the two kernels that earn this today (#1189/#1192). + * + * [KernelDispatch.mappedServableEncodings] derives the encodings actually served from the + * kernels registered with this marker (#1193) — the thing `StorageCapabilities.MAPPED_SERVABLE_DEFAULT` + * (`skainet-lang-core`) and `KernelSupportMatrixTest.mappedTiers()`'s `native-jni-direct` row still + * declare by hand, because neither can depend on this module (`skainet-backend-api`) without a + * dependency cycle. `KernelSupportMatrixTest.generate_and_gate_support_matrix()` installs the + * JVM-reachable mapped kernels and asserts the derived set against those hand-kept declarations, + * so a kernel gaining or losing this marker without updating them fails CI instead of drifting + * silently. + */ +@ExperimentalMemoryApi +public interface MappedCapableKernel diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBufferPackedKernels.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBufferPackedKernels.kt index 5cf9bce5..c212671b 100644 --- a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBufferPackedKernels.kt +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBufferPackedKernels.kt @@ -4,6 +4,7 @@ import java.nio.ByteBuffer import sk.ainet.backend.api.kernel.KernelDispatch import sk.ainet.backend.api.kernel.KernelKey import sk.ainet.backend.api.kernel.LayoutClass +import sk.ainet.backend.api.kernel.MappedCapableKernel import sk.ainet.backend.api.kernel.OperandKey import sk.ainet.backend.api.kernel.ReferenceMatmulKernel import sk.ainet.backend.api.kernel.ViewKernel @@ -51,7 +52,7 @@ public class JniRowMajorMatmulKernel( inputDim: Int, outputDim: Int, output: FloatArray, outputOffset: Int, ) -> Unit, -) : ViewKernel { +) : ViewKernel, MappedCapableKernel { override val name: String = "jni-rowmajor-$encodingName" diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/FfmRowMajorKernels.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/FfmRowMajorKernels.kt index 31d27611..1129de03 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/FfmRowMajorKernels.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/FfmRowMajorKernels.kt @@ -9,6 +9,7 @@ import java.lang.invoke.MethodHandle import sk.ainet.backend.api.kernel.KernelDispatch import sk.ainet.backend.api.kernel.KernelKey import sk.ainet.backend.api.kernel.LayoutClass +import sk.ainet.backend.api.kernel.MappedCapableKernel import sk.ainet.backend.api.kernel.OperandKey import sk.ainet.backend.api.kernel.ReferenceMatmulKernel import sk.ainet.backend.api.kernel.ViewKernel @@ -44,7 +45,7 @@ public class FfmRowMajorMatmulKernel internal constructor( encodingName: String, override val key: KernelKey, private val handle: MethodHandle, -) : ViewKernel { +) : ViewKernel, MappedCapableKernel { override val name: String = "ffm-rowmajor-$encodingName" diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt index 5b6c5fd9..c7da44da 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt @@ -1,9 +1,14 @@ package sk.ainet.exec.kernel import java.io.File +import kotlin.test.AfterTest import kotlin.test.Test import kotlin.test.assertEquals +import sk.ainet.backend.api.kernel.KernelDispatch import sk.ainet.backend.api.kernel.KernelProvider +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.StorageCapabilities +import sk.ainet.lang.tensor.storage.TensorEncoding /** * Emits `kernel-support.json` (introspected from the registered `KernelProvider`s) and @@ -19,8 +24,12 @@ import sk.ainet.backend.api.kernel.KernelProvider * probes the JDK incubator module / the loaded `.so`), so their *capability* is declared * here (the single place to edit when a provider gains a kernel) rather than probed. */ +@OptIn(ExperimentalMemoryApi::class) class KernelSupportMatrixTest { + @AfterTest + fun reset() = KernelDispatch.clearForTesting() + private val formats = listOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0") // platform key (display) -> the set of providers (by source-set) reaching it. @@ -59,15 +68,22 @@ class KernelSupportMatrixTest { /** * Mapped serving (#1189): kernels that read the weight in canonical row-major GGUF file * order straight from off-heap bytes (mmap/direct buffer) — the `_rm` symbols behind - * `JniBufferPackedMatmulKernel` on Android. K/N and remaining formats: #1192 follow-ups. Must stay in lockstep with - * `StorageCapabilities.MAPPED_SERVABLE_ENCODINGS` (dense F32 is mapped there too, but as - * element-view serving, not a matmul kernel — it has no row here on purpose). + * `JniBufferPackedMatmulKernel` on Android. K/N and remaining formats: #1192 follow-ups. + * + * The JVM row is *derived* (#1193), not hand-kept: `generate_and_gate_support_matrix` + * installs `FfmRowMajorKernelPack` and reads `KernelDispatch.mappedServableEncodings()`, + * which collects every registered [sk.ainet.backend.api.kernel.MappedCapableKernel]'s + * `BLOCKED_ROW_MAJOR` operand. That same test asserts the derived set against + * `StorageCapabilities.MAPPED_SERVABLE_DEFAULT` (`skainet-lang-core`), which cannot depend on + * this module to derive it directly — so this is the guard that catches the two drifting. + * The Android row stays hand-declared (nothing on a JVM test run can probe whether the JNI + * `.so` loads on a device); `FfmRowMajorKernelPack` and `JniMappedKernelPack` register the + * identical symbol list on purpose, so keep this row equal to the derived one by hand. */ private fun mappedTiers(): List = listOf( Tier("native-jni-direct", 100, setOf("Android"), setOf("Q4_K", "Q6_K", "Q5_K", "Q8_0", "Q4_0", "Q5_0", "Q5_1")), - Tier("ffm-rowmajor", 100, setOf("JVM"), - setOf("Q4_K", "Q6_K", "Q5_K", "Q8_0", "Q4_0", "Q5_0", "Q5_1")), + Tier("ffm-rowmajor", 100, setOf("JVM"), KernelDispatch.mappedServableEncodings().map { it.name }.toSet()), ) private fun best(fmt: String, platform: String, tiers: List): String? = @@ -119,6 +135,20 @@ class KernelSupportMatrixTest { val providers: List = listOf(ScalarKernelProvider, PanamaVectorKernelProvider, NativeKernelProvider) assertEquals(3, providers.size) + // Drift gate on mapped serving (#1193): install the JVM-reachable mapped-capable + // kernels and assert what they actually serve against StorageCapabilities' + // hand-kept mirror — the two lists can't independently drift without failing this test. + // Dense F32 is excluded: StorageCapabilities maps it too, but as element-view serving, + // not a matmul kernel, so it has no row in KernelDispatch.mappedServableEncodings(). + FfmRowMajorKernelPack.install() + assertEquals( + StorageCapabilities.MAPPED_SERVABLE_DEFAULT - TensorEncoding.Dense(4), + KernelDispatch.mappedServableEncodings(), + "KernelDispatch.mappedServableEncodings() drifted from " + + "StorageCapabilities.MAPPED_SERVABLE_DEFAULT — update both, and mappedTiers()'s " + + "native-jni-direct row", + ) + val jsonText = renderJson(tiers) val outDir = File("build/generated/kernel-support").apply { mkdirs() } File(outDir, "kernel-support.json").writeText(jsonText) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt index 21b506b7..ce61e1a4 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt @@ -27,9 +27,19 @@ public data class StorageCapabilities( val mappedServableEncodings: Set = MAPPED_SERVABLE_DEFAULT, ) { public companion object { - /** What the loaders can serve from mapped pages today: dense FP32 (#921) and the GGML + /** + * What the loaders can serve from mapped pages today: dense FP32 (#921) and the GGML * block formats (#1189 Q4_K/Q6_K, #1192 the rest). Ternary formats need a load-time - * repack and stay heap until the repack cache lands. */ + * repack and stay heap until the repack cache lands. + * + * This module cannot depend on `skainet-backend-api` (that dependency runs the other + * way, to avoid a cycle), so this constant can't be derived from + * `KernelDispatch.mappedServableEncodings()` directly (#1193). It is instead + * cross-checked there: `KernelSupportMatrixTest.generate_and_gate_support_matrix()` + * (`skainet-backend-native-cpu`) installs the mapped-capable kernel packs and asserts + * their derived encodings equal this set minus [TensorEncoding.Dense] (which is mapped + * as an element view, not served by a matmul kernel) — update both together. + */ public val MAPPED_SERVABLE_DEFAULT: Set = setOf( TensorEncoding.Dense(4), TensorEncoding.Q4_K, TensorEncoding.Q6_K, TensorEncoding.Q5_K,