rlx_ir/ops/io.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//! Graph I/O builders: inputs, parameters (plan #53).
6
7use crate::{Graph, NodeId, Op, Shape};
8
9impl Graph {
10 /// Graph input (runtime-provided tensor).
11 pub fn input(&mut self, name: impl Into<String>, shape: Shape) -> NodeId {
12 let name: String = name.into();
13 self.push(Op::Input { name: name.clone() }, vec![], shape, Some(name))
14 }
15
16 /// Model parameter (weight loaded at init).
17 pub fn param(&mut self, name: impl Into<String>, shape: Shape) -> NodeId {
18 let name: String = name.into();
19 self.push(Op::Param { name: name.clone() }, vec![], shape, Some(name))
20 }
21
22 /// Generic node constructor for custom ops.
23 pub fn add_node(&mut self, op: Op, inputs: Vec<NodeId>, shape: Shape) -> NodeId {
24 self.push(op, inputs, shape, None)
25 }
26
27 /// Materialize a row-major f32 `Op::Constant` tensor of the given shape.
28 /// Shared helper for the signal-processing graph builders (`vq`,
29 /// `spectral`, `upsample`, `dsp`) that bake windows / interpolation
30 /// matrices / filter taps into the graph.
31 pub(crate) fn const_f32_tensor(&mut self, data: Vec<f32>, dims: &[usize]) -> NodeId {
32 debug_assert_eq!(
33 data.len(),
34 dims.iter().product::<usize>(),
35 "const_f32_tensor: data length does not match dims"
36 );
37 let mut bytes = Vec::with_capacity(data.len() * 4);
38 for v in &data {
39 bytes.extend_from_slice(&v.to_le_bytes());
40 }
41 self.add_node(
42 Op::Constant { data: bytes },
43 vec![],
44 Shape::new(dims, crate::DType::F32),
45 )
46 }
47
48 /// Build an `Op::Custom` node, dispatching shape inference through
49 /// the global op registry. The named op must already be registered
50 /// via [`crate::register_op`]; `attrs` is forwarded verbatim to
51 /// the impl's `infer_shape` (and later, at execution time, to its
52 /// per-backend kernel).
53 ///
54 /// Panics if `name` is not registered or if `inputs.len()` does
55 /// not match the registered `num_inputs()` — both are programmer
56 /// errors that should fail loudly at graph-build time, not silently
57 /// at execution.
58 pub fn custom_op(
59 &mut self,
60 name: impl Into<String>,
61 attrs: Vec<u8>,
62 inputs: Vec<NodeId>,
63 ) -> NodeId {
64 let name: String = name.into();
65 let ext = crate::lookup_op(&name)
66 .unwrap_or_else(|| panic!("custom_op: '{name}' is not registered in the op registry"));
67 assert_eq!(
68 ext.num_inputs(),
69 inputs.len(),
70 "custom_op '{name}': registered op expects {} inputs, got {}",
71 ext.num_inputs(),
72 inputs.len(),
73 );
74 let in_shapes: Vec<&Shape> = inputs.iter().map(|id| self.shape(*id)).collect();
75 let out_shape = ext.infer_shape(&in_shapes, &attrs);
76 let num_inputs = ext.num_inputs() as u32;
77 self.push(
78 Op::Custom {
79 name,
80 num_inputs,
81 attrs,
82 },
83 inputs,
84 out_shape,
85 None,
86 )
87 }
88
89 /// Build an `Op::Custom` node with a caller-supplied output shape,
90 /// **bypassing** the registry's `infer_shape`. Use this for ops
91 /// whose output shape can't be determined by static input shapes
92 /// alone — most importantly, ops with multiple logical outputs
93 /// packed into one buffer.
94 ///
95 /// The canonical multi-output pattern:
96 ///
97 /// ```ignore
98 /// // Sparse-LU returns L_values + U_values packed end-to-end.
99 /// // Caller knows nnz_L and nnz_U from the symbolic factor.
100 /// let lu = g.custom_op_packed(
101 /// "sparse_lu",
102 /// attrs,
103 /// vec![A, b],
104 /// Shape::new(&[nnz_L + nnz_U], DType::F64),
105 /// );
106 /// let l_vals = g.narrow_(lu, 0, 0, nnz_L);
107 /// let u_vals = g.narrow_(lu, 0, nnz_L, nnz_U);
108 /// ```
109 ///
110 /// The op must still be registered (so `num_inputs` validation
111 /// and autodiff routing still work); only the shape is overridden.
112 pub fn custom_op_packed(
113 &mut self,
114 name: impl Into<String>,
115 attrs: Vec<u8>,
116 inputs: Vec<NodeId>,
117 out_shape: Shape,
118 ) -> NodeId {
119 let name: String = name.into();
120 let ext = crate::lookup_op(&name).unwrap_or_else(|| {
121 panic!("custom_op_packed: '{name}' is not registered in the op registry")
122 });
123 assert_eq!(
124 ext.num_inputs(),
125 inputs.len(),
126 "custom_op_packed '{name}': registered op expects {} inputs, got {}",
127 ext.num_inputs(),
128 inputs.len(),
129 );
130 let num_inputs = ext.num_inputs() as u32;
131 self.push(
132 Op::Custom {
133 name,
134 num_inputs,
135 attrs,
136 },
137 inputs,
138 out_shape,
139 None,
140 )
141 }
142
143 /// Non-panicking [`custom_op`](Self::custom_op): returns a
144 /// [`CustomOpError`](crate::CustomOpError) instead of panicking when the op
145 /// isn't registered or the input arity mismatches. Prefer this in library
146 /// code that builds graphs from user/config input.
147 pub fn try_custom_op(
148 &mut self,
149 name: impl Into<String>,
150 attrs: Vec<u8>,
151 inputs: Vec<NodeId>,
152 ) -> Result<NodeId, crate::CustomOpError> {
153 let name: String = name.into();
154 let ext = crate::lookup_op(&name)
155 .ok_or_else(|| crate::CustomOpError::NotRegistered(name.clone()))?;
156 if ext.num_inputs() != inputs.len() {
157 return Err(crate::CustomOpError::ArityMismatch {
158 name,
159 expected: ext.num_inputs(),
160 got: inputs.len(),
161 });
162 }
163 let in_shapes: Vec<&Shape> = inputs.iter().map(|id| self.shape(*id)).collect();
164 let out_shape = ext.infer_shape(&in_shapes, &attrs);
165 let num_inputs = ext.num_inputs() as u32;
166 Ok(self.push(
167 Op::Custom {
168 name,
169 num_inputs,
170 attrs,
171 },
172 inputs,
173 out_shape,
174 None,
175 ))
176 }
177
178 /// 1D FFT along the last axis.
179 ///
180 /// * **F32 / F64** — 2N real-block layout: last axis is `[re…, im…]`.
181 /// * **C64** — interleaved `[re, im]` pairs per complex element.
182 ///
183 /// Output shape matches input. Radix-2 when `N` is a power of two,
184 /// Bluestein otherwise. Default normalization is unnormalized
185 /// (`FftNorm::Backward`; `ifft(fft(x)) = N·x`).
186 pub fn fft(&mut self, x: NodeId, inverse: bool) -> NodeId {
187 self.fft_norm(x, inverse, crate::fft::FftNorm::Backward)
188 }
189
190 /// 1D FFT with explicit normalization mode.
191 pub fn fft_norm(&mut self, x: NodeId, inverse: bool, norm: crate::fft::FftNorm) -> NodeId {
192 let s = self.shape(x).clone();
193 crate::fft::fft_meta(&s);
194 self.push(Op::Fft { inverse, norm }, vec![x], s, None)
195 }
196
197 /// Fixed-point 1D FFT along the last axis (`Op::FftQ`).
198 ///
199 /// `I32` data in the 2N real-block layout, radix-2 only. `scale` decides
200 /// how the datapath stays in range and is a precision decision — see
201 /// [`crate::fft::FftQScale`], whose `bits_lost` and `headroom_needed`
202 /// spell out the trade for a given length.
203 pub fn fft_q(
204 &mut self,
205 x: NodeId,
206 inverse: bool,
207 norm: crate::fft::FftNorm,
208 scale: crate::fft::FftQScale,
209 ) -> NodeId {
210 let s = self.shape(x).clone();
211 let meta = crate::fft::fft_meta(&s);
212 assert!(
213 meta.n_complex.is_power_of_two(),
214 "Op::FftQ length {} is not a power of two; only radix-2 is supported",
215 meta.n_complex
216 );
217 self.push(
218 Op::FftQ {
219 inverse,
220 norm,
221 scale,
222 },
223 vec![x],
224 s,
225 None,
226 )
227 }
228
229 /// Ternary pruned radix-2 butterfly stage — see [`Op::FftButterflyStage`].
230 pub fn fft_butterfly_stage(
231 &mut self,
232 state: NodeId,
233 gate: NodeId,
234 rev: NodeId,
235 tw_re: NodeId,
236 tw_im: NodeId,
237 stage: u32,
238 n_fft: u32,
239 ) -> NodeId {
240 let s = self.shape(state).clone();
241 self.push(
242 Op::FftButterflyStage { stage, n_fft },
243 vec![state, gate, rev, tw_re, tw_im],
244 s,
245 None,
246 )
247 }
248
249 /// 1D FFT along an arbitrary axis. Lowers to
250 /// `Transpose(axis ↔ last) → Fft(last) → Transpose(last ↔ axis)`.
251 ///
252 /// AD is free: both `Op::Transpose` and `Op::Fft` have VJP/JVP rules.
253 pub fn fft_axis(&mut self, x: NodeId, axis: usize, inverse: bool) -> NodeId {
254 use crate::infer::GraphExt as _;
255 let rank = self.shape(x).rank();
256 assert!(
257 axis < rank,
258 "fft_axis: axis {axis} out of range for rank-{rank} tensor"
259 );
260 let last = rank - 1;
261 if axis == last {
262 return self.fft(x, inverse);
263 }
264 let mut perm: Vec<usize> = (0..rank).collect();
265 perm.swap(axis, last);
266
267 let x_t = self.transpose_(x, perm.clone());
268 let y_t = self.fft(x_t, inverse);
269 self.transpose_(y_t, perm)
270 }
271
272 /// N-dimensional FFT along `axes` (NumPy `fftn` semantics).
273 ///
274 /// Applies a 1D FFT along each listed axis in ascending order.
275 /// Empty `axes` is a no-op. For multi-axis transforms on tensors
276 /// with more than one spatial dimension, use `DType::C64`; the
277 /// F32/F64 2N-block layout only describes a single complex axis.
278 pub fn fftn(&mut self, x: NodeId, axes: &[usize], inverse: bool) -> NodeId {
279 let rank = self.shape(x).rank();
280 let axes = crate::fft::normalize_fftn_axes(rank, axes);
281 if axes.is_empty() {
282 return x;
283 }
284 if axes.len() > 1 && !self.shape(x).dtype().is_complex() {
285 panic!(
286 "fftn: multi-axis FFT on {:?} requires DType::C64; \
287 the F32/F64 2N real-block layout supports only one complex axis — \
288 call fft_axis for a single transform",
289 self.shape(x).dtype()
290 );
291 }
292 let mut y = x;
293 for axis in axes {
294 y = self.fft_axis(y, axis, inverse);
295 }
296 y
297 }
298
299 /// Inverse N-dimensional FFT — alias for `fftn(..., inverse: true)`.
300 pub fn ifftn(&mut self, x: NodeId, axes: &[usize]) -> NodeId {
301 self.fftn(x, axes, true)
302 }
303}