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 if !matches!(node.shape.dtype(), DType::F32 | DType::F16 | DType::BF16) {
553 // An op that *produces* a non-float produces it under every
554 // precision policy — integer indices, masks, and fixed-point
555 // transforms (`Op::FftQ`) are bit patterns, not AMP floats.
556 // This is the same rule the input loop above applies; it was
557 // only being enforced for boundaries, so an integer-valued
558 // compute op was relabeled F32 and its backend then rejected
559 // it on dtype.
560 node.shape.clone()
561 } else {
562 node.shape.clone().with_dtype(target.dtype())
563 };
564 let new_id = new_graph.add_node(node.op.clone(), new_inputs, new_shape);
565 id_map.insert(node.id, new_id);
566 node_precision.insert(node.id, target);
567 }
568
569 // Outputs always stay at F32 — cast back if needed.
570 let new_outputs: Vec<NodeId> = graph
571 .outputs
572 .iter()
573 .map(|&out_id| {
574 let src_new_id = id_map[&out_id];
575 let src_prec = node_precision
576 .get(&out_id)
577 .copied()
578 .unwrap_or(Precision::F32);
579 if src_prec == Precision::F32 {
580 src_new_id
581 } else {
582 let shape = new_graph
583 .node(src_new_id)
584 .shape
585 .clone()
586 .with_dtype(DType::F32);
587 new_graph.add_node(Op::Cast { to: DType::F32 }, vec![src_new_id], shape)
588 }
589 })
590 .collect();
591 new_graph.set_outputs(new_outputs);
592
593 new_graph
594 }
595}
596
597#[cfg(test)]
598mod tests {
599 use super::*;
600
601 #[test]
602 fn always_f32_is_noop() {
603 let mut g = Graph::new("test");
604 let x = g.input("x", Shape::new(&[2, 4], DType::F32));
605 let w = g.param("w", Shape::new(&[4, 3], DType::F32));
606 let mm = g.matmul(x, w, Shape::new(&[2, 3], DType::F32));
607 g.set_outputs(vec![mm]);
608
609 let pass = AutoMixedPrecision::new(PrecisionPolicy::AlwaysF32);
610 let out = pass.run(g);
611 assert_eq!(out.len(), 3); // input, param, matmul — no casts
612 }
613
614 #[test]
615 fn auto_mixed_inserts_casts_at_boundary() {
616 let mut g = Graph::new("test");
617 let x = g.input("x", Shape::new(&[2, 4], DType::F32));
618 let w = g.param("w", Shape::new(&[4, 3], DType::F32));
619 let mm = g.matmul(x, w, Shape::new(&[2, 3], DType::F32));
620 // Under `AutoMixed`, MatMul (Compute) stays F32; the F16 region begins at
621 // an elementwise op, so the f32↔f16 boundary casts appear around the relu.
622 let act = g.relu(mm);
623 g.set_outputs(vec![act]);
624
625 let pass = AutoMixedPrecision::new(PrecisionPolicy::AutoMixed);
626 let out = pass.run(g);
627
628 // input(f32), param(f32), matmul(f32), cast(f32→f16) into relu,
629 // relu(f16), cast(f16→f32) at the output boundary = 6 nodes, and the
630 // final output node is a Cast back to F32.
631 assert!(out.len() >= 6);
632 let final_node = out.node(out.outputs[0]);
633 assert!(matches!(final_node.op, Op::Cast { to: DType::F32 }));
634 }
635
636 #[test]
637 fn auto_mixed_preserves_u8_quant_weights() {
638 use rlx_ir::quant::QuantScheme;
639 let mut g = Graph::new("q1_amp");
640 let x = g.input("x", Shape::new(&[1, 128], DType::F32));
641 let w = g.param("w", Shape::new(&[2304], DType::U8));
642 let y = g.add_node(
643 Op::DequantMatMul {
644 scheme: QuantScheme::GgufQ1_0,
645 },
646 vec![x, w],
647 Shape::new(&[1, 64], DType::F32),
648 );
649 g.set_outputs(vec![y]);
650
651 let out = AutoMixedPrecision::new(PrecisionPolicy::AutoMixed).run(g);
652 let params: Vec<_> = out
653 .nodes()
654 .iter()
655 .filter(|n| matches!(n.op, Op::Param { .. }))
656 .collect();
657 assert_eq!(params.len(), 1);
658 assert_eq!(params[0].shape.dtype(), DType::U8);
659 for n in out.nodes() {
660 if let Op::Cast { to } = &n.op {
661 let src = out.node(n.inputs[0]);
662 assert!(
663 !matches!(src.op, Op::Param { .. }),
664 "AMP must not cast U8 Param to {to:?}"
665 );
666 }
667 }
668 }
669}