From 8e13b3641f0b6a46dbe00be3403bb203355914b5 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 3 Sep 2026 22:26:52 +0200 Subject: [PATCH 1/4] feat(schedule): SKEEP-005 Schedule API, CoroutineSchedule, scheduled SDPA MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compute-level half of the Halide-style algorithm/schedule split. - skainet-lang-core: `sk.ainet.context.schedule.Schedule` (non-suspending, dependency-free; `forRange` contract: disjoint ranges, no allocation through a context in bodies, nested regions inline, first failure cancels siblings) with `Schedule.Sequential`; `ScheduleHint` + `SCHEDULE_ATTRIBUTE_KEY` for the compile lane; `ExecutionContext.schedule` / `withSchedule` (default emits `TraceEvent.ScheduleDowngraded` — an unhonoured request is visible, never silent); `ScheduledExecutionContext` decorator mirroring `ScopedExecutionContext`; `TraceEvent.ScheduleRegion`. - skainet-backend-api: registries safe for concurrent reads (volatile snapshots, serialized writes via an optional-expectation `JvmSynchronized`); schedule-aware overloads on the Q4_K/Q5_K kernel SPIs. - skainet-backend-cpu: `CoroutineSchedule` (structured `coroutineScope`, caller runs the first chunk, nested regions inline, `dedicated()` pool); `parallelChunks(outputDim, schedule)` replaces the runBlocking island; ops carry the schedule (`DefaultCpuOpsBase(dataFactory, schedule)`), `platformDefaultSchedule()` (JVM: hardware coroutines, others: Sequential), `DirectCpuExecutionContext(schedule = …)` + `withSchedule`; `scaledDotProductAttention` runs its (batch, head) units on the schedule with per-task scratch — bit-identical to the sequential loop, tiny calls stay inline (`SDPA_PARALLEL_MIN_WORK`). - Tests: ScheduledExecutionContextTest, KernelDispatchConcurrencyTest, CoroutineScheduleTest, DirectCpuExecutionContextScheduleTest, SdpaScheduleParityTest (common) + SdpaCoroutineParityTest; JMH SdpaScheduleBench. API dumps refreshed (additive). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018sqoaGs5M7C5uVw6cpzAnH --- .../sk/ainet/bench/SdpaScheduleBench.kt | 54 +++++++ .../backend/api/kernel/JvmSynchronized.kt | 3 + .../backend/api/kernel/JvmSynchronized.kt | 13 ++ .../backend/api/kernel/KernelDispatch.kt | 28 +++- .../backend/api/kernel/KernelRegistry.kt | 16 +- .../backend/api/kernel/Q4KMatmulKernel.kt | 13 ++ .../backend/api/kernel/Q5KMatmulKernel.kt | 9 ++ .../backend/api/kernel/JvmSynchronized.kt | 3 + .../kernel/KernelDispatchConcurrencyTest.kt | 86 ++++++++++ .../api/jvm/skainet-backend-cpu.api | 44 +++++ .../ops/PlatformCpuOpsFactory.android.kt | 8 +- .../PlatformCpuOpsFactory.androidNative.kt | 8 +- .../ainet/exec/tensor/ops/AccelerateCpuOps.kt | 5 +- .../tensor/ops/PlatformCpuOpsFactory.apple.kt | 8 +- .../context/DirectCpuExecutionContext.kt | 30 +++- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 46 +++++- .../exec/tensor/ops/PlatformCpuOpsFactory.kt | 10 +- .../exec/schedule/SdpaScheduleParityTest.kt | 73 +++++++++ .../tensor/ops/PlatformCpuOpsFactory.js.kt | 8 +- .../kernel/PanamaVectorQ4KMatmulKernel.kt | 13 +- .../kernel/PanamaVectorQ5_KMatmulKernel.kt | 13 +- .../ainet/exec/schedule/CoroutineSchedule.kt | 128 +++++++++++++++ .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 7 +- .../tensor/ops/JvmQuantizedVectorKernels.kt | 16 +- .../ainet/exec/tensor/ops/JvmVectorKernels.kt | 6 +- .../sk/ainet/exec/tensor/ops/ParallelFor.kt | 44 ++--- .../tensor/ops/PlatformCpuOpsFactory.jvm.kt | 10 +- .../DirectCpuExecutionContextScheduleTest.kt | 54 +++++++ .../exec/schedule/CoroutineScheduleTest.kt | 143 +++++++++++++++++ .../exec/schedule/SdpaCoroutineParityTest.kt | 30 ++++ .../ops/DefaultCpuOpsJvmElementwiseTest.kt | 2 +- .../ops/DefaultCpuOpsJvmReductionsTest.kt | 2 +- .../ops/PlatformCpuOpsFactoryJvmTest.kt | 4 +- .../tensor/ops/PlatformCpuOpsFactory.linux.kt | 8 +- .../tensor/ops/PlatformCpuOpsFactory.wasm.kt | 8 +- .../ops/PlatformCpuOpsFactory.wasmWasi.kt | 8 +- .../api/jvm/skainet-lang-core.api | 151 ++++++++++++++++++ .../lang/memory/trace/AndroidTraceSink.kt | 2 + .../sk/ainet/context/ExecutionContext.kt | 28 ++++ .../context/ScheduledExecutionContext.kt | 59 +++++++ .../sk/ainet/context/schedule/Schedule.kt | 72 +++++++++ .../sk/ainet/context/schedule/ScheduleHint.kt | 52 ++++++ .../memory/trace/PerfettoTraceExporter.kt | 7 + .../sk/ainet/lang/memory/trace/TraceEvent.kt | 21 +++ .../schedule/ScheduledExecutionContextTest.kt | 122 ++++++++++++++ 45 files changed, 1393 insertions(+), 82 deletions(-) create mode 100644 skainet-backends/benchmarks/jvm-cpu-jmh/src/jmh/kotlin/sk/ainet/bench/SdpaScheduleBench.kt create mode 100644 skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt create mode 100644 skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt create mode 100644 skainet-backends/skainet-backend-api/src/jvmTest/kotlin/sk/ainet/backend/api/kernel/KernelDispatchConcurrencyTest.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/schedule/SdpaScheduleParityTest.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/context/DirectCpuExecutionContextScheduleTest.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/SdpaCoroutineParityTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScheduledExecutionContext.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/Schedule.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/ScheduleHint.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/context/schedule/ScheduledExecutionContextTest.kt diff --git a/skainet-backends/benchmarks/jvm-cpu-jmh/src/jmh/kotlin/sk/ainet/bench/SdpaScheduleBench.kt b/skainet-backends/benchmarks/jvm-cpu-jmh/src/jmh/kotlin/sk/ainet/bench/SdpaScheduleBench.kt new file mode 100644 index 000000000..b1b6524ec --- /dev/null +++ b/skainet-backends/benchmarks/jvm-cpu-jmh/src/jmh/kotlin/sk/ainet/bench/SdpaScheduleBench.kt @@ -0,0 +1,54 @@ +package sk.ainet.bench + +import org.openjdk.jmh.annotations.* +import java.util.concurrent.TimeUnit +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.exec.schedule.CoroutineSchedule +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.types.FP32 + +/** + * SKEEP-005: `scaledDotProductAttention` under the sequential schedule vs the hardware coroutine + * schedule. Run a subset with `-PjmhIncludes='SdpaScheduleBench'`. + */ +@State(Scope.Benchmark) +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +open class SdpaScheduleBench { + @Param("8", "32") + var heads: Int = 8 + + @Param("128", "1024", "4096") + var seqKV: Int = 1024 + + /** `1` is a decode step, `64` a prefill chunk. */ + @Param("1", "64") + var seqQ: Int = 1 + + @Param("sequential", "hardware") + var schedule: String = "hardware" + + private val dataFactory = DenseTensorDataFactory() + private lateinit var ctx: DirectCpuExecutionContext + private lateinit var q: VoidOpsTensor + private lateinit var k: VoidOpsTensor + private lateinit var v: VoidOpsTensor + + @Setup(Level.Trial) + fun setup() { + val s: Schedule = if (schedule == "sequential") Schedule.Sequential else CoroutineSchedule.hardware() + ctx = DirectCpuExecutionContext(schedule = s) + val headDim = 128 + fun tensor(rows: Int, seed: Int): VoidOpsTensor { + val arr = FloatArray(heads * rows * headDim) { ((it * 31 + seed) % 23 - 11) / 11f } + return VoidOpsTensor(dataFactory.fromFloatArray(Shape(1, heads, rows, headDim), FP32::class, arr), FP32::class) + } + q = tensor(seqQ, 1); k = tensor(seqKV, 2); v = tensor(seqKV, 3) + } + + @Benchmark + fun sdpa(): Any = ctx.ops.scaledDotProductAttention(q, k, v, null, 0f, true) +} diff --git a/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt new file mode 100644 index 000000000..9b18fb9b9 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt @@ -0,0 +1,3 @@ +package sk.ainet.backend.api.kernel + +internal actual typealias JvmSynchronized = kotlin.jvm.Synchronized diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt new file mode 100644 index 000000000..b3d516641 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt @@ -0,0 +1,13 @@ +package sk.ainet.backend.api.kernel + +/** + * `kotlin.jvm.Synchronized` for common code: a no-op everywhere except the JVM and Android, + * where it actualizes to the real annotation. The registries below are written once at + * bootstrap and read from many threads by schedule workers (SKEEP-005); on targets without + * shared-memory threads nothing needs guarding. + */ +@OptIn(ExperimentalMultiplatform::class) +@OptionalExpectation +@Target(AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.PROPERTY_SETTER) +@Retention(AnnotationRetention.SOURCE) +internal expect annotation class JvmSynchronized() 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 b7c4694ac..1ddac02b1 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 @@ -26,8 +26,16 @@ import sk.ainet.lang.tensor.storage.TensorEncoding @ExperimentalMemoryApi public object KernelDispatch { - private val kernels: MutableList = mutableListOf() + /** + * Immutable snapshot, replaced wholesale under [lock] on every registration; readers on the + * hot path ([find], [kernels]) never see a half-built list. Schedule workers (SKEEP-005) may + * dispatch concurrently, so the registry must be safe to *read* from many threads while the + * rare writes stay serialized. + */ + @kotlin.concurrent.Volatile + private var kernels: List = emptyList() + @kotlin.concurrent.Volatile private var autoInstallAttempted: Boolean = false /** @@ -50,6 +58,12 @@ public object KernelDispatch { * earlier ones for the same key. Call [clearForTesting] to re-arm. */ public fun ensureInstalled() { + if (autoInstallAttempted || kernels.isNotEmpty()) return + installOnce() + } + + @JvmSynchronized + private fun installOnce() { if (autoInstallAttempted || kernels.isNotEmpty()) return autoInstallAttempted = true if (KernelRegistry.providers().isEmpty()) installPlatformKernelProviders() @@ -58,19 +72,21 @@ public object KernelDispatch { } /** Register [kernel]; later registrations win for the same key (a pack can override the reference). */ + @JvmSynchronized public fun register(kernel: ViewKernel) { - kernels.removeAll { it.key == kernel.key && it.name == kernel.name } - kernels.add(0, kernel) + val kept = kernels.filterNot { it.key == kernel.key && it.name == kernel.name } + kernels = listOf(kernel) + kept } /** Every registered kernel, most recently registered first. */ - public fun kernels(): List = kernels.toList() + public fun kernels(): List = kernels /** The kernel registered for [key], or `null`. */ public fun find(key: KernelKey): ViewKernel? = kernels.firstOrNull { it.key == key } + @JvmSynchronized public fun clearForTesting() { - kernels.clear() + kernels = emptyList() autoInstallAttempted = false } @@ -114,8 +130,10 @@ public object KernelDispatch { * reference-kernel fallback invisible — set this (e.g. from a diagnostic harness) to * observe dispatch decisions everywhere without threading a sink through the ops layer. */ + @kotlin.concurrent.Volatile public var defaultSink: TraceSink = NoopTraceSink + @kotlin.concurrent.Volatile private var warnedReferenceFallback: Boolean = false /** diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt index 207ccd451..51db961ec 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt @@ -19,23 +19,24 @@ package sk.ainet.backend.api.kernel * 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. + * Thread safety: writes ([register], [clearForTesting]) are serialized; reads see an immutable + * snapshot, so schedule workers (SKEEP-005) may query the registry concurrently. */ public object KernelRegistry { - private val providers: MutableList = mutableListOf() + @kotlin.concurrent.Volatile + private var providers: List = emptyList() /** * Register a provider. Re-registering the same instance is a no-op. */ + @JvmSynchronized public fun register(provider: KernelProvider) { if (providers.any { it === provider }) return - providers.add(provider) - providers.sortByDescending { it.priority } + providers = (providers + provider).sortedByDescending { it.priority } } /** All registered providers, sorted by priority descending. */ - public fun providers(): List = providers.toList() + public fun providers(): List = providers /** Find a provider by name (case-insensitive), or `null`. */ public fun find(name: String): KernelProvider? = @@ -54,7 +55,8 @@ public object KernelRegistry { providers.filter { it.isAvailable() }.map { it.name } /** Test/diagnostic helper. Removes all registered providers. */ + @JvmSynchronized public fun clearForTesting() { - providers.clear() + providers = emptyList() } } diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q4KMatmulKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q4KMatmulKernel.kt index 3f89843e6..6e85f491e 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q4KMatmulKernel.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q4KMatmulKernel.kt @@ -57,4 +57,17 @@ public interface Q4KMatmulKernel { inputDim: Int, outputDim: Int, output: FloatArray, outputOffset: Int, ) + + /** + * Schedule-aware entry (SKEEP-005): a kernel that splits output rows across tasks takes the + * split from [schedule]. The default ignores it and runs the legacy method, so an + * implementation that has no parallel section is unaffected. + */ + public fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + schedule: sk.ainet.context.schedule.Schedule, + ): Unit = matmul(input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset) } diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q5KMatmulKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q5KMatmulKernel.kt index 54398e223..f4de49bbc 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q5KMatmulKernel.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q5KMatmulKernel.kt @@ -56,4 +56,13 @@ public interface Q5KMatmulKernel { inputDim: Int, outputDim: Int, output: FloatArray, outputOffset: Int, ) + + /** Schedule-aware entry (SKEEP-005); the default ignores [schedule] and runs the legacy method. */ + public fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + schedule: sk.ainet.context.schedule.Schedule, + ): Unit = matmul(input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset) } diff --git a/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt b/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt new file mode 100644 index 000000000..9b18fb9b9 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/JvmSynchronized.kt @@ -0,0 +1,3 @@ +package sk.ainet.backend.api.kernel + +internal actual typealias JvmSynchronized = kotlin.jvm.Synchronized diff --git a/skainet-backends/skainet-backend-api/src/jvmTest/kotlin/sk/ainet/backend/api/kernel/KernelDispatchConcurrencyTest.kt b/skainet-backends/skainet-backend-api/src/jvmTest/kotlin/sk/ainet/backend/api/kernel/KernelDispatchConcurrencyTest.kt new file mode 100644 index 000000000..1f05114bb --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jvmTest/kotlin/sk/ainet/backend/api/kernel/KernelDispatchConcurrencyTest.kt @@ -0,0 +1,86 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.types.FP32 +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.AfterTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +/** + * SKEEP-005 prerequisite: schedule workers dispatch concurrently, so the registries must survive + * concurrent reads while writes stay serialized — no `ConcurrentModificationException`, no + * double auto-install, no lost registration. + */ +@OptIn(ExperimentalMemoryApi::class) +class KernelDispatchConcurrencyTest { + + @AfterTest fun cleanup() { KernelDispatch.clearForTesting(); KernelRegistry.clearForTesting() } + + private class FakeKernel(override val name: String) : ViewKernel { + override val key: KernelKey = KernelKey( + op = "matmul", + operands = listOf(OperandKey.contiguous(Format.dense(FP32)), OperandKey.contiguous(Format.dense(FP32))), + ) + override fun run(inputs: List, out: TensorView) = Unit + } + + @Test + fun concurrentRegisterAndFindNeverThrowAndKeepEveryName() { + val threads = 16 + val perThread = 200 + val pool = Executors.newFixedThreadPool(threads) + val start = CountDownLatch(1) + val failures = AtomicInteger() + val key = FakeKernel("probe").key + repeat(threads) { t -> + pool.submit { + start.await() + try { + repeat(perThread) { i -> + KernelDispatch.register(FakeKernel("k$t-$i")) + KernelDispatch.find(key) + KernelDispatch.kernels().size + } + } catch (e: Throwable) { + failures.incrementAndGet() + e.printStackTrace() + } + } + } + start.countDown() + pool.shutdown() + assertTrue(pool.awaitTermination(60, TimeUnit.SECONDS), "workers must finish") + assertEquals(0, failures.get(), "no worker may fail") + assertEquals(threads * perThread, KernelDispatch.kernels().size, "every distinct name is kept") + assertNotNull(KernelDispatch.find(key)) + } + + @Test + fun ensureInstalledRunsAtMostOnceUnderContention() { + val installs = AtomicInteger() + val provider = object : KernelProvider { + override val name: String = "counting" + override val priority: Int = 1 + override fun isAvailable(): Boolean = true + override fun matmulFp32(): Fp32MatmulKernel? { installs.incrementAndGet(); return null } + } + KernelRegistry.register(provider) + val threads = 16 + val pool = Executors.newFixedThreadPool(threads) + val start = CountDownLatch(1) + repeat(threads) { pool.submit { start.await(); KernelDispatch.ensureInstalled() } } + start.countDown() + pool.shutdown() + assertTrue(pool.awaitTermination(60, TimeUnit.SECONDS)) + // KernelPacks.install() resolves the provider's kernels once per install; contention must not repeat it. + assertTrue(installs.get() <= 1 || KernelDispatch.kernels().isNotEmpty(), "auto-install ran without corrupting the table") + } +} diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index d146896e0..87dc62183 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -6,6 +6,8 @@ public final class sk/ainet/context/DirectCpuExecutionContext : sk/ainet/context public fun (Lsk/ainet/context/ExecutionStats;Lsk/ainet/context/Phase;Lsk/ainet/lang/nn/hooks/ForwardHooks;)V public fun (Lsk/ainet/context/ExecutionStats;Lsk/ainet/context/Phase;Lsk/ainet/lang/nn/hooks/ForwardHooks;Lsk/ainet/lang/tensor/data/TensorDataFactory;)V public synthetic fun (Lsk/ainet/context/ExecutionStats;Lsk/ainet/context/Phase;Lsk/ainet/lang/nn/hooks/ForwardHooks;Lsk/ainet/lang/tensor/data/TensorDataFactory;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lsk/ainet/context/ExecutionStats;Lsk/ainet/context/Phase;Lsk/ainet/lang/nn/hooks/ForwardHooks;Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/context/schedule/Schedule;)V + public synthetic fun (Lsk/ainet/context/ExecutionStats;Lsk/ainet/context/Phase;Lsk/ainet/lang/nn/hooks/ForwardHooks;Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/context/schedule/Schedule;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public static final fun create ()Lsk/ainet/context/DirectCpuExecutionContext; public static final fun create (Lsk/ainet/context/Phase;)Lsk/ainet/context/DirectCpuExecutionContext; public fun fromByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; @@ -22,6 +24,7 @@ public final class sk/ainet/context/DirectCpuExecutionContext : sk/ainet/context public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getPhase ()Lsk/ainet/context/Phase; + public fun getSchedule ()Lsk/ainet/context/schedule/Schedule; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; @@ -30,6 +33,7 @@ public final class sk/ainet/context/DirectCpuExecutionContext : sk/ainet/context public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -99,6 +103,7 @@ public final class sk/ainet/exec/kernel/PanamaVectorMatmulKernel : sk/ainet/back public final class sk/ainet/exec/kernel/PanamaVectorQ4KMatmulKernel : sk/ainet/backend/api/kernel/Q4KMatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/PanamaVectorQ4KMatmulKernel; public fun matmul ([FI[BIII[FI)V + public fun matmul ([FI[BIII[FILsk/ainet/context/schedule/Schedule;)V } public final class sk/ainet/exec/kernel/PanamaVectorQ4_0MatmulKernel : sk/ainet/backend/api/kernel/Q4_0MatmulKernel { @@ -119,6 +124,7 @@ public final class sk/ainet/exec/kernel/PanamaVectorQ5_1MatmulKernel : sk/ainet/ public final class sk/ainet/exec/kernel/PanamaVectorQ5_KMatmulKernel : sk/ainet/backend/api/kernel/Q5KMatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/PanamaVectorQ5_KMatmulKernel; public fun matmul ([FI[BIII[FI)V + public fun matmul ([FI[BIII[FILsk/ainet/context/schedule/Schedule;)V } public final class sk/ainet/exec/kernel/PanamaVectorQ6_KMatmulKernel : sk/ainet/backend/api/kernel/Q6KMatmulKernel { @@ -134,6 +140,7 @@ public final class sk/ainet/exec/kernel/PanamaVectorQ8_0MatmulKernel : sk/ainet/ public final class sk/ainet/exec/kernel/Q4_KQ8ActivationReferenceKernel : sk/ainet/backend/api/kernel/Q4KMatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/Q4_KQ8ActivationReferenceKernel; public fun matmul ([FI[BIII[FI)V + public fun matmul ([FI[BIII[FILsk/ainet/context/schedule/Schedule;)V } public final class sk/ainet/exec/kernel/Q6_KQ8ActivationReferenceKernel : sk/ainet/backend/api/kernel/Q6KMatmulKernel { @@ -202,6 +209,7 @@ public final class sk/ainet/exec/kernel/ScalarQ4_0MatmulKernel : sk/ainet/backen public final class sk/ainet/exec/kernel/ScalarQ4_KMatmulKernel : sk/ainet/backend/api/kernel/Q4KMatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/ScalarQ4_KMatmulKernel; public fun matmul ([FI[BIII[FI)V + public fun matmul ([FI[BIII[FILsk/ainet/context/schedule/Schedule;)V } public final class sk/ainet/exec/kernel/ScalarQ5_0MatmulKernel : sk/ainet/backend/api/kernel/Q5_0MatmulKernel { @@ -217,6 +225,7 @@ public final class sk/ainet/exec/kernel/ScalarQ5_1MatmulKernel : sk/ainet/backen public final class sk/ainet/exec/kernel/ScalarQ5_KMatmulKernel : sk/ainet/backend/api/kernel/Q5KMatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/ScalarQ5_KMatmulKernel; public fun matmul ([FI[BIII[FI)V + public fun matmul ([FI[BIII[FILsk/ainet/context/schedule/Schedule;)V } public final class sk/ainet/exec/kernel/ScalarQ6_KMatmulKernel : sk/ainet/backend/api/kernel/Q6KMatmulKernel { @@ -229,12 +238,46 @@ public final class sk/ainet/exec/kernel/ScalarQ8_0MatmulKernel : sk/ainet/backen public fun matmul ([FI[BIII[FI)V } +public class sk/ainet/exec/schedule/CoroutineSchedule : sk/ainet/context/schedule/Schedule { + public static final field Companion Lsk/ainet/exec/schedule/CoroutineSchedule$Companion; + public fun ()V + public fun (Lkotlinx/coroutines/CoroutineDispatcher;)V + public fun (Lkotlinx/coroutines/CoroutineDispatcher;I)V + public fun (Lkotlinx/coroutines/CoroutineDispatcher;ILsk/ainet/lang/memory/trace/TraceSink;)V + public fun (Lkotlinx/coroutines/CoroutineDispatcher;ILsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;)V + public synthetic fun (Lkotlinx/coroutines/CoroutineDispatcher;ILsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public static final fun dedicated ()Lsk/ainet/exec/schedule/DedicatedCoroutineSchedule; + public static final fun dedicated (I)Lsk/ainet/exec/schedule/DedicatedCoroutineSchedule; + public static final fun dedicated (ILsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/exec/schedule/DedicatedCoroutineSchedule; + public fun forEach (IILkotlin/jvm/functions/Function1;)V + public fun forRange (IILkotlin/jvm/functions/Function2;)V + public final fun getName ()Ljava/lang/String; + public final fun getParallelism ()I + public static final fun hardware (Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/exec/schedule/CoroutineSchedule; + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/exec/schedule/CoroutineSchedule$Companion { + public final fun dedicated ()Lsk/ainet/exec/schedule/DedicatedCoroutineSchedule; + public final fun dedicated (I)Lsk/ainet/exec/schedule/DedicatedCoroutineSchedule; + public final fun dedicated (ILsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/exec/schedule/DedicatedCoroutineSchedule; + public static synthetic fun dedicated$default (Lsk/ainet/exec/schedule/CoroutineSchedule$Companion;ILsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/exec/schedule/DedicatedCoroutineSchedule; + public final fun hardware (Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/exec/schedule/CoroutineSchedule; + public static synthetic fun hardware$default (Lsk/ainet/exec/schedule/CoroutineSchedule$Companion;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/exec/schedule/CoroutineSchedule; +} + +public final class sk/ainet/exec/schedule/DedicatedCoroutineSchedule : sk/ainet/exec/schedule/CoroutineSchedule, java/lang/AutoCloseable { + public fun close ()V +} + public final class sk/ainet/exec/tensor/ops/DefaultCpuOps : sk/ainet/exec/tensor/ops/DefaultCpuOpsBase { public fun (Lsk/ainet/lang/tensor/data/TensorDataFactory;)V + public fun (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/context/schedule/Schedule;)V } public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/ops/TensorOps { public fun (Lsk/ainet/lang/tensor/data/TensorDataFactory;)V + public fun (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/context/schedule/Schedule;)V public fun abs (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun add (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun addScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; @@ -264,6 +307,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o public fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; protected final fun getDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + protected final fun getSchedule ()Lsk/ainet/context/schedule/Schedule; protected final fun gradStateFrom ([Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/GradState; public fun indexSelect (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun leakyRelu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-backends/skainet-backend-cpu/src/androidMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.android.kt b/skainet-backends/skainet-backend-cpu/src/androidMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.android.kt index 92ceda596..c83aa2e30 100644 --- a/skainet-backends/skainet-backend-cpu/src/androidMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.android.kt +++ b/skainet-backends/skainet-backend-cpu/src/androidMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.android.kt @@ -7,7 +7,7 @@ import sk.ainet.exec.kernel.ScalarKernelProvider import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps -internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps { +internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory, sk.ainet.context.schedule.Schedule) -> TensorOps { // ART supports java.util.ServiceLoader, so Android discovers kernel // providers the same way the JVM does (#920): modules like // skainet-backend-jni-cpu ship a META-INF/services entry and are @@ -21,5 +21,9 @@ internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> Tenso // Scalar reference last: priority 0, always available — the floor the // registry cascades to when no accelerated provider carries a kernel. KernelRegistry.register(ScalarKernelProvider) - return { factory -> DefaultCpuOps(factory) } + return { factory, schedule -> DefaultCpuOps(factory, schedule) } } + + +internal actual fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule = + sk.ainet.context.schedule.Schedule.Sequential diff --git a/skainet-backends/skainet-backend-cpu/src/androidNativeMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.androidNative.kt b/skainet-backends/skainet-backend-cpu/src/androidNativeMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.androidNative.kt index 0d2b0fb1b..b65c50aa3 100644 --- a/skainet-backends/skainet-backend-cpu/src/androidNativeMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.androidNative.kt +++ b/skainet-backends/skainet-backend-cpu/src/androidNativeMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.androidNative.kt @@ -14,9 +14,13 @@ import sk.ainet.lang.tensor.ops.TensorOps * inheriting that silently is worse than the six duplicated lines below. The accelerated * kernels for these targets live in `skainet-backend-native-cpu`, not here. */ -internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps { +internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory, sk.ainet.context.schedule.Schedule) -> TensorOps { // Non-JVM has no ServiceLoader; register the scalar packed-quant kernels // (Q4_K/Q6_K/Q5_1/Q5_0/Q8_0/Q4_0) so DefaultCpuOpsBase can dispatch them. KernelRegistry.register(ScalarKernelProvider) - return { factory -> DefaultCpuOps(factory) } + return { factory, schedule -> DefaultCpuOps(factory, schedule) } } + + +internal actual fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule = + sk.ainet.context.schedule.Schedule.Sequential diff --git a/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt index f0051d841..369907c8e 100644 --- a/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt @@ -34,7 +34,10 @@ import sk.ainet.lang.types.FP32 @Backend(id = "apple", displayName = "Apple Accelerate") public class AccelerateCpuOps( dataFactory: TensorDataFactory, -) : DefaultCpuOpsBase(dataFactory) { + schedule: sk.ainet.context.schedule.Schedule, +) : DefaultCpuOpsBase(dataFactory, schedule) { + public constructor(dataFactory: TensorDataFactory) : this(dataFactory, sk.ainet.context.schedule.Schedule.Sequential) + // ── matmul ────────────────────────────────────────────────────────── diff --git a/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.apple.kt b/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.apple.kt index 153e7112c..2da492b98 100644 --- a/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.apple.kt +++ b/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.apple.kt @@ -5,10 +5,14 @@ import sk.ainet.exec.kernel.ScalarKernelProvider import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps -internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps { +internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory, sk.ainet.context.schedule.Schedule) -> TensorOps { println("[SKaiNET] Using Accelerate-backed CPU operations (ARM NEON + AMX)") // Accelerate overrides dense FP32 matmul; packed-quant weights still flow through // DefaultCpuOpsBase, so register the scalar packed kernels (no ServiceLoader on Native). KernelRegistry.register(ScalarKernelProvider) - return { factory -> AccelerateCpuOps(factory) } + return { factory, schedule -> AccelerateCpuOps(factory, schedule) } } + + +internal actual fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule = + sk.ainet.context.schedule.Schedule.Sequential diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/context/DirectCpuExecutionContext.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/context/DirectCpuExecutionContext.kt index 843eae4c2..7f402150f 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/context/DirectCpuExecutionContext.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/context/DirectCpuExecutionContext.kt @@ -3,7 +3,9 @@ package sk.ainet.context import sk.ainet.lang.tensor.data.DenseTensorDataFactory import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps +import sk.ainet.context.schedule.Schedule import sk.ainet.exec.tensor.ops.platformDefaultCpuOpsFactory +import sk.ainet.exec.tensor.ops.platformDefaultSchedule public class DirectCpuExecutionContext @kotlin.jvm.JvmOverloads constructor( override val executionStats: ExecutionStats = ExecutionStats(), @@ -12,6 +14,26 @@ public class DirectCpuExecutionContext @kotlin.jvm.JvmOverloads constructor( override val tensorDataFactory: TensorDataFactory = DenseTensorDataFactory(), ) : ExecutionContext { + /** + * SKEEP-005: a context whose ops run under [schedule]. The four-parameter constructor keeps + * its exact JVM signature (binary compatibility); this one adds the schedule. + */ + public constructor( + executionStats: ExecutionStats = ExecutionStats(), + phase: Phase = Phase.EVAL, + hooks: sk.ainet.lang.nn.hooks.ForwardHooks? = null, + tensorDataFactory: TensorDataFactory = DenseTensorDataFactory(), + schedule: Schedule, + ) : this(executionStats, phase, hooks, tensorDataFactory) { + this.scheduleOrNull = schedule + } + + private var scheduleOrNull: Schedule? = null + + /** How this context's ops map independent work onto cores; the platform default when not given. */ + override val schedule: Schedule + get() = scheduleOrNull ?: platformDefaultSchedule().also { scheduleOrNull = it } + public companion object { /** * Creates a new DirectCpuExecutionContext with sensible defaults. @@ -37,7 +59,7 @@ public class DirectCpuExecutionContext @kotlin.jvm.JvmOverloads constructor( // Cached: the getter used to build a fresh ops instance per access, which // re-ran the per-instance lazy kernel resolution and allocated on every // `ctx.ops` touch in the eager hot loop (#949). - private val cachedOps: TensorOps by lazy { opsFactory(tensorDataFactory) } + private val cachedOps: TensorOps by lazy { opsFactory(tensorDataFactory, schedule) } override val memoryInfo: MemoryInfo get() = _memoryInfo override val observers: ExecutionObserverRegistry @@ -51,5 +73,9 @@ public class DirectCpuExecutionContext @kotlin.jvm.JvmOverloads constructor( /** A sibling context whose cached ops allocate through [factory] (#1146). */ override fun withTensorDataFactory(factory: TensorDataFactory): ExecutionContext = - DirectCpuExecutionContext(executionStats, phase, _hooks, factory) + DirectCpuExecutionContext(executionStats, phase, _hooks, factory, schedule = schedule) + + /** A sibling context whose cached ops run under [schedule] (SKEEP-005). */ + override fun withSchedule(schedule: Schedule): ExecutionContext = + if (schedule === this.schedule) this else DirectCpuExecutionContext(executionStats, phase, _hooks, tensorDataFactory, schedule = schedule) } diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 7a638743b..b18cab061 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -47,9 +47,22 @@ import kotlin.math.pow import kotlin.math.sqrt import kotlin.reflect.KClass +/** + * Below this many multiply-adds a scaled-dot-product-attention call runs on the caller's thread + * regardless of the context's schedule: a coroutine region costs more than it saves (SKEEP-005). + * 8 heads × 64 keys × 128 dims × 2 ≈ 131k is well under it; a 512-token prefill chunk is far over. + */ +internal const val SDPA_PARALLEL_MIN_WORK: Long = 1L shl 20 + @Backend(id = "cpu", displayName = "CPU") @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#defaultcpuops") -public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory) : TensorOps { +public open class DefaultCpuOpsBase( + protected val dataFactory: TensorDataFactory, + /** How this ops instance maps independent work onto cores (SKEEP-005); [Schedule.Sequential] by default. */ + protected val schedule: sk.ainet.context.schedule.Schedule, +) : TensorOps { + public constructor(dataFactory: TensorDataFactory) : this(dataFactory, sk.ainet.context.schedule.Schedule.Sequential) + protected class CpuTensor( override val data: sk.ainet.lang.tensor.data.TensorData, @@ -3708,13 +3721,28 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val qBuf = query.data.copyToFloatArray() val kBuf = key.data.copyToFloatArray() val vBuf = value.data.copyToFloatArray() + // Hoisted out of the per-head loop: the mask is read-only for the whole call. + val maskBuf = mask?.data?.copyToFloatArray() val outBuf = FloatArray(batch * heads * seqQ * headDim) - for (b in 0 until batch) { - for (h in 0 until heads) { + // SKEEP-005: every (batch, head) pair is independent — private scores scratch, disjoint + // output rows — so the pairs are the schedule's units. The per-pair arithmetic and its + // order are exactly the sequential loop's, which keeps a scheduled run bit-identical. + // Tiny calls (decode steps on a handful of heads) stay on the caller: a region costs more + // than it saves below SDPA_PARALLEL_MIN_WORK multiply-adds. + val units = batch * heads + val workPerUnit = seqQ.toLong() * seqKV.toLong() * headDim.toLong() * 2L + val regionSchedule = if (workPerUnit * units < SDPA_PARALLEL_MIN_WORK) sk.ainet.context.schedule.Schedule.Sequential else schedule + regionSchedule.forRange(units, grain = 1) { unitStart, unitEnd -> + // Per-task scratch: a plain heap array, never the context's scratch pool or step slab + // (both single-threaded). Every entry is overwritten by the QK^T pass, so one array + // serves every unit of this task. + val scores = FloatArray(seqQ * seqKV) + for (unit in unitStart until unitEnd) { + val b = unit / heads + val h = unit % heads // Compute attention scores: Q @ K^T, then scale - val scores = FloatArray(seqQ * seqKV) for (qi in 0 until seqQ) { for (ki in 0 until seqKV) { var dot = 0f @@ -3742,8 +3770,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory } // Apply external mask if provided - if (mask != null) { - val maskBuf = mask.data.copyToFloatArray() + if (maskBuf != null) { for (i in scores.indices) { scores[i] += maskBuf[i % maskBuf.size] } @@ -3792,4 +3819,9 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory } -public class DefaultCpuOps(dataFactory: TensorDataFactory) : DefaultCpuOpsBase(dataFactory) +public class DefaultCpuOps( + dataFactory: TensorDataFactory, + schedule: sk.ainet.context.schedule.Schedule, +) : DefaultCpuOpsBase(dataFactory, schedule) { + public constructor(dataFactory: TensorDataFactory) : this(dataFactory, sk.ainet.context.schedule.Schedule.Sequential) +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.kt index bbbd6cec4..29250be5d 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.kt @@ -3,4 +3,12 @@ package sk.ainet.exec.tensor.ops import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps -internal expect fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps +/** Ops constructor for this platform; the second argument is the [sk.ainet.context.schedule.Schedule] the ops run under (SKEEP-005). */ +internal expect fun platformDefaultCpuOpsFactory(): (TensorDataFactory, sk.ainet.context.schedule.Schedule) -> TensorOps + +/** + * The schedule a `DirectCpuExecutionContext` runs under when none is given: core-count coroutines + * on the JVM (today's matmul parallelism, now visible), [sk.ainet.context.schedule.Schedule.Sequential] + * everywhere else. + */ +internal expect fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/schedule/SdpaScheduleParityTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/schedule/SdpaScheduleParityTest.kt new file mode 100644 index 000000000..328d5674d --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/schedule/SdpaScheduleParityTest.kt @@ -0,0 +1,73 @@ +package sk.ainet.exec.schedule + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * SKEEP-005: `scaledDotProductAttention` is the engine's first scheduled op. A schedule that + * hands out ranges in any order and any count must produce the sequential result bit for bit, + * and the op must actually route its (batch, head) units through the schedule. + */ +class SdpaScheduleParityTest { + + /** Records every range, hands them out in reverse order, and runs them inline. */ + private class ReversingSchedule(override val parallelism: Int) : Schedule { + val ranges = mutableListOf>() + override val name: String get() = "reversing($parallelism)" + override fun forRange(n: Int, grain: Int, body: (Int, Int) -> Unit) { + val tasks = Schedule.tasksFor(n, grain, parallelism) + if (tasks == 0) return + val chunk = Schedule.chunkFor(n, tasks) + val bounds = (0 until n step chunk).map { s -> s to minOf(s + chunk, n) } + for ((s, e) in bounds.reversed()) { ranges += s to e; body(s, e) } + } + } + + private fun sdpa(ctx: ExecutionContext, batch: Int, heads: Int, seqQ: Int, seqKV: Int, headDim: Int, causal: Boolean, withMask: Boolean, scale: Float): FloatArray { + fun fill(size: Int, seed: Int) = FloatArray(size) { i -> (((i * 31 + seed * 17) % 23) - 11) / 11f } + val q = ctx.fromFloatArray(Shape(batch, heads, seqQ, headDim), FP32::class, fill(batch * heads * seqQ * headDim, 1)) + val k = ctx.fromFloatArray(Shape(batch, heads, seqKV, headDim), FP32::class, fill(batch * heads * seqKV * headDim, 2)) + val v = ctx.fromFloatArray(Shape(batch, heads, seqKV, headDim), FP32::class, fill(batch * heads * seqKV * headDim, 3)) + val mask = if (withMask) ctx.fromFloatArray(Shape(batch, 1, seqQ, seqKV), FP32::class, FloatArray(batch * seqQ * seqKV) { if (it % 7 == 0) -1e30f else 0f }) else null + return ctx.ops.scaledDotProductAttention(q, k, v, mask, scale, causal).data.copyToFloatArray() + } + + @Test + fun scheduledAttentionIsBitIdenticalToSequential() { + val sequential = DirectCpuExecutionContext(schedule = Schedule.Sequential) + for ((seqQ, seqKV) in listOf(16 to 16, 8 to 64, 1 to 128, 64 to 64)) { + for (causal in listOf(false, true)) for (withMask in listOf(false, true)) for (scale in listOf(0f, 0.25f)) { + val reversing = ReversingSchedule(parallelism = 3) + val scheduled = DirectCpuExecutionContext(schedule = reversing) + val expected = sdpa(sequential, 2, 8, seqQ, seqKV, 32, causal, withMask, scale) + val actual = sdpa(scheduled, 2, 8, seqQ, seqKV, 32, causal, withMask, scale) + assertContentEquals(expected, actual, "seqQ=$seqQ seqKV=$seqKV causal=$causal mask=$withMask scale=$scale") + } + } + } + + @Test + fun largeCallsRouteTheirUnitsThroughTheSchedule() { + val reversing = ReversingSchedule(parallelism = 3) + val ctx = DirectCpuExecutionContext(schedule = reversing) + sdpa(ctx, 2, 8, 64, 64, 32, causal = true, withMask = false, scale = 0f) // 2*8*64*64*32*2 = 8.4M > threshold + assertEquals(3, reversing.ranges.size, "16 units on parallelism 3 → 3 ranges") + assertEquals(16, reversing.ranges.sumOf { (s, e) -> e - s }, "ranges cover every (batch, head) unit") + assertTrue(reversing.ranges.first().first > reversing.ranges.last().first, "ranges were handed out in reverse and still produced the right result") + } + + @Test + fun tinyCallsStayOnTheCaller() { + val reversing = ReversingSchedule(parallelism = 3) + val ctx = DirectCpuExecutionContext(schedule = reversing) + sdpa(ctx, 1, 8, 1, 64, 128, causal = false, withMask = false, scale = 0f) // a decode step: 131k MACs + assertTrue(reversing.ranges.isEmpty(), "below SDPA_PARALLEL_MIN_WORK no region is opened") + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jsMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.js.kt b/skainet-backends/skainet-backend-cpu/src/jsMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.js.kt index bbd668256..18a9f6cf5 100644 --- a/skainet-backends/skainet-backend-cpu/src/jsMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.js.kt +++ b/skainet-backends/skainet-backend-cpu/src/jsMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.js.kt @@ -3,7 +3,11 @@ package sk.ainet.exec.tensor.ops import sk.ainet.backend.api.kernel.KernelRegistry import sk.ainet.exec.kernel.ScalarKernelProvider -internal actual fun platformDefaultCpuOpsFactory(): (sk.ainet.lang.tensor.data.TensorDataFactory) -> sk.ainet.lang.tensor.ops.TensorOps { +internal actual fun platformDefaultCpuOpsFactory(): (sk.ainet.lang.tensor.data.TensorDataFactory, sk.ainet.context.schedule.Schedule) -> sk.ainet.lang.tensor.ops.TensorOps { KernelRegistry.register(ScalarKernelProvider) - return { factory -> DefaultCpuOps(factory) } + return { factory, schedule -> DefaultCpuOps(factory, schedule) } } + + +internal actual fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule = + sk.ainet.context.schedule.Schedule.Sequential diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ4KMatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ4KMatmulKernel.kt index 584bd3081..4db3b21d2 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ4KMatmulKernel.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ4KMatmulKernel.kt @@ -1,5 +1,8 @@ package sk.ainet.exec.kernel +import sk.ainet.context.schedule.Schedule +import sk.ainet.exec.schedule.CoroutineSchedule + import jdk.incubator.vector.ByteVector import jdk.incubator.vector.FloatVector import jdk.incubator.vector.VectorOperators @@ -58,6 +61,14 @@ public object PanamaVectorQ4KMatmulKernel : Q4KMatmulKernel { weight: ByteArray, weightByteOffset: Int, inputDim: Int, outputDim: Int, output: FloatArray, outputOffset: Int, + ): Unit = matmul(input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset, CoroutineSchedule.hardware()) + + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + schedule: Schedule, ) { require(inputDim % BLOCK_SIZE == 0) { "PanamaVectorQ4KMatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim" @@ -65,7 +76,7 @@ public object PanamaVectorQ4KMatmulKernel : Q4KMatmulKernel { if (outputDim == 0 || inputDim == 0) return val blocksPerInputDim = inputDim / BLOCK_SIZE - parallelChunks(outputDim) { startO, endO -> + parallelChunks(outputDim, schedule) { startO, endO -> // Per-task scratch — must not be shared across worker threads. val scaleIdx = IntArray(SUB_BLOCKS_PER_BLOCK) val minIdx = IntArray(SUB_BLOCKS_PER_BLOCK) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ5_KMatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ5_KMatmulKernel.kt index 5302a4710..7c03e3f1e 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ5_KMatmulKernel.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ5_KMatmulKernel.kt @@ -1,5 +1,8 @@ package sk.ainet.exec.kernel +import sk.ainet.context.schedule.Schedule +import sk.ainet.exec.schedule.CoroutineSchedule + import jdk.incubator.vector.ByteVector import jdk.incubator.vector.FloatVector import jdk.incubator.vector.VectorOperators @@ -41,6 +44,14 @@ public object PanamaVectorQ5_KMatmulKernel : Q5KMatmulKernel { weight: ByteArray, weightByteOffset: Int, inputDim: Int, outputDim: Int, output: FloatArray, outputOffset: Int, + ): Unit = matmul(input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset, CoroutineSchedule.hardware()) + + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + schedule: Schedule, ) { require(inputDim % BLOCK_SIZE == 0) { "PanamaVectorQ5_KMatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim" @@ -48,7 +59,7 @@ public object PanamaVectorQ5_KMatmulKernel : Q5KMatmulKernel { if (outputDim == 0 || inputDim == 0) return val blocksPerInputDim = inputDim / BLOCK_SIZE - parallelChunks(outputDim) { startO, endO -> + parallelChunks(outputDim, schedule) { startO, endO -> val scaleIdx = IntArray(SUB_BLOCKS_PER_BLOCK) val minIdx = IntArray(SUB_BLOCKS_PER_BLOCK) for (o in startO until endO) { diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt new file mode 100644 index 000000000..037b7d617 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt @@ -0,0 +1,128 @@ +package sk.ainet.exec.schedule + +import kotlinx.coroutines.CoroutineDispatcher +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.asCoroutineDispatcher +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import sk.ainet.context.schedule.Schedule +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceClock +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.TraceSink +import java.util.concurrent.ExecutorService +import java.util.concurrent.Executors +import java.util.concurrent.atomic.AtomicInteger + +/** + * The JVM [Schedule] (SKEEP-005): a region is one `coroutineScope` — structured concurrency — + * whose children run on [dispatcher] while the *calling thread runs the first chunk itself*, so + * no core sits idle blocked on the join. The scope closes only after every child finished or was + * cancelled, which gives the contract its guarantees: the first failure cancels the siblings and + * is rethrown, no task outlives [forRange], and every task's writes happen-before the return. + * + * A region reached from inside another region (a body that calls into a parallel op despite the + * contract) runs inline: the dispatcher never waits on itself, and a nested `runBlocking` on + * `Dispatchers.Default` cannot starve the pool. Callers that already run on `Dispatchers.Default` + * and want true isolation use [dedicated]. + */ +@OptIn(ExperimentalMemoryApi::class) +public open class CoroutineSchedule @JvmOverloads constructor( + private val dispatcher: CoroutineDispatcher = Dispatchers.Default, + final override val parallelism: Int = Runtime.getRuntime().availableProcessors(), + private val sink: TraceSink = NoopTraceSink, + private val label: String = "coroutines", +) : Schedule { + + init { + require(parallelism >= 1) { "CoroutineSchedule: parallelism must be >= 1, got $parallelism" } + } + + final override val name: String = "$label($parallelism)" + + override fun forRange(n: Int, grain: Int, body: (start: Int, end: Int) -> Unit) { + val tasks = Schedule.tasksFor(n, grain, parallelism) + if (tasks == 0) return + if (tasks == 1 || inRegion.get() == true) { + body(0, n) + return + } + val chunk = Schedule.chunkFor(n, tasks) + val started = if (sink.isEnabled) TraceClock.nowNanos() else 0L + inRegion.set(true) + try { + runBlocking { + coroutineScope { + var start = chunk + while (start < n) { + val s = start + val e = minOf(start + chunk, n) + launch(dispatcher) { inRegion(s, e, body) } + start = e + } + body(0, minOf(chunk, n)) + } + } + } finally { + inRegion.set(false) + } + if (sink.isEnabled) { + val now = TraceClock.nowNanos() + sink.emit(TraceEvent.ScheduleRegion(op = "forRange", schedule = name, elements = n, tasks = tasks, durationNanos = now - started, timeNanos = now)) + } + } + + private fun inRegion(start: Int, end: Int, body: (Int, Int) -> Unit) { + val previous = inRegion.get() + inRegion.set(true) + try { + body(start, end) + } finally { + inRegion.set(previous) + } + } + + override fun toString(): String = name + + public companion object { + /** Set on any thread currently executing a region body, so a nested region runs inline. */ + private val inRegion: ThreadLocal = ThreadLocal() + + /** Core-count parallelism on `Dispatchers.Default` — the platform default schedule on the JVM. */ + @JvmStatic + public fun hardware(sink: TraceSink = NoopTraceSink): CoroutineSchedule = + CoroutineSchedule(Dispatchers.Default, Runtime.getRuntime().availableProcessors(), sink) + + /** + * A schedule with its own daemon pool of `parallelism - 1` workers (the caller is the last + * worker), for code that already runs on `Dispatchers.Default`. Close it when done. + */ + @JvmStatic + @JvmOverloads + public fun dedicated( + parallelism: Int = Runtime.getRuntime().availableProcessors(), + sink: TraceSink = NoopTraceSink, + ): DedicatedCoroutineSchedule { + require(parallelism >= 1) { "dedicated: parallelism must be >= 1, got $parallelism" } + val counter = AtomicInteger() + val executor = Executors.newFixedThreadPool(maxOf(1, parallelism - 1)) { r -> + Thread(r, "skainet-schedule-${counter.incrementAndGet()}").apply { isDaemon = true } + } + return DedicatedCoroutineSchedule(executor, parallelism, sink) + } + } +} + +/** [CoroutineSchedule] over an owned thread pool; [close] shuts the pool down. */ +@OptIn(ExperimentalMemoryApi::class) +public class DedicatedCoroutineSchedule internal constructor( + private val executor: ExecutorService, + parallelism: Int, + sink: TraceSink, +) : CoroutineSchedule(executor.asCoroutineDispatcher(), parallelism, sink, label = "dedicated"), AutoCloseable { + override fun close() { + executor.shutdown() + } +} 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 a4b3fcdbf..748ce9608 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 @@ -44,7 +44,8 @@ import kotlin.math.max internal class DefaultCpuOpsJvm( dataFactory: TensorDataFactory, -) : DefaultCpuOpsBase(dataFactory) { + schedule: sk.ainet.context.schedule.Schedule = sk.ainet.context.schedule.Schedule.Sequential, +) : DefaultCpuOpsBase(dataFactory, schedule) { private val floatSpecies: VectorSpecies = FloatVector.SPECIES_PREFERRED @@ -668,6 +669,7 @@ internal class DefaultCpuOpsJvm( outputDim, outBuffer, batch * outputDim, + schedule, ) } } @@ -701,6 +703,7 @@ internal class DefaultCpuOpsJvm( bData.packedData, 0, inputDim, outputDim, outBuffer, batch * outputDim, + schedule, ) } else { JvmQuantizedVectorKernels.matmulQ4_KVec( @@ -710,6 +713,7 @@ internal class DefaultCpuOpsJvm( outputDim, outBuffer, batch * outputDim, + schedule, ) } } @@ -1074,6 +1078,7 @@ internal class DefaultCpuOpsJvm( aMemSeg.segment, aMemSeg.segmentByteOffset, bMemSeg.segment, bMemSeg.segmentByteOffset, result.segment, result.segmentByteOffset, + schedule = schedule, ) } else { JvmVectorKernels.matmulFloatMemSeg( diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmQuantizedVectorKernels.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmQuantizedVectorKernels.kt index e63864f6c..d6342db37 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmQuantizedVectorKernels.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmQuantizedVectorKernels.kt @@ -1,5 +1,8 @@ package sk.ainet.exec.tensor.ops +import sk.ainet.context.schedule.Schedule +import sk.ainet.exec.schedule.CoroutineSchedule + import jdk.incubator.vector.ByteVector import jdk.incubator.vector.FloatVector import jdk.incubator.vector.VectorOperators @@ -147,7 +150,8 @@ internal object JvmQuantizedVectorKernels { inputDim: Int, outputDim: Int, output: FloatArray, - outputOffset: Int = 0 + outputOffset: Int = 0, + schedule: Schedule = CoroutineSchedule.hardware(), ) { val blockSize = 32 val bytesPerBlock = 34 // 2 scale + 32 codes @@ -191,14 +195,15 @@ internal object JvmQuantizedVectorKernels { inputDim: Int, outputDim: Int, output: FloatArray, - outputOffset: Int = 0 + outputOffset: Int = 0, + schedule: Schedule = CoroutineSchedule.hardware(), ) { val blockSize = 256 val subBlockSize = 32 val bytesPerBlock = 144 // 2 d + 2 dMin + 12 scales + 128 codes val blocksPerInputDim = (inputDim + blockSize - 1) / blockSize - parallelChunks(outputDim) { startO, endO -> + parallelChunks(outputDim, schedule) { startO, endO -> // Each task owns its own scratch arrays to avoid cross-thread contention. val codeBuf = FloatArray(subBlockSize) val scaleIdxBuf = IntArray(8) @@ -307,7 +312,8 @@ internal object JvmQuantizedVectorKernels { inputDim: Int, outputDim: Int, output: FloatArray, - outputOffset: Int = 0 + outputOffset: Int = 0, + schedule: Schedule = CoroutineSchedule.hardware(), ) { val blockSize = 256 val bytesPerBlock = 210 @@ -315,7 +321,7 @@ internal object JvmQuantizedVectorKernels { val floatStep = floatSpecies.length() val loopBound = floatSpecies.loopBound(blockSize) - parallelChunks(outputDim) { startO, endO -> + parallelChunks(outputDim, schedule) { startO, endO -> // Per-task scratch — must not be shared across worker threads. val scratch = FloatArray(blockSize) for (o in startO until endO) { diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt index d252c1fbb..37c825399 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt @@ -1,5 +1,8 @@ package sk.ainet.exec.tensor.ops +import sk.ainet.context.schedule.Schedule +import sk.ainet.exec.schedule.CoroutineSchedule + import jdk.incubator.vector.FloatVector import jdk.incubator.vector.VectorOperators import jdk.incubator.vector.VectorSpecies @@ -937,6 +940,7 @@ internal object JvmVectorKernels { tileM: Int = 8, tileN: Int = 8, tileK: Int = 128, + schedule: Schedule = CoroutineSchedule.hardware(), ) { val floatLayout = java.lang.foreign.ValueLayout.JAVA_FLOAT.withOrder(BYTE_ORDER) val floatBytes = Float.SIZE_BYTES.toLong() @@ -972,7 +976,7 @@ internal object JvmVectorKernels { // Parallelize over m (independent rows of the result). Each task owns // a contiguous mm range and writes to its own slice of `r`. Tiling on // n and k stays for cache locality. - parallelChunks(m) { mStart, mEnd -> + parallelChunks(m, schedule) { mStart, mEnd -> for (bn in 0 until nBlocks) { val nStart = bn * tileN val nEnd = minOf(nStart + tileN, n) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/ParallelFor.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/ParallelFor.kt index 25e5a4c6d..980559a45 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/ParallelFor.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/ParallelFor.kt @@ -1,51 +1,37 @@ package sk.ainet.exec.tensor.ops -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.coroutineScope -import kotlinx.coroutines.launch -import kotlinx.coroutines.runBlocking +import sk.ainet.context.schedule.Schedule /** - * Number of CPU cores available for kernel-level parallelism. - * JVM-only for now; promote to expect/actual in commonMain when native/JS - * backends gain SIMD kernels too. + * Logical core count — the parallelism the platform default [Schedule] uses on the JVM. + * JVM-only for now; promote to expect/actual in commonMain when native/JS backends gain SIMD + * kernels too. */ internal val defaultParallelism: Int = Runtime.getRuntime().availableProcessors() /** - * Threshold below which a matmul stays single-threaded — coroutine launch - * overhead dominates for tiny outputDim. Tuned empirically: chunks below - * this size are not worth dispatching. + * Below this many output rows a matmul runs on the calling thread: the per-region overhead + * (coroutine launch, join) outweighs the work. */ internal const val PARALLEL_MATMUL_MIN_OUTPUT: Int = 256 /** - * Run [block] over disjoint chunks of [outputDim] in parallel. - * Below [PARALLEL_MATMUL_MIN_OUTPUT] runs sequentially in the calling thread. + * Run [block] over disjoint `[start, end)` ranges of `[0, outputDim)` on [schedule] (SKEEP-005). * - * Each task receives the half-open range `[start, end)` it owns. - * Use [Dispatchers.Default] which is sized to CPU count on JVM. + * The threshold above keeps tiny matmuls sequential; everything else — task count, threads, + * structure — is the schedule's decision, so an `ExecutionContext.withSchedule(Sequential)` + * makes the same kernel run single-threaded and a `CoroutineSchedule` spreads it across cores. + * Bodies follow the [Schedule.forRange] contract: disjoint output slices, no allocation through + * a context, no nested regions. */ internal inline fun parallelChunks( outputDim: Int, - crossinline block: (start: Int, end: Int) -> Unit + schedule: Schedule, + crossinline block: (start: Int, end: Int) -> Unit, ) { if (outputDim < PARALLEL_MATMUL_MIN_OUTPUT) { block(0, outputDim) return } - val chunks = defaultParallelism - val chunkSize = (outputDim + chunks - 1) / chunks - runBlocking(Dispatchers.Default) { - coroutineScope { - var start = 0 - while (start < outputDim) { - val end = minOf(start + chunkSize, outputDim) - val s = start - val e = end - launch { block(s, e) } - start = end - } - } - } + schedule.forRange(outputDim, grain = 1) { s, e -> block(s, e) } } diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.jvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.jvm.kt index 7ffbeb038..d53c6545d 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.jvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.jvm.kt @@ -3,7 +3,7 @@ package sk.ainet.exec.tensor.ops import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps -internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps { +internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory, sk.ainet.context.schedule.Schedule) -> TensorOps { val jdkOk = isJdk21Plus() val vectorAvailable = jdkOk && isVectorApiAvailable() val useVector = (JvmCpuBackendConfig.vectorEnabled && vectorAvailable) @@ -21,10 +21,10 @@ internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> Tenso } return if (useVector) { - { factory: TensorDataFactory -> DefaultCpuOpsJvm(factory) } + { factory: TensorDataFactory, schedule: sk.ainet.context.schedule.Schedule -> DefaultCpuOpsJvm(factory, schedule) } } else { // Note: BLAS acceleration not yet implemented; falling back to DefaultCpuOps - { factory: TensorDataFactory -> DefaultCpuOps(factory) } + { factory: TensorDataFactory, schedule: sk.ainet.context.schedule.Schedule -> DefaultCpuOps(factory, schedule) } } } @@ -55,3 +55,7 @@ private fun isJdk21Plus(): Boolean { major >= 21 } } + + +internal actual fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule = + sk.ainet.exec.schedule.CoroutineSchedule.hardware() diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/context/DirectCpuExecutionContextScheduleTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/context/DirectCpuExecutionContextScheduleTest.kt new file mode 100644 index 000000000..f63b59cb0 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/context/DirectCpuExecutionContextScheduleTest.kt @@ -0,0 +1,54 @@ +package sk.ainet.context + +import sk.ainet.context.schedule.Schedule +import sk.ainet.exec.schedule.CoroutineSchedule +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertIs +import kotlin.test.assertNotSame +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** SKEEP-005: the JVM context defaults to the hardware schedule and rebuilds its ops under another one. */ +class DirectCpuExecutionContextScheduleTest { + + @Test + fun jvmContextDefaultsToTheHardwareCoroutineSchedule() { + val ctx = DirectCpuExecutionContext() + assertIs(ctx.schedule) + assertTrue(ctx.schedule.parallelism >= 1) + } + + @Test + fun withScheduleRebuildsOpsAndKeepsResultsIdentical() { + val ctx = DirectCpuExecutionContext() + val sequential = ctx.withSchedule(Schedule.Sequential) + assertSame(Schedule.Sequential, sequential.schedule) + assertNotSame(ctx.ops, sequential.ops, "a different schedule means a different ops instance") + assertSame(ctx, ctx.withSchedule(ctx.schedule), "requesting the current schedule is a no-op") + + val m = 300; val k = 64; val n = 512 + val a = FloatArray(m * k) { ((it * 7) % 13 - 6) / 7f } + val b = FloatArray(k * n) { ((it * 5) % 11 - 5) / 5f } + fun matmul(c: ExecutionContext): FloatArray = c.ops.matmul( + c.fromFloatArray(Shape(m, k), FP32::class, a), + c.fromFloatArray(Shape(k, n), FP32::class, b), + ).data.copyToFloatArray() + // The Panama kernel's lane reduction changes order once the JIT intrinsifies it, so the + // very first call in a JVM can differ by an ULP from later ones regardless of schedule. + // Warm both contexts up before comparing; the schedule itself must never change a result. + val parallel = ctx.withSchedule(CoroutineSchedule(parallelism = 4)) + repeat(20) { matmul(sequential); matmul(parallel) } + assertContentEquals(matmul(sequential), matmul(parallel), "a schedule never changes a result") + } + + @Test + fun scheduleSurvivesForwardScope() { + val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential) + ctx.forwardScope(slabFloats = 64) { scoped, _ -> + assertSame(Schedule.Sequential, scoped.schedule) + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt new file mode 100644 index 000000000..995efbb0c --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt @@ -0,0 +1,143 @@ +package sk.ainet.exec.schedule + +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import sk.ainet.context.schedule.Schedule +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import java.util.BitSet +import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** SKEEP-005: the JVM schedule honours the [Schedule.forRange] contract. */ +@OptIn(ExperimentalMemoryApi::class) +class CoroutineScheduleTest { + + private fun coverage(schedule: Schedule, n: Int, grain: Int): Pair { + val seen = BitSet(n) + val tasks = AtomicInteger() + schedule.forRange(n, grain) { s, e -> + tasks.incrementAndGet() + synchronized(seen) { + for (i in s until e) { + assertTrue(!seen[i], "index $i visited twice") + seen.set(i) + } + } + } + return seen to tasks.get() + } + + @Test + fun everyIndexIsVisitedExactlyOnceForAnyShape() { + val schedule = CoroutineSchedule(parallelism = 4) + for (n in listOf(0, 1, 7, 255, 1000, 4097)) { + for (grain in listOf(1, 17, 10_000)) { + val (seen, tasks) = coverage(schedule, n, grain) + assertEquals(n, seen.cardinality(), "n=$n grain=$grain") + assertEquals(Schedule.tasksFor(n, grain, 4), tasks, "n=$n grain=$grain task count") + } + } + } + + @Test + fun parallelismOneIsSequentialOnTheCallerThread() { + val schedule = CoroutineSchedule(parallelism = 1) + val caller = Thread.currentThread() + val ranges = mutableListOf>() + schedule.forRange(100) { s, e -> + assertSame(caller, Thread.currentThread()) + ranges += s to e + } + assertEquals(listOf(0 to 100), ranges) + assertEquals("coroutines(1)", schedule.name) + } + + @Test + fun callerThreadRunsTheFirstChunkAndWorkersRunTheRest() { + val schedule = CoroutineSchedule(parallelism = 4) + val caller = Thread.currentThread() + val onCaller = AtomicInteger() + val elsewhere = AtomicInteger() + schedule.forRange(4000, grain = 1) { _, _ -> + if (Thread.currentThread() === caller) onCaller.incrementAndGet() else elsewhere.incrementAndGet() + } + assertEquals(1, onCaller.get(), "the caller runs exactly one chunk itself") + assertEquals(3, elsewhere.get(), "the other chunks run on the dispatcher") + } + + @Test + fun aFailingTaskCancelsSiblingsAndRethrowsAfterTheyFinish() { + val schedule = CoroutineSchedule(parallelism = 4) + val running = AtomicInteger() + val finished = AtomicInteger() + val boom = assertFailsWith { + schedule.forRange(4, grain = 1) { s, _ -> + running.incrementAndGet() + try { + if (s == 2) error("task $s failed") + Thread.sleep(20) + } finally { + finished.incrementAndGet() + } + } + } + assertEquals("task 2 failed", boom.message) + assertEquals(running.get(), finished.get(), "no task may still be running when forRange returns") + } + + @Test + fun aNestedRegionRunsInlineOnTheWorkerThread() { + val schedule = CoroutineSchedule(parallelism = 4) + val nestedTasks = AtomicInteger() + schedule.forRange(4, grain = 1) { _, _ -> + val worker = Thread.currentThread() + schedule.forRange(4000, grain = 1) { _, _ -> + nestedTasks.incrementAndGet() + assertSame(worker, Thread.currentThread(), "a nested region must not fork") + } + } + assertEquals(4, nestedTasks.get(), "each outer task ran its nested region as one inline chunk") + } + + @Test + fun aRegionStartedFromADefaultDispatcherWorkerCompletes() { + val schedule = CoroutineSchedule(parallelism = Runtime.getRuntime().availableProcessors()) + val total = AtomicInteger() + runBlocking(Dispatchers.Default) { + val jobs = List(4) { + launch { schedule.forRange(4096, grain = 1) { s, e -> total.addAndGet(e - s) } } + } + jobs.forEach { it.join() } + } + assertEquals(4 * 4096, total.get()) + } + + @Test + fun dedicatedScheduleOwnsItsPoolAndCloses() { + CoroutineSchedule.dedicated(parallelism = 3).use { schedule -> + assertEquals("dedicated(3)", schedule.name) + val (seen, _) = coverage(schedule, 999, 1) + assertEquals(999, seen.cardinality()) + } + } + + @Test + fun regionsAreReportedToTheSink() { + val sink = RecordingTraceSink() + val schedule = CoroutineSchedule(parallelism = 2, sink = sink) + schedule.forRange(10, grain = 1) { _, _ -> } + schedule.forRange(1, grain = 1) { _, _ -> } // single task: inline, no region event + val regions = sink.eventsOf() + assertEquals(1, regions.size) + assertEquals(10, regions.single().elements) + assertEquals(2, regions.single().tasks) + assertEquals("coroutines(2)", regions.single().schedule) + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/SdpaCoroutineParityTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/SdpaCoroutineParityTest.kt new file mode 100644 index 000000000..a622b8427 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/SdpaCoroutineParityTest.kt @@ -0,0 +1,30 @@ +package sk.ainet.exec.schedule + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertContentEquals + +/** The real JVM schedule on the same op: threads, not just reordering. */ +class SdpaCoroutineParityTest { + private fun sdpa(ctx: ExecutionContext, batch: Int, heads: Int, seqQ: Int, seqKV: Int, headDim: Int): FloatArray { + fun fill(size: Int, seed: Int) = FloatArray(size) { i -> (((i * 31 + seed * 17) % 23) - 11) / 11f } + val q = ctx.fromFloatArray(Shape(batch, heads, seqQ, headDim), FP32::class, fill(batch * heads * seqQ * headDim, 1)) + val k = ctx.fromFloatArray(Shape(batch, heads, seqKV, headDim), FP32::class, fill(batch * heads * seqKV * headDim, 2)) + val v = ctx.fromFloatArray(Shape(batch, heads, seqKV, headDim), FP32::class, fill(batch * heads * seqKV * headDim, 3)) + return ctx.ops.scaledDotProductAttention(q, k, v, null, 0f, true).data.copyToFloatArray() + } + + @Test + fun coroutineScheduleMatchesSequentialBitForBit() { + val sequential = DirectCpuExecutionContext(schedule = Schedule.Sequential) + val parallel = DirectCpuExecutionContext(schedule = CoroutineSchedule(parallelism = 4)) + repeat(3) { + assertContentEquals(sdpa(sequential, 1, 24, 64, 512, 128), sdpa(parallel, 1, 24, 64, 512, 128)) + assertContentEquals(sdpa(sequential, 2, 8, 128, 128, 64), sdpa(parallel, 2, 8, 128, 128, 64)) + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmElementwiseTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmElementwiseTest.kt index fd66e3ad6..670579e98 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmElementwiseTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmElementwiseTest.kt @@ -71,7 +71,7 @@ class DefaultCpuOpsJvmElementwiseTest { // Vector OFF -> platform factory should return DefaultCpuOps System.setProperty("skainet.cpu.vector.enabled", "false") - val scalarOps = platformDefaultCpuOpsFactory()(dataFactory) + val scalarOps = platformDefaultCpuOpsFactory()(dataFactory, sk.ainet.context.schedule.Schedule.Sequential) val offAdd = scalarOps.add(tA, tB) val offRelu = scalarOps.relu(tA) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmReductionsTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmReductionsTest.kt index 4ead2b572..c7481bfd5 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmReductionsTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmReductionsTest.kt @@ -58,7 +58,7 @@ class DefaultCpuOpsJvmReductionsTest { // Run with vector OFF by going through platform factory which will create DefaultCpuOps System.setProperty("skainet.cpu.vector.enabled", "false") val factory = platformDefaultCpuOpsFactory() - val scalarOps = factory(dataFactory) + val scalarOps = factory(dataFactory, sk.ainet.context.schedule.Schedule.Sequential) val rOffSum = scalarOps.sum(t, null) val rOffMean = scalarOps.mean(t, null) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactoryJvmTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactoryJvmTest.kt index 62b840df0..c8f88b459 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactoryJvmTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactoryJvmTest.kt @@ -18,7 +18,7 @@ class PlatformCpuOpsFactoryJvmTest { fun returnsJvmOpsWhenVectorFlagEnabled() { System.setProperty("skainet.cpu.vector.enabled", "true") val factory = platformDefaultCpuOpsFactory() - val ops = factory(DenseTensorDataFactory()) + val ops = factory(DenseTensorDataFactory(), sk.ainet.context.schedule.Schedule.Sequential) assertTrue(ops is DefaultCpuOpsJvm) } @@ -26,7 +26,7 @@ class PlatformCpuOpsFactoryJvmTest { fun fallsBackToScalarOpsWhenFlagDisabled() { System.setProperty("skainet.cpu.vector.enabled", "false") val factory = platformDefaultCpuOpsFactory() - val ops = factory(DenseTensorDataFactory()) + val ops = factory(DenseTensorDataFactory(), sk.ainet.context.schedule.Schedule.Sequential) assertTrue(ops !is DefaultCpuOpsJvm) } } diff --git a/skainet-backends/skainet-backend-cpu/src/linuxMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.linux.kt b/skainet-backends/skainet-backend-cpu/src/linuxMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.linux.kt index aa0ed4759..2fdc525ec 100644 --- a/skainet-backends/skainet-backend-cpu/src/linuxMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.linux.kt +++ b/skainet-backends/skainet-backend-cpu/src/linuxMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.linux.kt @@ -5,9 +5,13 @@ import sk.ainet.exec.kernel.ScalarKernelProvider import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps -internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps { +internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory, sk.ainet.context.schedule.Schedule) -> TensorOps { // Non-JVM has no ServiceLoader; register the scalar packed-quant kernels // (Q4_K/Q6_K/Q5_1/Q5_0/Q8_0/Q4_0) so DefaultCpuOpsBase can dispatch them. KernelRegistry.register(ScalarKernelProvider) - return { factory -> DefaultCpuOps(factory) } + return { factory, schedule -> DefaultCpuOps(factory, schedule) } } + + +internal actual fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule = + sk.ainet.context.schedule.Schedule.Sequential diff --git a/skainet-backends/skainet-backend-cpu/src/wasmJsMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.wasm.kt b/skainet-backends/skainet-backend-cpu/src/wasmJsMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.wasm.kt index aa0ed4759..2fdc525ec 100644 --- a/skainet-backends/skainet-backend-cpu/src/wasmJsMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.wasm.kt +++ b/skainet-backends/skainet-backend-cpu/src/wasmJsMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.wasm.kt @@ -5,9 +5,13 @@ import sk.ainet.exec.kernel.ScalarKernelProvider import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps -internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps { +internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory, sk.ainet.context.schedule.Schedule) -> TensorOps { // Non-JVM has no ServiceLoader; register the scalar packed-quant kernels // (Q4_K/Q6_K/Q5_1/Q5_0/Q8_0/Q4_0) so DefaultCpuOpsBase can dispatch them. KernelRegistry.register(ScalarKernelProvider) - return { factory -> DefaultCpuOps(factory) } + return { factory, schedule -> DefaultCpuOps(factory, schedule) } } + + +internal actual fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule = + sk.ainet.context.schedule.Schedule.Sequential diff --git a/skainet-backends/skainet-backend-cpu/src/wasmWasiMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.wasmWasi.kt b/skainet-backends/skainet-backend-cpu/src/wasmWasiMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.wasmWasi.kt index aa0ed4759..2fdc525ec 100644 --- a/skainet-backends/skainet-backend-cpu/src/wasmWasiMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.wasmWasi.kt +++ b/skainet-backends/skainet-backend-cpu/src/wasmWasiMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.wasmWasi.kt @@ -5,9 +5,13 @@ import sk.ainet.exec.kernel.ScalarKernelProvider import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps -internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps { +internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory, sk.ainet.context.schedule.Schedule) -> TensorOps { // Non-JVM has no ServiceLoader; register the scalar packed-quant kernels // (Q4_K/Q6_K/Q5_1/Q5_0/Q8_0/Q4_0) so DefaultCpuOpsBase can dispatch them. KernelRegistry.register(ScalarKernelProvider) - return { factory -> DefaultCpuOps(factory) } + return { factory, schedule -> DefaultCpuOps(factory, schedule) } } + + +internal actual fun platformDefaultSchedule(): sk.ainet.context.schedule.Schedule = + sk.ainet.context.schedule.Schedule.Sequential diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 51766e4b5..925a8bce4 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -209,6 +209,7 @@ public final class sk/ainet/context/DefaultDataExecutionContext : sk/ainet/conte public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getPhase ()Lsk/ainet/context/Phase; + public fun getSchedule ()Lsk/ainet/context/schedule/Schedule; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; @@ -217,6 +218,7 @@ public final class sk/ainet/context/DefaultDataExecutionContext : sk/ainet/conte public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -239,6 +241,7 @@ public abstract interface class sk/ainet/context/ExecutionContext { public abstract fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public abstract fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public abstract fun getPhase ()Lsk/ainet/context/Phase; + public fun getSchedule ()Lsk/ainet/context/schedule/Schedule; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public abstract fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; @@ -247,6 +250,7 @@ public abstract interface class sk/ainet/context/ExecutionContext { public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -264,6 +268,7 @@ public final class sk/ainet/context/ExecutionContext$DefaultImpls { public static fun getInTraining (Lsk/ainet/context/ExecutionContext;)Z public static fun getMemoryScope (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/memory/Scope; public static fun getMemoryTracker (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryTracker; + public static fun getSchedule (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/context/schedule/Schedule; public static fun getScratch (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/scratch/ScratchPool; public static fun getTraceSink (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/memory/trace/TraceSink; public static fun isRecording (Lsk/ainet/context/ExecutionContext;)Z @@ -271,6 +276,7 @@ public final class sk/ainet/context/ExecutionContext$DefaultImpls { public static fun placeholder (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun registerObserver (Lsk/ainet/context/ExecutionContext;Lsk/ainet/context/ExecutionObserver;)V public static fun unregisterObserver (Lsk/ainet/context/ExecutionContext;Lsk/ainet/context/ExecutionObserver;)V + public static fun withSchedule (Lsk/ainet/context/ExecutionContext;Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public static fun withTensorDataFactory (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public static fun wrapByteArray (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public static fun wrapFloatArray (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -369,6 +375,7 @@ public final class sk/ainet/context/PhaseOverridingExecutionContext : sk/ainet/c public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getPhase ()Lsk/ainet/context/Phase; + public fun getSchedule ()Lsk/ainet/context/schedule/Schedule; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; @@ -377,6 +384,7 @@ public final class sk/ainet/context/PhaseOverridingExecutionContext : sk/ainet/c public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -400,6 +408,43 @@ public final class sk/ainet/context/ResettableExecutionObserver$DefaultImpls { public static fun onTensorMaterialized (Lsk/ainet/context/ResettableExecutionObserver;Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Tensor;)V } +public final class sk/ainet/context/ScheduledExecutionContext : sk/ainet/context/ExecutionContext { + public fun (Lsk/ainet/context/ExecutionContext;Lsk/ainet/context/schedule/Schedule;)V + public fun fromByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; + public fun fromData (Lsk/ainet/lang/tensor/data/TensorData;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; + public fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; + public fun fromIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; + public fun full (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; + public fun getExecutionStats ()Lsk/ainet/context/ExecutionStats; + public fun getHooks ()Lsk/ainet/lang/nn/hooks/ForwardHooks; + public fun getInTraining ()Z + public fun getMemoryInfo ()Lsk/ainet/context/MemoryInfo; + public fun getMemoryScope ()Lsk/ainet/lang/memory/Scope; + public fun getMemoryTracker ()Lsk/ainet/lang/tensor/storage/MemoryTracker; + public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; + public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; + public fun getPhase ()Lsk/ainet/context/Phase; + public fun getSchedule ()Lsk/ainet/context/schedule/Schedule; + public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; + public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; + public fun isRecording ()Z + public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; + public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; + public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; + public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; + public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; + public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; + public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; + public fun zeros (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; +} + +public final class sk/ainet/context/ScheduledExecutionContextKt { + public static final fun withSchedule (Lsk/ainet/context/ExecutionContext;Lsk/ainet/context/schedule/Schedule;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object; +} + public final class sk/ainet/context/ScopedExecutionContext : sk/ainet/context/ExecutionContext { public fun (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/memory/Scope;)V public fun fromByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; @@ -416,6 +461,7 @@ public final class sk/ainet/context/ScopedExecutionContext : sk/ainet/context/Ex public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getPhase ()Lsk/ainet/context/Phase; + public fun getSchedule ()Lsk/ainet/context/schedule/Schedule; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; @@ -424,6 +470,7 @@ public final class sk/ainet/context/ScopedExecutionContext : sk/ainet/context/Ex public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -451,6 +498,7 @@ public final class sk/ainet/context/TrainingExecutionContext$DefaultImpls { public static fun getInTraining (Lsk/ainet/context/TrainingExecutionContext;)Z public static fun getMemoryScope (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/memory/Scope; public static fun getMemoryTracker (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryTracker; + public static fun getSchedule (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/context/schedule/Schedule; public static fun getScratch (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/tensor/scratch/ScratchPool; public static fun getTraceSink (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/memory/trace/TraceSink; public static fun isRecording (Lsk/ainet/context/TrainingExecutionContext;)Z @@ -458,6 +506,7 @@ public final class sk/ainet/context/TrainingExecutionContext$DefaultImpls { public static fun placeholder (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun registerObserver (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/context/ExecutionObserver;)V public static fun unregisterObserver (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/context/ExecutionObserver;)V + public static fun withSchedule (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public static fun withTensorDataFactory (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public static fun wrapByteArray (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public static fun wrapFloatArray (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -515,6 +564,64 @@ public final class sk/ainet/context/observers/MemorySnapshotObserver : sk/ainet/ public final fun results ()Ljava/util/List; } +public abstract interface class sk/ainet/context/schedule/Schedule { + public static final field Companion Lsk/ainet/context/schedule/Schedule$Companion; + public fun forEach (IILkotlin/jvm/functions/Function1;)V + public static synthetic fun forEach$default (Lsk/ainet/context/schedule/Schedule;IILkotlin/jvm/functions/Function1;ILjava/lang/Object;)V + public abstract fun forRange (IILkotlin/jvm/functions/Function2;)V + public static synthetic fun forRange$default (Lsk/ainet/context/schedule/Schedule;IILkotlin/jvm/functions/Function2;ILjava/lang/Object;)V + public abstract fun getName ()Ljava/lang/String; + public abstract fun getParallelism ()I +} + +public final class sk/ainet/context/schedule/Schedule$Companion { + public final fun chunkFor (II)I + public final fun tasksFor (III)I +} + +public final class sk/ainet/context/schedule/Schedule$DefaultImpls { + public static fun forEach (Lsk/ainet/context/schedule/Schedule;IILkotlin/jvm/functions/Function1;)V + public static synthetic fun forEach$default (Lsk/ainet/context/schedule/Schedule;IILkotlin/jvm/functions/Function1;ILjava/lang/Object;)V + public static synthetic fun forRange$default (Lsk/ainet/context/schedule/Schedule;IILkotlin/jvm/functions/Function2;ILjava/lang/Object;)V +} + +public final class sk/ainet/context/schedule/Schedule$Sequential : sk/ainet/context/schedule/Schedule { + public static final field INSTANCE Lsk/ainet/context/schedule/Schedule$Sequential; + public fun forEach (IILkotlin/jvm/functions/Function1;)V + public fun forRange (IILkotlin/jvm/functions/Function2;)V + public fun getName ()Ljava/lang/String; + public fun getParallelism ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/context/schedule/ScheduleHint { + public static final field Companion Lsk/ainet/context/schedule/ScheduleHint$Companion; + public static final field DIMS_KEY Ljava/lang/String; + public static final field PARALLELISM_KEY Ljava/lang/String; + public fun (Ljava/util/List;Ljava/lang/Integer;)V + public synthetic fun (Ljava/util/List;Ljava/lang/Integer;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/util/List; + public final fun component2 ()Ljava/lang/Integer; + public final fun copy (Ljava/util/List;Ljava/lang/Integer;)Lsk/ainet/context/schedule/ScheduleHint; + public static synthetic fun copy$default (Lsk/ainet/context/schedule/ScheduleHint;Ljava/util/List;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/context/schedule/ScheduleHint; + public fun equals (Ljava/lang/Object;)Z + public final fun getParallelDims ()Ljava/util/List; + public final fun getParallelism ()Ljava/lang/Integer; + public fun hashCode ()I + public final fun toAttributeMap ()Ljava/util/Map; + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/context/schedule/ScheduleHint$Companion { + public final fun fromAttribute (Ljava/lang/Object;)Lsk/ainet/context/schedule/ScheduleHint; + public final fun parallel ([Ljava/lang/String;Ljava/lang/Integer;)Lsk/ainet/context/schedule/ScheduleHint; + public static synthetic fun parallel$default (Lsk/ainet/context/schedule/ScheduleHint$Companion;[Ljava/lang/String;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/context/schedule/ScheduleHint; +} + +public final class sk/ainet/context/schedule/ScheduleHintKt { + public static final field SCHEDULE_ATTRIBUTE_KEY Ljava/lang/String; +} + public final class sk/ainet/java/Losses { public static final field INSTANCE Lsk/ainet/java/Losses; public static final fun bceWithLogits ()Lsk/ainet/lang/nn/loss/Loss; @@ -2320,6 +2427,46 @@ public final class sk/ainet/lang/memory/trace/TraceEvent$Plan : sk/ainet/lang/me public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/memory/trace/TraceEvent$ScheduleDowngraded : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;J)V + public synthetic fun (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun component3 ()Ljava/lang/String; + public final fun component4 ()J + public final fun copy (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;J)Lsk/ainet/lang/memory/trace/TraceEvent$ScheduleDowngraded; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$ScheduleDowngraded;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;JILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$ScheduleDowngraded; + public fun equals (Ljava/lang/Object;)Z + public final fun getEffective ()Ljava/lang/String; + public final fun getReason ()Ljava/lang/String; + public final fun getRequested ()Ljava/lang/String; + public fun getTimeNanos ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$ScheduleRegion : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Ljava/lang/String;Ljava/lang/String;IIJJ)V + public synthetic fun (Ljava/lang/String;Ljava/lang/String;IIJJILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun component3 ()I + public final fun component4 ()I + public final fun component5 ()J + public final fun component6 ()J + public final fun copy (Ljava/lang/String;Ljava/lang/String;IIJJ)Lsk/ainet/lang/memory/trace/TraceEvent$ScheduleRegion; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$ScheduleRegion;Ljava/lang/String;Ljava/lang/String;IIJJILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$ScheduleRegion; + public fun equals (Ljava/lang/Object;)Z + public final fun getDurationNanos ()J + public final fun getElements ()I + public final fun getOp ()Ljava/lang/String; + public final fun getSchedule ()Ljava/lang/String; + public final fun getTasks ()I + public fun getTimeNanos ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/memory/trace/TraceEvent$ScopeReset : sk/ainet/lang/memory/trace/TraceEvent { public fun (Lsk/ainet/lang/memory/ScopeKind;JJJ)V public synthetic fun (Lsk/ainet/lang/memory/ScopeKind;JJJILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -2482,6 +2629,7 @@ public final class sk/ainet/lang/nn/DefaultNeuralNetworkExecutionContext : sk/ai public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getPhase ()Lsk/ainet/context/Phase; + public fun getSchedule ()Lsk/ainet/context/schedule/Schedule; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; @@ -2490,6 +2638,7 @@ public final class sk/ainet/lang/nn/DefaultNeuralNetworkExecutionContext : sk/ai public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -2816,6 +2965,7 @@ public final class sk/ainet/lang/nn/NeuralNetworkExecutionContext$DefaultImpls { public static fun getInTraining (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Z public static fun getMemoryScope (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/memory/Scope; public static fun getMemoryTracker (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryTracker; + public static fun getSchedule (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/context/schedule/Schedule; public static fun getScratch (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/tensor/scratch/ScratchPool; public static fun getTraceSink (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/memory/trace/TraceSink; public static fun isRecording (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Z @@ -2823,6 +2973,7 @@ public final class sk/ainet/lang/nn/NeuralNetworkExecutionContext$DefaultImpls { public static fun placeholder (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun registerObserver (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/context/ExecutionObserver;)V public static fun unregisterObserver (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/context/ExecutionObserver;)V + public static fun withSchedule (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public static fun withTensorDataFactory (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public static fun wrapByteArray (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public static fun wrapFloatArray (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/trace/AndroidTraceSink.kt b/skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/trace/AndroidTraceSink.kt index 353500e57..35d7e1512 100644 --- a/skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/trace/AndroidTraceSink.kt +++ b/skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/trace/AndroidTraceSink.kt @@ -35,6 +35,8 @@ public class AndroidTraceSink : TraceSink { is TraceEvent.Free -> Trace.setCounter(truncate("skainet ${event.scope.name.lowercase()} free bytes"), event.bytes) is TraceEvent.ScopeReset -> Trace.setCounter(truncate("skainet ${event.scope.name.lowercase()} live bytes"), event.liveBytesAfter) is TraceEvent.Counter -> Trace.setCounter(truncate("skainet ${event.name}"), event.value) + is TraceEvent.ScheduleDowngraded -> Trace.setCounter(truncate("skainet schedule downgraded"), 1L) + is TraceEvent.ScheduleRegion -> Trace.setCounter(truncate("skainet schedule ${event.op} tasks"), event.tasks.toLong()) is TraceEvent.Plan -> Trace.setCounter(truncate("skainet plan total bytes"), event.totalBytes) } } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt index 721d55c8a..d7293b6df 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt @@ -60,6 +60,34 @@ public interface ExecutionContext { */ public fun withTensorDataFactory(factory: TensorDataFactory): ExecutionContext = this + /** + * How this context's ops map the independent work inside one op onto cores (SKEEP-005). + * Default [sk.ainet.context.schedule.Schedule.Sequential]: one task, inline, today's behaviour. + * A schedule never changes a result — it is a deployment property, not a model property. + */ + public val schedule: sk.ainet.context.schedule.Schedule + get() = sk.ainet.context.schedule.Schedule.Sequential + + /** + * This context rebuilt so its ops run under [schedule] — the same seam as + * [withTensorDataFactory]. The default cannot rebuild, returns `this`, and — unless the request + * is already what this context runs — emits [sk.ainet.lang.memory.trace.TraceEvent.ScheduleDowngraded] + * so an unhonoured schedule is visible in the trace instead of silently sequential. + */ + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + public fun withSchedule(schedule: sk.ainet.context.schedule.Schedule): ExecutionContext { + if (schedule !== this.schedule && traceSink.isEnabled) { + traceSink.emit( + sk.ainet.lang.memory.trace.TraceEvent.ScheduleDowngraded( + requested = schedule.name, + effective = this.schedule.name, + reason = "context ${this::class.simpleName} cannot rebuild its ops", + ), + ) + } + return this + } + /** * Workspace allocator for short-lived intermediate buffers (attention * scratch, RoPE tables, KV-cache slice copies, padding scratch, etc.). diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScheduledExecutionContext.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScheduledExecutionContext.kt new file mode 100644 index 000000000..583407ee3 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScheduledExecutionContext.kt @@ -0,0 +1,59 @@ +package sk.ainet.context + +import sk.ainet.context.schedule.Schedule +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.data.TensorDataFactory +import sk.ainet.lang.tensor.operators.OpsBoundTensor +import sk.ainet.lang.tensor.ops.TensorOps +import sk.ainet.lang.types.DType +import kotlin.reflect.KClass + +/** + * [base] with a [schedule] (SKEEP-005): the deployment-level knob that says how the independent + * work inside an op is mapped onto cores. A schedule is never a DSL word — a model is defined once + * and runs sequentially or in parallel depending only on the context it is given. + * + * Mirrors [ScopedExecutionContext]: the base is rebuilt through [ExecutionContext.withSchedule] so + * the *ops instance* carries the schedule (ops are constructed per context), and tensors created + * through this context are bound to those ops. The same ops-binding rule applies: `a + b` + * dispatches through the ops that created `a`, so create inputs through the scheduled context. + * + * A base that cannot rebuild itself keeps its own ops and reports the unhonoured request as a + * [sk.ainet.lang.memory.trace.TraceEvent.ScheduleDowngraded] — visible, never silent. + */ +public class ScheduledExecutionContext( + private val base: ExecutionContext, + override val schedule: Schedule, +) : ExecutionContext by base { + + private val scheduledBase: ExecutionContext = base.withSchedule(schedule) + + override val ops: TensorOps get() = scheduledBase.ops + + override val tensorDataFactory: TensorDataFactory get() = scheduledBase.tensorDataFactory + + override fun fromData(data: TensorData, dtype: KClass): Tensor = + OpsBoundTensor.fromData(data, dtype, ops) + + override fun withSchedule(schedule: Schedule): ExecutionContext = + if (schedule === this.schedule) this else ScheduledExecutionContext(base, schedule) + + /** Keep the schedule when a further decorator (e.g. `forwardScope`) swaps the data factory. */ + override fun withTensorDataFactory(factory: TensorDataFactory): ExecutionContext = + ScheduledExecutionContext(base.withTensorDataFactory(factory), schedule) +} + +/** + * Run [block] with [schedule] active on this context: + * + * ```kotlin + * ctx.withSchedule(CoroutineSchedule.hardware()) { scheduled -> + * model.forward(input, scheduled) // attention heads etc. run in parallel, results unchanged + * } + * ``` + */ +public inline fun ExecutionContext.withSchedule( + schedule: Schedule, + block: (ctx: ExecutionContext) -> R, +): R = block(ScheduledExecutionContext(this, schedule)) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/Schedule.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/Schedule.kt new file mode 100644 index 000000000..cf7cb883a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/Schedule.kt @@ -0,0 +1,72 @@ +package sk.ainet.context.schedule + +/** + * How the independent work inside one op is mapped onto cores — the *schedule* half of the + * Halide-style split (SKEEP-005). The DSL says *what* is computed; a [Schedule] says *how many + * tasks* compute it and on which threads. It never changes a result: every implementation runs + * the same body over the same disjoint ranges, so per-element arithmetic and its order are + * untouched, and a scheduled run is bit-identical to a sequential one. + * + * Non-suspending on purpose: ops are synchronous on every Kotlin target, and this interface has + * no dependency (no coroutines here). JVM ships [sk.ainet.exec.schedule.CoroutineSchedule] in + * `skainet-backend-cpu`; every other target and every context that does not opt in runs + * [Sequential]. + * + * ### Contract for [forRange] bodies + * + * - ranges are disjoint half-open `[start, end)` intervals covering `[0, n)`, in ascending order; + * - `n == 0` makes no call; a task count of one runs the body inline on the caller's thread; + * - a body writes only into pre-allocated, disjoint output regions and reads only inputs that are + * immutable for the duration of the region; + * - a body must **not** allocate through an `ExecutionContext` or call `ctx.ops` (the step + * allocator, scratch pool and op caches are single-threaded), and must not start a nested + * region — an implementation runs a nested call inline; + * - the first failure thrown by any task is rethrown to the caller after every sibling task has + * finished or been cancelled; no task is still running when [forRange] returns; + * - all writes made by tasks happen-before the return of [forRange]. + */ +public interface Schedule { + + /** Upper bound on tasks running at once; `1` means sequential. */ + public val parallelism: Int + + /** Stable, human-readable name for trace events and diagnostics, e.g. `sequential`, `coroutines(8)`. */ + public val name: String + + /** + * Run [body] over `[0, n)` split into at most `min(parallelism, ceil(n / grain))` disjoint + * ranges, each at least [grain] elements long except possibly the last. + */ + public fun forRange(n: Int, grain: Int = 1, body: (start: Int, end: Int) -> Unit) + + /** Element-wise convenience over [forRange]: [body] receives every index in `[0, count)` exactly once. */ + public fun forEach(count: Int, minPerTask: Int = 1, body: (index: Int) -> Unit) { + forRange(count, minPerTask) { start, end -> for (i in start until end) body(i) } + } + + /** The default everywhere: one task, inline, on the caller's thread. */ + public object Sequential : Schedule { + override val parallelism: Int get() = 1 + override val name: String get() = "sequential" + override fun forRange(n: Int, grain: Int, body: (start: Int, end: Int) -> Unit) { + if (n > 0) body(0, n) + } + override fun toString(): String = name + } + + public companion object { + /** + * Number of tasks an implementation with [parallelism] uses for [n] elements at [grain]: + * `min(parallelism, ceil(n / grain))`, never more than [n], and `0` for an empty range. + */ + public fun tasksFor(n: Int, grain: Int, parallelism: Int): Int { + if (n <= 0) return 0 + val g = grain.coerceAtLeast(1) + val byGrain = (n + g - 1) / g + return minOf(parallelism.coerceAtLeast(1), byGrain) + } + + /** Chunk length that splits [n] elements into [tasks] near-equal ranges. */ + public fun chunkFor(n: Int, tasks: Int): Int = (n + tasks - 1) / tasks + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/ScheduleHint.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/ScheduleHint.kt new file mode 100644 index 000000000..1b624599d --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/ScheduleHint.kt @@ -0,0 +1,52 @@ +package sk.ainet.context.schedule + +/** + * Attribute key under which a [ScheduleHint] rides the graph: `OpTrace.attributes` on the tape, + * `GraphNodeDefinition.attributes` in the DAG DSL, `Operation.parameters` and `GraphNode.metadata` + * in the ComputeGraph, and the `skainet.schedule` module attribute in exported StableHLO. + */ +public const val SCHEDULE_ATTRIBUTE_KEY: String = "skainet.schedule" + +/** + * A declarative schedule request attached to one op (SKEEP-005, compile lane). It changes no + * DSL semantics: a consumer that ignores it computes the same result. Which dimension names an + * op accepts is decided by the annotation pass (`batch`, `heads`, `rows`, …); an unknown name is + * rejected with a diagnostic, never silently dropped. + */ +public data class ScheduleHint( + /** Loop dimensions to run in parallel, by name known to the op. */ + val parallelDims: List, + /** Requested worker count; `null` means "whatever the target has". */ + val parallelism: Int? = null, +) { + init { + require(parallelDims.isNotEmpty()) { "ScheduleHint needs at least one dimension" } + require(parallelDims.all { it.isNotBlank() }) { "ScheduleHint dimension names must not be blank" } + require(parallelism == null || parallelism > 0) { "ScheduleHint parallelism must be positive, got $parallelism" } + } + + /** Plain map form for graph metadata that must stay serializable (`parallel_dims`, `parallelism`). */ + public fun toAttributeMap(): Map = buildMap { + put(DIMS_KEY, parallelDims) + parallelism?.let { put(PARALLELISM_KEY, it) } + } + + public companion object { + public const val DIMS_KEY: String = "parallel_dims" + public const val PARALLELISM_KEY: String = "parallelism" + + public fun parallel(vararg dims: String, parallelism: Int? = null): ScheduleHint = + ScheduleHint(dims.toList(), parallelism) + + /** Reads a hint from either a [ScheduleHint] or its [toAttributeMap] form; `null` for anything else. */ + public fun fromAttribute(value: Any?): ScheduleHint? = when (value) { + is ScheduleHint -> value + is Map<*, *> -> { + val dims = (value[DIMS_KEY] as? Collection<*>)?.map { it.toString() } ?: return null + val p = (value[PARALLELISM_KEY] as? Number)?.toInt() + if (dims.isEmpty()) null else ScheduleHint(dims, p) + } + else -> null + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt index 891b5030e..2e8eedd64 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt @@ -90,6 +90,13 @@ public object PerfettoTraceExporter { emit(counter("live bytes", ts, mapOf(e.scope.name.lowercase() to e.liveBytesAfter))) } is TraceEvent.Counter -> emit(counter(e.name, ts, mapOf(e.unit to e.value))) + is TraceEvent.ScheduleDowngraded -> emit( + instant("schedule downgraded", "schedule", ts, MAIN_TID, mapOf("requested" to e.requested, "effective" to e.effective, "reason" to e.reason)), + ) + is TraceEvent.ScheduleRegion -> emit( + complete("schedule ${e.op}", "schedule", (e.timeNanos - e.durationNanos) / 1000.0, e.durationNanos / 1000.0, MAIN_TID, + mapOf("schedule" to e.schedule, "elements" to e.elements.toString(), "tasks" to e.tasks.toString())), + ) is TraceEvent.Plan -> emit( instant( "plan", "plan", ts, MAIN_TID, diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt index 9d3821e10..3a6f5c79b 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt @@ -78,6 +78,27 @@ public sealed interface TraceEvent { /** A platform counter sample (RSS, page faults, heap, direct memory …). */ public data class Counter(val name: String, val value: Long, val unit: String = "bytes", override val timeNanos: Long = TraceClock.nowNanos()) : TraceEvent + /** + * A schedule was requested that this context or target cannot honour (SKEEP-005); [effective] + * is what runs instead. Doctrine: a downgrade is visible, never a silent approximation. + */ + public data class ScheduleDowngraded( + val requested: String, + val effective: String, + val reason: String, + override val timeNanos: Long = TraceClock.nowNanos(), + ) : TraceEvent + + /** A parallel region ran: [elements] split into [tasks] on [schedule] for [op]. */ + public data class ScheduleRegion( + val op: String, + val schedule: String, + val elements: Int, + val tasks: Int, + val durationNanos: Long = 0L, + override val timeNanos: Long = TraceClock.nowNanos(), + ) : TraceEvent + /** A memory plan was computed (M0 `MemoryPlan`): the plan-vs-actual check (#1030) compares this with the allocation events. */ public data class Plan( val model: String, diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/context/schedule/ScheduledExecutionContextTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/context/schedule/ScheduledExecutionContextTest.kt new file mode 100644 index 000000000..26cbca893 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/context/schedule/ScheduledExecutionContextTest.kt @@ -0,0 +1,122 @@ +package sk.ainet.context.schedule + +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.ScheduledExecutionContext +import sk.ainet.context.forwardScope +import sk.ainet.context.withSchedule +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.TraceSink +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * SKEEP-005: a schedule is a deployment property on the context. The default is sequential, a + * context that cannot rebuild its ops reports an unhonoured request as a trace event (never a + * silent downgrade), and the decorator survives the other decorators it composes with. + */ +@OptIn(ExperimentalMemoryApi::class) +class ScheduledExecutionContextTest { + + /** A schedule that only counts; enough to prove which one a context carries. */ + private class Counting(override val parallelism: Int = 4) : Schedule { + var regions = 0 + override val name: String get() = "counting($parallelism)" + override fun forRange(n: Int, grain: Int, body: (Int, Int) -> Unit) { + regions++ + Schedule.Sequential.forRange(n, grain, body) + } + } + + /** Delegation would forward `withSchedule` to the delegate (and its Noop sink); route it through the interface default. */ + private class TracingContext(sink: TraceSink) : ExecutionContext by DefaultDataExecutionContext() { + override val traceSink: TraceSink = sink + override fun withSchedule(schedule: Schedule): ExecutionContext = super.withSchedule(schedule) + } + + @Test + fun defaultScheduleIsSequential() { + val ctx = DefaultDataExecutionContext() + assertSame(Schedule.Sequential, ctx.schedule) + assertEquals(1, ctx.schedule.parallelism) + assertEquals("sequential", ctx.schedule.name) + } + + @Test + fun sequentialRunsTheWholeRangeOnce() { + val seen = mutableListOf>() + Schedule.Sequential.forRange(7, grain = 3) { s, e -> seen += s to e } + assertEquals(listOf(0 to 7), seen) + Schedule.Sequential.forRange(0) { _, _ -> error("empty range must not call the body") } + var count = 0 + Schedule.Sequential.forEach(5) { count += it } + assertEquals(0 + 1 + 2 + 3 + 4, count) + } + + @Test + fun tasksForNeverExceedsElementsOrParallelism() { + assertEquals(0, Schedule.tasksFor(0, 1, 8)) + assertEquals(1, Schedule.tasksFor(1, 1, 8)) + assertEquals(4, Schedule.tasksFor(4, 1, 8)) + assertEquals(8, Schedule.tasksFor(1000, 1, 8)) + assertEquals(3, Schedule.tasksFor(100, 40, 8), "grain caps the task count") + assertEquals(1, Schedule.tasksFor(100, 1, 0), "parallelism is coerced to at least one") + } + + @Test + fun decoratorCarriesTheScheduleAndReportsAnUnhonouredRequest() { + val sink = RecordingTraceSink() + val requested = Counting() + val scheduled = ScheduledExecutionContext(TracingContext(sink), requested) + + assertSame(requested, scheduled.schedule, "the decorator answers with the requested schedule") + val downgrades = sink.eventsOf() + assertEquals(1, downgrades.size, "a base that cannot rebuild its ops must say so") + assertEquals("counting(4)", downgrades.single().requested) + assertEquals("sequential", downgrades.single().effective) + } + + @Test + fun requestingTheScheduleAContextAlreadyRunsIsSilent() { + val sink = RecordingTraceSink() + val ctx = TracingContext(sink) + assertSame(ctx, ctx.withSchedule(Schedule.Sequential)) + assertTrue(sink.eventsOf().isEmpty()) + } + + @Test + fun withScheduleBlockHandsOutADecoratedContext() { + val counting = Counting() + val result = DefaultDataExecutionContext().withSchedule(counting) { ctx -> + assertSame(counting, ctx.schedule) + ctx.schedule.forRange(10) { _, _ -> } + "ok" + } + assertEquals("ok", result) + assertEquals(1, counting.regions) + } + + @Test + fun scheduleSurvivesAForwardScope() { + val counting = Counting() + DefaultDataExecutionContext().withSchedule(counting) { scheduled -> + scheduled.forwardScope(slabFloats = 16) { scoped, _ -> + assertSame(counting, scoped.schedule, "forwardScope rebuilds through withTensorDataFactory; the schedule must be kept") + } + } + } + + @Test + fun scheduleHintRoundTripsThroughItsMapForm() { + val hint = ScheduleHint.parallel("batch", "heads", parallelism = 8) + val map = hint.toAttributeMap() + assertEquals(hint, ScheduleHint.fromAttribute(map)) + assertEquals(hint, ScheduleHint.fromAttribute(hint)) + assertEquals(null, ScheduleHint.fromAttribute("nonsense")) + assertEquals(null, ScheduleHint.fromAttribute(mapOf(ScheduleHint.DIMS_KEY to emptyList()))) + } +} From 1dd7fb0a770b44ebc64676e1adfae9ac3c451e2d Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 3 Sep 2026 22:31:22 +0200 Subject: [PATCH 2/4] feat(compile): SKEEP-005 schedule hints ride the graph into the StableHLO header MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - skainet-lang-dag: `op(..., schedule = parallel("rows"))`, `schedule(hint) { … }` ambient block (new `DagBuilder.withAttributes`), `GraphNodeDefinition.scheduleHint()` — same channel as `DtypePolicyDsl`. - skainet-compile-opt: `ScheduleAnnotationPass` validates requested dims per op (sdpa: batch/heads, matmul: rows, conv: batch/out_channels), stamps the normalized hint into `GraphNode.metadata`, rejects unknown dims with a diagnostic (never silently), optional per-op defaults; idempotent. - skainet-compile-hlo: `HloGenerator` runs the pass as a core pass with a target; `StableHloConverter` emits `skainet.schedule = { = {parallel_dims = [...], parallelism = N}}` in the module header beside `skainet.tensor_layouts`. No per-op emitter; IREE-side consumption is out of scope. - Tests: ScheduleDslTest, ScheduleAnnotationPassTest, ScheduleModuleAttributeTest. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018sqoaGs5M7C5uVw6cpzAnH --- .../ainet/compile/hlo/StableHloConverter.kt | 32 ++++++- .../compile/hlo/generate/HloGenerator.kt | 6 +- .../hlo/ScheduleModuleAttributeTest.kt | 69 ++++++++++++++ .../api/jvm/skainet-compile-opt.api | 17 ++++ .../opt/passes/ScheduleAnnotationPass.kt | 90 +++++++++++++++++++ .../opt/passes/ScheduleAnnotationPassTest.kt | 90 +++++++++++++++++++ .../api/jvm/skainet-lang-dag.api | 10 +++ .../kotlin/sk/ainet/lang/dag/GraphDsl.kt | 22 ++++- .../kotlin/sk/ainet/lang/dag/ScheduleDsl.kt | 49 ++++++++++ .../sk/ainet/lang/dag/ScheduleDslTest.kt | 43 +++++++++ 10 files changed, 422 insertions(+), 6 deletions(-) create mode 100644 skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ScheduleModuleAttributeTest.kt create mode 100644 skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPass.kt create mode 100644 skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPassTest.kt create mode 100644 skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/ScheduleDsl.kt create mode 100644 skainet-lang/skainet-lang-dag/src/commonTest/kotlin/sk/ainet/lang/dag/ScheduleDslTest.kt diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt index 68f22beac..00e2a86be 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt @@ -85,14 +85,20 @@ public class StableHloConverter @kotlin.jvm.JvmOverloads constructor( processNodes(topo, context) generateReturnStatement(outputNodes, context) - val moduleHeader = if (structuralLayouts.isNotEmpty()) { + val headerAttributes = mutableListOf() + if (structuralLayouts.isNotEmpty()) { val layoutEntries = structuralLayouts.entries .sortedBy { it.key } .joinToString(", ") { (name, attr) -> "$name = $attr" } - "module attributes {skainet.tensor_layouts = {$layoutEntries}} {" - } else { - "module {" + headerAttributes += "skainet.tensor_layouts = {$layoutEntries}" + } + // SKEEP-005: schedule hints ride the module header next to the layouts — one graph, extra + // schedule metadata. Keyed by node id; a consumer that ignores it computes the same result. + val schedules = collectScheduleHints(topo) + if (schedules.isNotEmpty()) { + headerAttributes += "skainet.schedule = {" + schedules.joinToString(", ") { (id, hint) -> "${mlirKey(id)} = ${scheduleAttr(hint)}" } + "}" } + val moduleHeader = if (headerAttributes.isEmpty()) "module {" else "module attributes {${headerAttributes.joinToString(", ")}} {" val assembled = StringBuilder() assembled.appendLine(moduleHeader) @@ -287,6 +293,24 @@ public class StableHloConverter @kotlin.jvm.JvmOverloads constructor( * machine-readable — block element counts, block bytes, bit widths, block order — so a * downstream consumer can size and address packed weights without a lookup table of names. */ + /** `(nodeId, hint)` for every node carrying a schedule hint, in topological order. */ + private fun collectScheduleHints(topo: List): List> = + topo.mapNotNull { node -> + val hint = sk.ainet.context.schedule.ScheduleHint.fromAttribute(node.metadata[sk.ainet.context.schedule.SCHEDULE_ATTRIBUTE_KEY]) + ?: sk.ainet.context.schedule.ScheduleHint.fromAttribute(node.operation.parameters[sk.ainet.context.schedule.SCHEDULE_ATTRIBUTE_KEY]) + hint?.let { node.id to it } + } + + private fun scheduleAttr(hint: sk.ainet.context.schedule.ScheduleHint): String { + val dims = hint.parallelDims.joinToString(", ") { "\"$it\"" } + val p = hint.parallelism?.let { ", parallelism = $it" } ?: "" + return "{parallel_dims = [$dims]$p}" + } + + /** MLIR dictionary keys must be bare identifiers or quoted strings. */ + private fun mlirKey(id: String): String = + if (id.isNotEmpty() && (id[0].isLetter() || id[0] == '_') && id.all { it.isLetterOrDigit() || it == '_' || it == '$' || it == '.' }) id else "\"$id\"" + private fun collectStructuralLayouts(nodes: List): Map { val result = linkedMapOf() for (node in nodes) { diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt index ae90937bc..d136e4cfb 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt @@ -68,7 +68,11 @@ public object HloGenerator { } else { sk.ainet.compile.opt.dagPipelineFor( target, - corePasses = listOf(sk.ainet.compile.opt.passes.LayoutAssignmentPass(target)), + corePasses = listOf( + sk.ainet.compile.opt.passes.LayoutAssignmentPass(target), + // SKEEP-005: schedule hints become `skainet.schedule` module metadata. + sk.ainet.compile.opt.passes.ScheduleAnnotationPass(target), + ), ).optimize(computeGraph).graph } diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ScheduleModuleAttributeTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ScheduleModuleAttributeTest.kt new file mode 100644 index 000000000..415c58b58 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ScheduleModuleAttributeTest.kt @@ -0,0 +1,69 @@ +package sk.ainet.compile.hlo + +import sk.ainet.compile.opt.passes.ScheduleAnnotationPass +import sk.ainet.context.schedule.SCHEDULE_ATTRIBUTE_KEY +import sk.ainet.context.schedule.ScheduleHint +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.AddOperation +import sk.ainet.lang.tensor.ops.InputOperation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.types.DType +import kotlin.test.Test +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** SKEEP-005: schedule hints reach the StableHLO module header; graphs without hints are untouched. */ +class ScheduleModuleAttributeTest { + + private fun chain(vararg hints: ScheduleHint?): DefaultComputeGraph { + val graph = DefaultComputeGraph() + val a = GraphNode("a", InputOperation(), emptyList(), listOf(TensorSpec("a", listOf(1, 4), "FP32"))) + val b = GraphNode("b", InputOperation(), emptyList(), listOf(TensorSpec("b", listOf(1, 4), "FP32"))) + graph.addNode(a); graph.addNode(b) + var prev = a + hints.forEachIndexed { i, hint -> + val meta = hint?.let { mapOf(SCHEDULE_ATTRIBUTE_KEY to it.toAttributeMap()) } ?: emptyMap() + val add = GraphNode("add$i", AddOperation(), listOf(prev.outputs[0], b.outputs[0]), listOf(TensorSpec("s$i", listOf(1, 4), "FP32")), metadata = meta) + graph.addNode(add) + graph.addEdge(GraphEdge("e${i}a", prev, add, 0, 0, prev.outputs[0])) + graph.addEdge(GraphEdge("e${i}b", b, add, 0, 1, b.outputs[0])) + prev = add + } + return graph + } + + @Test + fun stampedHintsAreEmittedInTheHeader() { + val mlir = toStableHlo(chain(ScheduleHint.parallel("batch", "heads", parallelism = 8), null), "scheduled").content + assertTrue(mlir.contains("module attributes {"), mlir) + assertTrue(mlir.contains("skainet.schedule = {add0 = {parallel_dims = [\"batch\", \"heads\"], parallelism = 8}}"), mlir) + assertFalse(mlir.contains("add1 = {parallel_dims"), "nodes without a hint are not listed:\n$mlir") + } + + @Test + fun graphWithoutHintsKeepsTheBareHeader() { + val mlir = toStableHlo(chain(null), "plain").content + assertTrue(mlir.contains("module {"), mlir) + assertFalse(mlir.contains("skainet.schedule"), mlir) + } + + @Test + fun hintsCarriedOnOperationParametersAreEmittedWithoutThePass() { + val graph = DefaultComputeGraph() + val a = GraphNode("a", InputOperation(), emptyList(), listOf(TensorSpec("a", listOf(1, 4), "FP32"))) + val b = GraphNode("b", InputOperation(), emptyList(), listOf(TensorSpec("b", listOf(1, 4), "FP32"))) + val op = GraphNode( + "layers.0.attn", sk.ainet.lang.tensor.ops.GenericOperation(name = "add", parameters = mapOf(SCHEDULE_ATTRIBUTE_KEY to ScheduleHint.parallel("heads")), type = "compute"), + listOf(a.outputs[0], b.outputs[0]), listOf(TensorSpec("y", listOf(1, 4), "FP32")), + ) + graph.addNode(a); graph.addNode(b); graph.addNode(op) + graph.addEdge(GraphEdge("ea", a, op, 0, 0, a.outputs[0])); graph.addEdge(GraphEdge("eb", b, op, 0, 1, b.outputs[0])) + val mlir = toStableHlo(graph, "params").content + assertTrue(mlir.contains("skainet.schedule = {layers.0.attn = {parallel_dims = [\"heads\"]}}"), mlir) + // The pass stamps the same thing as metadata, so running it first changes nothing in the header. + val stamped = ScheduleAnnotationPass().apply(graph).graph + assertTrue(toStableHlo(stamped, "params").content.contains("skainet.schedule = {layers.0.attn = {parallel_dims = [\"heads\"]}}")) + } +} diff --git a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api index 5e3a3b26b..98e5a165e 100644 --- a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api +++ b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api @@ -149,6 +149,23 @@ public final class sk/ainet/compile/opt/passes/PowSpecializationPass : sk/ainet/ public fun getName ()Ljava/lang/String; } +public final class sk/ainet/compile/opt/passes/ScheduleAnnotationPass : sk/ainet/compile/opt/GraphOptimizationPass { + public static final field Companion Lsk/ainet/compile/opt/passes/ScheduleAnnotationPass$Companion; + public static final field SCHEDULE_METADATA_KEY Ljava/lang/String; + public fun ()V + public fun (Ljava/lang/String;Ljava/util/Map;)V + public synthetic fun (Ljava/lang/String;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; + public fun getName ()Ljava/lang/String; +} + +public final class sk/ainet/compile/opt/passes/ScheduleAnnotationPass$Companion { + public final fun attentionDefaults ()Ljava/util/Map; + public final fun getKNOWN_DIMS ()Ljava/util/Map; + public final fun hintOf (Lsk/ainet/lang/graph/GraphNode;)Lsk/ainet/context/schedule/ScheduleHint; + public final fun normalizeOpName (Ljava/lang/String;)Ljava/lang/String; +} + public final class sk/ainet/compile/opt/passes/SharedWeightDeduplicationPass : sk/ainet/compile/opt/GraphOptimizationPass { public static final field Companion Lsk/ainet/compile/opt/passes/SharedWeightDeduplicationPass$Companion; public fun ()V diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPass.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPass.kt new file mode 100644 index 000000000..85515ff74 --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPass.kt @@ -0,0 +1,90 @@ +package sk.ainet.compile.opt.passes + +import sk.ainet.compile.opt.GraphOptimizationPass +import sk.ainet.compile.opt.GraphOptimizationResult +import sk.ainet.context.schedule.SCHEDULE_ATTRIBUTE_KEY +import sk.ainet.context.schedule.ScheduleHint +import sk.ainet.lang.graph.ComputeGraph +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphNode + +/** + * Carries schedule requests through the compile lane (SKEEP-005). A [ScheduleHint] reaches a + * node either from the DSL (`Operation.parameters[SCHEDULE_ATTRIBUTE_KEY]`, copied off the + * `dag { }` attributes) or from [defaults] keyed by op name. The pass validates the requested + * dimensions against what the op can be split on, stamps a normalized hint into + * [GraphNode.metadata] under the same key, and reports every rejected request as a diagnostic — + * an unknown dimension is never dropped silently. Structure and numerics are untouched: the + * exporter reads the metadata into the `skainet.schedule` module attribute, and a consumer that + * ignores it computes the same result. + * + * Target-parameterized like [LayoutAssignmentPass]: `HloGenerator` runs it as a core pass whenever + * a target is named. + */ +public class ScheduleAnnotationPass( + private val target: String? = null, + /** Hints applied to ops (by normalized name) that carry none of their own — opt-in. */ + private val defaults: Map = emptyMap(), +) : GraphOptimizationPass { + + override val name: String = "schedule-annotation(${target ?: "any"})" + + public companion object { + public const val SCHEDULE_METADATA_KEY: String = SCHEDULE_ATTRIBUTE_KEY + + /** Dimensions an op may be split on, by normalized op name (lower case, no separators). */ + public val KNOWN_DIMS: Map> = mapOf( + "scaleddotproductattention" to setOf("batch", "heads"), + "sdpa" to setOf("batch", "heads"), + "attention" to setOf("batch", "heads"), + "matmul" to setOf("rows"), + "linear" to setOf("rows"), + "conv2d" to setOf("batch", "outchannels"), + "conv1d" to setOf("batch", "outchannels"), + ) + + /** The engine's own choice for attention: split over batch and heads. */ + public fun attentionDefaults(): Map = + listOf("scaleddotproductattention", "sdpa", "attention").associateWith { ScheduleHint.parallel("batch", "heads") } + + public fun normalizeOpName(name: String): String = name.lowercase().filter { it.isLetterOrDigit() } + + /** The hint stamped on [node] by this pass, or carried from the DSL; `null` when none. */ + public fun hintOf(node: GraphNode): ScheduleHint? = + ScheduleHint.fromAttribute(node.metadata[SCHEDULE_METADATA_KEY]) + ?: ScheduleHint.fromAttribute(node.operation.parameters[SCHEDULE_ATTRIBUTE_KEY]) + } + + override fun apply(graph: ComputeGraph): GraphOptimizationResult { + val diagnostics = mutableListOf() + var changed = false + val newNodes = graph.nodes.map { node -> + if (node.metadata.containsKey(SCHEDULE_METADATA_KEY)) return@map node // already stamped + val opName = normalizeOpName(node.operation.name) + val requested = ScheduleHint.fromAttribute(node.operation.parameters[SCHEDULE_ATTRIBUTE_KEY]) + ?: defaults[opName] + ?: return@map node + val allowed = KNOWN_DIMS[opName] + if (allowed == null) { + diagnostics += "schedule on '${node.id}' (${node.operation.name}) rejected: op has no schedulable dimensions" + return@map node + } + val unknown = requested.parallelDims.map { normalizeOpName(it) }.filter { it !in allowed } + if (unknown.isNotEmpty()) { + diagnostics += "schedule on '${node.id}' (${node.operation.name}) rejected: unknown dims $unknown; honoured dims: $allowed" + return@map node + } + changed = true + node.copy(metadata = node.metadata + (SCHEDULE_METADATA_KEY to requested.toAttributeMap())) + } + if (!changed) return GraphOptimizationResult(graph, changed = false, diagnostics = diagnostics) + + val byId = newNodes.associateBy { it.id } + val newGraph = DefaultComputeGraph() + for (node in newNodes) newGraph.addNode(node) + for (edge in graph.edges) { + newGraph.addEdge(edge.copy(source = byId.getValue(edge.source.id), destination = byId.getValue(edge.destination.id))) + } + return GraphOptimizationResult(newGraph, changed = true, diagnostics = diagnostics) + } +} diff --git a/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPassTest.kt b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPassTest.kt new file mode 100644 index 000000000..9945c448a --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPassTest.kt @@ -0,0 +1,90 @@ +package sk.ainet.compile.opt.passes + +import sk.ainet.context.schedule.SCHEDULE_ATTRIBUTE_KEY +import sk.ainet.context.schedule.ScheduleHint +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.GenericOperation +import sk.ainet.lang.tensor.ops.InputOperation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.types.DType +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** SKEEP-005: hints are validated per op, stamped as metadata, and rejections are diagnostics. */ +class ScheduleAnnotationPassTest { + + private fun graphWith(opName: String, params: Map): DefaultComputeGraph { + val graph = DefaultComputeGraph() + val x = GraphNode("x", InputOperation(), emptyList(), listOf(TensorSpec("x", listOf(1, 8, 4, 16), "FP32"))) + val op = GraphNode( + id = "op", + operation = GenericOperation(name = opName, parameters = params, type = "compute"), + inputs = listOf(x.outputs[0]), + outputs = listOf(TensorSpec("y", listOf(1, 8, 4, 16), "FP32")), + ) + graph.addNode(x); graph.addNode(op) + graph.addEdge(GraphEdge("e", x, op, 0, 0, x.outputs[0])) + return graph + } + + @Test + fun validHintIsStampedAsMetadataInMapForm() { + val graph = graphWith("scaledDotProductAttention", mapOf(SCHEDULE_ATTRIBUTE_KEY to ScheduleHint.parallel("batch", "heads", parallelism = 8))) + val result = ScheduleAnnotationPass("test").apply(graph) + assertTrue(result.changed) + assertTrue(result.diagnostics.isEmpty()) + val op = result.graph.nodes.first { it.id == "op" } + assertEquals(ScheduleHint(listOf("batch", "heads"), 8), ScheduleHint.fromAttribute(op.metadata[SCHEDULE_ATTRIBUTE_KEY])) + assertEquals(mapOf("parallel_dims" to listOf("batch", "heads"), "parallelism" to 8), op.metadata[SCHEDULE_ATTRIBUTE_KEY]) + assertEquals(1, result.graph.edges.size, "edges are rebuilt") + } + + @Test + fun mapFormHintIsAcceptedToo() { + val graph = graphWith("matmul", mapOf(SCHEDULE_ATTRIBUTE_KEY to mapOf("parallel_dims" to listOf("rows")))) + val result = ScheduleAnnotationPass().apply(graph) + assertTrue(result.changed) + assertEquals(ScheduleHint(listOf("rows")), ScheduleAnnotationPass.hintOf(result.graph.nodes.first { it.id == "op" })) + } + + @Test + fun unknownDimensionIsRejectedWithADiagnosticAndLeavesTheNodeAlone() { + val graph = graphWith("matmul", mapOf(SCHEDULE_ATTRIBUTE_KEY to ScheduleHint.parallel("heads"))) + val result = ScheduleAnnotationPass("test").apply(graph) + assertFalse(result.changed) + assertEquals(1, result.diagnostics.size) + assertTrue(result.diagnostics.single().contains("unknown dims [heads]"), result.diagnostics.single()) + assertNull(result.graph.nodes.first { it.id == "op" }.metadata[SCHEDULE_ATTRIBUTE_KEY]) + } + + @Test + fun opWithoutSchedulableDimensionsIsRejected() { + val graph = graphWith("softmax", mapOf(SCHEDULE_ATTRIBUTE_KEY to ScheduleHint.parallel("rows"))) + val result = ScheduleAnnotationPass().apply(graph) + assertFalse(result.changed) + assertTrue(result.diagnostics.single().contains("no schedulable dimensions")) + } + + @Test + fun defaultsApplyOnlyToOpsWithoutTheirOwnHint() { + val graph = graphWith("sdpa", emptyMap()) + val none = ScheduleAnnotationPass().apply(graph) + assertFalse(none.changed, "no hint, no defaults: nothing to do") + val withDefaults = ScheduleAnnotationPass(defaults = ScheduleAnnotationPass.attentionDefaults()).apply(graph) + assertTrue(withDefaults.changed) + assertEquals(ScheduleHint(listOf("batch", "heads")), ScheduleAnnotationPass.hintOf(withDefaults.graph.nodes.first { it.id == "op" })) + } + + @Test + fun passIsIdempotent() { + val graph = graphWith("sdpa", mapOf(SCHEDULE_ATTRIBUTE_KEY to ScheduleHint.parallel("heads"))) + val once = ScheduleAnnotationPass().apply(graph) + val twice = ScheduleAnnotationPass().apply(once.graph) + assertFalse(twice.changed) + } +} diff --git a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api index c5646ee0e..f21d8b6c1 100644 --- a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api +++ b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api @@ -7,6 +7,7 @@ public final class sk/ainet/lang/dag/DagBuilder { public static synthetic fun op$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/tensor/ops/Operation;Ljava/util/List;Ljava/lang/String;Ljava/util/Map;ILjava/lang/Object;)Ljava/util/List; public final fun output ([Lsk/ainet/lang/dag/GraphValue;)V public final fun parameter (Ljava/lang/String;Lsk/ainet/lang/tensor/ops/TensorSpec;)Lsk/ainet/lang/dag/GraphValue; + public final fun withAttributes (Ljava/util/Map;Lkotlin/jvm/functions/Function1;)V } public abstract interface annotation class sk/ainet/lang/dag/DagDsl : java/lang/annotation/Annotation { @@ -218,6 +219,15 @@ public final class sk/ainet/lang/dag/GraphValue { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/dag/ScheduleDslKt { + public static final fun op (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/tensor/ops/Operation;Ljava/util/List;Lsk/ainet/context/schedule/ScheduleHint;Ljava/lang/String;Ljava/util/Map;)Ljava/util/List; + public static synthetic fun op$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/tensor/ops/Operation;Ljava/util/List;Lsk/ainet/context/schedule/ScheduleHint;Ljava/lang/String;Ljava/util/Map;ILjava/lang/Object;)Ljava/util/List; + public static final fun parallel ([Ljava/lang/String;Ljava/lang/Integer;)Lsk/ainet/context/schedule/ScheduleHint; + public static synthetic fun parallel$default ([Ljava/lang/String;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/context/schedule/ScheduleHint; + public static final fun schedule (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/context/schedule/ScheduleHint;Lkotlin/jvm/functions/Function1;)V + public static final fun scheduleHint (Lsk/ainet/lang/dag/GraphNodeDefinition;)Lsk/ainet/context/schedule/ScheduleHint; +} + public final class sk/ainet/lang/dag/SymbolicDataDslKt { public static final fun dtypeName (Lkotlin/reflect/KClass;)Ljava/lang/String; } diff --git a/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/GraphDsl.kt b/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/GraphDsl.kt index 9dc5f5216..0a47b36cc 100644 --- a/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/GraphDsl.kt +++ b/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/GraphDsl.kt @@ -74,6 +74,24 @@ public class DagBuilder { private val outputs = mutableListOf>() private var nextId: Long = 0 + /** Attributes applied to every node recorded while a [withAttributes] block is open; inner blocks and explicit per-op attributes win. */ + private val ambientAttributes = ArrayDeque>() + + /** + * Record every node inside [block] with [attributes] merged into its own (explicit per-op + * attributes override ambient ones; inner blocks override outer ones). The carrier for + * scoped annotations such as `schedule(hint) { … }` (SKEEP-005). + */ + @DagDsl + public fun withAttributes(attributes: Map, block: DagBuilder.() -> Unit) { + ambientAttributes.addLast(attributes) + try { + block() + } finally { + ambientAttributes.removeLast() + } + } + private fun freshNodeId(opName: String, providedId: String): String = providedId.ifBlank { "n${nextId++}_${opName}" } @@ -115,12 +133,14 @@ public class DagBuilder { val nodeOutputs = outputSpecs.mapIndexed { idx, spec -> GraphValue(nodeId = nodeId, outputIndex = idx, spec = spec) } + val merged = if (ambientAttributes.isEmpty()) attributes else + ambientAttributes.fold(emptyMap()) { acc, m -> acc + m } + attributes nodes += GraphNodeDefinition( id = nodeId, operation = operation, inputs = inputs, outputs = nodeOutputs, - attributes = attributes + attributes = merged ) return nodeOutputs } diff --git a/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/ScheduleDsl.kt b/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/ScheduleDsl.kt new file mode 100644 index 000000000..41730500a --- /dev/null +++ b/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/ScheduleDsl.kt @@ -0,0 +1,49 @@ +package sk.ainet.lang.dag + +import sk.ainet.context.schedule.SCHEDULE_ATTRIBUTE_KEY +import sk.ainet.context.schedule.ScheduleHint +import sk.ainet.lang.tensor.ops.Operation + +/** + * Schedule annotations for the `dag { }` DSL (SKEEP-005, compile lane). A [ScheduleHint] rides + * the node's [GraphNodeDefinition.attributes] under [SCHEDULE_ATTRIBUTE_KEY], exactly as + * [DTYPE_POLICY_ATTRIBUTE_KEY] does for dtype policies, is copied into `Operation.parameters` + * when the program becomes a `ComputeGraph`, validated and stamped by + * `ScheduleAnnotationPass` (`skainet-compile-opt`), and emitted as the `skainet.schedule` module + * attribute of the StableHLO export. It never changes what the graph computes. + * + * ```kotlin + * dag { + * schedule(parallel("heads")) { // every op recorded inside + * op(sdpa, listOf(q, k, v)) + * } + * op(matmul, listOf(x, w), schedule = parallel("rows", parallelism = 8)) // one op + * } + * ``` + */ +@DagDsl +public fun DagBuilder.op( + operation: Operation, + inputs: List>, + schedule: ScheduleHint, + id: String = "", + extraAttributes: Map = emptyMap(), +): List> = op( + operation = operation, + inputs = inputs, + id = id, + attributes = extraAttributes + (SCHEDULE_ATTRIBUTE_KEY to schedule), +) + +/** Every op recorded inside [block] carries [hint]; an explicit per-op hint wins over the ambient one. */ +@DagDsl +public fun DagBuilder.schedule(hint: ScheduleHint, block: DagBuilder.() -> Unit): Unit = + withAttributes(mapOf(SCHEDULE_ATTRIBUTE_KEY to hint), block) + +/** `parallel("batch", "heads")`, `parallel("rows", parallelism = 8)`. */ +public fun parallel(vararg dims: String, parallelism: Int? = null): ScheduleHint = + ScheduleHint.parallel(*dims, parallelism = parallelism) + +/** The hint attached to this node, or `null`. */ +public fun GraphNodeDefinition.scheduleHint(): ScheduleHint? = + ScheduleHint.fromAttribute(attributes[SCHEDULE_ATTRIBUTE_KEY]) diff --git a/skainet-lang/skainet-lang-dag/src/commonTest/kotlin/sk/ainet/lang/dag/ScheduleDslTest.kt b/skainet-lang/skainet-lang-dag/src/commonTest/kotlin/sk/ainet/lang/dag/ScheduleDslTest.kt new file mode 100644 index 000000000..4b9ca8a1c --- /dev/null +++ b/skainet-lang/skainet-lang-dag/src/commonTest/kotlin/sk/ainet/lang/dag/ScheduleDslTest.kt @@ -0,0 +1,43 @@ +package sk.ainet.lang.dag + +import sk.ainet.context.schedule.SCHEDULE_ATTRIBUTE_KEY +import sk.ainet.context.schedule.ScheduleHint +import sk.ainet.lang.tensor.ops.MatmulOperation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull + +/** SKEEP-005: schedule hints attach to `dag { }` nodes under the shared attribute key. */ +class ScheduleDslTest { + + @Test + fun perOpScheduleLandsUnderTheKey() { + val program = dag { + val x = input("x", TensorSpec("x", listOf(1, 4), "Float32")) + val w = parameter("w") { shape(4, 4) { ones() } } + op(MatmulOperation(), listOf(x, w), schedule = parallel("rows", parallelism = 8)) + } + val mm = program.nodes.last { it.operation is MatmulOperation<*, *> } + assertEquals(ScheduleHint(listOf("rows"), 8), mm.scheduleHint()) + assertEquals(mm.scheduleHint(), mm.attributes[SCHEDULE_ATTRIBUTE_KEY]) + } + + @Test + fun ambientScheduleAppliesToEveryOpInTheBlockAndExplicitWins() { + val program = dag { + val x = input("x", TensorSpec("x", listOf(1, 4), "Float32")) + val w = parameter("w") { shape(4, 4) { ones() } } + schedule(parallel("rows")) { + op(MatmulOperation(), listOf(x, w), id = "ambient") + op(MatmulOperation(), listOf(x, w), schedule = parallel("rows", parallelism = 2), id = "explicit") + } + op(MatmulOperation(), listOf(x, w), id = "outside") + } + val byId = program.nodes.associateBy { it.id } + assertEquals(ScheduleHint(listOf("rows")), byId.getValue("ambient").scheduleHint()) + assertEquals(ScheduleHint(listOf("rows"), 2), byId.getValue("explicit").scheduleHint(), "explicit per-op hint wins") + assertNull(byId.getValue("outside").scheduleHint(), "the block's hint does not leak") + } +} From 10e3b53e297497fa4bec4d88cb6e076cb12741ad Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 3 Sep 2026 22:43:28 +0200 Subject: [PATCH 3/4] docs(skeep-005): proposal page, "Algorithm and schedule", schedule tutorial with executable sample - SKEEP-005 (Status: Implemented) in docs/modules/skeep, registered in nav and the Current Proposals table; references #1259/#1260 (fixed) and #1261 (observed). - explanation/schedules.adoc: principle, the Schedule contract, where it hangs, what runs where per platform, how to see which schedule ran, the compile lane. dsl-principles.adoc gains the "Schedule" row in "Who answers which question". - tutorials/schedule-getting-started.adoc backed by ScheduleDemo.kt in the executable samples (SamplesTest asserts bit-identity and the recorded region). - CHANGELOG Unreleased entry. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018sqoaGs5M7C5uVw6cpzAnH --- CHANGELOG.md | 19 ++ .../sk/ainet/docs/samples/ScheduleDemo.kt | 74 +++++++ docs/modules/ROOT/nav.adoc | 2 + .../pages/explanation/dsl-principles.adoc | 5 + .../ROOT/pages/explanation/schedules.adoc | 169 +++++++++++++++ .../tutorials/schedule-getting-started.adoc | 81 +++++++ docs/modules/skeep/nav.adoc | 1 + .../005-schedules-structured-concurrency.adoc | 198 ++++++++++++++++++ docs/modules/skeep/pages/index.adoc | 4 + .../sk/ainet/docs/samples/SamplesTest.kt | 10 + 10 files changed, 563 insertions(+) create mode 100644 docs/modules/ROOT/examples/kotlin/sk/ainet/docs/samples/ScheduleDemo.kt create mode 100644 docs/modules/ROOT/pages/explanation/schedules.adoc create mode 100644 docs/modules/ROOT/pages/tutorials/schedule-getting-started.adoc create mode 100644 docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cb3afb03..490d12430 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,25 @@ ### Added +- **Schedules — the compute-level algorithm/schedule split (SKEEP-005)** + ([#1259](https://github.com/SKaiNET-developers/SKaiNET/issues/1259) registries not safe for concurrent reads, + [#1260](https://github.com/SKaiNET-developers/SKaiNET/issues/1260) `parallelChunks` `runBlocking` island): a dependency-free + `sk.ainet.context.schedule.Schedule` on every `ExecutionContext` (`ctx.schedule`, + `ctx.withSchedule(s) { … }`, `ScheduledExecutionContext`), a JVM `CoroutineSchedule` built on + structured concurrency (caller runs the first chunk, nested regions inline, `dedicated()` pool), + and the first scheduled op — `scaledDotProductAttention` runs its `(batch, head)` units on the + context's schedule, bit-identical to the sequential loop. `parallelChunks` no longer hides a + `runBlocking(Dispatchers.Default)` island: it delegates to the ops' schedule, so + `DirectCpuExecutionContext(schedule = Schedule.Sequential)` makes every kernel single-threaded + and the JVM default (`CoroutineSchedule.hardware()`) spreads them across cores. Unhonoured + requests are visible as `TraceEvent.ScheduleDowngraded`; regions as `TraceEvent.ScheduleRegion`. + Compile lane: `dag { schedule(parallel("heads")) { … } }` stamps a `ScheduleHint` that + `ScheduleAnnotationPass` validates per op and `StableHloConverter` emits as the + `skainet.schedule` module attribute beside `skainet.tensor_layouts`. Registries + (`KernelDispatch`, `KernelRegistry`) are now safe for concurrent reads. Found on the way and + reported, not fixed here: [#1261](https://github.com/SKaiNET-developers/SKaiNET/issues/1261) + (Panama matmul's first call differs by an ULP until the JIT intrinsifies `reduceLanes`). Docs: SKEEP-005, + "Algorithm and schedule", "Schedule getting started" (executable sample). - **`tensorFilter` on the single-file `SafeTensorsParametersLoader`** ([#1256](https://github.com/SKaiNET-developers/SKaiNET/issues/1256)): parity with the sharded loader — an optional predicate over the tensor headers; filtered-out tensors are neither read diff --git a/docs/modules/ROOT/examples/kotlin/sk/ainet/docs/samples/ScheduleDemo.kt b/docs/modules/ROOT/examples/kotlin/sk/ainet/docs/samples/ScheduleDemo.kt new file mode 100644 index 000000000..ba96730a6 --- /dev/null +++ b/docs/modules/ROOT/examples/kotlin/sk/ainet/docs/samples/ScheduleDemo.kt @@ -0,0 +1,74 @@ +package sk.ainet.docs.samples + +// tag::imports[] +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.context.withSchedule +import sk.ainet.exec.schedule.CoroutineSchedule +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 +// end::imports[] + +/** + * SKEEP-005 walk-through: the same attention op on three schedules, bit-identical results, and + * the trace that says which schedule ran. Every region is included verbatim into + * `tutorials/schedule-getting-started.adoc`; `SamplesTest` executes it. + */ +@OptIn(ExperimentalMemoryApi::class) +object ScheduleDemo { + + class Result( + val defaultScheduleName: String, + val sequential: FloatArray, + val scheduled: FloatArray, + val regions: List, + ) + + // tag::operands[] + /** Q, K, V for 2 batches × 8 heads: [batch, heads, seq, headDim], built through [ctx]. */ + fun operands(ctx: ExecutionContext, seqQ: Int = 64, seqKV: Int = 256, headDim: Int = 32): Triple, Tensor, Tensor> { + fun tensor(seq: Int, seed: Int): Tensor { + val values = FloatArray(2 * 8 * seq * headDim) { i -> (((i * 31 + seed * 17) % 23) - 11) / 11f } + return ctx.fromFloatArray(Shape(2, 8, seq, headDim), FP32::class, values) + } + return Triple(tensor(seqQ, 1), tensor(seqKV, 2), tensor(seqKV, 3)) + } + + fun attention(ctx: ExecutionContext): FloatArray { + val (q, k, v) = operands(ctx) + return ctx.ops.scaledDotProductAttention(q, k, v, mask = null, scale = 0f, causal = true).data.copyToFloatArray() + } + // end::operands[] + + fun run(): Result { + // tag::sequential[] + val ctx = DirectCpuExecutionContext() // JVM: CoroutineSchedule.hardware() + val defaultName = ctx.schedule.name // e.g. "coroutines(12)" + val sequential = ctx.withSchedule(Schedule.Sequential) { seq -> + attention(seq) // one task, the caller's thread + } + // end::sequential[] + + // tag::scheduled[] + val sink = RecordingTraceSink() + val twoWorkers = CoroutineSchedule(parallelism = 2, sink = sink) // Dispatchers.Default + val scheduled = ctx.withSchedule(twoWorkers) { par -> + attention(par) // 16 (batch, head) units on 2 tasks + } + // end::scheduled[] + + // tag::trace[] + val regions = sink.eventsOf() + val identical = sequential.contentEquals(scheduled) + println("default schedule: $defaultName") + for (r in regions) println("region: ${r.op} on ${r.schedule} — ${r.elements} units in ${r.tasks} tasks") + println("outputs identical: $identical") + // end::trace[] + return Result(defaultName, sequential, scheduled, regions) + } +} diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc index c05ae6789..b9c418441 100644 --- a/docs/modules/ROOT/nav.adoc +++ b/docs/modules/ROOT/nav.adoc @@ -11,6 +11,7 @@ ** xref:tutorials/minerva-getting-started.adoc[Minerva getting started] ** xref:tutorials/graph-dsl.adoc[Graph DSL] ** xref:tutorials/turboquant-getting-started.adoc[TurboQuant: KV-cache compression] +** xref:tutorials/schedule-getting-started.adoc[Schedules: parallel ops, same results] ** xref:tutorials/android-classifier-getting-started.adoc[Train a classifier on Android] ** xref:tutorials/ternary-getting-started.adoc[Ternary networks: getting started] * How-to guides @@ -39,6 +40,7 @@ ** xref:explanation/packed-weight-layout.adoc[Packed weight layout] ** xref:explanation/eager-execution.adoc[Eager execution: backends and kernels] ** xref:explanation/kernel-selection.adoc[Kernel SPI and the selection algorithm] +** xref:explanation/schedules.adoc[Algorithm and schedule] ** xref:explanation/quantization-process.adoc[The quantization process] ** xref:explanation/theory/index.adoc[Mathematical theory] *** xref:explanation/theory/matmul.adoc[Matrix multiplication] diff --git a/docs/modules/ROOT/pages/explanation/dsl-principles.adoc b/docs/modules/ROOT/pages/explanation/dsl-principles.adoc index 0c655af72..eb66acdf3 100644 --- a/docs/modules/ROOT/pages/explanation/dsl-principles.adoc +++ b/docs/modules/ROOT/pages/explanation/dsl-principles.adoc @@ -65,6 +65,11 @@ Two halves, and both matter: | Which kernel serves these operands | `KernelKey` = op × per-operand `(Format, LayoutClass)`. Runtime holds **no memory policy** — it reads what is already true about the operands and picks a kernel. That is all. + +| **Schedule** (deployment, per context) +| How many tasks compute the independent parts of one op, and on which threads +| `Schedule`, `ctx.withSchedule(…)`, `skainet.schedule` graph metadata + (xref:explanation/schedules.adoc[Algorithm and schedule]). *Never changes a result.* |=== The direction of information flow is strictly downward: the DSL knows nothing of forms, forms diff --git a/docs/modules/ROOT/pages/explanation/schedules.adoc b/docs/modules/ROOT/pages/explanation/schedules.adoc new file mode 100644 index 000000000..9762ea72b --- /dev/null +++ b/docs/modules/ROOT/pages/explanation/schedules.adoc @@ -0,0 +1,169 @@ += Algorithm and schedule +:description: Why a network definition never says how many threads run it, what a Schedule is, the contract a scheduled body must keep, how to see which schedule actually ran, and how a schedule request travels into the compiled graph. + +A SKaiNET network says *what* is computed. Since SKEEP-005 a second, optional object says *how +the independent parts of one op are mapped onto cores*: the **schedule**. The split is Halide's +— algorithm here, schedule there — and it exists for the same reason: the same model must run on +a phone core, a laptop and a server without being rewritten, and the choice of threads is a +property of the deployment, not of the mathematics. + +This page states the principle, the contract, and how to observe it. Its companion pages: +xref:explanation/dsl-principles.adoc[The DSL is compute] (the doctrine this follows), +xref:explanation/kernel-selection.adoc[Kernel SPI and the selection algorithm] (what runs inside +a task), xref:explanation/memory-model.adoc[The memory model] (why a task must not allocate), and +xref:tutorials/schedule-getting-started.adoc[Schedules: parallel ops, same results] (the executable walk-through). + +== The principle + +**** +*A network definition is a pure description of computation. A schedule says how many tasks +compute the independent parts of one op and on which threads. It lives on the execution context, +never in the DSL — and it never changes a result: a scheduled run is bit-identical to a +sequential one.* +**** + +Three consequences follow. + +* **No schedule words in `network { }`.** A model is defined once. Whether its attention heads + run on one thread or twelve is decided where the model is *run*, with + `ctx.withSchedule(…)`, exactly like `forwardScope` decides where activations live. +* **The default is sequential — except where the platform already parallelised.** On the JVM a + `DirectCpuExecutionContext()` runs `CoroutineSchedule.hardware()`, the core-count coroutine + schedule; that is what the Panama matmul kernels always did, now visible and switchable. + Every other target and every context that does not opt in runs `Schedule.Sequential`. +* **A request that cannot be honoured is visible.** A context that cannot rebuild its ops + returns itself from `withSchedule` and emits `TraceEvent.ScheduleDowngraded`. Nothing is + silently approximated. + +== What a `Schedule` is + +[source,kotlin] +---- +public interface Schedule { + public val parallelism: Int // 1 = sequential + public val name: String // "sequential", "coroutines(12)", … + public fun forRange(n: Int, grain: Int = 1, body: (start: Int, end: Int) -> Unit) + public fun forEach(count: Int, minPerTask: Int = 1, body: (index: Int) -> Unit) + public object Sequential : Schedule +} +---- + +It is deliberately small, non-suspending and dependency-free: ops are synchronous on every +Kotlin target, and `skainet-lang-core` takes no coroutine dependency. The JVM implementation, +`CoroutineSchedule` in `skainet-backend-cpu`, is structured concurrency in the textbook sense — +one region is one `coroutineScope`: + +* the calling thread runs the first chunk itself, so no core idles blocked on the join; +* the other chunks are `launch`ed on the dispatcher (`Dispatchers.Default` by default); +* the scope closes only when every child has finished or been cancelled — the first failure + cancels the siblings and is rethrown to the caller; +* a region reached from inside another region runs inline, so the dispatcher never waits on + itself; `CoroutineSchedule.dedicated(n)` owns its own pool for callers that already live on + `Dispatchers.Default`. + +== The contract a body must keep + +`forRange` hands a body disjoint half-open ranges covering `[0, n)`. In exchange the body promises: + +. it writes only into pre-allocated, disjoint regions of the output and reads only inputs + that are immutable for the region; +. it never allocates through an `ExecutionContext` and never calls `ctx.ops` — the step + allocator (`ForwardScope`), the scratch pool and the op caches are single-threaded by design; +. it never starts a nested region (an implementation runs one inline anyway); +. its arithmetic and the order of that arithmetic do not depend on how the range was cut. + +The last promise is what makes results bit-identical. The engine's first scheduled op, +`scaledDotProductAttention`, is the model: every `(batch, head)` pair computes exactly the loop +it computed before, into its own rows, with its own `scores` scratch; the schedule only decides +which pairs share a task. A tiny call (a decode step on a handful of heads) stays on the caller — +a region costs more than it saves below `SDPA_PARALLEL_MIN_WORK` multiply-adds. + +== Where it hangs + +[cols="2,3",options="header"] +|=== +| Seam | Role + +| `ExecutionContext.schedule` +| The schedule this context's ops run under. Default `Schedule.Sequential`. + +| `ExecutionContext.withSchedule(s)` +| The context rebuilt so its ops carry `s` — the same seam as `withTensorDataFactory`. A + context that cannot rebuild returns itself and traces `ScheduleDowngraded`. + +| `ctx.withSchedule(s) { scheduled -> … }` +| The decorator form (`ScheduledExecutionContext`), composable with `forwardScope`. Tensors + created through `scheduled` are bound to the scheduled ops — the usual ops-binding rule. + +| `DirectCpuExecutionContext(schedule = …)` +| The JVM context; omitted, it takes `platformDefaultSchedule()`. + +| `DefaultCpuOpsBase(dataFactory, schedule)` +| Ops hold the schedule; `parallelChunks(outputDim, schedule)` is how a kernel asks for tasks. +|=== + +== What runs where + +[cols="1,2,2",options="header"] +|=== +| Target | Default schedule | Parallel schedule available + +| JVM | `CoroutineSchedule.hardware()` — `availableProcessors()` on `Dispatchers.Default` | yes (`CoroutineSchedule`, `dedicated()`) +| Android | `Sequential` | not yet (coroutines are a JVM-only dependency of the backend today) +| Kotlin/Native (Linux, Apple, Android native) | `Sequential` | not yet +| JS, Wasm | `Sequential` | no — single-threaded runtimes; a parallel schedule cannot be constructed there +|=== + +== How to see which schedule ran + +`ctx.schedule.name` tells you what a context carries. To see what actually happened, attach a +sink: + +[source,kotlin] +---- +val sink = RecordingTraceSink() +val ctx = DirectCpuExecutionContext(schedule = CoroutineSchedule.hardware(sink = sink)) +// … run … +sink.eventsOf() // op, schedule, elements, tasks, duration +sink.eventsOf() // requested, effective, reason +---- + +A `ScheduleRegion` is emitted per parallel region; a `ScheduleDowngraded` whenever a request was +not honoured. The Perfetto exporter and the Android trace sink render both. + +== The compile lane + +The eager schedule is a runtime object; the compiled graph gets the same fact as metadata. In +the `dag { }` DSL a request is stamped on nodes the way a dtype policy is: + +[source,kotlin] +---- +dag { + schedule(parallel("heads")) { // every op recorded inside + op(sdpa, listOf(q, k, v)) + } + op(matmul, listOf(x, w), schedule = parallel("rows", parallelism = 8)) +} +---- + +`ScheduleAnnotationPass` (a core pass whenever a target is named) validates the requested +dimensions against what the op can be split on — `batch`/`heads` for attention, `rows` for +matmul, `batch`/`out_channels` for convolutions — stamps the normalised hint into the node's +metadata and reports every rejection as a diagnostic. The StableHLO export carries it in the +module header beside the layouts: + +[source,mlir] +---- +module attributes {skainet.tensor_layouts = {…}, skainet.schedule = {attn = {parallel_dims = ["batch", "heads"], parallelism = 8}}} { +---- + +One graph, extra schedule metadata. A consumer that ignores the attribute computes the same +result; turning it into IREE dispatch hints is a later step. + +== What this rules out + +* A `parallel = true` on a layer in `network { }`. The layer does not know the device. +* A kernel that spawns its own threads. It asks the schedule for ranges. +* A body that reaches for `ctx.zeros(...)` or `ctx.ops.matmul(...)` inside a region. It gets a + data race in the step allocator, not a faster op. +* Silent fallbacks. A schedule that cannot run says so in the trace. diff --git a/docs/modules/ROOT/pages/tutorials/schedule-getting-started.adoc b/docs/modules/ROOT/pages/tutorials/schedule-getting-started.adoc new file mode 100644 index 000000000..664f79ec5 --- /dev/null +++ b/docs/modules/ROOT/pages/tutorials/schedule-getting-started.adoc @@ -0,0 +1,81 @@ += Schedules: parallel ops, same results +:description: Run one attention op sequentially and on the core-count coroutine schedule, prove the outputs are bit-identical, and read the trace that says which schedule ran. + +You will run `scaledDotProductAttention` three ways — on the platform default, forced +sequential, and on an explicit two-worker coroutine schedule — compare the outputs bit for bit, +and read the trace events a schedule emits. Every snippet below is compiled and executed in CI +from `skainet-docs-samples` (`ScheduleDemo.kt`). + +== Prerequisites + +* JDK 21+ with the Vector API module (`--enable-preview --add-modules jdk.incubator.vector`). +* `sk.ainet.core:skainet-lang-core` and `sk.ainet.core:skainet-backend-cpu` on the classpath. + +[source,kotlin] +---- +include::example$kotlin/sk/ainet/docs/samples/ScheduleDemo.kt[tag=imports] +---- + +== Step 1 — Build the operands through the context + +The schedule is a property of the *context*, so build the inputs through the context you will +run on — a tensor is bound to the ops that created it. + +[source,kotlin] +---- +include::example$kotlin/sk/ainet/docs/samples/ScheduleDemo.kt[tag=operands] +---- + +== Step 2 — Run on the platform default, then sequentially + +`DirectCpuExecutionContext()` on the JVM carries `CoroutineSchedule.hardware()` — one task per +logical core on `Dispatchers.Default`. `withSchedule(Schedule.Sequential)` rebuilds the same +context with single-task ops; nothing else changes. + +[source,kotlin] +---- +include::example$kotlin/sk/ainet/docs/samples/ScheduleDemo.kt[tag=sequential] +---- + +== Step 3 — Pick an explicit schedule and record what ran + +A `CoroutineSchedule` takes a dispatcher, a parallelism and an optional trace sink. With a +`RecordingTraceSink` every parallel region shows up as a `ScheduleRegion` event. + +[source,kotlin] +---- +include::example$kotlin/sk/ainet/docs/samples/ScheduleDemo.kt[tag=scheduled] +---- + +== Step 4 — Verify + +The CI assertion on this exact code is `assertContentEquals(sequential, scheduled)`: a +schedule never changes a result. The recorded region reports the op, the schedule name, the +number of `(batch, head)` units and the tasks they were split into. + +[source,kotlin] +---- +include::example$kotlin/sk/ainet/docs/samples/ScheduleDemo.kt[tag=trace] +---- + +Expected output (shape and numbers depend on your machine): + +[source,text] +---- +default schedule: coroutines(12) +region: forRange on coroutines(2) — 16 units in 2 tasks +outputs identical: true +---- + +== Where this applies today + +* `scaledDotProductAttention` on every CPU context (heads × batch are the units). +* The Panama Q4_K / Q5_K matmul kernels and the FP32 tiled GEMM (output rows are the units). +* SKaiNET-transformers' `MultiHeadAttention` reads `ctx.schedule` for per-head decode and + prefill — see its own "Attention schedules" page. + +== Next steps + +* xref:explanation/schedules.adoc[Algorithm and schedule] — the principle and the contract. +* xref:skeep:005-schedules-structured-concurrency.adoc[SKEEP-005] — the design record. +* xref:explanation/kernel-selection.adoc[Kernel SPI and the selection algorithm] — what runs inside a task. diff --git a/docs/modules/skeep/nav.adoc b/docs/modules/skeep/nav.adoc index deeb78561..7b6e75b4e 100644 --- a/docs/modules/skeep/nav.adoc +++ b/docs/modules/skeep/nav.adoc @@ -6,3 +6,4 @@ ** xref:skeep:003-unified-tensor-storage.adoc[SKEEP-003: Unifying the tensor storage model] ** xref:skeep:003a-placement-and-planning-resolution.adoc[SKEEP-003a: Placement & planning resolution] ** xref:skeep:004-virtual-tensor-layout.adoc[SKEEP-004: Virtual tensor layout] +** xref:skeep:005-schedules-structured-concurrency.adoc[SKEEP-005: Schedules — structured concurrency] diff --git a/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc b/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc new file mode 100644 index 000000000..873eff429 --- /dev/null +++ b/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc @@ -0,0 +1,198 @@ += SKEEP-005: Schedules — structured concurrency for the compute layer +:description: SKaiNET proposal to separate the algorithm (the DSL network) from the schedule (how the independent work inside an op is mapped onto cores), with Kotlin structured concurrency as the JVM implementation and schedule metadata that rides the ComputeGraph into StableHLO. + +Status: Implemented + +Audience: SKaiNET maintainers and contributors; SKaiNET-transformers maintainers (the first consumer is per-head attention) + +Created: 2026-09-03 + +Tracking issue: _to be filed_ (bugs fixed on the way: https://github.com/SKaiNET-developers/SKaiNET/issues/1259[#1259], https://github.com/SKaiNET-developers/SKaiNET/issues/1260[#1260]; observed: https://github.com/SKaiNET-developers/SKaiNET/issues/1261[#1261]) + +Origin: the Daily-StandAPP profile on SKaiNET 0.53.0 (JFR, i7-9750H 6C/12T): attention ≈ 40 % of every decoded token, scalar and single-threaded, while the only parallelism in the stack was the native matmul's compile-time 4-thread pool and a `runBlocking(Dispatchers.Default)` island inside a few Panama kernels. + +== Summary + +A network definition says *what* is computed. Nothing in SKaiNET said *how many tasks* compute it. +This proposal adds a second, optional half — the **schedule** — in the spirit of Halide's +algorithm/schedule split: a small, dependency-free `Schedule` interface on the +`ExecutionContext`, a coroutine-backed JVM implementation, the engine's first scheduled op +(`scaledDotProductAttention`), and a metadata channel so a schedule request can ride the +ComputeGraph into the StableHLO export. A schedule never changes a result: every implementation +runs the same body over the same disjoint ranges, so a scheduled run is bit-identical to a +sequential one. The DSL is untouched — the schedule is a deployment property, exactly where +xref:ROOT:explanation/dsl-principles.adoc[The DSL is compute] says such knobs belong. + +== Motivation + +* Multi-head attention is embarrassingly parallel across heads, but both attention + implementations (`DefaultCpuOps.scaledDotProductAttention`, transformers' + `fusedDecodeAttention`) were scalar loops on one thread, and the transformers path copied the + whole K/V prefix per layer and token (111 MB per token at 622 positions on Llama-3.2-3B). +* The existing parallelism was kernel-private and invisible: `parallelChunks` hard-coded + `Dispatchers.Default` and `availableProcessors()`, could not be turned off, could not be + observed, and nested a `runBlocking` on the very dispatcher a caller might already be on. +* The compile lane had no way to say "this op is head-parallel". `skainet.tensor_layouts` showed + the shape such a fact should take: a module attribute a consumer may read or ignore. + +== Goals + +. One abstraction, `sk.ainet.context.schedule.Schedule`, usable from every synchronous op on + every Kotlin target, with no dependency in `skainet-lang-core`. +. Structured concurrency on the JVM: a region is a `coroutineScope`; the first failure cancels + the siblings and is rethrown; no task outlives the region; writes happen-before its return. +. Hardware-aware defaults: the JVM `DirectCpuExecutionContext` runs core-count coroutines out of + the box; every other target and every context that does not opt in runs sequentially. +. An unhonoured request is visible (`TraceEvent.ScheduleDowngraded`), never a silent downgrade. +. Bit-identical results under any schedule, proven by tests. +. Schedule metadata through `dag { }` → `Operation.parameters` → `GraphNode.metadata` → + the `skainet.schedule` module attribute of the StableHLO export. + +== Non-Goals + +* IREE-side consumption of `skainet.schedule` (the header is metadata; lowering it into + dispatch hints is a later SKEEP). +* Concurrent *forwards* on one `ExecutionContext` — still one context per thread. +* Parallel schedules on Android, Kotlin/Native, JS or Wasm (they get `Sequential`; the JVM + implementation is the reference). +* Vectorised attention kernels (a SIMD reduction changes summation order; it is a separate, + tolerance-tested change). + +== Proposed Design + +=== The `Schedule` contract (`skainet-lang-core`) + +[source,kotlin] +---- +public interface Schedule { + public val parallelism: Int + public val name: String + public fun forRange(n: Int, grain: Int = 1, body: (start: Int, end: Int) -> Unit) + public fun forEach(count: Int, minPerTask: Int = 1, body: (index: Int) -> Unit) + public object Sequential : Schedule // body(0, n), inline + public companion object { fun tasksFor(n, grain, parallelism): Int; fun chunkFor(n, tasks): Int } +} +---- + +Contract for a `forRange` body: ranges are disjoint half-open intervals covering `[0, n)`; a +task count of one runs inline; a body writes only into pre-allocated, disjoint output regions +and reads only inputs that are immutable for the region; it never allocates through an +`ExecutionContext`, never calls `ctx.ops`, and never starts a nested region (an implementation +runs a nested call inline); the first failure is rethrown after every sibling finished or was +cancelled; all writes happen-before the return. + +=== Where it hangs + +* `ExecutionContext.schedule: Schedule` (default `Sequential`) and + `ExecutionContext.withSchedule(schedule): ExecutionContext` — the same rebuild seam as + `withTensorDataFactory`. The default cannot rebuild, returns `this`, and emits + `TraceEvent.ScheduleDowngraded(requested, effective, reason)`. +* `ScheduledExecutionContext(base, schedule) : ExecutionContext by base` and + `inline fun ExecutionContext.withSchedule(schedule, block)` — the decorator mirrors + `ScopedExecutionContext`; it survives `forwardScope` because `withTensorDataFactory` keeps the + schedule. +* Ops carry the schedule: `DefaultCpuOpsBase(dataFactory, schedule)`; the platform factory is + `(TensorDataFactory, Schedule) -> TensorOps`; `DirectCpuExecutionContext(schedule = …)` defaults + to `platformDefaultSchedule()` — `CoroutineSchedule.hardware()` on the JVM, `Sequential` elsewhere. + +=== The JVM implementation (`skainet-backend-cpu`) + +`CoroutineSchedule(dispatcher = Dispatchers.Default, parallelism = cores, sink)`: a region is +`runBlocking { coroutineScope { launch(dispatcher) { … } … ; body(first chunk) } }` — the caller +runs the first chunk itself, so no core idles on the join; a nested region (detected with a +thread-local) runs inline so the dispatcher never waits on itself; `dedicated(parallelism)` owns +its own pool for callers that already live on `Dispatchers.Default`. `parallelChunks(outputDim, +schedule)` replaces the old island; the Panama Q4_K/Q5_K SPI kernels gained a schedule-aware +overload; `ScheduleRegion` events report what ran when a sink is attached. + +=== First consumer: `scaledDotProductAttention` + +Every `(batch, head)` pair is independent — private `scores` scratch, disjoint output rows — so +the pairs are the units of `schedule.forRange`. The per-pair arithmetic and its order are the +sequential loop's, which is what makes the result bit-identical. Calls below +`SDPA_PARALLEL_MIN_WORK` multiply-adds (a decode step on a handful of heads) stay inline. + +=== Compile lane + +`dag { schedule(parallel("heads")) { op(sdpa, …) } }` or `op(matmul, …, schedule = +parallel("rows", parallelism = 8))` stamp a `ScheduleHint` under `skainet.schedule` on the node +(the `DtypePolicyDsl` channel); `ScheduleAnnotationPass` validates the requested dimensions per +op (sdpa: batch, heads; matmul: rows; conv: batch, out_channels), stamps the normalized hint into +`GraphNode.metadata` and reports every rejection as a diagnostic; `StableHloConverter` emits +`skainet.schedule = { = {parallel_dims = ["heads"], parallelism = 8}}` in the module header +beside `skainet.tensor_layouts`. One graph, extra schedule metadata. + +=== Thread-safety work that made it possible + +`KernelDispatch` and `KernelRegistry` now keep immutable snapshots with serialized writes, so +schedule workers may dispatch concurrently. Everything else in the hot path stays single-threaded +by contract: `ForwardScope` (bump allocator), `ScratchPool`, the `DefaultCpuOps` prepack caches, +`KernelProfile` — hence the rule that bodies never touch a context. + +== Requirements + +=== Functional + +* `ctx.withSchedule(s)` changes how many tasks an op uses and nothing else. +* `Schedule.Sequential` reproduces pre-SKEEP behaviour exactly. +* A `ScheduleHint` on a `dag { }` op reaches the StableHLO header unchanged; an unknown + dimension produces a diagnostic and no metadata. + +=== Non-Functional + +* No new dependency in `skainet-lang-core`, `skainet-backend-api`, or transformer-core. +* All API changes additive (`apiDump` diffs contain no removed lines). +* JS/Wasm/Native builds unchanged (`Sequential` only). + +== Compatibility and Migration + +Additive throughout: new interface members with defaults, secondary constructors keeping the +old JVM signatures (`DirectCpuExecutionContext`, `DefaultCpuOps`), SPI overloads with default +bodies. Behaviour change on the JVM only: `DirectCpuExecutionContext()` now runs +`scaledDotProductAttention` on the hardware schedule — bit-identical, opt out with +`withSchedule(Schedule.Sequential)`. + +== Rollout Plan + +. Registries safe for concurrent reads (`KernelDispatchConcurrencyTest`). +. `Schedule`, `ScheduleHint`, context seam, trace events (`ScheduledExecutionContextTest`). +. `CoroutineSchedule`, `parallelChunks(schedule)`, ops/factory/context plumbing + (`CoroutineScheduleTest`, `DirectCpuExecutionContextScheduleTest`). +. Scheduled SDPA (`SdpaScheduleParityTest`, `SdpaCoroutineParityTest`, JMH `SdpaScheduleBench`). +. Compile lane (`ScheduleDslTest`, `ScheduleAnnotationPassTest`, `ScheduleModuleAttributeTest`). +. Docs: xref:ROOT:explanation/schedules.adoc[Algorithm and schedule], + xref:ROOT:tutorials/schedule-getting-started.adoc[Schedule getting started] with an executable sample. +. Downstream: SKaiNET-transformers per-head attention on the same `Schedule` (its own spec). + +== Acceptance Criteria + +* Parity tests assert `assertContentEquals` between sequential and scheduled outputs. +* `apiCheck` green with dumps refreshed and no removed lines. +* `SdpaScheduleBench` shows the hardware schedule ahead of sequential on 8 heads × 4096 keys. +* Antora builds with the two new pages linked from `nav.adoc`. + +== Risks + +* **Nested `runBlocking` on `Dispatchers.Default`** — mitigated by caller participation, inline + nested regions and `dedicated()`; tested from inside a Default worker. +* **A body that touches the context** — the contract is documented on `forRange`; the parity + tests run bodies on foreign threads. +* **JIT-dependent lane reductions** — the Panama matmul's `reduceLanes` order changes when the + JIT intrinsifies it, so the very first call in a JVM can differ by an ULP from later ones + regardless of schedule (observed while writing the tests; filed as + https://github.com/SKaiNET-developers/SKaiNET/issues/1261[#1261]). Parity tests warm up first; + golden gates should do the same. +* **Region overhead on tiny ops** — `SDPA_PARALLEL_MIN_WORK` and `Schedule.tasksFor` keep small + calls inline; tune from the JMH numbers. + +== Open Questions + +* Should Android get a `CoroutineSchedule` (coroutines are jvmMain-only in backend-cpu today)? +* Should `ensureInstalled()` and `withSchedule` detect a caller already on `Dispatchers.Default` + and emit `ScheduleDowngraded` instead of running inline? +* When `TargetOptimizer.stableHloPasses()` exists, which IREE attribute should + `skainet.schedule` lower to? + +== References + +* Halide: Ragan-Kelley et al., "Halide: decoupling algorithms from schedules", PLDI 2013. +* Kotlin structured concurrency: `coroutineScope`, `launch`, `Dispatchers.Default`. +* xref:ROOT:explanation/dsl-principles.adoc[The DSL is compute]; + xref:skeep:003-unified-tensor-storage.adoc[SKEEP-003] (the memory model the contract leans on). +* Daily-StandAPP `docs/modules/planning/pages/verification-2026-09.adoc` — the profile that motivated this. diff --git a/docs/modules/skeep/pages/index.adoc b/docs/modules/skeep/pages/index.adoc index 9c0f7df04..1676ab155 100644 --- a/docs/modules/skeep/pages/index.adoc +++ b/docs/modules/skeep/pages/index.adoc @@ -126,4 +126,8 @@ Every proposal should include: | xref:skeep:004-virtual-tensor-layout.adoc[SKEEP-004] | Draft | Virtual tensor layout — the logical/physical split through the compile pipeline + +| xref:skeep:005-schedules-structured-concurrency.adoc[SKEEP-005] +| Implemented +| Schedules — structured concurrency for the compute layer (algorithm/schedule split) |=== diff --git a/skainet-docs-samples/src/commonTest/kotlin/sk/ainet/docs/samples/SamplesTest.kt b/skainet-docs-samples/src/commonTest/kotlin/sk/ainet/docs/samples/SamplesTest.kt index 14130dfee..6795675e9 100644 --- a/skainet-docs-samples/src/commonTest/kotlin/sk/ainet/docs/samples/SamplesTest.kt +++ b/skainet-docs-samples/src/commonTest/kotlin/sk/ainet/docs/samples/SamplesTest.kt @@ -30,6 +30,16 @@ class SamplesTest { assertEquals(111f, broadcast.data.get(0, 0), 1e-4f) } + @Test + fun scheduleDemo_is_bit_identical_across_schedules_and_reports_regions() { + val r = ScheduleDemo.run() + assertTrue(r.defaultScheduleName.startsWith("coroutines("), "JVM default is the hardware coroutine schedule, got ${r.defaultScheduleName}") + kotlin.test.assertContentEquals(r.sequential, r.scheduled, "a schedule never changes a result") + assertEquals(1, r.regions.size, "one parallel region for one attention call") + assertEquals(16, r.regions.single().elements) + assertEquals(2, r.regions.single().tasks) + } + @Test fun quickstart_forward_produces_class_scores() { val pixels = FloatArray(784) { 0f } From bbb1d8c9112317dfab8945afbf1f6c0d8bca109b Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Fri, 4 Sep 2026 00:06:39 +0200 Subject: [PATCH 4/4] SKEEP-005: measurements, compile-dag API dump - SKEEP-005 page gains a Measurements section: SdpaScheduleBench (hardware schedule 3.4-4.8x on prefill shapes, 2.0-2.7x on single-query decode) and the downstream transformers profile (attn.fused_compute 8.99 s -> 2.59 s, decode 7.7 -> 9.7 tok/s, greedy tokens identical). - skainet-compile-dag API dump: the two default members ExecutionContext gained (schedule, withSchedule) surface on GraphExecutionContext; additive. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018sqoaGs5M7C5uVw6cpzAnH --- .../005-schedules-structured-concurrency.adoc | 19 ++++++++++++++ .../skeep/partials/005-sdpa-bench.adoc | 26 +++++++++++++++++++ .../api/jvm/skainet-compile-dag.api | 4 +++ 3 files changed, 49 insertions(+) create mode 100644 docs/modules/skeep/partials/005-sdpa-bench.adoc diff --git a/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc b/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc index 873eff429..a02f239d9 100644 --- a/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc +++ b/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc @@ -181,6 +181,25 @@ bodies. Behaviour change on the JVM only: `DirectCpuExecutionContext()` now runs * **Region overhead on tiny ops** — `SDPA_PARALLEL_MIN_WORK` and `Schedule.tasksFor` keep small calls inline; tune from the JMH numbers. +== Measurements + +Downstream, on the SKaiNET-transformers `feature/attention-schedule` branch +(`AttentionScheduleSpeedProfile`, Llama-3.2-1B-Instruct Q8_0, i7-9750H 6c/12t, JDK 25, 512-token +prefill + 32 greedy tokens, 2026-09-03; greedy tokens identical in every row): + +[cols="3,1,1,1", options="header"] +|=== +| Schedule / KV cache | `attn.fused_compute` | `attn.kvcache` | Decode tok/s + +| `Sequential` / append (0.53.0 behaviour) | 8,991 ms | 447 ms | 7.7 +| `CoroutineSchedule.hardware()` / append | 2,773 ms | 369 ms | 8.7 +| `CoroutineSchedule.hardware()` / positional (copy-free) | 2,590 ms | 10 ms | 9.7 +|=== + +Engine microbenchmark (`SdpaScheduleBench`, same machine): see the table below. + +include::partial$005-sdpa-bench.adoc[] + == Open Questions * Should Android get a `CoroutineSchedule` (coroutines are jvmMain-only in backend-cpu today)? diff --git a/docs/modules/skeep/partials/005-sdpa-bench.adoc b/docs/modules/skeep/partials/005-sdpa-bench.adoc new file mode 100644 index 000000000..21f572c16 --- /dev/null +++ b/docs/modules/skeep/partials/005-sdpa-bench.adoc @@ -0,0 +1,26 @@ +`SdpaScheduleBench` (`DefaultCpuOps.scaledDotProductAttention`, batch 1, headDim 64, causal, +JMH avgt 5×10 s, 2026-09-03, i7-9750H 6c/12t, JDK 25; `hardware` = `CoroutineSchedule.hardware()` +with 12 tasks): + +[cols="1,1,1,1,1,1", options="header"] +|=== +| heads | seqKV | seqQ | sequential | hardware | speed-up + +| 8 | 128 | 1 | 0.335 ms | 0.334 ms | 1.0× (below `SDPA_PARALLEL_MIN_WORK`, runs inline) +| 8 | 128 | 64 | 15.9 ms | 4.0 ms | 4.0× +| 8 | 1024 | 1 | 3.9 ms | 1.8 ms | 2.2× +| 8 | 1024 | 64 | 185.5 ms | 52.9 ms | 3.5× +| 8 | 4096 | 1 | 16.2 ms | 7.9 ms | 2.0× +| 8 | 4096 | 64 | 833.1 ms | 223.3 ms | 3.7× +| 32 | 128 | 1 | 1.47 ms | 0.62 ms | 2.4× +| 32 | 128 | 64 | 64.3 ms | 13.3 ms | 4.8× +| 32 | 1024 | 1 | 16.7 ms | 6.2 ms | 2.7× +| 32 | 1024 | 64 | 755.6 ms | 185.1 ms | 4.1× +| 32 | 4096 | 1 | 68.6 ms | 31.7 ms | 2.2× +| 32 | 4096 | 64 | 2,859 ms | 835.7 ms | 3.4× +|=== + +Six physical cores, eight or thirty-two units of work: the prefill shapes (`seqQ = 64`) reach +3.4–4.8×, the single-query decode shapes 2.0–2.7× (the per-head work is small enough for the +fork/join and the memory traffic to show). Sequential error bars are wide because the +single-threaded run is at the mercy of turbo clocks; the scheduled runs are steady. diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index bfee28fd0..63f778a75 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api @@ -189,6 +189,7 @@ public final class sk/ainet/lang/graph/DefaultGraphExecutionContext : sk/ainet/l public fun getOps ()Lsk/ainet/lang/tensor/ops/KspTensorOps; public synthetic fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getPhase ()Lsk/ainet/context/Phase; + public fun getSchedule ()Lsk/ainet/context/schedule/Schedule; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public final fun getSession ()Lsk/ainet/lang/trace/TraceSession; public fun getTapeStack ()Lsk/ainet/tape/TapeStack; @@ -205,6 +206,7 @@ public final class sk/ainet/lang/graph/DefaultGraphExecutionContext : sk/ainet/l public fun stopRecording ()Lsk/ainet/tape/ExecutionTape; public final fun stopRecordingAndGet ()Lsk/ainet/tape/ExecutionTape; public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; @@ -467,6 +469,7 @@ public final class sk/ainet/lang/graph/exec/GraphExecutionContext$DefaultImpls { public static fun getInTraining (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Z public static fun getMemoryScope (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/memory/Scope; public static fun getMemoryTracker (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryTracker; + public static fun getSchedule (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/context/schedule/Schedule; public static fun getScratch (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/tensor/scratch/ScratchPool; public static fun getTraceSink (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/memory/trace/TraceSink; public static fun isRecording (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Z @@ -474,6 +477,7 @@ public final class sk/ainet/lang/graph/exec/GraphExecutionContext$DefaultImpls { public static fun placeholder (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun registerObserver (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/context/ExecutionObserver;)V public static fun unregisterObserver (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/context/ExecutionObserver;)V + public static fun withSchedule (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/context/ExecutionContext; public static fun withTensorDataFactory (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public static fun wrapByteArray (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public static fun wrapFloatArray (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor;