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