rlx_ir/ops/
upsample.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//! 1-D upsampling helpers for time-series U-Net decoders and diffusion
17//! samplers (U-Sleep, EEGDM, …).
18//!
19//! `conv_transpose1d` wraps the existing 2-D transposed convolution with a
20//! singleton height axis. `interpolate1d` resamples the last axis by an
21//! arbitrary factor using a host-built `[L_in, L_out]` interpolation matrix
22//! applied as a matmul — so both lower entirely to existing kernels and run
23//! on every backend.
24
25use crate::infer::GraphExt as _;
26use crate::{Graph, NodeId};
27
28/// Interpolation mode for [`Graph::interpolate1d`].
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub enum InterpMode {
31    /// Nearest-neighbor (round to closest source sample).
32    Nearest,
33    /// Linear interpolation between the two neighboring source samples.
34    Linear,
35}
36
37impl Graph {
38    /// Transposed 1-D convolution (a.k.a. deconvolution / fractionally-strided
39    /// conv), the learned upsampler in Conv1d U-Net decoders.
40    ///
41    /// * `input`  — `[N, C_in, L]`.
42    /// * `weight` — `[C_in, C_out/groups, K]` (PyTorch `ConvTranspose1d` layout).
43    ///
44    /// Returns `[N, C_out, L_out]` with
45    /// `L_out = (L−1)·stride − 2·padding + dilation·(K−1) + output_padding + 1`.
46    #[allow(clippy::too_many_arguments)]
47    pub fn conv_transpose1d(
48        &mut self,
49        input: NodeId,
50        weight: NodeId,
51        kernel: usize,
52        stride: usize,
53        padding: usize,
54        dilation: usize,
55        output_padding: usize,
56        groups: usize,
57    ) -> NodeId {
58        let in_s = self.shape(input).clone();
59        let w_s = self.shape(weight).clone();
60        assert_eq!(
61            in_s.rank(),
62            3,
63            "conv_transpose1d: input must be [N, C_in, L]"
64        );
65        assert_eq!(
66            w_s.rank(),
67            3,
68            "conv_transpose1d: weight must be [C_in, C_out/groups, K]"
69        );
70        let n = in_s.dim(0).unwrap_static();
71        let c_in = in_s.dim(1).unwrap_static();
72        let l = in_s.dim(2).unwrap_static();
73        let c_out_g = w_s.dim(1).unwrap_static();
74
75        // Lift to NCHW with a singleton height: [N, C_in, 1, L], weight
76        // [C_in, C_out/g, 1, K], then run the 2-D transposed conv.
77        let input4 = self.reshape_(input, vec![n as i64, c_in as i64, 1, l as i64]);
78        let weight4 = self.reshape_(weight, vec![c_in as i64, c_out_g as i64, 1, kernel as i64]);
79        let out4 = self.conv_transpose2d(
80            input4,
81            weight4,
82            [1, kernel],
83            [1, stride],
84            [0, padding],
85            [1, dilation],
86            [0, output_padding],
87            groups,
88        );
89        let o = self.shape(out4).clone();
90        let c_out = o.dim(1).unwrap_static();
91        let l_out = o.dim(3).unwrap_static();
92        self.reshape_(out4, vec![n as i64, c_out as i64, l_out as i64])
93    }
94
95    /// Resample the last axis of `x` from `L_in` to `l_out` samples.
96    ///
97    /// Works for any rank ≥ 1; leading axes are treated as batch. Uses
98    /// align-corners endpoint mapping (`pos = j·(L_in−1)/(l_out−1)`), matching
99    /// `torch.nn.functional.interpolate(..., align_corners=True)`.
100    pub fn interpolate1d(&mut self, x: NodeId, l_out: usize, mode: InterpMode) -> NodeId {
101        assert!(l_out > 0, "interpolate1d: l_out must be positive");
102        let xs = self.shape(x).clone();
103        let rank = xs.rank();
104        let l_in = xs.dim(rank - 1).unwrap_static();
105        if l_in == l_out {
106            return x;
107        }
108        let batch: usize = (0..rank - 1).map(|i| xs.dim(i).unwrap_static()).product();
109
110        // Host-build the [L_in, L_out] interpolation matrix W, so that
111        // out[.., j] = Σ_i x[.., i]·W[i, j].
112        let mut w = vec![0f32; l_in * l_out];
113        for j in 0..l_out {
114            let pos = if l_out == 1 {
115                0.0
116            } else {
117                j as f32 * (l_in as f32 - 1.0) / (l_out as f32 - 1.0)
118            };
119            match mode {
120                InterpMode::Nearest => {
121                    let i = (pos + 0.5) as usize;
122                    let i = i.min(l_in - 1);
123                    w[i * l_out + j] = 1.0;
124                }
125                InterpMode::Linear => {
126                    let i0 = pos.floor() as usize;
127                    let i0 = i0.min(l_in - 1);
128                    let i1 = (i0 + 1).min(l_in - 1);
129                    let frac = pos - i0 as f32;
130                    w[i0 * l_out + j] += 1.0 - frac;
131                    w[i1 * l_out + j] += frac;
132                }
133            }
134        }
135        let w_node = self.const_f32_tensor(w, &[l_in, l_out]);
136
137        // Flatten batch → [batch, L_in], matmul, restore shape.
138        let x2 = self.reshape_(x, vec![batch as i64, l_in as i64]);
139        let y2 = self.mm(x2, w_node); // [batch, l_out]
140        let mut out_dims: Vec<i64> = (0..rank - 1)
141            .map(|i| xs.dim(i).unwrap_static() as i64)
142            .collect();
143        out_dims.push(l_out as i64);
144        self.reshape_(y2, out_dims)
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use crate::{DType, Shape};
152
153    fn dims(g: &Graph, id: NodeId) -> Vec<usize> {
154        g.shape(id)
155            .dims()
156            .iter()
157            .map(|d| d.unwrap_static())
158            .collect()
159    }
160    fn input(g: &mut Graph, shape: &[usize]) -> NodeId {
161        g.input("x", Shape::new(shape, DType::F32))
162    }
163
164    #[test]
165    fn conv_transpose1d_upsamples_by_stride() {
166        let mut g = Graph::new("ct1d");
167        let x = input(&mut g, &[2, 3, 8]); // N,C_in,L
168        let w = g.param("w", Shape::new(&[3, 5, 4], DType::F32)); // C_in,C_out,K
169        let y = g.conv_transpose1d(x, w, 4, 2, 1, 1, 0, 1);
170        // L_out = (8-1)*2 - 2*1 + 1*(4-1) + 0 + 1 = 14 - 2 + 3 + 1 = 16
171        assert_eq!(dims(&g, y), vec![2, 5, 16]);
172    }
173
174    #[test]
175    fn interpolate1d_linear_shape() {
176        let mut g = Graph::new("interp");
177        let x = input(&mut g, &[2, 4, 10]);
178        let y = g.interpolate1d(x, 25, InterpMode::Linear);
179        assert_eq!(dims(&g, y), vec![2, 4, 25]);
180    }
181
182    #[test]
183    fn interpolate1d_nearest_identity() {
184        let mut g = Graph::new("interp_id");
185        let x = input(&mut g, &[3, 7]);
186        let y = g.interpolate1d(x, 7, InterpMode::Nearest);
187        assert_eq!(y, x); // no-op when l_in == l_out
188    }
189}
rlx_ir/ops/upsample.rs

rlx_ir/ops/
upsample.rs
