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