diff --git a/CHANGELOG.md b/CHANGELOG.md index 14b13a3e..74efe019 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/docs/modules/ROOT/pages/explanation/mixed-precision.adoc b/docs/modules/ROOT/pages/explanation/mixed-precision.adoc index 85442a54..749c43dd 100644 --- a/docs/modules/ROOT/pages/explanation/mixed-precision.adoc +++ b/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. diff --git a/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc b/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc index 378ec988..2640b29d 100644 --- a/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc +++ b/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 diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index a5e205e0..f70295d6 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/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; diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt index 668880be..a96f25a8 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt +++ b/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 = trace.attributes.filterValues { it != null } as Map + override val parameters: Map = + trace.attributes.filter { (key, value) -> value != null && key !in TRACE_METADATA_KEYS } as Map override fun execute(inputs: List>): List> = emptyList() override fun validateInputs(inputs: List): sk.ainet.lang.tensor.ops.ValidationResult = sk.ainet.lang.tensor.ops.ValidationResult.Valid override fun inferOutputs(inputs: List): List = outputs diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultGraphExecutionContext.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultGraphExecutionContext.kt index 081f57b8..e5eb2670 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultGraphExecutionContext.kt +++ b/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() diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/PrecisionBoundaryScope.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/PrecisionBoundaryScope.kt new file mode 100644 index 00000000..67570282 --- /dev/null +++ b/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 = 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() + + 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), + ) + } +} diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt index 230f1018..6692d907 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt +++ b/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 + val parameters = trace.attributes + .filter { (key, value) -> value != null && key !in sk.ainet.lang.graph.TRACE_METADATA_KEYS } + .toMutableMap() as MutableMap + // 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) diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/MixedPrecisionExportTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/MixedPrecisionExportTest.kt index 921b9bf0..af97d71c 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/MixedPrecisionExportTest.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/MixedPrecisionExportTest.kt @@ -1,12 +1,16 @@ package sk.ainet.compile.hlo +import sk.ainet.compile.opt.GraphOptimizationPipeline import sk.ainet.compile.opt.passes.CastFoldingPass +import sk.ainet.compile.opt.passes.DTypeConstraintResolutionPass import sk.ainet.compile.opt.passes.DtypeForwardPropagationPass import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.dag.DTYPE_POLICY_ATTRIBUTE_KEY import sk.ainet.lang.graph.ComputeGraph import sk.ainet.lang.graph.DefaultComputeGraph import sk.ainet.lang.graph.GraphEdge import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.graph.PRECISION_BOUNDARY_METADATA_KEY import sk.ainet.lang.graph.exec.ComputeGraphExecutor import sk.ainet.lang.graph.DefaultGraphExecutionContext import sk.ainet.lang.nn.Module @@ -18,6 +22,7 @@ import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.ops.GenericOperation import sk.ainet.lang.tensor.ops.TensorSpec import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DTypePolicy import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 import kotlin.test.Test @@ -202,6 +207,97 @@ class MixedPrecisionExportTest { assertTrue(expected != values.map { it.toRawBits() }, "the chain must round, or it proves nothing") } + // ── declared dtype as a checked contract (#1495) ────────────────────────────────────────── + + private fun declared(node: GraphNode): Pair? { + val boundary = node.metadata[PRECISION_BOUNDARY_METADATA_KEY] as? String ?: return null + return boundary to (node.metadata[DTYPE_POLICY_ATTRIBUTE_KEY] as DTypePolicy) + } + + @Test + fun trace_marksOpsInsideABoundaryWithItsDeclaredDtype() { + val nodes = trace(::mixedNetwork).getTopologicalOrder() + + val marked = nodes.mapNotNull { node -> declared(node)?.let { node.operationName to it } } + assertEquals( + setOf( + "bf16Dense" to DTypePolicy.Prefer(BF16), + "head" to DTypePolicy.Prefer(FP16), + ), + marked.map { it.second }.toSet(), + ) + // Three dense layers: a matmul and a bias add each, plus the weight transposes. + assertEquals(3, marked.count { it.first == "matmul" }) + assertEquals(3, marked.count { it.first == "add" }) + + // The boundary's own casts have a different dtype on one side by design; inputs and + // weights are not ops of the region. + for (node in nodes) { + if (node.operationName in setOf("convert", "input", "weight")) { + assertEquals(null, declared(node), "${node.id} must not carry a boundary policy") + } + // Where an op ran is metadata; its parameters stay what it computes. + assertFalse(DTYPE_POLICY_ATTRIBUTE_KEY in node.operation.parameters, node.id) + assertFalse(PRECISION_BOUNDARY_METADATA_KEY in node.operation.parameters, node.id) + } + } + + @Test + fun trace_withoutBoundaries_carriesNoPolicy() { + val graph = trace { ctx -> + definition { + network(ctx) { + input(4) + dense(3, "h1") + dense(2, "h2") + } + } + } + assertTrue(graph.nodes.all { it.metadata.isEmpty() }, graph.nodes.map { it.id to it.metadata }.toString()) + } + + @Test + fun constraintPass_acceptsAnUntouchedNetwork_andReportsARetypedBoundary() { + val untouched = DTypeConstraintResolutionPass().apply(trace(::mixedNetwork)) + assertTrue(untouched.diagnostics.isEmpty(), untouched.diagnostics.toString()) + + // Coercing every float source to BF16 turns the FP16 stage into a BF16 one. The + // user declared FP16 there, and the pass now says so instead of staying silent. + val coerced = DtypeForwardPropagationPass(targetFloatDtype = "BF16").apply(trace(::mixedNetwork)).graph + val diagnostics = DTypeConstraintResolutionPass().apply(coerced).diagnostics + + assertTrue(diagnostics.isNotEmpty(), "the retyped FP16 stage must be reported") + assertTrue(diagnostics.all { "precision boundary 'head'" in it && "Float16" in it }, diagnostics.toString()) + // The BF16 dense still computes in BF16: nothing to report for it. + assertTrue(diagnostics.none { "bf16Dense" in it }, diagnostics.toString()) + } + + @Test + fun defaultPipeline_keepsATracedMixedPrecisionGraphConnected() { + // Every op inside a boundary now carries a policy, so the constraint pass visits it. + val graph = trace(::mixedNetwork) + val optimized = GraphOptimizationPipeline( + passes = listOf(DTypeConstraintResolutionPass(), CastFoldingPass()), + ).optimize(graph).graph + + // One cast pair folded: one node and one edge fewer, nothing else lost. + assertEquals(graph.nodes.size - 1, optimized.nodes.size) + assertEquals(graph.edges.size - 1, optimized.edges.size) + val mlir = StableHloConverterFactory.createExtended().convert(optimized, "mixed").content + assertFalse("Unsupported" in mlir, mlir) + assertEquals(3, convertLines(mlir).size, mlir) + assertTrue("(%arg0: tensor<1x4xf32>) -> (tensor<1x2xf32>)" in mlir, mlir) + } + + @Test + fun castFolding_keepsTheBoundaryPolicies() { + val before = trace(::mixedNetwork) + val after = CastFoldingPass().apply(trace(::mixedNetwork)).graph + + fun policies(g: ComputeGraph) = g.getTopologicalOrder().mapNotNull { n -> declared(n)?.let { n.operationName to it } } + assertEquals(policies(before), policies(after)) + } + @Test fun dtypePropagation_keepsTheBoundaries() { // The pass makes edges dtype-consistent and must treat convert as intentionally diff --git a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api index 656be1b4..faef1bec 100644 --- a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api +++ b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api @@ -120,6 +120,7 @@ public final class sk/ainet/compile/opt/passes/ConstantFoldingPass : sk/ainet/co } public final class sk/ainet/compile/opt/passes/DTypeConstraintResolutionPass : sk/ainet/compile/opt/GraphOptimizationPass { + public static final field BOUNDARY_KEY Ljava/lang/String; public static final field Companion Lsk/ainet/compile/opt/passes/DTypeConstraintResolutionPass$Companion; public static final field POLICY_KEY Ljava/lang/String; public static final field RESOLVED_KEY Ljava/lang/String; diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/DTypeConstraintResolutionPass.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/DTypeConstraintResolutionPass.kt index 8002523f..6eee6c75 100644 --- a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/DTypeConstraintResolutionPass.kt +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/DTypeConstraintResolutionPass.kt @@ -3,6 +3,7 @@ package sk.ainet.compile.opt.passes import sk.ainet.compile.opt.GraphOptimizationPass import sk.ainet.compile.opt.GraphOptimizationResult import sk.ainet.lang.graph.ComputeGraph +import sk.ainet.lang.graph.DefaultComputeGraph import sk.ainet.lang.graph.GraphNode import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType @@ -52,11 +53,20 @@ public class DTypeConstraintResolutionPass : GraphOptimizationPass { override fun apply(graph: ComputeGraph): GraphOptimizationResult { val diagnostics = mutableListOf() - var changed = false + val resolved = LinkedHashMap() for (node in graph.nodes) { val policy = node.metadata[POLICY_KEY] as? DTypePolicy ?: continue + // A policy that comes from a precision boundary is about the float format the + // region computes in; an index or mask input inside it is not a violation. + val boundary = node.metadata[BOUNDARY_KEY] as? String val inputDtypes = node.inputs.map { it.dtype } + .filter { boundary == null || isFloatDtypeString(it) } + val subject = if (boundary == null) { + "Node '${node.id}' (${node.operationName})" + } else { + "Node '${node.id}' (${node.operationName}) in precision boundary '$boundary'" + } when (policy) { DTypePolicy.Any -> { /* permissive; no-op */ } @@ -80,7 +90,7 @@ public class DTypeConstraintResolutionPass : GraphOptimizationPass { val targetName = policy.target.name for ((i, dtypeStr) in inputDtypes.withIndex()) { if (!dtypeStringMatches(dtypeStr, policy.target)) { - diagnostics += "Node '${node.id}' (${node.operationName}) prefers " + + diagnostics += "$subject prefers " + "$targetName but input $i has dtype '$dtypeStr' — using the existing dtype." } } @@ -103,13 +113,24 @@ public class DTypeConstraintResolutionPass : GraphOptimizationPass { // Mark the node as resolved by this pass. Use copy to keep // the immutable-copy convention the other passes follow. - val resolved = node.copy(metadata = node.metadata + (RESOLVED_KEY to true)) - graph.removeNode(node) - graph.addNode(resolved) - changed = true + if (node.metadata[RESOLVED_KEY] != true) { + resolved[node.id] = node.copy(metadata = node.metadata + (RESOLVED_KEY to true)) + } } - return GraphOptimizationResult(graph, changed = changed, diagnostics = diagnostics) + if (resolved.isEmpty()) return GraphOptimizationResult(graph, changed = false, diagnostics = diagnostics) + + // Rebuild so the edges point at the marked nodes. Replacing a node in place + // (removeNode + addNode) would drop every edge connected to it. + val newGraph = DefaultComputeGraph() + val added = HashMap() + for (node in graph.nodes) added[node.id] = newGraph.addNode(resolved[node.id] ?: node) + for (edge in graph.edges) { + val source = added[edge.source.id] ?: continue + val destination = added[edge.destination.id] ?: continue + newGraph.addEdge(edge.copy(source = source, destination = destination)) + } + return GraphOptimizationResult(newGraph, changed = true, diagnostics = diagnostics) } /** @@ -134,10 +155,22 @@ public class DTypeConstraintResolutionPass : GraphOptimizationPass { } } + private fun isFloatDtypeString(dtypeStr: String): Boolean = + FLOAT_DTYPES.any { dtypeStringMatches(dtypeStr, it) } + public companion object { /** Attribute key shared with the DSL extension (W6). */ public const val POLICY_KEY: String = "dtype_policy" + /** + * Metadata key naming the `PrecisionBoundary` a traced node was recorded inside + * (`sk.ainet.lang.graph.PRECISION_BOUNDARY_METADATA_KEY`, #1495). Such a node's + * policy applies to its floating-point inputs only. + */ + public const val BOUNDARY_KEY: String = "precision_boundary" + + private val FLOAT_DTYPES: List = listOf(FP16, BF16, FP32, FP64) + /** Marker the pass writes onto every node it visits. */ public const val RESOLVED_KEY: String = "dtype_resolved" } diff --git a/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/DTypeConstraintResolutionPassTest.kt b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/DTypeConstraintResolutionPassTest.kt index e2b52b3b..419b45c0 100644 --- a/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/DTypeConstraintResolutionPassTest.kt +++ b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/DTypeConstraintResolutionPassTest.kt @@ -8,6 +8,7 @@ import kotlin.test.assertTrue import sk.ainet.compile.opt.passes.DTypeConstraintResolutionPass import sk.ainet.compile.opt.passes.DtypeConstraintViolationException import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge import sk.ainet.lang.graph.GraphNode import sk.ainet.lang.tensor.ops.GenericOperation import sk.ainet.lang.tensor.ops.TensorSpec @@ -117,4 +118,64 @@ class DTypeConstraintResolutionPassTest { assertTrue(msg.contains("OneOf"), msg) assertTrue(msg.contains("Float32"), msg) } + + @Test + fun resolved_nodes_keep_their_edges() { + // Marking a node used to replace it in place, which dropped every edge connected to it. + val g = DefaultComputeGraph() + val a = g.addNode(node("a", opName = "input")) + val b = g.addNode(node("b", policy = DTypePolicy.Prefer(FP32))) + val c = g.addNode(node("c", opName = "relu")) + g.addEdge(GraphEdge("ab", a, b, tensorSpec = a.outputs.single())) + g.addEdge(GraphEdge("bc", b, c, tensorSpec = b.outputs.single())) + + val result = DTypeConstraintResolutionPass().apply(g) + + assertTrue(result.changed) + assertEquals(listOf("a" to "b", "b" to "c"), result.graph.edges.map { it.source.id to it.destination.id }) + assertEquals(listOf("a", "b", "c"), result.graph.getTopologicalOrder().map { it.id }) + val marked = result.graph.nodes.single { it.id == "b" } + assertEquals(true, marked.metadata[DTypeConstraintResolutionPass.RESOLVED_KEY]) + assertTrue(result.graph.edges.all { it.source in result.graph.nodes && it.destination in result.graph.nodes }) + } + + @Test + fun a_second_run_changes_nothing() { + val g = DefaultComputeGraph() + g.addNode(node("n0", policy = DTypePolicy.Require(FP32))) + + val first = DTypeConstraintResolutionPass().apply(g) + val second = DTypeConstraintResolutionPass().apply(first.graph) + + assertTrue(first.changed) + assertFalse(second.changed, "an already resolved graph must let a pipeline converge") + } + + @Test + fun boundary_policy_applies_to_float_inputs_only() { + // An op recorded inside a precision boundary: its float input is checked against the + // declared dtype, an index input is not. + fun boundaryNode(floatDtype: String) = GraphNode( + id = "g", + operation = GenericOperation("gather"), + inputs = listOf( + TensorSpec(name = "table", shape = listOf(8, 4), dtype = floatDtype), + TensorSpec(name = "ids", shape = listOf(2), dtype = "Int32"), + ), + outputs = listOf(TensorSpec(name = "rows", shape = listOf(2, 4), dtype = floatDtype)), + metadata = mapOf( + DTypeConstraintResolutionPass.POLICY_KEY to DTypePolicy.Prefer(BF16), + DTypeConstraintResolutionPass.BOUNDARY_KEY to "encoder", + ), + ) + + val kept = DefaultComputeGraph().also { it.addNode(boundaryNode("BF16")) } + assertEquals(emptyList(), DTypeConstraintResolutionPass().apply(kept).diagnostics) + + val retyped = DefaultComputeGraph().also { it.addNode(boundaryNode("FP32")) } + val diagnostics = DTypeConstraintResolutionPass().apply(retyped).diagnostics + assertEquals(1, diagnostics.size, diagnostics.toString()) + assertTrue("in precision boundary 'encoder'" in diagnostics.single(), diagnostics.single()) + assertTrue("input 0" in diagnostics.single() && "FP32" in diagnostics.single(), diagnostics.single()) + } } diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 8136cd13..b305d8ca 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -2937,8 +2937,10 @@ public final class sk/ainet/lang/nn/NeuralNetworkExecutionContextKt { public final class sk/ainet/lang/nn/PrecisionBoundary : sk/ainet/lang/nn/Module { public fun (Ljava/lang/String;Lsk/ainet/lang/nn/Module;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)V + public final fun getInnerType ()Lsk/ainet/lang/types/DType; public fun getModules ()Ljava/util/List; public fun getName ()Ljava/lang/String; + public final fun getOuterType ()Lsk/ainet/lang/types/DType; } public final class sk/ainet/lang/nn/TrainingRunnerKt { diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt index b875f69f..99f4c2fd 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt @@ -15,6 +15,10 @@ import sk.ainet.lang.types.DType * the tracing wrappers and lowers to `stablehlo.convert`, so the boundary stays visible * in traces and exports. * + * A tracing context marks every op recorded inside the wrapped module with + * `DTypePolicy.Prefer(innerType)`, so `DTypeConstraintResolutionPass` reports a boundary + * whose ops a later pass retyped. + * * `convert` is differentiable between floating-point dtypes, so gradients cross the * boundary and the wrapped layers train. The gradient is cast like the values: it is * rounded to [innerType] where it enters the wrapped module and widened back to @@ -24,8 +28,10 @@ import sk.ainet.lang.types.DType public class PrecisionBoundary( override val name: String, private val inner: Module, - private val innerType: TInner, - private val outerType: TOuter, + /** The dtype the wrapped module's weights and computation use. */ + public val innerType: TInner, + /** The dtype of the surrounding network, which the boundary's input and output have. */ + public val outerType: TOuter, ) : Module() { @Suppress("UNCHECKED_CAST")