P7: AllocationResolver — one resolver decides domain and scope

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -1360,6 +1360,17 @@ public final class sk/ainet/lang/memory/plan/ActualMemory$Companion {
 	public final fun from (Lsk/ainet/lang/memory/trace/RecordingTraceSink;)Lsk/ainet/lang/memory/plan/ActualMemory;
 }
 
+public final class sk/ainet/lang/memory/plan/AllocationResolver {
+	public static final field INSTANCE Lsk/ainet/lang/memory/plan/AllocationResolver;
+	public final fun explain (Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;)Ljava/lang/String;
+	public static synthetic fun explain$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Ljava/lang/String;
+	public final fun residentFormat (Lsk/ainet/lang/memory/plan/PlanTensor;)Lsk/ainet/lang/memory/Format;
+	public final fun resolve (Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;)Lsk/ainet/lang/memory/AllocationSpec;
+	public static synthetic fun resolve$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec;
+	public final fun resolveTransient (Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;Lsk/ainet/lang/memory/ScopeKind;)Lsk/ainet/lang/memory/AllocationSpec;
+	public static synthetic fun resolveTransient$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;Lsk/ainet/lang/memory/ScopeKind;ILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec;
+}
+
 public final class sk/ainet/lang/memory/plan/Budget {
 	public static final field Companion Lsk/ainet/lang/memory/plan/Budget$Companion;
 	public static final field RESERVE_ANDROID_JVM J
@@ -1649,6 +1660,8 @@ public final class sk/ainet/lang/memory/plan/PlanLine {
 public final class sk/ainet/lang/memory/plan/PlanTensor {
 	public fun <init> (Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJLsk/ainet/lang/memory/plan/WeightForm;)V
 	public synthetic fun <init> (Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJLsk/ainet/lang/memory/plan/WeightForm;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public final fun allocation (Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;)Lsk/ainet/lang/memory/AllocationSpec;
+	public static synthetic fun allocation$default (Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec;
 	public final fun component1 ()Ljava/lang/String;
 	public final fun component2 ()Lsk/ainet/lang/tensor/TensorId;
 	public final fun component3 ()Lsk/ainet/lang/memory/Format;
@@ -1658,7 +1671,6 @@ public final class sk/ainet/lang/memory/plan/PlanTensor {
 	public final fun copy (Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJLsk/ainet/lang/memory/plan/WeightForm;)Lsk/ainet/lang/memory/plan/PlanTensor;
 	public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanTensor;Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJLsk/ainet/lang/memory/plan/WeightForm;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanTensor;
 	public fun equals (Ljava/lang/Object;)Z
-	public final fun getAllocation ()Lsk/ainet/lang/memory/AllocationSpec;
 	public final fun getBytes ()J
 	public final fun getElementCount ()J
 	public final fun getForm ()Lsk/ainet/lang/memory/plan/WeightForm;
@@ -1795,6 +1807,27 @@ public final class sk/ainet/lang/memory/plan/ProfiledPlan {
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class sk/ainet/lang/memory/plan/StorageCapabilities {
+	public static final field Companion Lsk/ainet/lang/memory/plan/StorageCapabilities$Companion;
+	public fun <init> (ZZ)V
+	public synthetic fun <init> (ZZILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public final fun component1 ()Z
+	public final fun component2 ()Z
+	public final fun copy (ZZ)Lsk/ainet/lang/memory/plan/StorageCapabilities;
+	public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/StorageCapabilities;ZZILjava/lang/Object;)Lsk/ainet/lang/memory/plan/StorageCapabilities;
+	public fun equals (Ljava/lang/Object;)Z
+	public final fun getSupportsMappedFiles ()Z
+	public final fun getSupportsOffHeap ()Z
+	public fun hashCode ()I
+	public fun toString ()Ljava/lang/String;
+}
+
+public final class sk/ainet/lang/memory/plan/StorageCapabilities$Companion {
+	public final fun current ()Lsk/ainet/lang/memory/plan/StorageCapabilities;
+	public final fun getFULL ()Lsk/ainet/lang/memory/plan/StorageCapabilities;
+	public final fun getHEAP_ONLY ()Lsk/ainet/lang/memory/plan/StorageCapabilities;
+}
+
 public final class sk/ainet/lang/memory/plan/Suggestion {
 	public fun <init> (Ljava/lang/String;J)V
 	public final fun component1 ()Ljava/lang/String;

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt

@@ -0,0 +1,131 @@
+package sk.ainet.lang.memory.plan
+
+import sk.ainet.lang.memory.AllocationSpec
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Format
+import sk.ainet.lang.memory.PlatformStorage
+import sk.ainet.lang.memory.ScopeKind
+import sk.ainet.lang.tensor.storage.MemoryDomain
+
+/**
+ * What the running platform's storage can actually do — the third input of [AllocationResolver],
+ * separated from [PlatformStorage] so a test can resolve for a platform it is not running on.
+ */
+@ExperimentalMemoryApi
+public data class StorageCapabilities(
+    val supportsMappedFiles: Boolean,
+    val supportsOffHeap: Boolean = true,
+) {
+    public companion object {
+        /** The platform this code is running on. */
+        public fun current(): StorageCapabilities = StorageCapabilities(
+            supportsMappedFiles = PlatformStorage.supportsMappedFiles,
+            supportsOffHeap = PlatformStorage.supports(MemoryDomain.HOST_OFFHEAP),
+        )
+
+        /** A JVM/native-class platform: everything works. */
+        public val FULL: StorageCapabilities = StorageCapabilities(supportsMappedFiles = true)
+
+        /** A browser-class platform: heap only, no mmap, no off-heap. */
+        public val HEAP_ONLY: StorageCapabilities = StorageCapabilities(supportsMappedFiles = false, supportsOffHeap = false)
+    }
+}
+
+/**
+ * Decides where a tensor's bytes belong — the placement counterpart of [WeightFormResolver] (#1143,
+ * closing the question #1133 asked).
+ *
+ * Same contract as the form resolver: a pure function of *what will be held* (the resolved
+ * [WeightForm]), *what the profile says* ([PlannerProfile.domainFor], its thresholds) and *what the
+ * platform can do* ([StorageCapabilities]). Consumers — the plan, the loader, a context — carry the
+ * result; none of them decide. Nothing here allocates.
+ */
+@ExperimentalMemoryApi
+public object AllocationResolver {
+
+    /**
+     * The allocation a weight needs once [PlanTensor.form] is honoured.
+     *
+     * Served from file-backed pages only when every condition holds: the form asks for
+     * [WeightResidency.MAPPED], the platform can map, and the bytes really are the file's bytes —
+     * a re-encoded ([EncodingRequest.DequantizeTo]/[EncodingRequest.RequantizeTo]) or re-ordered
+     * ([WeightByteOrder.KERNEL_FEED]) weight is a load-time copy, and a copy cannot be paged from
+     * the file it no longer matches. Everything else falls to [PlannerProfile.domainFor] over the
+     * bytes actually held, so a dequantized giant goes off-heap and a small bias stays on it.
+     */
+    public fun resolve(
+        weight: PlanTensor,
+        profile: PlannerProfile,
+        platform: StorageCapabilities = StorageCapabilities.current(),
+    ): AllocationSpec {
+        val form = weight.form
+        val fileBytesAreTheBytes = form == null ||
+            (form.encoding == EncodingRequest.KeepAsStored && form.order == WeightByteOrder.AS_STORED)
+        val wantsMapped = form?.residency == WeightResidency.MAPPED
+        val mapped = wantsMapped && platform.supportsMappedFiles && fileBytesAreTheBytes
+        val domain = if (mapped) MemoryDomain.MMAP_FILE else fallbackDomain(weight.residentBytes, profile, platform)
+        return AllocationSpec(
+            format = residentFormat(weight),
+            elementCount = weight.elementCount,
+            domain = domain,
+            scope = ScopeKind.MODEL,
+            mutable = false,
+        )
+    }
+
+    /**
+     * The allocation a transient (activation/scratch) tensor needs: never mapped, never
+     * model-lifetime — only the profile's heap/off-heap threshold and the caller's scope.
+     */
+    public fun resolveTransient(
+        format: Format,
+        elementCount: Long,
+        profile: PlannerProfile,
+        platform: StorageCapabilities = StorageCapabilities.current(),
+        scope: ScopeKind = ScopeKind.FORWARD,
+    ): AllocationSpec {
+        val bytes = format.physicalBytes(elementCount) ?: (format.dtype.sizeInBytes.toLong() * elementCount)
+        return AllocationSpec(format, elementCount, fallbackDomain(bytes, profile, platform), scope, mutable = true)
+    }
+
+    /** The [Format] the weight holds once its form is honoured — dense after a dequantization. */
+    public fun residentFormat(weight: PlanTensor): Format = when (val request = weight.form?.encoding) {
+        null, EncodingRequest.KeepAsStored -> weight.format
+        is EncodingRequest.DequantizeTo -> Format.dense(request.dtype)
+        is EncodingRequest.RequantizeTo -> Format(weight.format.dtype, request.encoding)
+    }
+
+    /**
+     * One line saying where [weight] lands and *why* — the transparency counterpart of
+     * [resolve], for plan renders and load-time traces. The decision is recomputed, so the
+     * explanation can never drift from what the resolver actually did.
+     */
+    public fun explain(
+        weight: PlanTensor,
+        profile: PlannerProfile,
+        platform: StorageCapabilities = StorageCapabilities.current(),
+    ): String {
+        val spec = resolve(weight, profile, platform)
+        val form = weight.form
+        val why = when {
+            spec.domain == MemoryDomain.MMAP_FILE ->
+                "form asks MAPPED, platform can map, bytes are the file's bytes"
+            form?.residency == WeightResidency.MAPPED && !platform.supportsMappedFiles ->
+                "form asks MAPPED but this platform cannot map files"
+            form?.residency == WeightResidency.MAPPED && form.encoding != EncodingRequest.KeepAsStored ->
+                "form asks MAPPED but the weight is re-encoded at load — a copy cannot be paged from the file"
+            form?.residency == WeightResidency.MAPPED && form.order == WeightByteOrder.KERNEL_FEED ->
+                "form asks MAPPED but kernel-feed order is a load-time copy"
+            else ->
+                "resident ${MemoryPlans.formatBytes(weight.residentBytes)} vs off-heap threshold " +
+                    MemoryPlans.formatBytes(profile.offHeapThresholdBytes) +
+                    if (!platform.supportsOffHeap) " (no off-heap on this platform)" else ""
+        }
+        return "${weight.name}: ${spec.domain}/${spec.scope} — $why"
+    }
+
+    private fun fallbackDomain(bytes: Long, profile: PlannerProfile, platform: StorageCapabilities): MemoryDomain {
+        val preferred = profile.domainFor(bytes)
+        return if (preferred == MemoryDomain.HOST_OFFHEAP && !platform.supportsOffHeap) MemoryDomain.HOST_HEAP else preferred
+    }
+}

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt

@@ -3,9 +3,7 @@ package sk.ainet.lang.memory.plan
 import sk.ainet.lang.memory.AllocationSpec
 import sk.ainet.lang.memory.ExperimentalMemoryApi
 import sk.ainet.lang.memory.Format
-import sk.ainet.lang.memory.ScopeKind
 import sk.ainet.lang.tensor.TensorId
-import sk.ainet.lang.tensor.storage.MemoryDomain
 import sk.ainet.lang.tensor.storage.TensorEncoding
 
 /**
@@ -44,9 +42,15 @@ public data class PlanTensor(
                 Format(format.dtype, request.encoding).physicalBytes(elementCount) ?: bytes
         }
 
-    /** The allocation this weight needs: mapped, model-lifetime, read-only. */
-    val allocation: AllocationSpec
-        get() = AllocationSpec(format, elementCount, MemoryDomain.MMAP_FILE, ScopeKind.MODEL, mutable = false)
+    /**
+     * The allocation this weight needs — resolved, not assumed (#1143). The old form of this
+     * property hardcoded mapped/model-lifetime for every weight regardless of what [form] asked,
+     * what the profile said, or whether the platform could map at all.
+     */
+    public fun allocation(
+        profile: PlannerProfile,
+        platform: StorageCapabilities = StorageCapabilities.current(),
+    ): AllocationSpec = AllocationResolver.resolve(this, profile, platform)
 }
 
 /** The transformer geometry the KV-cache and forward-slab estimates need (from the GGUF header). */

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/AllocationResolverTest.kt

@@ -0,0 +1,149 @@
+package sk.ainet.lang.memory.plan
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Format
+import sk.ainet.lang.memory.ScopeKind
+import sk.ainet.lang.tensor.storage.MemoryDomain
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertNotEquals
+import kotlin.test.assertTrue
+
+/**
+ * #1143 / #1133: placement is a resolver decision — (what will be held × profile × platform) in,
+ * an [sk.ainet.lang.memory.AllocationSpec] out. These tests pin the rules the old
+ * `PlanTensor.allocation` hardcode ignored.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+class AllocationResolverTest {
+
+    private val q4k = Format(FP32, TensorEncoding.Q4_K)
+
+    private fun weight(
+        elements: Long = 1L shl 20, // 1Mi elements: Q4_K packed ≈ 576 KiB, dense FP32 = 4 MiB
+        form: WeightForm?,
+    ) = PlanTensor(
+        name = "blk.0.attn_q.weight",
+        id = null,
+        format = q4k,
+        elementCount = elements,
+        bytes = q4k.physicalBytes(elements)!!,
+        form = form,
+    )
+
+    @Test
+    fun mappedKeptAsStoredWeightIsServedFromTheFile() {
+        val spec = AllocationResolver.resolve(
+            weight(form = WeightForm(residency = WeightResidency.MAPPED)),
+            PlannerProfile.MOBILE_2GB,
+            StorageCapabilities.FULL,
+        )
+        assertEquals(MemoryDomain.MMAP_FILE, spec.domain)
+        assertEquals(ScopeKind.MODEL, spec.scope)
+        assertFalse(spec.mutable)
+        assertEquals(q4k, spec.format)
+    }
+
+    @Test
+    fun unmappablePlatformFallsBackByTheProfileThreshold() {
+        val spec = AllocationResolver.resolve(
+            weight(form = WeightForm(residency = WeightResidency.MAPPED)),
+            PlannerProfile.MOBILE_2GB,
+            StorageCapabilities(supportsMappedFiles = false),
+        )
+        // 576 KiB packed is over the 256 KiB off-heap threshold
+        assertEquals(MemoryDomain.HOST_OFFHEAP, spec.domain)
+        assertEquals(ScopeKind.MODEL, spec.scope)
+    }
+
+    @Test
+    fun heapOnlyPlatformEndsOnTheHeapNoMatterTheSize() {
+        val spec = AllocationResolver.resolve(
+            weight(form = WeightForm(residency = WeightResidency.MAPPED)),
+            PlannerProfile.MOBILE_2GB,
+            StorageCapabilities.HEAP_ONLY,
+        )
+        assertEquals(MemoryDomain.HOST_HEAP, spec.domain)
+    }
+
+    @Test
+    fun dequantizedWeightIsNeverMapped() {
+        val form = WeightForm(encoding = EncodingRequest.DequantizeTo(FP32), residency = WeightResidency.MAPPED)
+        val spec = AllocationResolver.resolve(weight(form = form), PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL)
+        assertNotEquals(MemoryDomain.MMAP_FILE, spec.domain)
+        // the resident bytes are the dense bytes, and they price the domain decision
+        assertEquals(Format.dense(FP32), spec.format)
+        assertEquals(MemoryDomain.HOST_OFFHEAP, spec.domain) // 4 MiB dense is far over threshold
+    }
+
+    @Test
+    fun kernelFeedOrderIsALoadTimeCopySoNotMapped() {
+        val form = WeightForm(order = WeightByteOrder.KERNEL_FEED, residency = WeightResidency.MAPPED)
+        val spec = AllocationResolver.resolve(weight(form = form), PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL)
+        assertNotEquals(MemoryDomain.MMAP_FILE, spec.domain)
+    }
+
+    @Test
+    fun smallHeapWeightStaysOnTheHeap() {
+        val dense = Format.dense(FP32)
+        val bias = PlanTensor("blk.0.attn_q.bias", null, dense, 768, 4L * 768, WeightForm())
+        val spec = AllocationResolver.resolve(bias, PlannerProfile.DESKTOP, StorageCapabilities.FULL)
+        assertEquals(MemoryDomain.HOST_HEAP, spec.domain)
+        assertEquals(ScopeKind.MODEL, spec.scope)
+    }
+
+    @Test
+    fun noFormMeansTheFileBytesByTheProfileRules() {
+        val spec = AllocationResolver.resolve(weight(form = null), PlannerProfile.DESKTOP, StorageCapabilities.FULL)
+        assertNotEquals(MemoryDomain.MMAP_FILE, spec.domain) // nothing asked for mapping
+        assertEquals(MemoryDomain.HOST_OFFHEAP, spec.domain)
+    }
+
+    @Test
+    fun transientAllocationsFollowScopeAndThreshold() {
+        val dense = Format.dense(FP32)
+        val small = AllocationResolver.resolveTransient(dense, 1024, PlannerProfile.DESKTOP, StorageCapabilities.FULL)
+        assertEquals(MemoryDomain.HOST_HEAP, small.domain)
+        assertEquals(ScopeKind.FORWARD, small.scope)
+        assertTrue(small.mutable)
+
+        val big = AllocationResolver.resolveTransient(
+            dense, 1L shl 20, PlannerProfile.DESKTOP, StorageCapabilities.FULL, scope = ScopeKind.AMBIENT
+        )
+        assertEquals(MemoryDomain.HOST_OFFHEAP, big.domain)
+        assertEquals(ScopeKind.AMBIENT, big.scope)
+    }
+
+    @Test
+    fun explainSaysWhereAndWhy() {
+        val mapped = AllocationResolver.explain(
+            weight(form = WeightForm(residency = WeightResidency.MAPPED)),
+            PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL,
+        )
+        assertTrue("MMAP_FILE" in mapped && "file's bytes" in mapped, mapped)
+
+        val noMmap = AllocationResolver.explain(
+            weight(form = WeightForm(residency = WeightResidency.MAPPED)),
+            PlannerProfile.MOBILE_2GB, StorageCapabilities(supportsMappedFiles = false),
+        )
+        assertTrue("cannot map" in noMmap, noMmap)
+
+        val dequant = AllocationResolver.explain(
+            weight(form = WeightForm(encoding = EncodingRequest.DequantizeTo(FP32), residency = WeightResidency.MAPPED)),
+            PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL,
+        )
+        assertTrue("re-encoded at load" in dequant, dequant)
+    }
+
+    @Test
+    fun planTensorAllocationDelegatesToTheResolver() {
+        val w = weight(form = WeightForm(residency = WeightResidency.MAPPED))
+        assertEquals(
+            AllocationResolver.resolve(w, PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL),
+            w.allocation(PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL),
+        )
+    }
+}