rlx_compile/
precision.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Precision policy + AutoMixedPrecision rewrite pass.
6//!
7//! The `PrecisionPolicy` is a high-level declarative spec that maps
8//! op kinds to numeric precisions. The `AutoMixedPrecision` pass
9//! consumes a policy and rewrites the graph: updates each node's
10//! shape dtype + inserts Cast nodes at precision boundaries.
11//!
12//! After this pass runs, the IR carries per-node precision info via
13//! `node.shape.dtype`, and the backend just reads it to pick the
14//! right kernel variant. Backends don't need any session-level
15//! precision flag.
16
17use rlx_fusion::pass::Pass;
18use rlx_ir::*;
19use std::collections::HashMap;
20
21/// Which numeric precision to use for an op.
22/// (Subset of DType — only the ones we currently dispatch on.)
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
24pub enum Precision {
25    F32,
26    F16,
27    BF16,
28}
29
30impl Precision {
31    pub fn dtype(self) -> DType {
32        match self {
33            Precision::F32 => DType::F32,
34            Precision::F16 => DType::F16,
35            Precision::BF16 => DType::BF16,
36        }
37    }
38}
39
40/// Cast configuration carried by ops that emit a typed output.
41///
42/// Inspired by TileKernels' `CastInputConfig` / `CastOutputConfig`: a single
43/// dataclass that flows from the layer down to the kernel selector, so adding
44/// new quantized formats (FP8 e4m3, FP4 e2m1, blocked scaling) becomes a
45/// matter of populating fields rather than threading new flags through call
46/// sites.
47///
48/// Today only `out_dtype` is consulted by backends — the scaling-factor
49/// fields are reserved for future quantization passes (FP8 / blocked SF).
50/// Constructed once by the precision pass and embedded in fused ops.
51#[derive(Debug, Clone, Copy, PartialEq)]
52pub struct CastConfig {
53    /// Destination dtype for the cast (fragment of the output tensor).
54    pub out_dtype: DType,
55    /// Scaling factor block size `(rows, cols)` for blocked quantization.
56    /// `None` means no scaling factor (plain cast).
57    pub sf_block: Option<(usize, usize)>,
58    /// Round scaling factors to powers of two (UE8M0 style).
59    pub round_sf: bool,
60}
61
62impl CastConfig {
63    /// Plain dtype cast with no scaling factor.
64    pub const fn plain(out_dtype: DType) -> Self {
65        Self {
66            out_dtype,
67            sf_block: None,
68            round_sf: false,
69        }
70    }
71    /// True when the cast does no work (out matches input dtype).
72    pub fn is_noop(&self, in_dtype: DType) -> bool {
73        self.out_dtype == in_dtype && self.sf_block.is_none()
74    }
75}
76
77/// High-level op categorization for precision policies.
78#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
79pub enum OpKind {
80    /// Matmul, FusedMatMulBiasAct, conv — compute-heavy ops that
81    /// benefit most from low precision.
82    Compute,
83    /// LayerNorm, RmsNorm, Softmax — reductions that need accuracy.
84    Reduction,
85    /// Add, Mul, GELU, SiLU — element-wise ops.
86    Elementwise,
87    /// Gather, Narrow, Reshape — data movement, no math.
88    DataMovement,
89    /// Inputs, parameters, outputs — user-facing.
90    Boundary,
91}
92
93fn op_kind(op: &Op) -> OpKind {
94    match op {
95        Op::MatMul
96        | Op::FusedMatMulBiasAct { .. }
97        | Op::FusedConvBiasAct { .. }
98        | Op::Conv { .. }
99        | Op::Im2Col { .. }
100        | Op::DotGeneral { .. }
101        | Op::DenseSolve
102        | Op::BatchedDenseSolve
103        | Op::Attention { .. }
104        | Op::FusedTransformerLayer { .. }
105        | Op::GroupedMatMul
106        | Op::DequantGroupedMatMul { .. }
107        | Op::DequantMoEWeights { .. }
108        | Op::LoraMatMul { .. }
109        | Op::DequantMatMul { .. }
110        | Op::SynthMatMul { .. }
111        | Op::ScaledMatMul { .. }
112        | Op::QMatMul { .. }
113        | Op::QConv2d { .. }
114        | Op::Conv2dBackwardInput { .. }
115        | Op::Conv2dBackwardWeight { .. }
116        | Op::AttentionBackward { .. } => OpKind::Compute,
117        Op::LayerNorm { .. }
118        | Op::RmsNorm { .. }
119        | Op::Softmax { .. }
120        | Op::FusedResidualLN { .. }
121        | Op::FusedResidualRmsNorm { .. }
122        | Op::Reduce { .. }
123        | Op::Cumsum { .. }
124        | Op::Sample { .. }
125        | Op::SelectiveScan { .. }
126        | Op::GatedDeltaNet { .. }
127        | Op::Lstm { .. }
128        | Op::Gru { .. }
129        | Op::Rnn { .. }
130        | Op::Mamba2 { .. }
131        | Op::SoftmaxCrossEntropy
132        | Op::SoftmaxCrossEntropyWithLogits
133        | Op::SoftmaxCrossEntropyBackward
134        | Op::LayerNormBackwardInput { .. }
135        | Op::LayerNormBackwardGamma { .. }
136        | Op::GroupNorm { .. }
137        // DiT adaLN — same reduction-accumulate story as LayerNorm.
138        | Op::AdaLayerNorm { .. }
139        | Op::AdaLayerNormBackward { .. } => OpKind::Reduction,
140        Op::Activation(_)
141        | Op::Binary(_)
142        | Op::FusedSwiGLU { .. }
143        | Op::Compare(_)
144        | Op::Where
145        | Op::ElementwiseRegion { .. }
146        | Op::Quantize { .. }
147        | Op::ScaledQuantize { .. }
148        | Op::ScaledQuantScale { .. }
149        | Op::ScaledDequantize { .. }
150        | Op::Dequantize { .. }
151        | Op::FakeQuantize { .. }
152        | Op::FakeQuantizeBackward { .. }
153        | Op::FakeQuantizeLSQ { .. }
154        | Op::FakeQuantizeLSQBackwardX { .. }
155        | Op::FakeQuantizeLSQBackwardScale { .. }
156        | Op::ReluBackward
157        | Op::ActivationBackward { .. }
158        | Op::ComplexNormSq
159        | Op::ComplexNormSqBackward
160        | Op::Conjugate
161        // DiT gate: x + y * sigmoid(gate) — elementwise fuse.
162        | Op::GatedResidual
163        | Op::GatedResidualBackward => OpKind::Elementwise,
164        Op::Gather { .. }
165        | Op::Narrow { .. }
166        | Op::Reshape { .. }
167        | Op::Transpose { .. }
168        | Op::Concat { .. }
169        | Op::Expand { .. }
170        | Op::Cast { .. }
171        | Op::Rope { .. }
172        | Op::Pool { .. }
173        | Op::FusedAttentionBlock { .. }
174        | Op::TopK { .. }
175        | Op::ScatterAdd
176        | Op::ScatterNd { .. }
177        | Op::ScatterElements { .. }
178        | Op::GatherNd { .. }
179        | Op::GatherElements { .. }
180        | Op::MaxPool2dBackward { .. }
181        | Op::ResizeNearest2x
182        | Op::AxialRope2d { .. } => OpKind::DataMovement,
183        Op::Input { .. } | Op::Param { .. } | Op::Constant { .. } => OpKind::Boundary,
184        // Control flow: treated as data movement (the inner sub-graph
185        // gets its own precision policy applied separately).
186        Op::If { .. } | Op::While { .. } => OpKind::DataMovement,
187        // Custom user-registered ops are opaque to the precision pass
188        // — classify as Compute by default; the registered op's own
189        // implementation decides what dtype it operates at.
190        Op::Custom { .. } => OpKind::Compute,
191        Op::Scan { .. } => OpKind::Compute,
192        Op::ScanBackward { .. } => OpKind::Compute,
193        Op::ScanBackwardXs { .. } => OpKind::Compute,
194        Op::CustomFn { .. } => OpKind::Compute,
195        Op::Fft { .. } => OpKind::Compute,
196        Op::FftButterflyStage { .. } => OpKind::Compute,
197        Op::LogMel => OpKind::Compute,
198        Op::LogMelBackward => OpKind::Compute,
199        // SPD / eigendecomposition — heavy Compute (host LAPACK or native
200        // solvers); never AutoMixed F16.
201        Op::BiMap
202        | Op::ReEig { .. }
203        | Op::LogEig { .. }
204        | Op::SpdBatchNorm { .. }
205        | Op::SpdKarcherMean { .. }
206        | Op::SpdKarcherMeanWeighted { .. }
207        | Op::SpdLogMap
208        | Op::SpdExpMap
209        | Op::SpdParallelTransport
210        | Op::SpdMatrixFnBatch { .. }
211        | Op::ReEigBackward { .. }
212        | Op::LogEigBackward { .. }
213        | Op::SpdBatchNormBackwardX { .. }
214        | Op::SpdBatchNormBackwardG { .. }
215        | Op::SpdLogMapBackward
216        | Op::SpdExpMapBackward
217        | Op::SpdParallelTransportBackward
218        | Op::SpdMatrixFnBatchBackward { .. }
219        | Op::Eigh
220        | Op::EighBackward
221        | Op::EighBatch
222        | Op::EighBatchBackward => OpKind::Compute,
223        _ => OpKind::Compute,
224    }
225}
226
227/// Declarative precision policy for graph compilation.
228#[derive(Debug, Clone, Default)]
229pub enum PrecisionPolicy {
230    /// All ops at F32. Default; safe; baseline accuracy.
231    #[default]
232    AlwaysF32,
233    /// All ops at F16. Maximum speed; may lose accuracy on reductions.
234    AlwaysF16,
235    /// Mixed precision, conservative variant. Forces F32 at every reduction
236    /// boundary, matching PyTorch's pre-2024 autocast and HuggingFace's
237    /// historical default. Accuracy is the highest of the AMP variants;
238    /// performance suffers from a Cast node before and after every
239    /// LayerNorm / Softmax in the graph.
240    ///   Compute → F16
241    ///   Reduction → F32  (← the cast tax — see AutoMixed for the fix)
242    ///   Elementwise → F16
243    ///   DataMovement → F16
244    ///   Boundary (input/param/output) → F32
245    AutoMixedConservative,
246    /// Mixed precision (Phase G — current default), tuned against Bonsai-27B
247    /// greedy tokens on Metal. Compute (matmul) and DataMovement stay F32:
248    /// F16 there diverged (residual-stream / Rope/Gather/Narrow paths) even
249    /// with `Cast`-`to` preserved. Elementwise and Reduction run F16 (reduction
250    /// kernels accumulate in f32 internally), which is where the win is without
251    /// the divergence. Boundaries stay user-visible F32.
252    ///   Compute → F32
253    ///   Reduction → F16  (kernel accumulates in f32 internally)
254    ///   Elementwise → F16
255    ///   DataMovement → F32
256    ///   Boundary (input/param/output) → F32
257    AutoMixed,
258    /// Mixed precision targeting BF16 on TPU/XLA. Same shape as
259    /// `AutoMixed` (compute + reduction + elementwise + data-movement
260    /// in the chosen low precision; boundaries stay F32) but the low
261    /// precision is BF16 instead of F16. BF16 is the native compute
262    /// dtype on TPU and recent GPUs; matches what JAX picks when
263    /// `jax.config.update("jax_default_dtype_bits", "bfloat16")`.
264    ///   Compute → BF16
265    ///   Reduction → BF16  (XLA's TPU codegen accumulates in f32)
266    ///   Elementwise → BF16
267    ///   DataMovement → BF16
268    ///   Boundary → F32
269    AutoMixedBf16,
270    /// Explicit per-op-kind override.
271    Custom(HashMap<OpKind, Precision>),
272}
273
274impl PrecisionPolicy {
275    /// Resolve the target precision for an op kind.
276    pub fn precision_for(&self, kind: OpKind) -> Precision {
277        match self {
278            PrecisionPolicy::AlwaysF32 => Precision::F32,
279            PrecisionPolicy::AlwaysF16 => match kind {
280                OpKind::Boundary => Precision::F32, // user-facing stays f32
281                _ => Precision::F16,
282            },
283            PrecisionPolicy::AutoMixed => match kind {
284                // Bonsai-27B Metal: Elementwise + Reduction F16 match F32
285                // greedy tokens. DataMovement→F16 still diverges (Rope/Gather/
286                // Narrow/Concat paths) even with Cast-`to` preserved — keep F32.
287                OpKind::Compute => Precision::F32,
288                OpKind::Reduction => Precision::F16,
289                OpKind::Elementwise => Precision::F16,
290                OpKind::DataMovement => Precision::F32,
291                OpKind::Boundary => Precision::F32,
292            },
293            PrecisionPolicy::AutoMixedConservative => match kind {
294                OpKind::Compute => Precision::F32,
295                OpKind::Reduction => Precision::F32,
296                OpKind::Elementwise => Precision::F16,
297                OpKind::DataMovement => Precision::F32,
298                OpKind::Boundary => Precision::F32,
299            },
300            PrecisionPolicy::AutoMixedBf16 => match kind {
301                OpKind::Compute => Precision::BF16,
302                OpKind::Reduction => Precision::BF16,
303                OpKind::Elementwise => Precision::BF16,
304                OpKind::DataMovement => Precision::BF16,
305                OpKind::Boundary => Precision::F32,
306            },
307            PrecisionPolicy::Custom(map) => map.get(&kind).copied().unwrap_or(Precision::F32),
308        }
309    }
310}
311
312/// Pass that rewrites a graph according to a `PrecisionPolicy`.
313///
314/// For each node:
315/// 1. Look up the target precision based on op kind.
316/// 2. Update `node.shape.dtype` to that precision.
317/// 3. If any input has a different dtype, insert a Cast node before it.
318///
319/// After this pass, every node knows its compute precision via its
320/// shape dtype. Backends dispatch kernels per-node.
321pub struct AutoMixedPrecision {
322    pub policy: PrecisionPolicy,
323}
324
325impl AutoMixedPrecision {
326    pub fn new(policy: PrecisionPolicy) -> Self {
327        Self { policy }
328    }
329}
330
331impl Pass for AutoMixedPrecision {
332    fn name(&self) -> &str {
333        "auto_mixed_precision"
334    }
335
336    fn run(&self, graph: Graph) -> Graph {
337        // Skip the pass entirely for AlwaysF32 — it's a no-op.
338        if matches!(self.policy, PrecisionPolicy::AlwaysF32) {
339            return graph;
340        }
341
342        let mut new_graph = Graph::new(&graph.name);
343        // Maps old NodeId → new NodeId at its post-rewrite precision.
344        let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();
345        // Tracks the precision each rewritten node ended up at.
346        let mut node_precision: HashMap<NodeId, Precision> = HashMap::new();
347        // Cast cache: avoid re-inserting identical Cast nodes.
348        // Key: (source new id, target precision)
349        let mut cast_cache: HashMap<(NodeId, Precision), NodeId> = HashMap::new();
350
351        for node in graph.nodes() {
352            // Native fp8 GEMM subgraph (from `scaled_quant_insert`): leave it
353            // intact. Its operands are U8 codes and its output is already f32 —
354            // relabeling the op's dtype or casting U8 inputs to f16 would
355            // corrupt it. Emit unchanged; if an f32 operand of a quantize op
356            // was lowered to f16 upstream, cast it back to f32 first.
357            if matches!(
358                node.op,
359                Op::ScaledMatMul { .. }
360                    | Op::ScaledQuantize { .. }
361                    | Op::ScaledQuantScale { .. }
362                    | Op::ScaledDequantize { .. }
363            ) {
364                let new_inputs: Vec<NodeId> = node
365                    .inputs
366                    .iter()
367                    .map(|&in_id| {
368                        let src = id_map[&in_id];
369                        let sd = new_graph.node(src).shape.dtype();
370                        if matches!(sd, DType::F16 | DType::BF16) {
371                            let shape = new_graph.node(src).shape.clone().with_dtype(DType::F32);
372                            new_graph.add_node(Op::Cast { to: DType::F32 }, vec![src], shape)
373                        } else {
374                            src
375                        }
376                    })
377                    .collect();
378                let new_id = new_graph.add_node(node.op.clone(), new_inputs, node.shape.clone());
379                id_map.insert(node.id, new_id);
380                // The matmul's output is f32; consumers cast it down as needed.
381                node_precision.insert(node.id, Precision::F32);
382                continue;
383            }
384
385            // Ops whose Metal kernels are still f32-only. Casting the node to
386            // F16 while the kernel reads/writes floats zeros the residual
387            // stream (Bonsai-27B AMP → all token-0 / "!!!!!!!!"). Keep these at
388            // F32; Cast boundaries reconnect the f16 residual around them.
389            let force_f32 = match &node.op {
390                Op::GatedDeltaNet { .. }
391                | Op::SelectiveScan { .. }
392                | Op::Conv { .. }
393                | Op::Conv3d { .. }
394                | Op::ConvTranspose2d { .. }
395                | Op::ConvTranspose3d { .. }
396                | Op::Im2Col { .. }
397                | Op::Pool { .. }
398                | Op::QConv2d { .. }
399                | Op::Conv2dBackwardInput { .. }
400                | Op::Conv2dBackwardWeight { .. } => true,
401                // Reduction / data-movement ops without Metal `_h` yet.
402                Op::Cumsum { .. }
403                | Op::Sample { .. }
404                | Op::Lstm { .. }
405                | Op::Gru { .. }
406                | Op::Rnn { .. }
407                | Op::Mamba2 { .. }
408                | Op::SoftmaxCrossEntropy
409                | Op::SoftmaxCrossEntropyWithLogits
410                | Op::SoftmaxCrossEntropyBackward
411                | Op::LayerNormBackwardInput { .. }
412                | Op::LayerNormBackwardGamma { .. }
413                | Op::GroupNorm { .. }
414                | Op::FusedAttentionBlock { .. }
415                | Op::TopK { .. }
416                | Op::ScatterAdd
417                | Op::ScatterNd { .. }
418                | Op::ScatterElements { .. }
419                | Op::GatherNd { .. }
420                | Op::GatherElements { .. }
421                | Op::Expand { .. }
422                | Op::MaxPool2dBackward { .. }
423                | Op::ResizeNearest2x
424                | Op::AxialRope2d { .. } => true,
425                // SPD / eigh — host LAPACK or f32 native solvers.
426                Op::BiMap
427                | Op::ReEig { .. }
428                | Op::LogEig { .. }
429                | Op::SpdBatchNorm { .. }
430                | Op::SpdKarcherMean { .. }
431                | Op::SpdKarcherMeanWeighted { .. }
432                | Op::SpdLogMap
433                | Op::SpdExpMap
434                | Op::SpdParallelTransport
435                | Op::SpdMatrixFnBatch { .. }
436                | Op::ReEigBackward { .. }
437                | Op::LogEigBackward { .. }
438                | Op::SpdBatchNormBackwardX { .. }
439                | Op::SpdBatchNormBackwardG { .. }
440                | Op::SpdLogMapBackward
441                | Op::SpdExpMapBackward
442                | Op::SpdParallelTransportBackward
443                | Op::SpdMatrixFnBatchBackward { .. }
444                | Op::Eigh
445                | Op::EighBackward
446                | Op::EighBatch
447                | Op::EighBatchBackward => true,
448                // Keep *all* packed DequantMatMul at F32 until the full-graph
449                // AMP path (dual-Q1 / fused MLP + f16 residual) is proven.
450                Op::DequantMatMul { .. } => true,
451                // Same policy for the codebook-synthesis matmul.
452                Op::SynthMatMul { .. } => true,
453                _ => false,
454            };
455            if force_f32 {
456                let new_inputs: Vec<NodeId> = node
457                    .inputs
458                    .iter()
459                    .map(|&in_id| {
460                        let src = id_map[&in_id];
461                        let sd = new_graph.node(src).shape.dtype();
462                        if matches!(sd, DType::F16 | DType::BF16) {
463                            let shape = new_graph.node(src).shape.clone().with_dtype(DType::F32);
464                            new_graph.add_node(Op::Cast { to: DType::F32 }, vec![src], shape)
465                        } else {
466                            src
467                        }
468                    })
469                    .collect();
470                let new_id = new_graph.add_node(
471                    node.op.clone(),
472                    new_inputs,
473                    node.shape.clone().with_dtype(DType::F32),
474                );
475                id_map.insert(node.id, new_id);
476                node_precision.insert(node.id, Precision::F32);
477                continue;
478            }
479
480            let kind = op_kind(&node.op);
481            let mut target = self.policy.precision_for(kind);
482
483            // Inputs / params keep their original dtype (they're external);
484            // outputs stay user-visible at F32.
485            if kind == OpKind::Boundary {
486                target = Precision::F32;
487            }
488
489            // `Op::Cast { to }` must keep `to` — DataMovement→F16 was
490            // relabeling Cast-to-F32 nodes as F16 and zeroing AMP graphs.
491            if let Op::Cast { to } = &node.op {
492                target = match to {
493                    DType::F16 => Precision::F16,
494                    DType::BF16 => Precision::BF16,
495                    DType::F32 => Precision::F32,
496                    _ => {
497                        // Bool / int casts: emit unchanged.
498                        let new_id = new_graph.add_node(
499                            node.op.clone(),
500                            vec![id_map[&node.inputs[0]]],
501                            node.shape.clone(),
502                        );
503                        id_map.insert(node.id, new_id);
504                        node_precision.insert(node.id, Precision::F32);
505                        continue;
506                    }
507                };
508            }
509
510            // Resolve each input: insert a Cast if precision differs.
511            // Never cast integer / bool / packed-quant (U8/I8) tensors — those
512            // are bit patterns (GGUF weights, indices, masks), not AMP floats.
513            // Casting U8 Q1 packs to F16 silently zeros Metal decode.
514            let mut new_inputs = Vec::with_capacity(node.inputs.len());
515            for &in_id in &node.inputs {
516                let src_new_id = id_map[&in_id];
517                let src_dtype = new_graph.node(src_new_id).shape.dtype();
518                if !matches!(src_dtype, DType::F32 | DType::F16 | DType::BF16) {
519                    new_inputs.push(src_new_id);
520                    continue;
521                }
522                let src_prec = node_precision
523                    .get(&in_id)
524                    .copied()
525                    .unwrap_or(Precision::F32);
526                if src_prec == target {
527                    new_inputs.push(src_new_id);
528                } else {
529                    // Insert (or reuse cached) cast
530                    let cast_id = *cast_cache.entry((src_new_id, target)).or_insert_with(|| {
531                        let shape = new_graph
532                            .node(src_new_id)
533                            .shape
534                            .clone()
535                            .with_dtype(target.dtype());
536                        new_graph.add_node(Op::Cast { to: target.dtype() }, vec![src_new_id], shape)
537                    });
538                    new_inputs.push(cast_id);
539                }
540            }
541
542            // Build the rewritten node with the target dtype on its shape.
543            // Boundaries (Input/Param/Constant) keep their original dtype —
544            // packed U8/I8 GGUF weights and integer indices must not be
545            // relabeled as F32 (arena would oversize and AMP would try to
546            // cast them into the float stream).
547            // Cast keeps `to` (already reflected in `target` above).
548            let new_shape = if kind == OpKind::Boundary {
549                node.shape.clone()
550            } else if let Op::Cast { to } = &node.op {
551                node.shape.clone().with_dtype(*to)
552            } else {
553                node.shape.clone().with_dtype(target.dtype())
554            };
555            let new_id = new_graph.add_node(node.op.clone(), new_inputs, new_shape);
556            id_map.insert(node.id, new_id);
557            node_precision.insert(node.id, target);
558        }
559
560        // Outputs always stay at F32 — cast back if needed.
561        let new_outputs: Vec<NodeId> = graph
562            .outputs
563            .iter()
564            .map(|&out_id| {
565                let src_new_id = id_map[&out_id];
566                let src_prec = node_precision
567                    .get(&out_id)
568                    .copied()
569                    .unwrap_or(Precision::F32);
570                if src_prec == Precision::F32 {
571                    src_new_id
572                } else {
573                    let shape = new_graph
574                        .node(src_new_id)
575                        .shape
576                        .clone()
577                        .with_dtype(DType::F32);
578                    new_graph.add_node(Op::Cast { to: DType::F32 }, vec![src_new_id], shape)
579                }
580            })
581            .collect();
582        new_graph.set_outputs(new_outputs);
583
584        new_graph
585    }
586}
587
588#[cfg(test)]
589mod tests {
590    use super::*;
591
592    #[test]
593    fn always_f32_is_noop() {
594        let mut g = Graph::new("test");
595        let x = g.input("x", Shape::new(&[2, 4], DType::F32));
596        let w = g.param("w", Shape::new(&[4, 3], DType::F32));
597        let mm = g.matmul(x, w, Shape::new(&[2, 3], DType::F32));
598        g.set_outputs(vec![mm]);
599
600        let pass = AutoMixedPrecision::new(PrecisionPolicy::AlwaysF32);
601        let out = pass.run(g);
602        assert_eq!(out.len(), 3); // input, param, matmul — no casts
603    }
604
605    #[test]
606    fn auto_mixed_inserts_casts_at_boundary() {
607        let mut g = Graph::new("test");
608        let x = g.input("x", Shape::new(&[2, 4], DType::F32));
609        let w = g.param("w", Shape::new(&[4, 3], DType::F32));
610        let mm = g.matmul(x, w, Shape::new(&[2, 3], DType::F32));
611        // Under `AutoMixed`, MatMul (Compute) stays F32; the F16 region begins at
612        // an elementwise op, so the f32↔f16 boundary casts appear around the relu.
613        let act = g.relu(mm);
614        g.set_outputs(vec![act]);
615
616        let pass = AutoMixedPrecision::new(PrecisionPolicy::AutoMixed);
617        let out = pass.run(g);
618
619        // input(f32), param(f32), matmul(f32), cast(f32→f16) into relu,
620        // relu(f16), cast(f16→f32) at the output boundary = 6 nodes, and the
621        // final output node is a Cast back to F32.
622        assert!(out.len() >= 6);
623        let final_node = out.node(out.outputs[0]);
624        assert!(matches!(final_node.op, Op::Cast { to: DType::F32 }));
625    }
626
627    #[test]
628    fn auto_mixed_preserves_u8_quant_weights() {
629        use rlx_ir::quant::QuantScheme;
630        let mut g = Graph::new("q1_amp");
631        let x = g.input("x", Shape::new(&[1, 128], DType::F32));
632        let w = g.param("w", Shape::new(&[2304], DType::U8));
633        let y = g.add_node(
634            Op::DequantMatMul {
635                scheme: QuantScheme::GgufQ1_0,
636            },
637            vec![x, w],
638            Shape::new(&[1, 64], DType::F32),
639        );
640        g.set_outputs(vec![y]);
641
642        let out = AutoMixedPrecision::new(PrecisionPolicy::AutoMixed).run(g);
643        let params: Vec<_> = out
644            .nodes()
645            .iter()
646            .filter(|n| matches!(n.op, Op::Param { .. }))
647            .collect();
648        assert_eq!(params.len(), 1);
649        assert_eq!(params[0].shape.dtype(), DType::U8);
650        for n in out.nodes() {
651            if let Op::Cast { to } = &n.op {
652                let src = out.node(n.inputs[0]);
653                assert!(
654                    !matches!(src.op, Op::Param { .. }),
655                    "AMP must not cast U8 Param to {to:?}"
656                );
657            }
658        }
659    }
660}
rlx_compile/precision.rs

rlx_compile/
precision.rs
