From 5a68a97bcfe862d0c346f7db425d90267e9fdbbf Mon Sep 17 00:00:00 2001
From: Michal Harakal <michal.harakal@gmail.com>
Date: Sat, 17 Oct 2026 18:36:01 +0200
Subject: [PATCH] perf(lang-core): bulk copyToFloatArray for
 MmapFloatTensorData
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

MmapFloatTensorData inherited TensorData.copyToFloatArray, which walks every element
through get(*indices) — an IntArray copy for the spread plus a boxed Float each.
RMSNormalization.eagerFusedRmsNorm copies its (mmap'd) gain every forward, so in a
BitNet 2B4T JVM decode this was ~10 GB of allocation per 45 s. Override with one
bulk FloatBuffer read (duplicate, so the shared view's position is untouched;
statements kept separate for Android's pre-Java-9 nio).

Refs #1357

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../ainet/lang/tensor/data/MmapTensorData.kt  | 14 +++++
 .../data/MmapFloatTensorDataCopyTest.kt       | 51 +++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/tensor/data/MmapFloatTensorDataCopyTest.kt

diff --git a/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/MmapTensorData.kt b/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/MmapTensorData.kt
index 1f66fef7..523b29b5 100644
--- a/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/MmapTensorData.kt
+++ b/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/tensor/data/MmapTensorData.kt
@@ -68,6 +68,20 @@ public class MmapFloatTensorData<T : DType>(
         floatBuffer.put(flatIndex, value)
     }
 
+    /**
+     * One bulk read of the mapped floats. The TensorData default walks every element through
+     * `get(*indices)` — an IntArray copy for the spread plus a boxed Float per element — which
+     * made per-forward copies of mmap'd weights (e.g. RMSNorm gains) the top allocation site.
+     */
+    override fun copyToFloatArray(): FloatArray {
+        val out = FloatArray(shape.volume)
+        // Separate statements: on Android's (pre-Java-9) nio, position() returns Buffer.
+        val view = floatBuffer.duplicate()
+        view.position(0)
+        view.get(out)
+        return out
+    }
+
     private fun calcFlatIndex(indices: IntArray): Int {
         require(indices.size == shape.dimensions.size) {
             "Number of indices (${indices.size}) must match tensor dimensions (${shape.dimensions.size})"
diff --git a/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/tensor/data/MmapFloatTensorDataCopyTest.kt b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/tensor/data/MmapFloatTensorDataCopyTest.kt
new file mode 100644
index 00000000..d22bbb9a
--- /dev/null
+++ b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/tensor/data/MmapFloatTensorDataCopyTest.kt
@@ -0,0 +1,51 @@
+package sk.ainet.lang.tensor.data
+
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.types.FP32
+import java.nio.ByteBuffer
+import java.nio.ByteOrder
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+
+class MmapFloatTensorDataCopyTest {
+
+    private fun leFloats(n: Int, f: (Int) -> Float): ByteBuffer {
+        val bb = ByteBuffer.allocateDirect(n * 4).order(ByteOrder.LITTLE_ENDIAN)
+        for (i in 0 until n) bb.putFloat(i * 4, f(i))
+        return bb
+    }
+
+    private fun elementwise(d: MmapFloatTensorData<FP32>): FloatArray {
+        val (a, b, c) = d.shape.dimensions
+        return FloatArray(d.shape.volume) { i -> d.get(i / (b * c), (i / c) % b, i % c) }
+    }
+
+    @Test
+    fun bulk_copy_matches_elementwise_get_in_row_major_order() {
+        val shape = Shape(2, 3, 4)
+        val d = MmapFloatTensorData<FP32>(shape, leFloats(24) { it * 0.5f - 3f })
+        assertContentEquals(elementwise(d), d.copyToFloatArray())
+    }
+
+    @Test
+    fun copies_only_the_tensor_volume_from_a_larger_buffer() {
+        val d = MmapFloatTensorData<FP32>(Shape(2, 2, 2), leFloats(16) { it.toFloat() })
+        assertContentEquals(FloatArray(8) { it.toFloat() }, d.copyToFloatArray())
+    }
+
+    @Test
+    fun slice_at_nonzero_offset_and_shared_buffer_position_untouched() {
+        // A tensor view at byte offset 40 of a larger region, the way MmapTensorSource slices it.
+        val region = leFloats(40) { 100f + it }
+        val dup = region.duplicate()
+        dup.position(40)
+        dup.limit(40 + 24 * 4)
+        val d = MmapFloatTensorData<FP32>(Shape(2, 3, 4), dup.slice())
+        val first = d.copyToFloatArray()
+        assertContentEquals(FloatArray(24) { 110f + it }, first)
+        assertEquals(0, d.floatBuffer.position())
+        assertContentEquals(first, d.copyToFloatArray())   // repeatable
+        assertEquals(110f, d.get(0, 0, 0))
+    }
+}