Follow-up of SKEEP-008 (#1433, closed), Open Question 2. Proposal: docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc.

🧠 D: DOCUMENT — Problem & Opportunity

A PrecisionBoundary casts its input into the inner dtype and its output back to the network dtype. Two boundaries in a row therefore meet in the network dtype, even when nothing runs there:

sequential<FP32, Float> {
    input(784)
    dense(256, dtype = BF16)
    stage("head", dtype = FP16) { dense(64); dense(10) }
}

exports as (pinned by MixedPrecisionExportTest):

%v13 = stablehlo.convert %v12 : (tensor<1x256xbf16>) -> tensor<1x256xf32>   // out of the BF16 dense
%v14 = stablehlo.convert %v13 : (tensor<1x256xf32>) -> tensor<1x256xf16>    // into the FP16 stage

The f32 tensor between the two casts has no other consumer. It costs one op and one full-width intermediate per adjacent pair, and on a target that wants a narrow dtype end to end (the reason DtypeForwardPropagationPass(targetFloatDtype = …) exists) it is an f32 island for no reason. With two boundaries of the same inner dtype the pair is f16 -> f32 -> f16, a no-op.

Today only the degenerate case is elided, at DSL time: a boundary whose inner dtype equals the network dtype records no cast (#1492).

Summary:

Add a DAG pass that folds a convert directly followed by another convert when the first one is a widening, so adjacent mixed-precision boundaries do not round-trip through the network dtype.

🔍 A: ASSESS — Feasibility & Impact

✔️ Feasibility

Straightforward, and value-preserving by construction. For a chain A -> B -> C where A -> B is exact (every value of A is representable in B: FP16 → FP32, BF16 → FP32):

• C == A: both casts go away (f16 -> f32 -> f16).
• C != A: one cast A -> C remains (bf16 -> f32 -> f16 becomes bf16 -> f16). Rounding an exactly widened value gives the same result as rounding the original.

The opposite order, f32 -> bf16 -> f32, is the boundary itself: the narrowing is the semantics and must never be folded by this pass.

The same pass should drop a convert whose input and output dtype are equal, which DtypeForwardPropagationPass(targetFloatDtype = …) can produce when it coerces a whole graph to one dtype.

It is a DAG-phase GraphOptimizationPass next to TransposeEliminationPass (same shape: producer/consumer maps, single-consumer check, rebuild the graph). It runs after tracing, so the autograd tape and convertBackward are not affected.

Why a compile pass and not the DSL. The builder could merge neighbours while it appends modules, but that changes the module tree the user wrote (names, modules, forward hooks, parameter paths), and it only covers the DSL: a hand-written convert chain or an imported graph would not benefit.

✔️ Should the folding decision be a callback / pluggable strategy?

Evaluated, recommendation: no callback; one constructor parameter at most.

• The default rule is not a judgement call. Whether A -> B is exact follows from the dtype lattice, so there is nothing for a callback to decide that the pass cannot decide itself.
• A lambda such as (CastChain) -> Decision is opaque to the pipeline: it cannot be printed in diagnostics, compared in tests, or carried in a serialized pipeline description. The codebase's idiom for this kind of choice is a small declared type on the pass or the exporter: ConstantMaterializationPolicy, ConversionErrorPolicy, PackedConstantHandling, OpGranularityPolicy.
• Per-target variation already has an extension point: TargetOptimizers.registerDagPasses(target) { … }. A target that wants different behaviour registers the pass with different arguments, or its own pass.
• The one strategy that is a real alternative is lossy: dropping a narrowing round trip (f32 -> bf16 -> f32) for a target that has no bf16 and should run the layer in f32. That changes numerics, so by the DSL principle that a lossy step is a visible, named call it should be an explicit opt-in, e.g. CastFoldingMode.EXACT (default) and CastFoldingMode.DROP_NARROWING, not a hook. It is also arguably a different feature ("ignore mixed precision on this target") and can wait until a target asks for it.

So: ship the exact rule with no knob, and add the enum only when the lossy case has a user.

✔️ Expected Impact

• One op and one f32 intermediate less per adjacent boundary pair; none for same-dtype neighbours.
• Narrow-dtype regions stay contiguous in exported modules, which is what NPU targets that require bf16 / f16 end to end need.

✔️ Risks / Constraints

• The intermediate tensor must have exactly one consumer and must not be a graph output; otherwise only the second cast may be rewritten to read from the original source.
• "Exact widening" needs one authoritative definition. Float → wider float is safe; integer ↔ float pairs are not automatically (Int32 → FP32 is not exact) and should be excluded in the first version.
• convert nodes come from tracing with the dtypes on the tensor specs and a targetType attribute; cast / to are aliases in the StableHLO converter and should be matched too.

✔️ Dependencies

skainet-compile-opt only. No third-party dependency.

📚 R: RESEARCH — What Must Be Understood First?

Research Tasks

• Decide where "A -> B is exact" lives: a helper next to DType.isConvertibleTo / the promotion lattice in sk.ainet.lang.types, or private to the pass
• Check how MLIR / StableHLO canonicalization folds convert(convert(x)), and mirror its conditions
• Decide which default pipelines in GraphOptimizationPipeline include the pass, and its position relative to DtypeForwardPropagationPass (after it, so coerced same-dtype casts are removed too)

Open Questions

Should the pass also fold a widening followed by a further widening (f16 -> f32 -> f64)? It is exact as well, but no DSL path produces it today.

🛠️ C: CODE — Implementation Plan

Development Tasks

• CastFoldingPass in skainet-compile-opt/.../passes, with diagnostics naming each folded pair
• Unit tests on hand-built graphs: f16 -> f32 -> f16 removed; bf16 -> f32 -> f16 becomes one cast; f32 -> bf16 -> f32 untouched; multi-consumer intermediate untouched; same-dtype convert removed
• Extend MixedPrecisionExportTest: the BF16-dense + FP16-stage network exports three stablehlo.convert instead of four after the pass, with identical output values on the CPU backend before and after
• Add the pass to the default DAG pipelines
• Update explanation/mixed-precision.adoc ("What an export looks like") and mark Open Question 2 of SKEEP-008 resolved
• CHANGELOG entry

Acceptance Criteria

• Adjacent boundaries no longer export a cast pair through the network dtype
• Executing the graph before and after the pass gives bit-identical results
• A narrowing round trip is never removed
• No regressions in the existing export and optimization tests