Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) { _ ->
Expand Down
Loading