P7: AllocationResolver + plan→load wiring (re-land #1153/#1154 onto develop)

skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt

@@ -0,0 +1,87 @@
+package sk.ainet.io.gguf
+
+import sk.ainet.io.RandomAccessSource
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.plan.AllocationResolver
+import sk.ainet.lang.memory.plan.KernelCapabilities
+import sk.ainet.lang.memory.plan.PlanInput
+import sk.ainet.lang.memory.plan.PlannerProfile
+import sk.ainet.lang.memory.plan.ProfiledPlan
+import sk.ainet.lang.memory.plan.StorageCapabilities
+import sk.ainet.lang.memory.plan.WeightForm
+import sk.ainet.lang.memory.plan.resolveWeightForms
+import sk.ainet.lang.memory.trace.NoopTraceSink
+import sk.ainet.lang.memory.trace.TraceSink
+
+/**
+ * Plan → load, wired (#1144): the loader consults what the resolvers decided, instead of the two
+ * pipelines sharing vocabulary and never talking.
+ *
+ * [resolve] reads the GGUF *header* (no payload), resolves every weight's [WeightForm] from
+ * file × [PlannerProfile] × [KernelCapabilities], applies the caller's per-tensor overrides,
+ * and returns a [Resolution]: a loader that will deliver exactly those forms, plus the resolved
+ * plan input so the same decisions are priceable ([Resolution.profiledPlan]) and explainable
+ * ([Resolution.explainPlacements]) *before a byte of payload is read*.
+ *
+ * ## Who decides — the precedence order
+ *
+ * **Your override > the resolver.** [overrideFormFor] outranks everything for the tensors it names,
+ * including the deliberately blunt `WeightForm(DequantizeTo(FP32), residency = HEAP)` — everything
+ * dense, on the managed heap. Tensors it returns `null` for get the resolver's answer. The plan and
+ * the explanations are computed *after* overrides are applied, so what you print is what you load.
+ */
+@ExperimentalMemoryApi
+public object ResolvedGguf {
+
+    /** What [resolve] decided: the loader that obeys it, and the resolved input that explains it. */
+    public data class Resolution(
+        val loader: StreamingGgufParametersLoader,
+        /** The header-derived plan input with every weight's resolved (and overridden) form. */
+        val input: PlanInput,
+        val profile: PlannerProfile,
+        val platform: StorageCapabilities,
+    ) {
+        /** The plan these forms cost against [availableBytes] — `requireFits()` refuses pre-load (M2-F6). */
+        public fun profiledPlan(availableBytes: Long): ProfiledPlan = profile.plan(input, availableBytes)
+
+        /** One line per weight: where it lands and why — [AllocationResolver.explain] over the resolved forms. */
+        public fun explainPlacements(): List<String> =
+            input.weights.map { AllocationResolver.explain(it, profile, platform) }
+    }
+
+    /**
+     * Resolve every weight's form from the header, apply [overrideFormFor], and build the loader
+     * that delivers it. Reads header only; costs a few kilobytes.
+     *
+     * @param capabilities what the backend's kernels can feed — pass the registry-backed
+     *   implementation from the backend in use, or [KernelCapabilities.DENSE_ONLY] to price the
+     *   worst case
+     * @param overrideFormFor the user-wins channel; `null` per tensor means "resolver decides"
+     * @throws IllegalStateException when the profile is strict and a weight would dequantize
+     */
+    public fun resolve(
+        sourceProvider: () -> RandomAccessSource,
+        profile: PlannerProfile,
+        capabilities: KernelCapabilities,
+        ctx: Int? = null,
+        platform: StorageCapabilities = StorageCapabilities.current(),
+        overrideFormFor: (tensorName: String) -> WeightForm? = { null },
+        onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> },
+        traceSink: TraceSink = NoopTraceSink,
+    ): Resolution {
+        val resolved = sourceProvider().use { source ->
+            StreamingGGUFReader.open(source).planInput(ctx)
+        }.resolveWeightForms(profile, capabilities)
+        val overridden = resolved.copy(
+            weights = resolved.weights.map { w -> overrideFormFor(w.name)?.let { w.copy(form = it) } ?: w },
+        )
+        val forms: Map<String, WeightForm?> = overridden.weights.associate { it.name to it.form }
+        val loader = StreamingGgufParametersLoader(
+            sourceProvider = sourceProvider,
+            onProgress = onProgress,
+            weightFormFor = { name -> forms[name] },
+            traceSink = traceSink,
+        )
+        return Resolution(loader, overridden, profile, platform)
+    }
+}

skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt

