perf(backend-cpu): dense FP32 x·Wᵀ runs under the context schedule (SKEEP-005)

CHANGELOG.md

@@ -2,6 +2,14 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- **`CoroutineSchedule.forRange` could call its body with an empty, inverted range** (`skainet-backend-cpu`): the chunk length is rounded up, so for some `n` and task counts the last chunk started past `n` (63 elements on 12 tasks: chunks of 6, the twelfth at 66..63). The engine's own bodies loop over the range and did nothing for it; a body that passes the range length to a kernel saw a negative size. The task count now follows the chunk length, and every range is non-empty and covers `[0, n)` as the `Schedule` contract states.
+
+### Changed
+
+- **Dense FP32 `x · Wᵀ` on the JVM native lane runs under the context schedule** (`skainet-backend-cpu`, SKEEP-005): `DefaultCpuOpsJvm` used to hand the whole product to `skainet_fp32_matmul_wt` in one call and leave threading to the C row pool, which is sized at compile time (four threads) and re-reads the weight for every activation row, so an encoder or prefill pass on a twelve-core host ran on about three cores at the shared cache's bandwidth. A product with several activation rows now splits the weight's rows across the `Schedule` — up to `parallelism` native calls over disjoint column slices, each slice narrow enough for a core's own cache — so `DirectCpuExecutionContext()` (the hardware schedule) uses every core and `withSchedule(Schedule.Sequential)` runs the old single call. Results are bit-identical either way: every output stays the same contiguous dot product in the same order. A single-row product (a decode step) is unchanged, as is a weight in a confined arena, which stays on the calling thread.
+
 ## [0.62.1] - 2026-10-21
 
 Headline: **the native Q8_0 / Q4_0 / Q5_0 / Q5_1 kernel bindings no longer zero the earlier rows of a

skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt

