diff --git a/CHANGELOG.md b/CHANGELOG.md index a84b7548..a4dd6a35 100644 --- a/CHANGELOG.md +++ b/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 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 index 94cecc58..f0918137 100644 --- 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 @@ -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) 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 4e592bba..aefd68ab 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 @@ -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 matmulDenseWeightTransposedNative(x: Tensor, weight: Tensor): Tensor? { @@ -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(outShape, x.dtype, out) as TensorData 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? { val bytes = floats * java.lang.Float.BYTES @@ -1402,3 +1439,12 @@ internal class DefaultCpuOpsJvm( return CpuTensor(outData as TensorData, 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 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 index 2c938b84..11aba77f 100644 --- 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 @@ -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>() + 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>() + 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) diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DenseWeightTransposedScheduleTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DenseWeightTransposedScheduleTest.kt new file mode 100644 index 00000000..b8669ae3 --- /dev/null +++ b/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 = ctx.fromFloatArray(Shape(m, k), FP32::class, fill(m * k, 1)) + val w: Tensor = 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 = ctx.fromFloatArray(Shape(m, k), FP32::class, fill(m * k, 1)) + val w: Tensor = 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") + } + } + } +}