feat(compile): a mixed-precision layer's dtype is a checked preference (#1495)

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Fixed
 
+- **`DTypeConstraintResolutionPass` no longer disconnects the nodes it resolves** (`skainet-compile-opt`, #1495): the pass marked a policy-carrying node by removing and re-adding it, which dropped every edge connected to that node, so a `dag { … dtypePolicy(…) }` graph left the pass with those nodes unwired. It now rebuilds the graph with the edges intact, and reports no change on an already resolved graph, so a multi-iteration pipeline converges.
 - **A same-dtype `PrecisionBoundary` leaves no cast in a trace** (`skainet-lang-core`, SKEEP-008 phase A2, #1433): `stage("body", dtype = FP32)` inside an FP32 network called `convert`, which returns its input on the CPU backend but is still recorded by the tracing wrapper, so the exported module carried two `stablehlo.convert` `f32 -> f32`. The boundary now passes a tensor of the target dtype through without calling `convert`.
 - **`concat` backward with inputs of different sizes** (`skainet-compile-dag`, #1490): the rule split the upstream gradient by one input's size and picked chunk `offset / size`, which only lines up when every input has the same extent along the concat axis; other shapes got another input's slice, or `null`, without an error. Each input now receives its own `narrow` of the upstream at the offset the forward gave it, for any sizes and a negative `dim`.
 - **`squeeze` backward with no axis** (`skainet-compile-dag`, #1490): `squeeze(x)` removes every size-1 axis, but the backward unsqueezed dim 0 alone, so `x: [1, 3, 1]` got a `[1, 3]` gradient. The gradient is now the upstream reshaped to the input's shape.
@@ -14,6 +15,7 @@
 
 ### Added
 
+- **A mixed-precision layer's dtype is a checked preference** (`skainet-compile-dag`, `skainet-compile-opt`, `skainet-lang-core`, #1495): a tracing context marks every op recorded inside a `PrecisionBoundary` with `DTypePolicy.Prefer(innerType)` and the boundary's name as graph-node metadata (`PRECISION_BOUNDARY_METADATA_KEY`); the boundary's own casts are not marked. `DTypeConstraintResolutionPass` checks that preference against the node's floating-point inputs, so a region retyped by a later pass, for example `DtypeForwardPropagationPass(targetFloatDtype = "BF16")` on an FP16 stage, is reported as a diagnostic instead of passing silently. `PrecisionBoundary.innerType` / `outerType` are public. Graphs traced without boundaries are unchanged.
 - **`CastFoldingPass`** (`skainet-compile-opt`, #1494): folds dtype casts that change no value. A cast that follows an exact widening (FP16 or BF16 to FP32, any of them to FP64) reads from the original tensor, so the `bf16 -> f32 -> f16` between two adjacent mixed-precision boundaries exports as one `stablehlo.convert` and `f16 -> f32 -> f16` as none; a cast to the dtype its input already has is removed. A narrowing round trip (`f32 -> bf16 -> f32`) and integer casts are left alone. The pass is part of the default, aggressive and LLM pipelines of `GraphOptimizationPipeline`.
 - **`ComputeGraphExecutor` runs `convert` / `cast` / `to` nodes** (`skainet-compile-dag`, #1494): the target is the node's output dtype; a graph traced from a mixed-precision network can now be executed.
 - **Training through a mixed-precision boundary** (`skainet-lang-core`, `skainet-compile-dag`, SKEEP-008 phase A3, #1433): `TensorOps.convert` is a `@Diff` op. Its backward rule treats the cast as the identity and casts the upstream gradient back to the source dtype, so a gradient entering an FP16 / BF16 region is rounded to that format and the layers behind `dense(dtype = …)` / `stage(id, dtype = …)` train; a cast from or to an integer dtype has no gradient. `MixedPrecisionBackwardTest` checks the rule, a finite-difference gradient through FP16 boundaries, parity with the all-FP32 network, and an SGD run. The inference-only notes are gone; the ops matrix Backward column reads 60/60.

docs/modules/ROOT/pages/explanation/mixed-precision.adoc

@@ -153,6 +153,27 @@ values as narrowing the original, and the two casts fold into one:
 Between two boundaries of the *same* inner dtype nothing is left at all. The pass never
 touches the other order, `f32 -> bf16 -> f32`: that pair is the rounding the layer asked for.
 
+== A declared dtype is a checked preference
+
+Tracing marks every op that runs inside a boundary with the dtype the code declared for it:
+the graph node carries `DTypePolicy.Prefer(dtype)` and the boundary's name as metadata. The
+boundary's own two casts are not marked, since their job is to have a different dtype on one
+side.
+
+`DTypeConstraintResolutionPass`, the first pass of the default pipelines, checks that
+preference against the node's floating-point inputs. Nothing is reported for an untouched
+network. If a later pass retypes the region — `DtypeForwardPropagationPass(targetFloatDtype =
+"BF16")` coerces a whole graph to one float dtype, which is what some NPU targets need — a
+second run of the constraint pass says so:
+
+[source,text]
+----
+Node 'n17_matmul' (matmul) in precision boundary 'head' prefers Float16 but input 0 has dtype 'BF16' — using the existing dtype.
+----
+
+It is a diagnostic, not an error: retyping a mixed-precision model for a single-dtype target
+is a legitimate choice, and it should be a visible one.
+
 == Training through a boundary
 
 `convert` is differentiable between floating-point dtypes, so a mixed-precision layer or
@@ -188,6 +209,8 @@ parameters are not affected.
 == Where to look
 
 * `sk.ainet.lang.nn.PrecisionBoundary` — the boundary module.
+* `DTypeConstraintResolutionPass` and `PRECISION_BOUNDARY_METADATA_KEY` — the declared dtype
+  as graph metadata, and the pass that checks it.
 * `NeuralNetworkDsl.dense(outputDimension, dtype, …)` and `stage(id, dtype, …)` — the DSL
   entry points.
 * `TensorOps.convert` — the cast, including FP16/BF16 rounding through the format codecs.

docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc

@@ -360,8 +360,14 @@ Each phase is a DARC feature with lane sub-issues, linked from the tracking issu
   *Resolved:* it is a member of `TensorStorage`, next to `repackTo`.
 . Does phase 2 need per-boundary `DTypePolicy` metadata, so that
   `DTypeConstraintResolutionPass` can see and validate the boundaries in a compiled graph? +
-  *Open.* A3 did not need it: the boundaries are ordinary `convert` nodes in the graph.
-  Tracked in https://github.com/SKaiNET-developers/SKaiNET/issues/1495[#1495].
+  *Resolved: yes, as metadata on the region*
+  (https://github.com/SKaiNET-developers/SKaiNET/issues/1495[#1495]). A tracing context marks
+  every op recorded inside a boundary with `DTypePolicy.Prefer(innerType)` and the boundary's
+  name; the boundary's own casts are not marked. `DTypeConstraintResolutionPass` checks the
+  preference against the node's floating-point inputs and reports a region that a later pass
+  retyped. `Prefer`, not `Require`: coercing a mixed-precision model to one dtype for a target
+  is allowed, but no longer silent. A cast strategy or callback on `PrecisionBoundary` was
+  considered and rejected: a lambda is invisible to the compile pipeline.
 
 == References
 

skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api

@@ -334,6 +334,10 @@ public final class sk/ainet/lang/graph/OutputDesignatedGraph : sk/ainet/lang/gra
 	public fun validate ()Lsk/ainet/lang/tensor/ops/ValidationResult;
 }
 
+public final class sk/ainet/lang/graph/PrecisionBoundaryScopeKt {
+	public static final field PRECISION_BOUNDARY_METADATA_KEY Ljava/lang/String;
+}
+
 public final class sk/ainet/lang/graph/ResolvedComputeGraph {
 	public static final field BACKEND_ASSIGNMENT_METADATA_KEY Ljava/lang/String;
 	public static final field Companion Lsk/ainet/lang/graph/ResolvedComputeGraph$Companion;

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt

@@ -103,7 +103,8 @@ public open class DefaultExecutionTape(
             val op = object : sk.ainet.lang.tensor.ops.Operation {
                 override val name: String = trace.opType
                 override val type: String = "trace"
-                override val parameters: Map<String, Any> = trace.attributes.filterValues { it != null } as Map<String, Any>
+                override val parameters: Map<String, Any> =
+                    trace.attributes.filter { (key, value) -> value != null && key !in TRACE_METADATA_KEYS } as Map<String, Any>
                 override fun <T : sk.ainet.lang.types.DType, V> execute(inputs: List<sk.ainet.lang.tensor.Tensor<T, V>>): List<sk.ainet.lang.tensor.Tensor<T, V>> = emptyList()
                 override fun validateInputs(inputs: List<TensorSpec>): sk.ainet.lang.tensor.ops.ValidationResult = sk.ainet.lang.tensor.ops.ValidationResult.Valid
                 override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> = outputs

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultGraphExecutionContext.kt

@@ -22,7 +22,11 @@ public class DefaultGraphExecutionContext(
     override val baseOps: TensorOps = VoidTensorOps(),
     override val phase: Phase = Phase.EVAL,
     override val tensorDataFactory: TensorDataFactory = DenseTensorDataFactory(),
-    override val hooks: ForwardHooks? = null,
+    /**
+     * Forward hooks of the caller. The context's own [hooks] wrap them to follow the module
+     * call stack (see [PrecisionBoundaryScope]) and call them unchanged.
+     */
+    hooks: ForwardHooks? = null,
     override val memoryInfo: MemoryInfo = MemoryInfo.getEmptyInfo(),
     override val executionStats: ExecutionStats = ExecutionStats(),
     override val createTapeFactory: (GraphExecutionContext) -> ExecutionTape =
@@ -37,6 +41,11 @@ public class DefaultGraphExecutionContext(
     private val baseSink: OpSink = NoOpSink,
     ) : GraphExecutionContext {
 
+    private val userHooks: ForwardHooks? = hooks
+    private val boundaryScope = PrecisionBoundaryScope(userHooks)
+
+    override val hooks: ForwardHooks get() = boundaryScope
+
     private val observerRegistry = ExecutionObserverRegistry()
 
     /**
@@ -62,7 +71,7 @@ public class DefaultGraphExecutionContext(
             baseOps = scheduled.withSchedule(schedule),
             phase = phase,
             tensorDataFactory = tensorDataFactory,
-            hooks = hooks,
+            hooks = userHooks,
             memoryInfo = memoryInfo,
             executionStats = executionStats,
             createTapeFactory = createTapeFactory,
@@ -121,11 +130,13 @@ public class DefaultGraphExecutionContext(
     private val _ops: KspTensorOps by lazy {
         val dynamicSink = object : OpSink {
             override fun onOpExecuted(trace: sk.ainet.lang.trace.OpTrace) {
+                // An op that runs inside a PrecisionBoundary carries the declared dtype (#1495).
+                val scoped = boundaryScope.annotate(trace)
                 val tape = tapeStack.currentTape
                 if (tapeStack.isRecording() && tape != null) {
-                    TapeSink(tape).onOpExecuted(trace)
+                    TapeSink(tape).onOpExecuted(scoped)
                 }
-                baseSink.onOpExecuted(trace)
+                baseSink.onOpExecuted(scoped)
             }
         }
         val sharedSession = sk.ainet.lang.trace.TraceSession()

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/PrecisionBoundaryScope.kt

@@ -0,0 +1,55 @@
+package sk.ainet.lang.graph
+
+import sk.ainet.lang.dag.DTYPE_POLICY_ATTRIBUTE_KEY
+import sk.ainet.lang.nn.PrecisionBoundary
+import sk.ainet.lang.nn.hooks.ForwardHooks
+import sk.ainet.lang.nn.topology.ModuleNode
+import sk.ainet.lang.trace.OpTrace
+import sk.ainet.lang.types.DTypePolicy
+
+/**
+ * Metadata key naming the [PrecisionBoundary] a graph node was recorded inside (#1495).
+ * Nodes that carry it also carry a [DTypePolicy.Prefer] of the boundary's inner dtype under
+ * [DTYPE_POLICY_ATTRIBUTE_KEY], which `DTypeConstraintResolutionPass` checks.
+ */
+public const val PRECISION_BOUNDARY_METADATA_KEY: String = "precision_boundary"
+
+/** Trace attributes that describe where an op ran, not what it does: node metadata, not parameters. */
+internal val TRACE_METADATA_KEYS: Set<String> = setOf(DTYPE_POLICY_ATTRIBUTE_KEY, PRECISION_BOUNDARY_METADATA_KEY)
+
+/**
+ * Follows the module call stack of a tracing context through its forward hooks, so that an
+ * op recorded inside a [PrecisionBoundary] can be marked with the dtype the user declared
+ * for it. The boundary's own two casts are not marked: they run in the boundary itself, not
+ * in the module it wraps, and their job is to have a different dtype on one side.
+ */
+internal class PrecisionBoundaryScope(private val delegate: ForwardHooks?) : ForwardHooks {
+
+    private val stack = ArrayList<ModuleNode>()
+
+    override fun onForwardBegin(module: ModuleNode, input: Any) {
+        stack.add(module)
+        delegate?.onForwardBegin(module, input)
+    }
+
+    override fun onForwardEnd(module: ModuleNode, input: Any, output: Any) {
+        delegate?.onForwardEnd(module, input, output)
+        // Pop down to this module: a forward that threw further in left its entries behind.
+        val index = stack.lastIndexOf(module)
+        if (index >= 0) while (stack.size > index) stack.removeAt(stack.size - 1)
+    }
+
+    /** [trace] with the declared dtype of the innermost enclosing boundary, or [trace] itself. */
+    fun annotate(trace: OpTrace): OpTrace {
+        if (stack.isEmpty()) return trace
+        val index = stack.indexOfLast { it is PrecisionBoundary<*, *, *> }
+        // Not inside a boundary, or in the boundary's own forward (its casts).
+        if (index < 0 || index == stack.size - 1) return trace
+        val boundary = stack[index] as PrecisionBoundary<*, *, *>
+        return trace.copy(
+            attributes = trace.attributes +
+                (DTYPE_POLICY_ATTRIBUTE_KEY to DTypePolicy.Prefer(boundary.innerType)) +
+                (PRECISION_BOUNDARY_METADATA_KEY to boundary.name),
+        )
+    }
+}

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt

@@ -60,7 +60,14 @@ public class TraceToGraphBuilder(
      * and registering the outputs as new producers.
      */
     public fun addTrace(trace: OpTrace) {
-        val parameters = trace.attributes.filterValues { it != null }.toMutableMap() as MutableMap<String, Any>
+        val parameters = trace.attributes
+            .filter { (key, value) -> value != null && key !in sk.ainet.lang.graph.TRACE_METADATA_KEYS }
+            .toMutableMap() as MutableMap<String, Any>
+        // Where the op ran (its precision boundary and declared dtype) is node metadata: the
+        // constraint pass reads it there, and operation parameters stay what the op computes.
+        val metadata = trace.attributes
+            .filter { (key, value) -> value != null && key in sk.ainet.lang.graph.TRACE_METADATA_KEYS }
+            .mapValues { it.value as Any }
 
         // If we have a session and weight embedding is enabled, try to resolve
         // constant inputs (weights/biases) for operations that need them during codegen.
@@ -190,7 +197,8 @@ public class TraceToGraphBuilder(
             id = nodeId,
             operation = op,
             inputs = inputSpecs,
-            outputs = outputSpecs
+            outputs = outputSpecs,
+            metadata = metadata,
         )
         graph.addNode(node)