feat(compile-opt): CastFoldingPass folds casts that change no value (#1494)

CHANGELOG.md

@@ -14,6 +14,8 @@
 
 ### Added
 
+- **`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.
 - **SKEEP-008 acceptance tests** (`skainet-lang-core`, `skainet-backend-cpu`, #1433): `repackTo(Dense(4))` is checked bit-exact for Q5_0, Q5_1, Q5_K, Q6_K and `BITNET_B1_58` too, and one FP32 network with both a BF16 `dense` and an FP16 `stage` runs against the FP32 reference.
 - **Mixed-precision export test** (`skainet-compile-hlo`, SKEEP-008 phase A2, #1433): `MixedPrecisionExportTest` traces an FP32 network with one BF16 `dense` and one FP16 `stage` and pins the result: one `convert` pair per boundary in the graph, inner layers typed in the narrow dtype, and `stablehlo.convert` `f32 -> bf16`, `bf16 -> f32`, `f32 -> f16`, `f16 -> f32` in the exported module.

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

@@ -107,8 +107,8 @@ hiding the conversion inside the layer:
   cast in a trace or an export.
 * *Visible in traces and exports.* The tracing wrappers record `convert` like any other op,
   and the StableHLO converter lowers it to `stablehlo.convert` — a compiled or exported graph
-  shows exactly where precision changes, and a later optimisation pass can fold adjacent
-  cast pairs.
+  shows exactly where precision changes, and `CastFoldingPass` folds the cast pair between
+  adjacent boundaries.
 * *Honest about cost.* Every boundary copies its input and its output. A network that
   switches dtype every layer can be *slower* than all-FP32, especially while narrow dtypes
   are float-backed. Few, coarse boundaries (a `stage` around a block) beat many fine ones.
@@ -137,9 +137,21 @@ module has this shape (abridged: constants, transposes and bias adds are left ou
 %v21 = stablehlo.convert %v20 : (tensor<1x10xf16>) -> tensor<1x10xf32>       // and back out
 ----
 
-The `bf16 -> f32` cast followed directly by `f32 -> f16` is the price of two adjacent
-boundaries with different inner dtypes: the network's own dtype is always restored in
-between. `MixedPrecisionExportTest` pins this output.
+The `bf16 -> f32` cast followed directly by `f32 -> f16` is what two adjacent boundaries
+with different inner dtypes record: each one restores the network's dtype, so the tensor is
+widened only to be narrowed again. `MixedPrecisionExportTest` pins this output.
+
+`CastFoldingPass`, part of the default optimization pipelines, removes that round trip. A
+widening from FP16 or BF16 to FP32 is exact, so narrowing the widened tensor gives the same
+values as narrowing the original, and the two casts fold into one:
+
+[source,mlir]
+----
+%v13 = stablehlo.convert %v12 : (tensor<1x256xbf16>) -> tensor<1x256xf16>    // BF16 dense straight into the FP16 stage
+----
+
+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.
 
 == Training through a boundary
 

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

@@ -200,7 +200,7 @@ public class PrecisionBoundary<TOuter : DType, TInner : DType, V>(
 A tensor that already has the target dtype is passed through without calling `convert` (the
 tracing wrapper would record the call even though the backend returns its input), so a boundary
 between equal dtypes costs nothing and leaves no cast in a trace. Two consecutive boundaries of the same inner dtype are a cast followed by its
-inverse; a later optimisation pass may fold them, and nothing in this design depends on that.
+inverse; `CastFoldingPass` (#1494) folds them, and nothing in this design depends on that.
 
 Weights of a `TLayer` dense are created in `TLayer` by the existing initialisers. User-supplied
 weight tensors must already be `TLayer`; the DSL does not silently cast them.
@@ -349,15 +349,19 @@ Each phase is a DARC feature with lane sub-issues, linked from the tracking issu
   *Resolved: no.* Only the value form shipped.
 . Should a stage boundary that ends in the same dtype as the next stage elide the cast pair at
   DSL time, or leave that to a compile pass? +
-  *Open.* The DSL elides only the degenerate case (inner dtype equal to the network dtype). Two
-  adjacent boundaries still export a `narrow -> f32` cast followed by `f32 -> narrow`; folding
-  that pair is left to a compile pass.
+  *Resolved: a compile pass.* The DSL elides only the degenerate case (inner dtype equal to
+  the network dtype). `CastFoldingPass`
+  (https://github.com/SKaiNET-developers/SKaiNET/issues/1494[#1494]) folds a cast that follows
+  an exact widening, so `narrow -> f32 -> narrow` between two adjacent boundaries becomes one
+  cast or none. Merging the boundaries in the DSL was rejected: it would change the module
+  tree the user wrote and would not help hand-written or imported graphs.
 . Is `TensorStorage` the right owner of `requantizeTo`, or should it be a free function in
   `sk.ainet.lang.tensor.data` next to `Requantizers`, keeping `TensorStorage` a descriptor? +
   *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].
 
 == References
 

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

@@ -352,6 +352,16 @@ public class ComputeGraphExecutor(
                 listOf(ops.gather(inputs[0], indices, dim))
             }
 
+            // Dtype cast. The target is the node's output dtype, which is what tracing
+            // records and what the StableHLO converter reads.
+            "convert", "cast", "to" -> {
+                val targetName = node.outputs.firstOrNull()?.dtype
+                    ?: error("cast node '${node.id}' declares no output dtype")
+                val target = dtypeNamed(targetName)
+                    ?: error("cast node '${node.id}' targets unknown dtype '$targetName'")
+                listOf(ops.convert(inputs[0], target) as Tensor<T, V>)
+            }
+
             // Pass-through / identity
             "identity", "input", "weight", "constant", "parameter" -> {
                 if (inputs.isNotEmpty()) inputs.toList() else listOf()
@@ -401,6 +411,18 @@ public class ComputeGraphExecutor(
         }
     }
 
+    /**
+     * The dtype a tensor spec names. Specs carry the short spelling tracing records
+     * (`FP16`, `BF16`, `FP32`) or a format's own (`f32`, `float32`), not [DType.name].
+     */
+    private fun dtypeNamed(name: String): DType? = when (val upper = name.uppercase()) {
+        "FP16", "F16", "FLOAT16" -> sk.ainet.lang.types.FP16
+        "BF16", "BFLOAT16" -> sk.ainet.lang.types.BF16
+        "FP32", "F32", "FLOAT32" -> sk.ainet.lang.types.FP32
+        "FP64", "F64", "FLOAT64" -> sk.ainet.lang.types.FP64
+        else -> DType.entries.firstOrNull { it.name.uppercase() == upper }
+    }
+
     private fun isInputNode(node: GraphNode): Boolean =
         node.operationName in INPUT_OPS
 

skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/MixedPrecisionExportTest.kt

@@ -1,15 +1,22 @@
 package sk.ainet.compile.hlo
 
+import sk.ainet.compile.opt.passes.CastFoldingPass
 import sk.ainet.compile.opt.passes.DtypeForwardPropagationPass
 import sk.ainet.context.DirectCpuExecutionContext
 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.exec.ComputeGraphExecutor
 import sk.ainet.lang.graph.DefaultGraphExecutionContext
 import sk.ainet.lang.nn.Module
 import sk.ainet.lang.nn.definition
 import sk.ainet.lang.nn.network
 import sk.ainet.lang.tape.toComputeGraph
 import sk.ainet.lang.tensor.Shape
 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.FP16
 import sk.ainet.lang.types.FP32
@@ -137,6 +144,64 @@ class MixedPrecisionExportTest {
         assertTrue(convertLines(mlir).isEmpty(), "no cast expected:\n$mlir")
     }
 
+    @Test
+    fun castFolding_dropsTheRoundTripBetweenAdjacentBoundaries() {
+        // The BF16 dense is directly followed by the FP16 stage: bf16 -> f32 -> f16 is one
+        // exact widening followed by a narrowing, and folds to bf16 -> f16 (#1494).
+        val folded = CastFoldingPass().apply(trace(::mixedNetwork))
+        assertTrue(folded.changed)
+
+        val mlir = StableHloConverterFactory.createExtended().convert(folded.graph, "mixed").content
+
+        assertFalse("Unsupported" in mlir, mlir)
+        assertEquals(
+            listOf(
+                "(tensor<1x4xf32>) -> tensor<1x4xbf16>",
+                "(tensor<1x3xbf16>) -> tensor<1x3xf16>",
+                "(tensor<1x2xf16>) -> tensor<1x2xf32>",
+            ),
+            convertLines(mlir).map { it.substringAfter(" : ") },
+            mlir,
+        )
+    }
+
+    @Test
+    fun castFolding_givesBitIdenticalResults() {
+        // The cast chain two adjacent boundaries produce, on values no 16-bit format holds:
+        // executing the folded graph must reproduce the unfolded one bit for bit.
+        fun cast(id: String, from: String, to: String) = GraphNode(
+            id = id,
+            operation = GenericOperation("convert"),
+            inputs = listOf(TensorSpec("t", listOf(1, 6), from)),
+            outputs = listOf(TensorSpec("t", listOf(1, 6), to)),
+        )
+        val graph = DefaultComputeGraph()
+        var previous = graph.addNode(
+            GraphNode("x", GenericOperation("input"), emptyList(), listOf(TensorSpec("x", listOf(1, 6), "FP32")))
+        )
+        val steps = listOf("FP32", "BF16", "FP32", "FP16", "FP32")
+        for (i in 1 until steps.size) {
+            val node = graph.addNode(cast("c$i", steps[i - 1], steps[i]))
+            graph.addEdge(GraphEdge("e$i", previous, node, tensorSpec = TensorSpec("t", listOf(1, 6), steps[i - 1])))
+            previous = node
+        }
+        val folded = CastFoldingPass().apply(graph).graph
+        assertEquals(3, folded.nodes.count { it.operationName == "convert" })
+
+        val cpu = DirectCpuExecutionContext()
+        val values = floatArrayOf(0.1234567f, -1.7654321f, 2.3456789f, 65000.5f, 1.0e-6f, -3.3333333f)
+        val x = cpu.fromFloatArray<FP32, Float>(Shape(1, 6), FP32::class, values)
+
+        fun run(g: ComputeGraph): List<Int> {
+            val out = ComputeGraphExecutor(g, cpu.ops).execute(mapOf("x" to x)).values.single()
+            return List(6) { out.data[0, it].toRawBits() }
+        }
+
+        val expected = run(graph)
+        assertEquals(expected, run(folded))
+        assertTrue(expected != values.map { it.toRawBits() }, "the chain must round, or it proves nothing")
+    }
+
     @Test
     fun dtypePropagation_keepsTheBoundaries() {
         // The pass makes edges dtype-consistent and must treat convert as intentionally

skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api

@@ -107,6 +107,12 @@ public final class sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass : sk/a
 public final class sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass$Companion {
 }
 
+public final class sk/ainet/compile/opt/passes/CastFoldingPass : sk/ainet/compile/opt/GraphOptimizationPass {
+	public fun <init> ()V
+	public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult;
+	public fun getName ()Ljava/lang/String;
+}
+
 public final class sk/ainet/compile/opt/passes/ConstantFoldingPass : sk/ainet/compile/opt/GraphOptimizationPass {
 	public fun <init> ()V
 	public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult;

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.CastFoldingPass
 import sk.ainet.compile.opt.passes.DTypeConstraintResolutionPass
 import sk.ainet.compile.opt.passes.DeadCodeEliminationPass
 import sk.ainet.compile.opt.passes.LLMFusionPass
@@ -81,6 +82,9 @@ public class GraphOptimizationPipeline(
                 // is the boundary where dtype problems surface — every
                 // later pass can assume dtype-validity.
                 DTypeConstraintResolutionPass(),
+                // Fold casts that change no value (a widening followed by another
+                // cast, a cast to the same dtype) while the graph is still unfused.
+                CastFoldingPass(),
                 // Rewrite pow(x, 2) to multiply(x, x) before fusion so
                 // the downstream passes see the multiply form. Runs after
                 // dtype resolution (still benefits from resolved dtypes)
@@ -98,6 +102,7 @@ public class GraphOptimizationPipeline(
         public fun createAggressive(): GraphOptimizationPipeline = GraphOptimizationPipeline(
             passes = listOf(
                 DTypeConstraintResolutionPass(),
+                CastFoldingPass(),
                 PowSpecializationPass(),
                 DeadCodeEliminationPass(),
                 ConstantFoldingPass(),
@@ -110,7 +115,8 @@ public class GraphOptimizationPipeline(
          * Creates an LLM-optimized pipeline with transformer-specific passes.
          *
          * Pass ordering:
-         * 0. DTypeConstraintResolution — resolve dtype policies before fusion
+         * 0. DTypeConstraintResolution — resolve dtype policies before fusion,
+         *    then CastFolding — drop casts that change no value
          * 1. TransposeElimination — fold transposes into matmuls
          * 2. SharedWeightDedup — deduplicate tied weights (e.g. token_embd ↔ output)
          * 3. LLMFusion — fuse RMSNorm, SwiGLU, QKV patterns
@@ -120,6 +126,7 @@ public class GraphOptimizationPipeline(
         public fun createLLM(): GraphOptimizationPipeline = GraphOptimizationPipeline(
             passes = listOf(
                 DTypeConstraintResolutionPass(),
+                CastFoldingPass(),
                 TransposeEliminationPass(),
                 SharedWeightDeduplicationPass(),
                 LLMFusionPass(),