SKEEP-005: schedules — structured concurrency for compute, parallel SDPA, compile-lane metadata

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

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<TraceEvent.ScheduleRegion>,
+    )
+
+    // 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<FP32, Float>, Tensor<FP32, Float>, Tensor<FP32, Float>> {
+        fun tensor(seq: Int, seed: Int): Tensor<FP32, Float> {
+            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<TraceEvent.ScheduleRegion>()
+        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)
+    }
+}

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]

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

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<TraceEvent.ScheduleRegion>()      // op, schedule, elements, tasks, duration
+sink.eventsOf<TraceEvent.ScheduleDowngraded>()  // 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.

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.

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]