Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
@@ -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 }
}
}
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
sk.ainet.exec.kernel.ScalarKernelProviderFactory
sk.ainet.exec.kernel.PanamaVectorKernelProviderFactory
Original file line number Diff line number Diff line change
@@ -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)
}
}
Loading