feat(storage): off-heap mmap tensor storage on Android — shared java.nio path + MappedGgufWeights (#921)

CHANGELOG.md

@@ -2,6 +2,26 @@
 
 ## [Unreleased]
 
+### Added
+
+- **Off-heap / mmap tensor storage on Android**
+  ([#921](https://github.com/SKaiNET-developers/SKaiNET/issues/921), SKEEP-002/SKEEP-003
+  slice): the java.nio memory-mapped storage that existed jvmMain-only is now shared
+  source between the JVM and Android compilations (`FileChannel.map` is API 1 — no JNI):
+  `MmapFloatTensorData`/`MmapTensorSource` (skainet-lang-core), `JvmMappedMemoryChunk`,
+  `MappedRandomAccessSource` and the `BufferHandle.FileBacked` resolver
+  `JvmFileBackedResolver` (skainet-io-core). New `MappedGgufWeights` (skainet-io-gguf,
+  JVM+Android) opens a GGUF once, maps it read-only, and serves dense F32 tensors as
+  zero-heap mapped views, any tensor as a `FileBacked` `TensorStorage` descriptor, and
+  packed payloads as heap bytes for the existing kernels. Weight bytes live in file-backed
+  pages the OS pages in/out — outside the hard ART heap cap that limited practical model
+  size on Android. Verified host-side (no device required): a 640 MB dense model loads and
+  reads through mapped views with **1.4 MB** of managed-heap allocation (0.0022x of the
+  dense size; per-thread allocation counters), and the Android compilation is exercised by
+  new `androidHostTest` suites (96 MB payload, ~240 KB used-heap growth). Files over 2 GB
+  are rejected fast (single-region mapping); windowed mapping is a follow-up under
+  SKEEP-003's IO pipeline improvement.
+
 ### Fixed
 
 - **GGUF `DEQUANTIZE_TO_FP32` no longer over-allocates**

build-logic/convention/src/main/kotlin/sk.ainet.dokka.gradle.kts

@@ -28,6 +28,17 @@ extensions.configure<DokkaExtension> {
             suppress.set(true)
         }
 
+        // Modules with a shared src/jvmAndroidMain source *directory* (#966)
+        // compile the same files into BOTH the jvm and android compilations.
+        // Dokka refuses files that belong to two source sets (dokka#3701), and
+        // the android pages would duplicate the jvm ones anyway — document the
+        // shared API once, via jvm, by suppressing the android source set.
+        if ((name == "androidMain" || name == "android") &&
+            projectDir.resolve("src/jvmAndroidMain").exists()
+        ) {
+            suppress.set(true)
+        }
+
         sourceLink {
             localDirectory.set(projectDir.resolve("src"))
             remoteUrl("https://github.com/SKaiNET-developers/skainet/tree/main/${project.path.replace(":", "/").removePrefix("/")}/src")

skainet-io/skainet-io-core/build.gradle.kts

@@ -90,6 +90,19 @@ kotlin {
             }
         }
 
+        // Memory-mapped IO (JvmMappedMemoryChunk, JvmFileBackedResolver,
+        // MappedRandomAccessSource), shared source between the JVM and Android
+        // compilations (#921): pure java.nio (FileChannel.map — API 1 on
+        // Android), so model weights can live in file-backed pages outside the
+        // ART heap. Shared directory rather than an intermediate source set —
+        // each compilation builds against its full platform classpath.
+        getByName("jvmMain") {
+            kotlin.srcDir("src/jvmAndroidMain/kotlin")
+        }
+        getByName("androidMain") {
+            kotlin.srcDir("src/jvmAndroidMain/kotlin")
+        }
+
         val commonTest by getting {
             dependencies {
                 implementation(libs.kotlin.test)

skainet-io/skainet-io-core/src/androidHostTest/kotlin/sk/ainet/io/MappedMemoryChunkAndroidHostTest.kt

@@ -0,0 +1,75 @@
+package sk.ainet.io
+
+import sk.ainet.lang.tensor.storage.BufferHandle
+import java.io.File
+import kotlin.test.AfterTest
+import kotlin.test.BeforeTest
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+
+/**
+ * Host-side test of the *Android compilation* of the shared mmap IO (#921):
+ * [JvmMappedMemoryChunk], [MappedRandomAccessSource] and
+ * [JvmFileBackedResolver] are java.nio-only (FileChannel.map is API 1) and
+ * are compiled into androidMain from the shared `jvmAndroidMain` source
+ * directory. This test compiles against the android variant and runs on the
+ * host JVM — no device/emulator required.
+ */
+class MappedMemoryChunkAndroidHostTest {
+
+    private val payload = ByteArray(64 * 1024) { ((it * 31) and 0xFF).toByte() }
+    private lateinit var file: File
+
+    @BeforeTest
+    fun setUp() {
+        file = File.createTempFile("android-mmap-chunk-", ".bin")
+        file.writeBytes(payload)
+    }
+
+    @AfterTest
+    fun tearDown() {
+        file.delete()
+    }
+
+    @Test
+    fun `mapped chunk reads, slices and offsets match the file`() {
+        JvmMappedMemoryChunk.open(file).use { chunk ->
+            assertEquals(payload.size.toLong(), chunk.size)
+            assertEquals(payload[0], chunk.readByte(0))
+            assertEquals(payload[12345], chunk.readByte(12345))
+            assertContentEquals(payload.copyOfRange(1000, 1256), chunk.readBytes(1000, 256))
+
+            val slice = chunk.slice(4096, 512)
+            assertContentEquals(payload.copyOfRange(4096, 4096 + 512), slice.readBytes(0, 512))
+        }
+    }
+
+    @Test
+    fun `mapped random access source serves positional reads`() {
+        MappedRandomAccessSource.open(file).use { source ->
+            assertEquals(payload.size.toLong(), source.size)
+            assertContentEquals(payload.copyOfRange(777, 777 + 64), source.readAt(777, 64))
+            val buf = ByteArray(128)
+            assertEquals(128, source.readAt(2048, buf, 0, 128))
+            assertContentEquals(payload.copyOfRange(2048, 2048 + 128), buf)
+        }
+    }
+
+    @Test
+    fun `FileBacked handles resolve to mmap-backed accessors`() {
+        val handle = BufferHandle.FileBacked(
+            path = file.absolutePath,
+            fileOffset = 8192,
+            sizeInBytes = 1024,
+        )
+        val accessor = JvmFileBackedResolver.resolveFileBacked(handle)
+        try {
+            assertEquals(1024, accessor.sizeInBytes)
+            assertEquals(payload[8192], accessor.readByte(0))
+            assertContentEquals(payload.copyOfRange(8192, 8192 + 1024), accessor.readBytes(0, 1024))
+        } finally {
+            accessor.close()
+        }
+    }
+}

skainet-io/skainet-io-gguf/build.gradle.kts

@@ -25,6 +25,9 @@ kotlin {
         compilerOptions {
             jvmTarget.set(JvmTarget.JVM_1_8)
         }
+        // Host-side (JVM) unit tests for the Android compilation — exercises
+        // the mmap-backed weight path (MappedGgufWeights) without a device (#921).
+        withHostTest {}
     }
 
     iosArm64()
@@ -63,6 +66,19 @@ kotlin {
                 implementation(libs.kotlin.test)
             }
         }
+
+        // Memory-mapped GGUF weight access (MappedGgufWeights), shared source
+        // between the JVM and Android compilations (#921): pure java.nio
+        // (FileChannel.map — API 1 on Android). Shared directory rather than an
+        // intermediate source set — each compilation builds against its full
+        // platform classpath.
+        getByName("jvmMain") {
+            kotlin.srcDir("src/jvmAndroidMain/kotlin")
+        }
+        getByName("androidMain") {
+            kotlin.srcDir("src/jvmAndroidMain/kotlin")
+        }
+
         val jvmTest by getting {
             dependencies {
                 implementation(libs.junit)

skainet-io/skainet-io-gguf/src/androidHostTest/kotlin/sk/ainet/io/gguf/MappedGgufWeightsAndroidHostTest.kt

@@ -0,0 +1,135 @@
+package sk.ainet.io.gguf
+
+import sk.ainet.io.JvmFileBackedResolver
+import sk.ainet.lang.tensor.storage.BufferHandle
+import sk.ainet.lang.types.FP32
+import java.io.File
+import java.io.RandomAccessFile
+import java.nio.ByteBuffer
+import java.nio.ByteOrder
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * Host-side test of the *Android compilation* of the mmap weight path (#921):
+ * this source set compiles against androidMain (android.jar nio API), so it
+ * proves `MappedGgufWeights`, `MmapFloatTensorData` and the shared
+ * `JvmFileBackedResolver` all build and behave on the Android variant —
+ * without a device or emulator (the precise allocation-counter budget gate
+ * lives in jvmTest's `MappedGgufHeapBudgetTest`, which exercises the
+ * byte-identical shared source).
+ *
+ * Includes a coarse heap check on a sparse 96 MB model: after load + reads,
+ * used-heap growth stays far under the payload size (lenient bound — host GC
+ * is not deterministic; the strict gate is the jvmTest allocation counter).
+ */
+class MappedGgufWeightsAndroidHostTest {
+
+    private fun writeSparseF32Gguf(file: File, tensorName: String, elements: Int, sentinels: Map<Int, Float>): File {
+        val head = ByteBuffer.allocate(4096).order(ByteOrder.LITTLE_ENDIAN)
+        head.putInt(0x46554747)
+        head.putInt(3)
+        head.putLong(1)
+        head.putLong(1)
+        val key = "general.architecture".encodeToByteArray()
+        head.putLong(key.size.toLong())
+        head.put(key)
+        head.putInt(GGUFValueType.STRING.value)
+        val value = "test".encodeToByteArray()
+        head.putLong(value.size.toLong())
+        head.put(value)
+        val nameBytes = tensorName.encodeToByteArray()
+        head.putLong(nameBytes.size.toLong())
+        head.put(nameBytes)
+        head.putInt(1)
+        head.putLong(elements.toLong())
+        head.putInt(GGMLQuantizationType.F32.value)
+        head.putLong(0)
+        val padding = (32 - (head.position() % 32)) % 32
+        repeat(padding) { head.put(0) }
+        val dataStart = head.position().toLong()
+
+        RandomAccessFile(file, "rw").use { raf ->
+            raf.write(head.array(), 0, head.position())
+            raf.setLength(dataStart + elements.toLong() * 4)
+            for ((idx, v) in sentinels) {
+                raf.seek(dataStart + idx.toLong() * 4)
+                val bits = v.toRawBits()
+                raf.write(
+                    byteArrayOf(
+                        (bits and 0xFF).toByte(),
+                        ((bits shr 8) and 0xFF).toByte(),
+                        ((bits shr 16) and 0xFF).toByte(),
+                        ((bits shr 24) and 0xFF).toByte(),
+                    )
+                )
+            }
+        }
+        return file
+    }
+
+    @Test
+    fun `android-compiled mapped weight path loads a 96 MB tensor without heap-sized allocation`() {
+        val elements = 24 * 1024 * 1024 // 96 MB dense F32
+        val sentinels = mapOf(0 to 1.5f, 12_345_678 to -2.25f, elements - 1 to 3.125f)
+        val file = File.createTempFile("android_mmap_", ".gguf")
+        file.deleteOnExit()
+        writeSparseF32Gguf(file, "big.weight", elements, sentinels)
+        try {
+            val rt = Runtime.getRuntime()
+            System.gc()
+            val usedBefore = rt.totalMemory() - rt.freeMemory()
+
+            var checksum = 0.0
+            MappedGgufWeights.open(file.absolutePath).use { weights ->
+                val tensor = weights.mappedFloatTensor<FP32>("big.weight")
+                assertEquals(elements, tensor.shape.volume)
+                var i = 0
+                while (i < elements) {
+                    checksum += tensor[i]
+                    i += 262_144
+                }
+                for ((idx, v) in sentinels) {
+                    assertEquals(v, tensor[idx], "big.weight[$idx]")
+                }
+
+                // FileBacked descriptor + the shared resolver, on the android variant.
+                val storage = weights.mappedStorage("big.weight")
+                assertTrue(storage.isFileBacked)
+                val accessor = JvmFileBackedResolver.resolveFileBacked(storage.buffer as BufferHandle.FileBacked)
+                try {
+                    // First 4 sentinel bytes come back identical through the mmap accessor.
+                    val viaAccessor = accessor.readBytes(0, 4)
+                    val expectedBits = 1.5f.toRawBits()
+                    assertContentEquals(
+                        byteArrayOf(
+                            (expectedBits and 0xFF).toByte(),
+                            ((expectedBits shr 8) and 0xFF).toByte(),
+                            ((expectedBits shr 16) and 0xFF).toByte(),
+                            ((expectedBits shr 24) and 0xFF).toByte(),
+                        ),
+                        viaAccessor,
+                    )
+                } finally {
+                    accessor.close()
+                }
+            }
+
+            System.gc()
+            val usedAfter = rt.totalMemory() - rt.freeMemory()
+            val growth = usedAfter - usedBefore
+            println("android host mapped load: payload=96 MB, used-heap growth=${growth / 1024} KB, checksum=$checksum")
+            // Lenient (GC nondeterminism): far below the 96 MB payload — a heap
+            // materialization would add >= 96 MB here.
+            assertTrue(
+                growth < 48L * 1024 * 1024,
+                "used heap grew by ${growth / (1024 * 1024)} MB for a 96 MB mapped payload — " +
+                    "weights are being materialized on the managed heap (#921)",
+            )
+        } finally {
+            file.delete()
+        }
+    }
+}