@@ -121,6 +121,20 @@ public class StreamingGgufParametersLoader(
      * silently resolved, so nobody loses a setting they thought they had.
      */
     private val weightForm: WeightForm? = null,
+    /**
+     * Per-tensor forms — the *user wins* channel (#1144).
+     *
+     * Precedence, explicit and documented: **this function > [weightForm] > the three legacy
+     * parameters > nothing**. Whatever you return for a tensor outranks every resolver and every
+     * profile — including the deliberately blunt `WeightForm(DequantizeTo(FP32), residency = HEAP)`
+     * ("everything dense, on the managed heap"). Return `null` for tensors you have no opinion on;
+     * they fall through to [weightForm] (or the legacy axes).
+     *
+     * The intended producer is `WeightFormResolver`/`resolveWeightForms` via [ResolvedGguf], which
+     * resolves per tensor from the file × profile × kernel capability — but the contract is the
+     * same for a hand-written lambda: the loader carries and obeys, it does not decide.
+     */
+    private val weightFormFor: ((tensorName: String) -> WeightForm?)? = null,
     /**
      * Where conversions are reported (#1117).
      *
@@ -184,47 +198,52 @@ public class StreamingGgufParametersLoader(
         },
     )
 
-    /** [QuantPolicy.DEQUANTIZE_TO_FP32] asked for through [form], whichever way it was set. */
-    private val dequantizeToDense: Boolean = form.encoding is EncodingRequest.DequantizeTo
-
-    /** [StagingPolicy.MAPPED] asked for through [form]. */
-    private val mapsTheFile: Boolean = form.residency == WeightResidency.MAPPED
-
-    /** [WeightOrientation.OUT_IN] asked for through [form]. */
-    private val reversesWeightShape: Boolean = form.shape == WeightShapeOrientation.OUT_IN
+    /**
+     * The form [tensorName] loads under — the precedence order of [weightFormFor], validated the
+     * same way the uniform [form] was at construction.
+     */
+    private fun formFor(tensorName: String): WeightForm {
+        val perTensor = weightFormFor?.invoke(tensorName) ?: return form
+        validateForm(perTensor, "weightFormFor('$tensorName')")
+        return perTensor
+    }
 
     init {
         require(quantPolicy != QuantPolicy.RAW_BYTES) {
             "StreamingGgufParametersLoader does not support QuantPolicy.RAW_BYTES — quantized " +
                 "tensors are preserved as packed block TensorData (NATIVE_OPTIMIZED) or " +
                 "dequantized to dense FP32 (DEQUANTIZE_TO_FP32)."
         }
-        if (weightForm != null) {
+        if (weightForm != null || weightFormFor != null) {
             require(
                 quantPolicy == QuantPolicy.NATIVE_OPTIMIZED &&
                     staging == StagingPolicy.HEAP &&
                     weightOrientation == WeightOrientation.AS_STORED,
             ) {
-                "a WeightForm and the quantPolicy/staging/weightOrientation parameters were both " +
-                    "set. They are the same three axes, so one of them would have been silently " +
-                    "ignored; pass only the form."
+                "a WeightForm (or weightFormFor) and the quantPolicy/staging/weightOrientation " +
+                    "parameters were both set. They are the same three axes, so one of them would " +
+                    "have been silently ignored; pass only the form."
             }
         }
+        validateForm(form, "weightForm")
+    }
+
+    private fun validateForm(form: WeightForm, where: String) {
         require(form.order == WeightByteOrder.AS_STORED || form.shape == WeightShapeOrientation.OUT_IN) {
-            "WeightByteOrder.KERNEL_FEED needs WeightShapeOrientation.OUT_IN: feed order is defined " +
-                "relative to a [out, in] weight — which block is 'block b of output row o' has no " +
-                "answer while the tensor is still labelled in the file's `ne` order."
+            "$where: WeightByteOrder.KERNEL_FEED needs WeightShapeOrientation.OUT_IN: feed order is " +
+                "defined relative to a [out, in] weight — which block is 'block b of output row o' " +
+                "has no answer while the tensor is still labelled in the file's `ne` order."
         }
         val requested = form.encoding
         require(requested !is EncodingRequest.RequantizeTo) {
-            "EncodingRequest.RequantizeTo(${requested.let { (it as? EncodingRequest.RequantizeTo)?.encoding?.name }}) " +
+            "$where: EncodingRequest.RequantizeTo(${requested.let { (it as? EncodingRequest.RequantizeTo)?.encoding?.name }}) " +
                 "is not supported by this loader: re-quantizing a weight the file does not already " +
                 "carry needs a quantizer per target encoding, and none of them exist here yet."
         }
         val dequantTarget = (requested as? EncodingRequest.DequantizeTo)?.dtype
         require(dequantTarget == null || dequantTarget == FP32) {
-            "EncodingRequest.DequantizeTo(${dequantTarget?.name}) is not supported: this loader " +
-                "dequantizes to FP32 only."
+            "$where: EncodingRequest.DequantizeTo(${dequantTarget?.name}) is not supported: this " +
+                "loader dequantizes to FP32 only."
         }
     }
 
@@ -233,9 +252,10 @@ public class StreamingGgufParametersLoader(
      * is reversed for `OUT_IN`, everything else is passed through as the file has it. Only 2-D
      * tensors are touched — a 1-D bias or norm has no orientation to get wrong.
      */
-    private fun shapeOf(tensorInfo: StreamingTensorInfo): Shape {
+    private fun shapeOf(tensorInfo: StreamingTensorInfo, form: WeightForm): Shape {
         val dims = tensorInfo.shape.map { it.toInt() }
-        val ordered = if (reversesWeightShape && dims.size == 2) dims.reversed() else dims
+        val reverses = form.shape == WeightShapeOrientation.OUT_IN
+        val ordered = if (reverses && dims.size == 2) dims.reversed() else dims
         return Shape(*ordered.toIntArray())
     }
 
@@ -248,7 +268,9 @@ public class StreamingGgufParametersLoader(
         val source = sourceProvider()
         // MAPPED staging needs a file to map; a Blob or an in-memory source has no path and
         // silently stays on the heap, which is the documented fallback rather than a failure.
-        val mapped = if (mapsTheFile) source.filePath?.let { openMappedFile(it) } else null
+        // With per-tensor forms the file is mapped whenever any tensor *might* ask for it.
+        val mightMap = form.residency == WeightResidency.MAPPED || weightFormFor != null
+        val mapped = if (mightMap) source.filePath?.let { openMappedFile(it) } else null
         try {
         StreamingGGUFReader.open(source).use { reader ->
             val tensors = reader.tensors
@@ -257,12 +279,15 @@ public class StreamingGgufParametersLoader(
             var current = 0L
 
             for (tensorInfo in tensors) {
-                val shape = shapeOf(tensorInfo)
+                val tensorForm = formFor(tensorInfo.name)
+                val shape = shapeOf(tensorInfo, tensorForm)
                 // A dense F32 tensor under MAPPED staging never reaches the heap: it is a view over
                 // file-backed pages. Everything else reads its bytes (out of the mapping when there
                 // is one — one page-cache copy instead of a channel read).
                 val mappedFloats: Tensor<T, V>? =
-                    if (mapped != null && tensorInfo.tensorType == GGMLQuantizationType.F32 && dtype == FP32::class) {
+                    if (mapped != null && tensorForm.residency == WeightResidency.MAPPED &&
+                        tensorInfo.tensorType == GGMLQuantizationType.F32 && dtype == FP32::class
+                    ) {
                         @Suppress("UNCHECKED_CAST")
                         ctx.fromData(
                             mapped.denseFloats<T>(tensorInfo.absoluteDataOffset, shape) as sk.ainet.lang.tensor.data.TensorData<T, V>,
@@ -340,7 +365,7 @@ public class StreamingGgufParametersLoader(
                     GGMLQuantizationType.Q5_0,
                     GGMLQuantizationType.Q5_1,
                     GGMLQuantizationType.TQ1_0,
-                    GGMLQuantizationType.TQ2_0 -> quantizedTensor(ctx, dtype, shape, tensorInfo, rawBytes)
+                    GGMLQuantizationType.TQ2_0 -> quantizedTensor(ctx, dtype, shape, tensorInfo, rawBytes, tensorForm)
 
                     else -> throw IllegalStateException(
                         "StreamingGgufParametersLoader: tensor '${tensorInfo.name}' of type " +
@@ -382,8 +407,9 @@ public class StreamingGgufParametersLoader(
         shape: Shape,
         tensorInfo: StreamingTensorInfo,
         rawBytes: ByteArray,
+        tensorForm: WeightForm,
     ): Tensor<T, V> {
-        if (dequantizeToDense &&
+        if (tensorForm.encoding is EncodingRequest.DequantizeTo &&
             (dtype == FP32::class || dtype == FP16::class)
         ) {
             val dest = DequantOps.dequantFromBytes(rawBytes, tensorInfo.tensorType, tensorInfo.nElements.toInt())
@@ -431,7 +457,7 @@ public class StreamingGgufParametersLoader(
                 "quantizedTensor called for non-quantized type ${tensorInfo.tensorType}"
             )
         }
-        val delivered = if (form.order == WeightByteOrder.KERNEL_FEED) feedOrdered(packed, tensorInfo) else packed
+        val delivered = if (tensorForm.order == WeightByteOrder.KERNEL_FEED) feedOrdered(packed, tensorInfo) else packed
         return ctx.fromData(delivered as sk.ainet.lang.tensor.data.TensorData<T, V>, dtype)
     }