rlx_ir/ops/
reduction.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//! Reduction builders: reduce, softmax, cumsum, sample
6//! (plan #53).
7
8use crate::op::ReduceOp;
9use crate::{Graph, NodeId, Op, Shape};
10
11impl Graph {
12    /// Reduce.
13    /// Reduce over `axes`.
14    ///
15    /// Non-adjacent axes are emitted as a chain of single-axis reductions,
16    /// highest index first, rather than one node.
17    ///
18    /// The backends only implement a reduction over a *contiguous* axis range
19    /// — the shape folds to `[outer, reduced, inner]` and the kernel walks
20    /// that. Handed `[1, 3]` the CPU backend emits `Thunk::Nop`, so the output
21    /// buffer keeps whatever it held and the reduction silently answers zeros.
22    /// That is not a hypothetical: a hand-written 3-D max pool built this way
23    /// looked correct because it had been written around the bug, and the
24    /// workaround cost 30.8 ms of a 34 ms CUDA forward pass.
25    ///
26    /// Chaining is exact for every `ReduceOp` this accepts. Sum/Max/Min/Prod
27    /// are associative and commutative; `Mean` works because each axis is
28    /// fully reduced, so the divisors multiply out to the same total count.
29    /// Going highest-first keeps the lower axis indices valid as the rank
30    /// drops.
31    pub fn reduce(
32        &mut self,
33        input: NodeId,
34        op: ReduceOp,
35        axes: Vec<usize>,
36        keep_dim: bool,
37        shape: Shape,
38    ) -> NodeId {
39        let mut sorted = axes.clone();
40        sorted.sort_unstable();
41        sorted.dedup();
42        let adjacent = sorted.windows(2).all(|w| w[1] == w[0] + 1);
43        if adjacent || sorted.len() < 2 {
44            return self.push(Op::Reduce { op, axes, keep_dim }, vec![input], shape, None);
45        }
46        let mut cur = input;
47        let mut dims: Vec<usize> = self
48            .shape(input)
49            .dims()
50            .iter()
51            .map(|d| d.unwrap_static())
52            .collect();
53        for &axis in sorted.iter().rev() {
54            if keep_dim {
55                dims[axis] = 1;
56            } else {
57                dims.remove(axis);
58            }
59            let step = Shape::new(&dims, shape.dtype());
60            cur = self.push(
61                Op::Reduce {
62                    op,
63                    axes: vec![axis],
64                    keep_dim,
65                },
66                vec![cur],
67                step,
68                None,
69            );
70        }
71        debug_assert_eq!(
72            dims,
73            shape
74                .dims()
75                .iter()
76                .map(|d| d.unwrap_static())
77                .collect::<Vec<_>>(),
78            "reduce: split over non-adjacent axes {sorted:?} did not land on the \
79             declared output shape"
80        );
81        cur
82    }
83
84    /// Softmax.
85    pub fn softmax(&mut self, input: NodeId, axis: i32, shape: Shape) -> NodeId {
86        self.push(Op::Softmax { axis }, vec![input], shape, None)
87    }
88
89    /// Histogram of `input` into `bins` equal-width buckets over `[min, max]`.
90    /// Output is a 1-D f32 tensor `[bins]` of counts. Elements outside the
91    /// range are dropped and `x == max` lands in the last bin (matches
92    /// `numpy.histogram`). Non-differentiable.
93    pub fn histogram(&mut self, input: NodeId, bins: usize, min: f32, max: f32) -> NodeId {
94        let out = Shape::new(&[bins], crate::DType::F32);
95        self.push(Op::Histogram { bins, min, max }, vec![input], out, None)
96    }
97
98    /// Cumulative sum along an axis (output shape == input shape).
99    pub fn cumsum(&mut self, input: NodeId, axis: i32, exclusive: bool, shape: Shape) -> NodeId {
100        self.push(Op::Cumsum { axis, exclusive }, vec![input], shape, None)
101    }
102
103    /// Cumulative product along an axis (output shape == input shape).
104    pub fn cumprod(&mut self, input: NodeId, axis: i32, exclusive: bool, shape: Shape) -> NodeId {
105        self.push(Op::CumProd { axis, exclusive }, vec![input], shape, None)
106    }
107
108    /// Cumulative maximum along an axis (output shape == input shape).
109    pub fn cummax(&mut self, input: NodeId, axis: i32, exclusive: bool, shape: Shape) -> NodeId {
110        self.push(Op::CumMax { axis, exclusive }, vec![input], shape, None)
111    }
112
113    /// Index of the max along `axis` (f32-encoded indices).
114    pub fn argmax(&mut self, input: NodeId, axis: usize, keep_dim: bool, shape: Shape) -> NodeId {
115        self.push(Op::ArgMax { axis, keep_dim }, vec![input], shape, None)
116    }
117
118    /// Index of the min along `axis` (f32-encoded indices).
119    pub fn argmin(&mut self, input: NodeId, axis: usize, keep_dim: bool, shape: Shape) -> NodeId {
120        self.push(Op::ArgMin { axis, keep_dim }, vec![input], shape, None)
121    }
122
123    /// Fused sample: logits → token id (one f32-encoded id per row).
124    /// `output_shape` should be `[batch]` (one id per logit row).
125    pub fn sample(
126        &mut self,
127        logits: NodeId,
128        top_k: usize,
129        top_p: f32,
130        temperature: f32,
131        seed: u64,
132        output_shape: Shape,
133    ) -> NodeId {
134        self.push(
135            Op::Sample {
136                top_k,
137                top_p,
138                temperature,
139                seed,
140            },
141            vec![logits],
142            output_shape,
143            None,
144        )
145    }
146}
rlx_ir/ops/reduction.rs

rlx_ir/ops/
reduction.rs
