From b05343eee3bd8715100d8038fc8b5fdecfed99b9 Mon Sep 17 00:00:00 2001
From: Michal Harakal <michal.harakal@gmail.com>
Date: Sat, 17 Oct 2026 18:34:12 +0200
Subject: [PATCH] perf(backend-cpu): dense-FP32 fast path for split
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

DefaultCpuOpsBase.split built every output element through dataFactory.init with
outIdx.copyOf() plus a vararg-spread get — two IntArray copies and a boxed value
per element. RoPE's split-half rotation calls it every step, which made it the top
decode allocation site once the Vector API fallback was gone (#1357). Dense FP32
(floatWindowOf: plain or slab-backed) now copies each chunk as outer contiguous runs
with copyInto and adopts the result through floatResult; other data keeps the
generic path.

Refs #1357

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 29 +++++++++++++++++++
 .../tensor/ops/DefaultCpuOpsShapeOpsTest.kt   | 29 +++++++++++++++++++
 2 files changed, 58 insertions(+)

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 db6da6530..9ef19c248 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
@@ -2417,6 +2417,35 @@ public open class DefaultCpuOpsBase(
         val total = tensor.shape[nd]
         val chunks = (total + splitSize - 1) / splitSize
         val result = ArrayList<Tensor<T, V>>(chunks)
+
+        // Dense-FP32 fast path: each chunk is `outer` contiguous runs of `size * inner` floats,
+        // copied with copyInto. The generic path below costs two IntArray copies (outIdx.copyOf
+        // + the vararg spread) and a boxed value per element — RoPE's per-step split made it a
+        // top decode allocation site (#1357).
+        val win = floatWindowOf(tensor)
+        if (win != null) {
+            val dims = tensor.shape.dimensions
+            var outer = 1
+            for (i in 0 until nd) outer *= dims[i]
+            var inner = 1
+            for (i in nd + 1 until rank) inner *= dims[i]
+            var start = 0
+            for (c in 0 until chunks) {
+                val size = minOf(splitSize, total - start)
+                val run = size * inner
+                val out = FloatArray(outer * run)
+                for (o in 0 until outer) {
+                    val from = win.off + (o * total + start) * inner
+                    win.arr.copyInto(out, destinationOffset = o * run, startIndex = from, endIndex = from + run)
+                }
+                val newDims = dims.copyOf()
+                newDims[nd] = size
+                result += floatResult(Shape(newDims), tensor.dtype, out, tensor)
+                start += size
+            }
+            return result
+        }
+
         var offset = 0
         for (c in 0 until chunks) {
             val size = minOf(splitSize, total - offset)
diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsShapeOpsTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsShapeOpsTest.kt
index 99fa60802..58c80a4d1 100644
--- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsShapeOpsTest.kt
+++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsShapeOpsTest.kt
@@ -98,6 +98,35 @@ class DefaultCpuOpsShapeOpsTest {
         }
     }
 
+    @Test
+    fun split_rank4_every_dim_matches_element_mapping() {
+        // Exercises the dense-FP32 copyInto path on every axis, with a remainder chunk:
+        // part p, element idx  ==  source element with idx[dim] + p * splitSize.
+        data(ctx) { _ ->
+            val dims = intArrayOf(2, 3, 5, 4)
+            val t = tensor<FP32, Float> {
+                shape(*dims) { init { (((it[0] * 3 + it[1]) * 5 + it[2]) * 4 + it[3]).toFloat() } }
+            }
+            for (dim in 0 until 4) {
+                val splitSize = 2
+                val parts = ops.split(t, splitSize = splitSize, dim = dim)
+                assertEquals((dims[dim] + splitSize - 1) / splitSize, parts.size)
+                for ((p, part) in parts.withIndex()) {
+                    val pd = part.shape.dimensions
+                    assertEquals(minOf(splitSize, dims[dim] - p * splitSize), pd[dim])
+                    for (a in 0 until pd[0]) for (b in 0 until pd[1]) for (c in 0 until pd[2]) for (d in 0 until pd[3]) {
+                        val src = intArrayOf(a, b, c, d).also { it[dim] += p * splitSize }
+                        assertEquals(t.data[src[0], src[1], src[2], src[3]], part.data[a, b, c, d], "dim=$dim part=$p [$a,$b,$c,$d]")
+                    }
+                }
+            }
+            // negative dim resolves like its positive twin
+            val neg = ops.split(t, splitSize = 3, dim = -1)
+            assertEquals(listOf(Shape(2, 3, 5, 3), Shape(2, 3, 5, 1)), neg.map { it.shape })
+            assertEquals(t.data[1, 2, 4, 3], neg[1].data[1, 2, 4, 0])
+        }
+    }
+
     @Test
     fun squeeze_and_unsqueeze_behavior() {
         data(ctx) { _ ->