@@ -53,13 +53,17 @@ public open class CoroutineSchedule @JvmOverloads constructor(
     private val executor: Executor = dispatcher.asExecutor()
 
     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) {
+        val wanted = Schedule.tasksFor(n, grain, parallelism)
+        if (wanted == 0) return
+        if (wanted == 1 || inRegion.get() == true) {
             body(0, n)
             return
         }
-        val chunk = Schedule.chunkFor(n, tasks)
+        val chunk = Schedule.chunkFor(n, wanted)
+        // Rounding the chunk up can leave the last of `wanted` chunks past the end (63 rows on 12
+        // tasks: chunks of 6, the twelfth starts at 66); the contract promises non-empty ranges that
+        // cover [0, n), so the task count follows the chunk, never the other way round.
+        val tasks = (n + chunk - 1) / chunk
         val started = if (sink.isEnabled) TraceClock.nowNanos() else 0L
         val region = Region(n, chunk, tasks, body)
         inRegion.set(true)

skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt

@@ -304,9 +304,20 @@ internal class DefaultCpuOpsJvm(
     /**
      * `x · Wᵀ` for a dense FP32 `[out, in]` weight through the native direct kernel — the JVM-side
      * half of SKEEP-007 §T0.2: no transposed copy of the weight is made, the weight is read where it
-     * lies (a segment-backed weight as is, a heap one after a single off-heap copy), and the C kernel
-     * threads over output rows. `null` when the kernel is not available or the operands are not
-     * dense FP32 the kernel can take, in which case the caller continues to the cached transpose.
+     * lies (a segment-backed weight as is, a heap one after a single off-heap copy). `null` when the
+     * kernel is not available or the operands are not dense FP32 the kernel can take, in which case
+     * the caller continues to the cached transpose.
+     *
+     * The output columns — the weight's rows — are the [schedule]'s units (SKEEP-005): a product with
+     * several activation rows (a prefill or encoder pass) runs as up to `parallelism` native calls,
+     * each over a disjoint slice of the weight's rows and all activation rows, into a private
+     * `[m, cols]` block copied into the result. The kernel reads its weight rows once per activation
+     * row, so a slice narrow enough for a core's own cache is what makes the pass scale; a slice is
+     * never wider than the kernel's own threading threshold, which keeps the C row pool out of a
+     * scheduled call. Every output element is the same contiguous dot product in the same `p` order
+     * whichever task computes it, so a scheduled run is bit-identical to a sequential one. A single
+     * row (a decode step) has nothing to amortise and goes to the kernel as one call, whose own pool
+     * threads over the columns.
      */
     @Suppress("UNCHECKED_CAST")
     private fun <T : DType, V> matmulDenseWeightTransposedNative(x: Tensor<T, V>, weight: Tensor<T, V>): Tensor<T, V>? {
@@ -319,7 +330,26 @@ internal class DefaultCpuOpsJvm(
         val activation = heapActivation(x) ?: return null
         val m = if (k == 0) 0 else x.shape.volume / k
         val out = FloatArray(m * n)
-        if (m > 0 && n > 0) kernel.matmulWt(activation.first, activation.second, resident.first, resident.second, out, 0, m, n, k)
+        if (m > 0 && n > 0) {
+            val (xArr, xOff) = activation
+            val (wSeg, wOff) = resident
+            val tasks = if (m >= 2 && m.toLong() * n >= PARALLEL_MATMUL_MIN_OUTPUT && isSharedSegment(wSeg)) {
+                sk.ainet.context.schedule.Schedule.tasksFor(n, WT_COLUMN_GRAIN, schedule.parallelism)
+            } else {
+                1
+            }
+            if (tasks <= 1) {
+                kernel.matmulWt(xArr, xOff, wSeg, wOff, out, 0, m, n, k)
+            } else {
+                val grain = maxOf(WT_COLUMN_GRAIN, (n + tasks - 1) / tasks)
+                schedule.forRange(n, grain) { o0, o1 ->
+                    val cols = o1 - o0
+                    val block = FloatArray(m * cols)
+                    kernel.matmulWt(xArr, xOff, wSeg, wOff + o0.toLong() * k * java.lang.Float.BYTES, block, 0, m, cols, k)
+                    for (i in 0 until m) System.arraycopy(block, i * cols, out, i * n + o0, cols)
+                }
+            }
+        }
         val outShape = if (x.rank == 1) Shape(n) else Shape(*(x.shape.dimensions.copyOf(x.rank - 1) + n))
         val outData = dataFactory.adoptFloatArray<T, Float>(outShape, x.dtype, out) as TensorData<T, V>
         return newTensor(outData, x.dtype, x, weight)
@@ -341,6 +371,13 @@ internal class DefaultCpuOpsJvm(
         else -> null
     }
 
+    /**
+     * Whether [segment] may be read from a thread other than the current one — a shared, automatic
+     * or global arena. A confined segment (another caller's `Arena.ofConfined()`) is not, and a
+     * product over it stays on the calling thread. The probe thread is never started.
+     */
+    private fun isSharedSegment(segment: java.lang.foreign.MemorySegment): Boolean = segment.isAccessibleBy(SEGMENT_PROBE_THREAD)
+
     /** The weight's `floats` floats as an off-heap segment and byte offset: its own segment, or the resident copy of its heap array. */
     private fun residentDenseWeight(weight: Tensor<*, *>, floats: Long): Pair<java.lang.foreign.MemorySegment, Long>? {
         val bytes = floats * java.lang.Float.BYTES
@@ -1402,3 +1439,12 @@ internal class DefaultCpuOpsJvm(
         return CpuTensor(outData as TensorData<T, V>, this, tensor.dtype)
     }
 }
+
+/** A thread that never runs: the probe for [java.lang.foreign.MemorySegment.isAccessibleBy] (a confined segment answers `false` for it). */
+private val SEGMENT_PROBE_THREAD: Thread = Thread { }
+
+/**
+ * Fewest weight rows one scheduled `x · Wᵀ` task takes: four is the kernel's row tile, sixteen
+ * keeps the per-call staging (an activation copy, an arena) a rounding error beside the dot products.
+ */
+private const val WT_COLUMN_GRAIN: Int = 16

skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt

@@ -26,6 +26,7 @@ class CoroutineScheduleTest {
         val tasks = AtomicInteger()
         schedule.forRange(n, grain) { s, e ->
             tasks.incrementAndGet()
+            assertTrue(s < e, "range [$s, $e) must be non-empty")
             synchronized(seen) {
                 for (i in s until e) {
                     assertTrue(!seen[i], "index $i visited twice")
@@ -36,18 +37,41 @@ class CoroutineScheduleTest {
         return seen to tasks.get()
     }
 
+    /** Ranges a schedule of [parallelism] cuts [n] elements at [grain] into: the chunk is rounded up, the count follows it. */
+    private fun expectedTasks(n: Int, grain: Int, parallelism: Int): Int {
+        val wanted = Schedule.tasksFor(n, grain, parallelism)
+        if (wanted == 0) return 0
+        val chunk = Schedule.chunkFor(n, wanted)
+        return (n + chunk - 1) / chunk
+    }
+
     @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")
+        for (parallelism in listOf(4, 12)) {
+            val schedule = CoroutineSchedule(parallelism = parallelism)
+            for (n in listOf(0, 1, 5, 7, 63, 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 parallelism=$parallelism")
+                    assertEquals(expectedTasks(n, grain, parallelism), tasks, "n=$n grain=$grain parallelism=$parallelism task count")
+                }
             }
         }
     }
 
+    @Test
+    fun aRoundedUpChunkNeverYieldsAnEmptyTrailingRange() {
+        // 63 on 12 tasks: chunks of 6 — eleven of them cover [0, 63), a twelfth would start at 66.
+        val ranges = mutableListOf<Pair<Int, Int>>()
+        CoroutineSchedule(parallelism = 12).forRange(63) { s, e -> synchronized(ranges) { ranges += s to e } }
+        assertEquals(11, ranges.size)
+        assertTrue(ranges.all { (s, e) -> s < e && e <= 63 })
+        // 5 on 4 tasks: chunks of 2 — three ranges, not four.
+        val small = mutableListOf<Pair<Int, Int>>()
+        CoroutineSchedule(parallelism = 4).forRange(5) { s, e -> synchronized(small) { small += s to e } }
+        assertEquals(setOf(0 to 2, 2 to 4, 4 to 5), small.toSet())
+    }
+
     @Test
     fun parallelismOneIsSequentialOnTheCallerThread() {
         val schedule = CoroutineSchedule(parallelism = 1)

skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DenseWeightTransposedScheduleTest.kt

@@ -0,0 +1,98 @@
+package sk.ainet.exec.tensor.ops
+
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.context.ExecutionContext
+import sk.ainet.context.schedule.Schedule
+import sk.ainet.exec.kernel.NativeFp32MatmulWtKernel
+import sk.ainet.exec.schedule.CoroutineSchedule
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.types.FP32
+import kotlin.math.abs
+import kotlin.test.BeforeTest
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertTrue
+
+/**
+ * The dense `x · Wᵀ` product through the native direct kernel under the context schedule
+ * (SKEEP-005): the activation rows are split across tasks, one native call each, and the result
+ * is bit-identical to the sequential schedule and to a plain row-dot reference. Shapes cross the
+ * kernel's own row-pool threshold (512 output columns) and stay under it, and include the
+ * single-row decode shape that has nothing to split.
+ */
+class DenseWeightTransposedScheduleTest {
+
+    @BeforeTest
+    fun checkAvailable() {
+        assertTrue(NativeFp32MatmulWtKernel.isAvailable(), "bundled libskainet_kernels must resolve skainet_fp32_matmul_wt")
+        // The direct path lives on the Vector API tier; without it the common ops serve and this test would check nothing.
+        assertTrue(DirectCpuExecutionContext().ops::class.simpleName == "DefaultCpuOpsJvm", "the test JVM must run with --add-modules jdk.incubator.vector")
+    }
+
+    private fun fill(size: Int, seed: Int) = FloatArray(size) { i -> (((i * 31L + seed * 17L) % 23L) - 11L) / 11f }
+
+    private fun product(ctx: ExecutionContext, m: Int, n: Int, k: Int): FloatArray {
+        val x: Tensor<FP32, Float> = ctx.fromFloatArray(Shape(m, k), FP32::class, fill(m * k, 1))
+        val w: Tensor<FP32, Float> = ctx.fromFloatArray(Shape(n, k), FP32::class, fill(n * k, 2))
+        val out = ctx.ops.matmulWeightTransposed(x, w)
+        assertContentEquals(intArrayOf(m, n), out.shape.dimensions)
+        return out.data.copyToFloatArray()
+    }
+
+    private fun reference(m: Int, n: Int, k: Int): FloatArray {
+        val x = fill(m * k, 1)
+        val w = fill(n * k, 2)
+        val out = FloatArray(m * n)
+        for (i in 0 until m) for (o in 0 until n) {
+            var acc = 0.0
+            for (p in 0 until k) acc += x[i * k + p].toDouble() * w[o * k + p]
+            out[i * n + o] = acc.toFloat()
+        }
+        return out
+    }
+
+    @Test
+    fun scheduledRowsMatchSequentialBitForBit() {
+        val sequential = DirectCpuExecutionContext(schedule = Schedule.Sequential)
+        val parallel = DirectCpuExecutionContext(schedule = CoroutineSchedule(parallelism = 4))
+        for ((m, n, k) in listOf(Triple(63, 1024, 1024), Triple(183, 640, 1024), Triple(5, 256, 64), Triple(2, 1024, 32), Triple(1, 1024, 1024))) {
+            repeat(2) {
+                val a = product(sequential, m, n, k)
+                val b = product(parallel, m, n, k)
+                assertContentEquals(a, b, "m=$m n=$n k=$k: the schedule must not change a single bit")
+                val ref = reference(m, n, k)
+                val tol = 1e-5f * k
+                for (i in ref.indices) assertTrue(abs(ref[i] - a[i]) <= tol * maxOf(1f, abs(ref[i])), "m=$m n=$n k=$k element $i: ${a[i]} vs reference ${ref[i]}")
+            }
+        }
+    }
+
+    /**
+     * Opt-in timing of the split: `-Dskainet.runBench=true`. Encoder-sized products (a 63- and a
+     * 183-token sequence against ModernBERT-large's 1024-wide projections and its 2624-wide
+     * feed-forward) on the sequential schedule and on every core.
+     */
+    @Test
+    fun scheduledRowsBench() {
+        if (System.getProperty("skainet.runBench") != "true") {
+            println("DenseWeightTransposedScheduleTest.scheduledRowsBench skipped — pass -Dskainet.runBench=true to enable.")
+            return
+        }
+        val hardware = CoroutineSchedule.hardware()
+        val contexts = listOf("sequential" to DirectCpuExecutionContext(schedule = Schedule.Sequential), hardware.name to DirectCpuExecutionContext(schedule = hardware))
+        for ((m, n, k) in listOf(Triple(63, 1024, 1024), Triple(63, 2624, 1024), Triple(63, 1024, 2624), Triple(183, 1024, 1024), Triple(1, 1024, 1024))) {
+            for ((label, ctx) in contexts) {
+                val x: Tensor<FP32, Float> = ctx.fromFloatArray(Shape(m, k), FP32::class, fill(m * k, 1))
+                val w: Tensor<FP32, Float> = ctx.fromFloatArray(Shape(n, k), FP32::class, fill(n * k, 2))
+                repeat(3) { ctx.ops.matmulWeightTransposed(x, w) }
+                val reps = 10
+                val t0 = System.nanoTime()
+                repeat(reps) { ctx.ops.matmulWeightTransposed(x, w) }
+                val perCall = (System.nanoTime() - t0) / 1e6 / reps
+                val gflops = 2.0 * m * n * k / (perCall * 1e6)
+                println("[wt-schedule] m=$m n=$n k=$k $label: ${"%.3f".format(perCall)} ms/call  ${"%.1f".format(gflops)} GFLOP/s")
+            }
+        }
+    }
+}