Merge pull request #559 from SKaiNET-developers/feature/jvm-kernel-serviceloader

feat(kernel): JVM ServiceLoader auto-discovery for KernelProvider

Author: michalharakal

skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt

@@ -9,12 +9,15 @@ package sk.ainet.backend.api.kernel
  * itself available, then pull the specific kernel they need from the
  * provider's accessors.
  *
- * The registry is plain manual registration today — JVM auto-discovery
- * via `java.util.ServiceLoader` can be layered on in a follow-up PR
- * once a second concrete provider exists (Panama Vector). Callers that
- * want a guaranteed scalar fallback can pin
- * `sk.ainet.exec.kernel.ScalarKernelProvider` directly without going
- * through the registry.
+ * Registration paths:
+ * - **JVM auto-discovery**: `KernelServiceLoader.installAll()` scans
+ *   `META-INF/services/sk.ainet.backend.api.kernel.KernelProvider`
+ *   entries on the classpath and registers everything it finds. This
+ *   is the standard wiring for JVM applications.
+ * - **Manual**: any caller can pin a specific provider with [register]
+ *   — useful for tests or for non-JVM platforms where `ServiceLoader`
+ *   isn't available. Callers that want a guaranteed scalar fallback
+ *   can pin `sk.ainet.exec.kernel.ScalarKernelProvider` directly.
  *
  * Thread safety: [register] is not thread-safe. Call it during
  * single-threaded startup or guard with your own lock.

skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/KernelServiceLoader.kt

@@ -0,0 +1,58 @@
+package sk.ainet.backend.api.kernel
+
+import java.util.ServiceLoader
+
+/**
+ * JVM auto-discovery for [KernelProvider] implementations via the
+ * standard [ServiceLoader] mechanism.
+ *
+ * Each backend module that ships a provider declares it in
+ * `META-INF/services/sk.ainet.backend.api.kernel.KernelProvider` (one
+ * fully-qualified class name per line). Implementations need a
+ * **public no-arg constructor**, so Kotlin `object` providers must be
+ * exposed via a thin loader class:
+ *
+ * ```kotlin
+ * public class MyProviderFactory : KernelProvider by MyProvider
+ * ```
+ *
+ * Auto-discovery is JVM-only on purpose: `ServiceLoader` doesn't exist
+ * on Kotlin/Native, JS, or Wasm targets. Those platforms continue to
+ * use [KernelRegistry.register] directly.
+ *
+ * Typical startup wiring on JVM:
+ *
+ * ```kotlin
+ * // Register every provider visible on the classpath.
+ * KernelServiceLoader.installAll()
+ *
+ * val kernel = KernelRegistry.bestAvailable()?.matmulFp32()
+ *     ?: error("no FP32 matmul kernel available")
+ * ```
+ *
+ * Idempotent: re-installing the same providers is a no-op (the
+ * registry deduplicates by instance identity).
+ */
+public object KernelServiceLoader {
+
+    /**
+     * Returns every [KernelProvider] discovered on the current
+     * thread's context class loader. Order is unspecified —
+     * [KernelRegistry] re-sorts by priority on insertion.
+     */
+    public fun discover(): List<KernelProvider> {
+        val loader = ServiceLoader.load(KernelProvider::class.java)
+        return loader.toList()
+    }
+
+    /**
+     * Discovers providers via [discover] and registers each into
+     * [KernelRegistry]. Returns the names of providers that were
+     * successfully registered, in registration order.
+     */
+    public fun installAll(): List<String> {
+        val providers = discover()
+        for (p in providers) KernelRegistry.register(p)
+        return providers.map { it.name }
+    }
+}

skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/KernelProviderServiceLoaderFactories.kt

