Add DTypeConstraintResolutionPass + register in pipeline (W7 of #615)

michalharakal · claude · michalharakal · commit faeebe47903d · 2026-05-17T21:34:46.000+02:00
New GraphOptimizationPass that walks the graph, reads each node's
attached DTypePolicy from `metadata["dtype_policy"]` (set by the
W6 DSL extension), and enforces the RFC's "fail before execution"
rule. Three policy arms:

- Require(target): every input edge MUST already have the target
  dtype. Mismatch throws DtypeConstraintViolationException with
  a message naming the failing edge, the required dtype, the
  actual dtype, and the suggested resolution (loader-side policy
  via SafeTensorsParametersLoader.withPolicy / StreamingGguf-
  ParametersLoader.withPolicy).
- Prefer(target): mismatch emits a diagnostic and falls through.
- OneOf(allowed): input must be in the allowed set; otherwise throw.
- Any: no-op (and nodes without an attached policy are passed
  through unvisited).

Visited nodes get `metadata["dtype_resolved"] = true` so the
W8 ResolvedComputeGraph wrapper can verify the pass has run.

dtypeStringMatches() handles both the canonical DType.name strings
("Float32", "BFloat16") and the short class-derived aliases the
DAG DSL produces ("FP32", "BF16", …) — the resolver works whether
the graph came from the DSL or from a TensorSpec-string-driven
source like the StableHLO converter.

NO new cast kernels in this PR — when a Require mismatches and no
cast kernel is registered, the pass fails fast. Cast-node
insertion is a follow-up that ships with concrete cast kernels.

Registered in all three pipeline factories (createDefault,
createAggressive, createLLM) BEFORE the fusion passes so fusion
sees dtype-resolved nodes. All existing pipeline tests still pass.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/GraphOptimizationPipeline.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/GraphOptimizationPipeline.kt
@@ -2,6 +2,7 @@ package sk.ainet.compile.opt
 
 import sk.ainet.lang.graph.ComputeGraph
 import sk.ainet.compile.opt.passes.ConstantFoldingPass
+import sk.ainet.compile.opt.passes.DTypeConstraintResolutionPass
 import sk.ainet.compile.opt.passes.DeadCodeEliminationPass
 import sk.ainet.compile.opt.passes.LLMFusionPass
 import sk.ainet.compile.opt.passes.OperationFusionPass
@@ -73,6 +74,12 @@ public class GraphOptimizationPipeline(
          */
         public fun createDefault(): GraphOptimizationPipeline = GraphOptimizationPipeline(
             passes = listOf(
+                // Resolve dtype constraints first so fusion / DCE / constant
+                // folding see the resolved-or-failed graph rather than a
+                // mix of policy-tagged and bare nodes. Per the RFC, this
+                // is the boundary where dtype problems surface — every
+                // later pass can assume dtype-validity.
+                DTypeConstraintResolutionPass(),
                 DeadCodeEliminationPass(),
                 ConstantFoldingPass(),
                 OperationFusionPass()
@@ -84,6 +91,7 @@ public class GraphOptimizationPipeline(
          */
         public fun createAggressive(): GraphOptimizationPipeline = GraphOptimizationPipeline(
             passes = listOf(
+                DTypeConstraintResolutionPass(),
                 DeadCodeEliminationPass(),
                 ConstantFoldingPass(),
                 OperationFusionPass()
@@ -95,6 +103,7 @@ public class GraphOptimizationPipeline(
          * Creates an LLM-optimized pipeline with transformer-specific passes.
          *
          * Pass ordering:
+         * 0. DTypeConstraintResolution — resolve dtype policies before fusion
          * 1. TransposeElimination — fold transposes into matmuls
          * 2. SharedWeightDedup — deduplicate tied weights (e.g. token_embd ↔ output)
          * 3. LLMFusion — fuse RMSNorm, SwiGLU, QKV patterns
@@ -103,6 +112,7 @@ public class GraphOptimizationPipeline(
          */
         public fun createLLM(): GraphOptimizationPipeline = GraphOptimizationPipeline(
             passes = listOf(
+                DTypeConstraintResolutionPass(),
                 TransposeEliminationPass(),
                 SharedWeightDeduplicationPass(),
                 LLMFusionPass(),
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
@@ -0,0 +1,153 @@
+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.GraphNode
+import sk.ainet.lang.types.BF16
+import sk.ainet.lang.types.DType
+import sk.ainet.lang.types.DTypePolicy
+import sk.ainet.lang.types.FP16
+import sk.ainet.lang.types.FP32
+import sk.ainet.lang.types.FP64
+import sk.ainet.lang.types.Int8
+import sk.ainet.lang.types.Int16
+import sk.ainet.lang.types.Int32
+import sk.ainet.lang.types.Int64
+
+/**
+ * Pass that enforces per-node [DTypePolicy] constraints attached to
+ * graph nodes (via the `dag { … dtypePolicy(…) }` DSL extension from
+ * W6 of #615). Implements the RFC's "fail before execution" rule —
+ * any [DTypePolicy.Require] that can't be satisfied raises
+ * [DtypeConstraintViolationException] *here*, at graph-prep time,
+ * not at forward execution.
+ *
+ * Policy semantics:
+ * - `Any`: never visited; nodes without an attached policy are
+ *   passed through.
+ * - `Require(target)`: every input edge to the node MUST already
+ *   have dtype matching `target`. Mismatch throws
+ *   [DtypeConstraintViolationException].
+ * - `Prefer(target)`: input dtype matching `target` is preferred;
+ *   mismatches emit a diagnostic but do not fail.
+ * - `OneOf(allowed)`: every input edge's dtype MUST already be in
+ *   `allowed`. Mismatch throws.
+ *
+ * **Scope intentionally narrow.** This pass does not insert cast
+ * nodes today — when a `Require` mismatches, it fails fast (which
+ * is the RFC's prescribed behaviour when no cast kernel exists).
+ * Cast-node insertion is a follow-up that ships alongside concrete
+ * cast kernels (Q4_K → Int8, FP32 → BF16, …). See the
+ * out-of-scope section of issue #615.
+ *
+ * Side effect on the graph: visited nodes get
+ * `metadata["dtype_resolved"] = true` so downstream passes (and the
+ * future `ResolvedComputeGraph` wrapper from W8) can confirm the
+ * pass has run.
+ */
+public class DTypeConstraintResolutionPass : GraphOptimizationPass {
+
+    override val name: String = "dtype-constraint-resolution"
+
+    override fun apply(graph: ComputeGraph): GraphOptimizationResult {
+        val diagnostics = mutableListOf<String>()
+        var changed = false
+
+        for (node in graph.nodes) {
+            val policy = node.metadata[POLICY_KEY] as? DTypePolicy ?: continue
+            val inputDtypes = node.inputs.map { it.dtype }
+
+            when (policy) {
+                DTypePolicy.Any -> { /* permissive; no-op */ }
+
+                is DTypePolicy.Require -> {
+                    val targetName = policy.target.name
+                    for ((i, dtypeStr) in inputDtypes.withIndex()) {
+                        if (!dtypeStringMatches(dtypeStr, policy.target)) {
+                            throw DtypeConstraintViolationException(
+                                "Node '${node.id}' (${node.operationName}) declares " +
+                                    "DTypePolicy.Require($targetName) but input $i has dtype '$dtypeStr'. " +
+                                    "Cast kernels are not registered for this conversion; resolve at the " +
+                                    "loader (e.g. SafeTensorsParametersLoader.withPolicy) or change the " +
+                                    "policy to Prefer/OneOf to permit fallback."
+                            )
+                        }
+                    }
+                }
+
+                is DTypePolicy.Prefer -> {
+                    val targetName = policy.target.name
+                    for ((i, dtypeStr) in inputDtypes.withIndex()) {
+                        if (!dtypeStringMatches(dtypeStr, policy.target)) {
+                            diagnostics += "Node '${node.id}' (${node.operationName}) prefers " +
+                                "$targetName but input $i has dtype '$dtypeStr' — using the existing dtype."
+                        }
+                    }
+                }
+
+                is DTypePolicy.OneOf -> {
+                    val allowedNames = policy.allowed.joinToString { it.name }
+                    for ((i, dtypeStr) in inputDtypes.withIndex()) {
+                        if (policy.allowed.none { dtypeStringMatches(dtypeStr, it) }) {
+                            throw DtypeConstraintViolationException(
+                                "Node '${node.id}' (${node.operationName}) declares " +
+                                    "DTypePolicy.OneOf($allowedNames) but input $i has dtype " +
+                                    "'$dtypeStr' which is outside the allowed set. Cast kernels " +
+                                    "are not registered; resolve at the loader."
+                            )
+                        }
+                    }
+                }
+            }
+
+            // 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
+        }
+
+        return GraphOptimizationResult(graph, changed = changed, diagnostics = diagnostics)
+    }
+
+    /**
+     * Matches the string form used by [sk.ainet.lang.tensor.ops.TensorSpec.dtype]
+     * against a typed [DType]. Handles both registry-canonical names
+     * (`"Float32"`, `"BFloat16"`) and the short class-derived
+     * aliases produced by the DAG DSL's `dtypeName()` helper (`"FP32"`,
+     * `"BF16"`, `"Int8"`, …).
+     */
+    internal fun dtypeStringMatches(dtypeStr: String, dtype: DType): Boolean {
+        if (dtypeStr == dtype.name) return true
+        return when (dtype) {
+            FP32 -> dtypeStr == "FP32" || dtypeStr == "F32"
+            FP16 -> dtypeStr == "FP16" || dtypeStr == "F16"
+            BF16 -> dtypeStr == "BF16"
+            FP64 -> dtypeStr == "FP64" || dtypeStr == "F64"
+            Int8 -> dtypeStr == "Int8" || dtypeStr == "I8"
+            Int16 -> dtypeStr == "Int16" || dtypeStr == "I16"
+            Int32 -> dtypeStr == "Int32" || dtypeStr == "I32"
+            Int64 -> dtypeStr == "Int64" || dtypeStr == "I64"
+            else -> false
+        }
+    }
+
+    public companion object {
+        /** Attribute key shared with the DSL extension (W6). */
+        public const val POLICY_KEY: String = "dtype_policy"
+
+        /** Marker the pass writes onto every node it visits. */
+        public const val RESOLVED_KEY: String = "dtype_resolved"
+    }
+}
+
+/**
+ * Raised when [DTypeConstraintResolutionPass] cannot satisfy a hard
+ * [DTypePolicy.Require] (or `OneOf` rejection) and no cast kernel
+ * is available to bridge the gap. Surfaces dtype problems at
+ * graph-prep time, before forward execution — exactly the RFC's
+ * "fail before execution" boundary.
+ */
+public class DtypeConstraintViolationException(message: String) : RuntimeException(message)
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
@@ -0,0 +1,120 @@
+package sk.ainet.compile.opt
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertFalse
+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.GraphNode
+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.FP32
+import sk.ainet.lang.types.Int8
+
+class DTypeConstraintResolutionPassTest {
+
+    private fun node(
+        id: String,
+        opName: String = "matmul",
+        inputDtype: String = "Float32",
+        policy: DTypePolicy? = null,
+    ): GraphNode {
+        val meta = if (policy != null) mapOf<String, Any>(DTypeConstraintResolutionPass.POLICY_KEY to policy) else emptyMap()
+        return GraphNode(
+            id = id,
+            operation = GenericOperation(opName),
+            inputs = listOf(TensorSpec(name = "$id-in", shape = listOf(4, 4), dtype = inputDtype)),
+            outputs = listOf(TensorSpec(name = "$id-out", shape = listOf(4, 4), dtype = inputDtype)),
+            metadata = meta,
+        )
+    }
+
+    @Test
+    fun nodes_without_policy_are_passed_through() {
+        val g = DefaultComputeGraph()
+        g.addNode(node("n0"))
+        g.addNode(node("n1"))
+        val result = DTypeConstraintResolutionPass().apply(g)
+        assertFalse(result.changed, "no policy = no work")
+        // Neither node should be marked resolved (only visited nodes get the marker).
+        assertEquals(emptyList(), result.graph.nodes.filter { it.metadata.containsKey(DTypeConstraintResolutionPass.RESOLVED_KEY) })
+    }
+
+    @Test
+    fun any_policy_passes_through() {
+        val g = DefaultComputeGraph()
+        g.addNode(node("n0", policy = DTypePolicy.Any))
+        val result = DTypeConstraintResolutionPass().apply(g)
+        assertTrue(result.changed, "the resolved-marker write counts as a change")
+        val n = result.graph.nodes.single()
+        assertTrue(n.metadata[DTypeConstraintResolutionPass.RESOLVED_KEY] == true)
+    }
+
+    @Test
+    fun require_matching_dtype_passes() {
+        val g = DefaultComputeGraph()
+        g.addNode(node("n0", inputDtype = "Float32", policy = DTypePolicy.Require(FP32)))
+        val result = DTypeConstraintResolutionPass().apply(g)
+        assertTrue(result.changed)
+    }
+
+    @Test
+    fun require_mismatched_dtype_fails_fast() {
+        val g = DefaultComputeGraph()
+        g.addNode(node("n0", inputDtype = "Float32", policy = DTypePolicy.Require(BF16)))
+        val ex = assertFailsWith<DtypeConstraintViolationException> {
+            DTypeConstraintResolutionPass().apply(g)
+        }
+        val msg = ex.message ?: ""
+        assertTrue(msg.contains("BFloat16"), "msg must name the required dtype: $msg")
+        assertTrue(msg.contains("Float32"), "msg must name the actual input dtype: $msg")
+        assertTrue(msg.contains("Cast kernels"), "msg must hint at the resolution path: $msg")
+    }
+
+    @Test
+    fun require_mismatched_dtype_with_short_alias_also_resolves() {
+        // DAG DSL emits dtype strings like "FP32" / "BF16" via dtypeName().
+        // The pass must handle both the registry canonical name and the short alias.
+        val g = DefaultComputeGraph()
+        g.addNode(node("n0", inputDtype = "FP32", policy = DTypePolicy.Require(FP32)))
+        val result = DTypeConstraintResolutionPass().apply(g)
+        assertTrue(result.changed, "alias 'FP32' must satisfy Require(FP32)")
+    }
+
+    @Test
+    fun prefer_mismatched_dtype_emits_diagnostic_no_throw() {
+        val g = DefaultComputeGraph()
+        g.addNode(node("n0", inputDtype = "Float32", policy = DTypePolicy.Prefer(BF16)))
+        val result = DTypeConstraintResolutionPass().apply(g)
+        assertTrue(result.changed)
+        assertTrue(
+            result.diagnostics.any { it.contains("prefers") && it.contains("BFloat16") },
+            "diagnostic must mention the preference: ${result.diagnostics}",
+        )
+    }
+
+    @Test
+    fun oneOf_in_set_passes() {
+        val g = DefaultComputeGraph()
+        g.addNode(node("n0", inputDtype = "Float32", policy = DTypePolicy.OneOf(setOf(FP32, BF16))))
+        val result = DTypeConstraintResolutionPass().apply(g)
+        assertTrue(result.changed)
+    }
+
+    @Test
+    fun oneOf_outside_set_fails_fast() {
+        val g = DefaultComputeGraph()
+        g.addNode(node("n0", inputDtype = "Float32", policy = DTypePolicy.OneOf(setOf(BF16, Int8))))
+        val ex = assertFailsWith<DtypeConstraintViolationException> {
+            DTypeConstraintResolutionPass().apply(g)
+        }
+        val msg = ex.message ?: ""
+        assertTrue(msg.contains("OneOf"), msg)
+        assertTrue(msg.contains("Float32"), msg)
+    }
+}
