diff --git a/CHANGELOG.md b/CHANGELOG.md index 3feadee3..14b13a3e 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/docs/modules/ROOT/pages/explanation/mixed-precision.adoc b/docs/modules/ROOT/pages/explanation/mixed-precision.adoc index 1b9164e0..85442a54 100644 --- a/docs/modules/ROOT/pages/explanation/mixed-precision.adoc +++ b/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 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 16e3469d..378ec988 100644 --- a/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc +++ b/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc @@ -200,7 +200,7 @@ public class PrecisionBoundary( 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 diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt index a99b6e99..30987423 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt +++ b/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) + } + // 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 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 be94a9f5..921b9bf0 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,8 +1,13 @@ 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 @@ -10,6 +15,8 @@ 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(Shape(1, 6), FP32::class, values) + + fun run(g: ComputeGraph): List { + 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 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 8b77e65a..656be1b4 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 @@ -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 ()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 ()V public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; 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 index d77c7222..d0967551 100644 --- 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.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(), diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/CastFoldingPass.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/CastFoldingPass.kt new file mode 100644 index 00000000..7fc7593f --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/CastFoldingPass.kt @@ -0,0 +1,160 @@ +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.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.TensorSpec + +/** + * Folds dtype casts (`convert` / `cast` / `to`) that change no value (#1494). + * + * Two rewrites, both value-preserving: + * + * 1. **A cast after an exact widening.** In `A -> B -> C`, when every value of `A` is + * representable in `B` (FP16 or BF16 to FP32, any of them to FP64), the second cast reads + * from the original `A` tensor instead: rounding an exactly widened value gives the same + * result as rounding the value itself. `bf16 -> f32 -> f16` becomes `bf16 -> f16`. The + * widening cast is removed once nothing else reads it. + * 2. **A cast to the dtype its input already has**, which rewrite 1 leaves behind for + * `f16 -> f32 -> f16` and which [DtypeForwardPropagationPass] can produce when it coerces a + * graph to one float dtype. The node is removed and its consumers read its input. + * + * This is what two adjacent mixed-precision boundaries need (SKEEP-008): each casts back to + * the network dtype, so between them a tensor is widened only to be narrowed again. + * + * A *narrowing* followed by a widening (`f32 -> bf16 -> f32`) is never touched: that pair is + * the rounding a mixed-precision layer asked for. Casts from or to integer dtypes are not + * treated as exact (Int32 to FP32 is not) and are left alone, except by rewrite 2. + */ +public class CastFoldingPass : GraphOptimizationPass { + + override val name: String = "cast-folding" + + private class Link( + val id: String, + var sourceId: String, + var sourceOutputIndex: Int, + val destinationId: String, + val destinationInputIndex: Int, + var tensorSpec: TensorSpec, + val metadata: Map, + ) + + override fun apply(graph: ComputeGraph): GraphOptimizationResult { + if (graph.nodes.none { isCast(it) }) return GraphOptimizationResult(graph, changed = false) + + val nodes = LinkedHashMap() + for (node in graph.nodes) nodes[node.id] = node + val links = graph.edges.mapTo(mutableListOf()) { + Link( + it.id, it.source.id, it.sourceOutputIndex, it.destination.id, + it.destinationInputIndex, it.tensorSpec, it.metadata, + ) + } + val diagnostics = mutableListOf() + + // One rewrite per round keeps the bookkeeping simple; a round is linear in the graph + // and the number of rounds is bounded by the number of cast nodes. + while (foldOne(nodes, links, diagnostics)) Unit + + if (diagnostics.isEmpty()) return GraphOptimizationResult(graph, changed = false) + + val newGraph = DefaultComputeGraph() + val added = HashMap() + for (node in nodes.values) added[node.id] = newGraph.addNode(node.copy()) + for (link in links) { + val source = added[link.sourceId] ?: continue + val destination = added[link.destinationId] ?: continue + newGraph.addEdge( + GraphEdge( + id = link.id, + source = source, + destination = destination, + sourceOutputIndex = link.sourceOutputIndex, + destinationInputIndex = link.destinationInputIndex, + tensorSpec = link.tensorSpec, + metadata = link.metadata, + ) + ) + } + return GraphOptimizationResult(newGraph, changed = true, diagnostics = diagnostics) + } + + /** Applies the first rewrite that matches; `false` when none does. */ + private fun foldOne( + nodes: MutableMap, + links: MutableList, + diagnostics: MutableList, + ): Boolean { + for (cast in nodes.values) { + if (!isCast(cast)) continue + val incoming = links.filter { it.destinationId == cast.id } + if (incoming.size != 1) continue + val input = incoming.single() + val producer = nodes[input.sourceId] ?: continue + val producerSpec = producer.outputs.getOrNull(input.sourceOutputIndex) ?: continue + val from = canonical(producerSpec.dtype) + val to = canonical(cast.outputs.firstOrNull()?.dtype ?: continue) + + // Rewrite 2: the cast changes nothing; its consumers read its input. + if (from == to) { + links.remove(input) + for (out in links) { + if (out.sourceId != cast.id) continue + out.sourceId = producer.id + out.sourceOutputIndex = input.sourceOutputIndex + out.tensorSpec = producerSpec + } + nodes.remove(cast.id) + diagnostics += "Removed ${cast.id}: cast from $from to the same dtype" + return true + } + + // Rewrite 1: the input is itself a cast, and an exact one; read through it. + if (!isCast(producer)) continue + val widenedFrom = links.filter { it.destinationId == producer.id }.singleOrNull() ?: continue + val origin = nodes[widenedFrom.sourceId] ?: continue + val originSpec = origin.outputs.getOrNull(widenedFrom.sourceOutputIndex) ?: continue + val originType = canonical(originSpec.dtype) + if (!isExactWidening(originType, from)) continue + + input.sourceId = origin.id + input.sourceOutputIndex = widenedFrom.sourceOutputIndex + input.tensorSpec = originSpec + nodes[cast.id] = cast.copy( + inputs = listOf(originSpec) + cast.inputs.drop(1), + ) + diagnostics += "Folded ${producer.id} ($originType -> $from) into ${cast.id}: now $originType -> $to" + if (links.none { it.sourceId == producer.id }) { + links.remove(widenedFrom) + nodes.remove(producer.id) + } + return true + } + return false + } + + private fun isCast(node: GraphNode): Boolean = node.operationName.lowercase() in CAST_OPS + + private fun canonical(dtype: String): String = when (val upper = dtype.uppercase()) { + "F16", "FLOAT16" -> "FP16" + "BFLOAT16" -> "BF16" + "F32", "FLOAT32" -> "FP32" + "F64", "FLOAT64" -> "FP64" + else -> upper + } + + /** Every value of [from] is representable in [to]. Floating-point dtypes only. */ + private fun isExactWidening(from: String, to: String): Boolean = when (to) { + "FP32" -> from == "FP16" || from == "BF16" + "FP64" -> from == "FP16" || from == "BF16" || from == "FP32" + else -> false + } + + private companion object { + val CAST_OPS: Set = setOf("convert", "cast", "to") + } +} diff --git a/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/CastFoldingPassTest.kt b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/CastFoldingPassTest.kt new file mode 100644 index 00000000..15b90216 --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/CastFoldingPassTest.kt @@ -0,0 +1,171 @@ +package sk.ainet.compile.opt + +import sk.ainet.compile.opt.passes.CastFoldingPass +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.tensor.ops.GenericOperation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.tensor.ops.ValidationResult +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertSame +import kotlin.test.assertTrue + +class CastFoldingPassTest { + + private fun spec(dtype: String) = TensorSpec(name = "t", shape = listOf(1, 4), dtype = dtype) + + private fun node(id: String, op: String, inDtype: String?, outDtype: String) = GraphNode( + id = id, + operation = GenericOperation(op), + inputs = listOfNotNull(inDtype?.let { spec(it) }), + outputs = listOf(spec(outDtype)), + ) + + /** A chain `x -> cast -> cast -> … -> relu`, one cast per consecutive dtype pair. */ + private fun chain(vararg dtypes: String, castOp: String = "convert"): ComputeGraph { + val graph = DefaultComputeGraph() + var previous = graph.addNode(node("x", "input", null, dtypes.first())) + for (i in 1 until dtypes.size) { + val cast = graph.addNode(node("c$i", castOp, dtypes[i - 1], dtypes[i])) + graph.addEdge(GraphEdge("e$i", previous, cast, tensorSpec = spec(dtypes[i - 1]))) + previous = cast + } + val sink = graph.addNode(node("out", "relu", dtypes.last(), dtypes.last())) + graph.addEdge(GraphEdge("eOut", previous, sink, tensorSpec = spec(dtypes.last()))) + return graph + } + + /** The casts of [graph] in execution order, as `from->to`, read off the edges. */ + private fun casts(graph: ComputeGraph): List = + graph.getTopologicalOrder() + .filter { it.operationName in setOf("convert", "cast", "to") } + .map { cast -> + val edge = graph.edges.single { it.destination.id == cast.id } + assertEquals(edge.tensorSpec.dtype, cast.inputs.single().dtype, "input spec of ${cast.id} must match its edge") + "${edge.source.outputs[edge.sourceOutputIndex].dtype}->${cast.outputs.single().dtype}" + } + + private fun fold(graph: ComputeGraph): GraphOptimizationResult = CastFoldingPass().apply(graph) + + @Test + fun wideningThenSameNarrowDtype_removesBothCasts() { + val result = fold(chain("FP16", "FP32", "FP16")) + + assertTrue(result.changed) + assertTrue(casts(result.graph).isEmpty()) + val sinkEdge = result.graph.edges.single() + assertEquals("x", sinkEdge.source.id) + assertEquals("out", sinkEdge.destination.id) + assertEquals("FP16", sinkEdge.tensorSpec.dtype) + assertEquals(ValidationResult.Valid, result.graph.validate()) + } + + @Test + fun wideningThenOtherNarrowDtype_becomesOneCast() { + val result = fold(chain("BF16", "FP32", "FP16")) + + assertTrue(result.changed) + assertEquals(listOf("BF16->FP16"), casts(result.graph)) + assertEquals(setOf("x", "c2", "out"), result.graph.nodes.map { it.id }.toSet()) + } + + @Test + fun adjacentBoundaries_keepTheirNarrowingAndLoseTheRoundTrip() { + // FP32 network, BF16 layer directly followed by an FP16 stage. + val result = fold(chain("FP32", "BF16", "FP32", "FP16", "FP32")) + + assertEquals(listOf("FP32->BF16", "BF16->FP16", "FP16->FP32"), casts(result.graph)) + } + + @Test + fun narrowingRoundTrip_isNeverFolded() { + // f32 -> bf16 -> f32 is the rounding a mixed-precision layer asked for. + val graph = chain("FP32", "BF16", "FP32") + + val result = fold(graph) + + assertFalse(result.changed) + assertSame(graph, result.graph) + } + + @Test + fun narrowToNarrow_isNotAnExactWidening() { + val graph = chain("FP16", "BF16", "FP16") + assertFalse(fold(graph).changed) + } + + @Test + fun integerCasts_areNotTreatedAsExact() { + // Int32 -> FP32 loses precision above 2^24. + val graph = chain("Int32", "FP32", "Int32") + assertFalse(fold(graph).changed) + } + + @Test + fun castToTheSameDtype_isRemoved() { + val result = fold(chain("BF16", "BF16")) + + assertTrue(result.changed) + assertTrue(casts(result.graph).isEmpty()) + assertEquals("x", result.graph.edges.single().source.id) + } + + @Test + fun chainOfWidenings_collapses() { + val result = fold(chain("FP16", "FP32", "FP64", "FP16")) + assertTrue(casts(result.graph).isEmpty()) + } + + @Test + fun wideningWithAnotherConsumer_staysAndTheSecondCastReadsTheSource() { + // x(BF16) -> widen(FP32) -> narrow(FP16) -> out + // \-> other (reads the FP32 tensor) + val graph = DefaultComputeGraph() + val x = graph.addNode(node("x", "input", null, "BF16")) + val widen = graph.addNode(node("widen", "convert", "BF16", "FP32")) + val narrow = graph.addNode(node("narrow", "convert", "FP32", "FP16")) + val out = graph.addNode(node("out", "relu", "FP16", "FP16")) + val other = graph.addNode(node("other", "relu", "FP32", "FP32")) + graph.addEdge(GraphEdge("e1", x, widen, tensorSpec = spec("BF16"))) + graph.addEdge(GraphEdge("e2", widen, narrow, tensorSpec = spec("FP32"))) + graph.addEdge(GraphEdge("e3", narrow, out, tensorSpec = spec("FP16"))) + graph.addEdge(GraphEdge("e4", widen, other, tensorSpec = spec("FP32"))) + + val result = fold(graph) + + assertTrue(result.changed) + assertEquals(setOf("BF16->FP32", "BF16->FP16"), casts(result.graph).toSet()) + assertEquals("widen", result.graph.edges.single { it.destination.id == "other" }.source.id) + assertEquals("x", result.graph.edges.single { it.destination.id == "narrow" }.source.id) + } + + @Test + fun castAliasesAndDtypeSpellings_areRecognised() { + val result = fold(chain("bf16", "float32", "f16", castOp = "cast")) + assertEquals(1, result.graph.nodes.count { it.operationName == "cast" }) + } + + @Test + fun graphWithoutCasts_isReturnedUnchanged() { + val graph = DefaultComputeGraph() + val x = graph.addNode(node("x", "input", null, "FP32")) + val out = graph.addNode(node("out", "relu", "FP32", "FP32")) + graph.addEdge(GraphEdge("e", x, out, tensorSpec = spec("FP32"))) + + val result = fold(graph) + + assertFalse(result.changed) + assertSame(graph, result.graph) + } + + @Test + fun diagnosticsNameEveryRewrite() { + val result = fold(chain("FP32", "BF16", "FP32", "FP16", "FP32")) + assertEquals(1, result.diagnostics.size, result.diagnostics.toString()) + assertTrue("BF16 -> FP16" in result.diagnostics.single(), result.diagnostics.single()) + } +}