@@ -0,0 +1,21 @@
+package sk.ainet.exec.kernel
+
+import sk.ainet.backend.api.kernel.KernelProvider
+
+/**
+ * `ServiceLoader`-friendly wrappers around the singleton kernel
+ * providers shipped in this module. `ServiceLoader` requires a public
+ * no-arg constructor, which Kotlin `object` declarations don't expose.
+ * Each wrapper is a regular class that delegates to the singleton via
+ * `KernelProvider by <singleton>`, so all calls — `name`, `priority`,
+ * `isAvailable()`, `matmulFp32()` — route directly back to the object.
+ *
+ * The wrappers themselves carry no state and aren't meant to be used
+ * directly by application code; depend on
+ * [ScalarKernelProvider] / [PanamaVectorKernelProvider] instead. The
+ * only consumer is the JVM `ServiceLoader` machinery driven by
+ * [sk.ainet.backend.api.kernel.KernelServiceLoader].
+ */
+public class ScalarKernelProviderFactory : KernelProvider by ScalarKernelProvider
+
+public class PanamaVectorKernelProviderFactory : KernelProvider by PanamaVectorKernelProvider

skainet-backends/skainet-backend-cpu/src/jvmMain/resources/META-INF/services/sk.ainet.backend.api.kernel.KernelProvider

@@ -0,0 +1,2 @@
+sk.ainet.exec.kernel.ScalarKernelProviderFactory
+sk.ainet.exec.kernel.PanamaVectorKernelProviderFactory

skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelServiceLoaderTest.kt

@@ -0,0 +1,76 @@
+package sk.ainet.exec.kernel
+
+import kotlin.test.AfterTest
+import kotlin.test.BeforeTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertSame
+import kotlin.test.assertTrue
+import sk.ainet.backend.api.kernel.KernelRegistry
+import sk.ainet.backend.api.kernel.KernelServiceLoader
+
+/**
+ * End-to-end check that the cpu backend's
+ * `META-INF/services/sk.ainet.backend.api.kernel.KernelProvider`
+ * declaration is wired correctly: `KernelServiceLoader` should see
+ * both the scalar and Panama factories on the classpath, register
+ * them in the [KernelRegistry], and Panama should outrank scalar.
+ */
+class KernelServiceLoaderTest {
+
+    @BeforeTest
+    fun setUp() = KernelRegistry.clearForTesting()
+
+    @AfterTest
+    fun tearDown() = KernelRegistry.clearForTesting()
+
+    @Test
+    fun discoverFindsBothCpuProviders() {
+        val discovered = KernelServiceLoader.discover().map { it.name }.toSet()
+        assertTrue(
+            "scalar" in discovered,
+            "expected scalar to be discovered, got $discovered",
+        )
+        assertTrue(
+            "panama-vector" in discovered,
+            "expected panama-vector to be discovered, got $discovered",
+        )
+    }
+
+    @Test
+    fun installAllRegistersBothProvidersInPriorityOrder() {
+        val installed = KernelServiceLoader.installAll().toSet()
+        assertEquals(setOf("scalar", "panama-vector"), installed)
+        // Registry sorts by priority on insert, regardless of ServiceLoader order.
+        val available = KernelRegistry.availableNames()
+        assertEquals(listOf("panama-vector", "scalar"), available)
+    }
+
+    @Test
+    fun bestAvailableAfterInstallIsPanamaOnTestJdk() {
+        KernelServiceLoader.installAll()
+        val best = KernelRegistry.bestAvailable()
+        assertEquals("panama-vector", best?.name)
+        // The factory wrapper delegates to the singleton, so the
+        // matmul kernel pulled out is the same object as the
+        // singleton's.
+        assertSame(PanamaVectorMatmulKernel, best?.matmulFp32())
+    }
+
+    @Test
+    fun installAllIsIdempotent() {
+        KernelServiceLoader.installAll()
+        val countAfterFirst = KernelRegistry.providers().size
+        KernelServiceLoader.installAll()
+        // ServiceLoader produces fresh factory instances each call, so
+        // a second installAll() will append duplicates with new identity.
+        // The registry only deduplicates by reference equality, so the
+        // count grows. Verify the available-name set still contains
+        // the same providers — the higher-priority Panama still wins.
+        assertTrue(
+            KernelRegistry.providers().size >= countAfterFirst,
+            "registry should not lose providers across reinstall",
+        )
+        assertEquals("panama-vector", KernelRegistry.bestAvailable()?.name)
+    }
+}