rlx_ir/ops/
splat.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//! 3D Gaussian splatting graph builders.
6
7use crate::infer::GraphExt;
8use crate::op::Op;
9use crate::{DType, Graph, NodeId, Shape};
10
11/// Packed scene + camera tensors for [`Op::GaussianSplatRender`].
12#[derive(Clone, Debug)]
13pub struct GaussianSplatInputs {
14    pub positions: NodeId,
15    pub scales: NodeId,
16    pub rotations: NodeId,
17    pub opacities: NodeId,
18    pub colors: NodeId,
19    pub sh_coeffs: NodeId,
20    pub meta: NodeId,
21}
22
23/// Render-parameter bundle (embedded in the op, not separate tensors).
24#[derive(Clone, Copy, Debug)]
25pub struct GaussianSplatRenderParams {
26    pub width: u32,
27    pub height: u32,
28    pub tile_size: u32,
29    pub radius_scale: f32,
30    pub alpha_cutoff: f32,
31    pub max_splat_steps: u32,
32    pub transmittance_threshold: f32,
33    pub max_list_entries: u32,
34}
35
36/// Training backward parameters for [`Op::GaussianSplatRenderBackward`].
37#[derive(Clone, Copy, Debug)]
38pub struct GaussianSplatBackwardParams {
39    pub render: GaussianSplatRenderParams,
40    pub loss_grad_clip: f32,
41    pub sh_band: u32,
42    pub max_anisotropy: f32,
43}
44
45impl Default for GaussianSplatBackwardParams {
46    fn default() -> Self {
47        Self {
48            render: GaussianSplatRenderParams::default(),
49            loss_grad_clip: 1.0,
50            sh_band: 0,
51            max_anisotropy: 10.0,
52        }
53    }
54}
55
56/// Trailing raster params in a packed prepare buffer (`width`, `height`, …).
57pub const GAUSSIAN_SPLAT_PREP_RASTER_PARAMS_FLOATS: usize = 11;
58
59/// Tile count for a framebuffer (matches `rlx_splat::prep_layout::tile_count`).
60pub fn gaussian_splat_tile_count(width: u32, height: u32, tile_size: u32) -> u32 {
61    let tw = width.div_ceil(tile_size);
62    let th = height.div_ceil(tile_size);
63    tw * th
64}
65
66/// Packed prepare-buffer length for `N` splats (must match `rlx_splat::prep_layout::pack_prepared`).
67pub fn gaussian_splat_prep_packed_len(
68    count: usize,
69    max_list_entries: u32,
70    width: u32,
71    height: u32,
72    tile_size: u32,
73) -> usize {
74    let n = count.max(1);
75    let max_list = max_list_entries as usize;
76    let tiles = gaussian_splat_tile_count(width, height, tile_size) as usize;
77    let pixels = (width as usize).saturating_mul(height as usize).max(1);
78    n * 4
79        + n
80        + n * 3
81        + n * 3
82        + n * 4
83        + max_list
84        + tiles * 2
85        + pixels * 3
86        + GAUSSIAN_SPLAT_PREP_RASTER_PARAMS_FLOATS
87}
88
89/// Packed scene gradient layout lengths for `N` splats and `sh_coeff_count` SH bands.
90pub fn gaussian_splat_packed_grad_len(count: usize, sh_coeff_count: usize) -> usize {
91    count * (3 + 3 + 4 + 1 + 3) + count * sh_coeff_count.max(1) * 3
92}
93
94/// Unpack [`Op::GaussianSplatRenderBackward`] output into per-parameter gradients.
95pub fn unpack_gaussian_splat_packed_grads(
96    g: &mut Graph,
97    packed: NodeId,
98    count: usize,
99    sh_coeff_count: usize,
100) -> GaussianSplatInputs {
101    let mut off = 0usize;
102    let mut take = |len: usize| -> NodeId {
103        let id = g.narrow_(packed, 0, off, len);
104        off += len;
105        id
106    };
107    let positions = take(count * 3);
108    let scales = take(count * 3);
109    let rotations = take(count * 4);
110    let opacities = take(count);
111    let colors = take(count * 3);
112    let sh_coeffs = take(count * sh_coeff_count.max(1) * 3);
113    let _ = off;
114    GaussianSplatInputs {
115        positions,
116        scales,
117        rotations,
118        opacities,
119        colors,
120        sh_coeffs,
121        meta: packed,
122    }
123}
124
125impl Default for GaussianSplatRenderParams {
126    fn default() -> Self {
127        Self {
128            width: 64,
129            height: 64,
130            tile_size: 16,
131            radius_scale: 1.6,
132            alpha_cutoff: 1.0 / 255.0,
133            max_splat_steps: 32,
134            transmittance_threshold: 0.01,
135            max_list_entries: 18 * 32,
136        }
137    }
138}
139
140impl Graph {
141    /// First-class CPU reference Gaussian splat forward render.
142    ///
143    /// See [`Op::GaussianSplatRender`] for the seven-input contract and
144    /// [`GaussianSplatRenderParams`] for framebuffer settings.
145    pub fn gaussian_splat_render(
146        &mut self,
147        inputs: GaussianSplatInputs,
148        params: GaussianSplatRenderParams,
149    ) -> NodeId {
150        let out_elems = (params.width as usize) * (params.height as usize) * 4;
151        let dtype = self.shape(inputs.positions).dtype();
152        let out_shape = Shape::new(&[out_elems], dtype);
153        self.push(
154            Op::GaussianSplatRender {
155                width: params.width,
156                height: params.height,
157                tile_size: params.tile_size,
158                radius_scale: params.radius_scale,
159                alpha_cutoff: params.alpha_cutoff,
160                max_splat_steps: params.max_splat_steps,
161                transmittance_threshold: params.transmittance_threshold,
162                max_list_entries: params.max_list_entries,
163            },
164            vec![
165                inputs.positions,
166                inputs.scales,
167                inputs.rotations,
168                inputs.opacities,
169                inputs.colors,
170                inputs.sh_coeffs,
171                inputs.meta,
172            ],
173            out_shape,
174            None,
175        )
176    }
177
178    /// Build the 23-float `meta` vector expected by [`Op::GaussianSplatRender`].
179    pub fn gaussian_splat_render_meta(
180        &mut self,
181        camera_position: [f32; 3],
182        camera_target: [f32; 3],
183        camera_up: [f32; 3],
184        fov_y_degrees: f32,
185        near: f32,
186        far: f32,
187        background: [f32; 3],
188        params: GaussianSplatRenderParams,
189    ) -> NodeId {
190        let values = vec![
191            camera_position[0],
192            camera_position[1],
193            camera_position[2],
194            camera_target[0],
195            camera_target[1],
196            camera_target[2],
197            camera_up[0],
198            camera_up[1],
199            camera_up[2],
200            fov_y_degrees,
201            near,
202            far,
203            background[0],
204            background[1],
205            background[2],
206            params.width as f32,
207            params.height as f32,
208            params.tile_size as f32,
209            params.radius_scale,
210            params.alpha_cutoff,
211            params.max_splat_steps as f32,
212            params.transmittance_threshold,
213            params.max_list_entries as f32,
214        ];
215        let bytes: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
216        self.add_node(
217            Op::Constant { data: bytes },
218            vec![],
219            Shape::new(&[23], DType::F32),
220        )
221    }
222
223    /// Strict IR stage 1: project + bin + sort + rays → packed prepare buffer.
224    pub fn gaussian_splat_prepare(
225        &mut self,
226        inputs: GaussianSplatInputs,
227        params: GaussianSplatRenderParams,
228    ) -> NodeId {
229        let count = self.shape(inputs.positions).num_elements().unwrap_or(0) / 3;
230        let packed_len = gaussian_splat_prep_packed_len(
231            count,
232            params.max_list_entries,
233            params.width,
234            params.height,
235            params.tile_size,
236        );
237        let dtype = self.shape(inputs.positions).dtype();
238        self.push(
239            Op::GaussianSplatPrepare {
240                width: params.width,
241                height: params.height,
242                tile_size: params.tile_size,
243                radius_scale: params.radius_scale,
244                alpha_cutoff: params.alpha_cutoff,
245                max_splat_steps: params.max_splat_steps,
246                transmittance_threshold: params.transmittance_threshold,
247                max_list_entries: params.max_list_entries,
248            },
249            vec![
250                inputs.positions,
251                inputs.scales,
252                inputs.rotations,
253                inputs.opacities,
254                inputs.colors,
255                inputs.sh_coeffs,
256                inputs.meta,
257            ],
258            Shape::new(&[packed_len], dtype),
259            None,
260        )
261    }
262
263    /// Strict IR stage 2: rasterize from prepare buffer + meta.
264    pub fn gaussian_splat_rasterize(
265        &mut self,
266        prep: NodeId,
267        meta: NodeId,
268        params: GaussianSplatRenderParams,
269    ) -> NodeId {
270        let out_elems = (params.width as usize) * (params.height as usize) * 4;
271        let dtype = self.shape(prep).dtype();
272        self.push(
273            Op::GaussianSplatRasterize {
274                width: params.width,
275                height: params.height,
276                tile_size: params.tile_size,
277                alpha_cutoff: params.alpha_cutoff,
278                max_splat_steps: params.max_splat_steps,
279                transmittance_threshold: params.transmittance_threshold,
280                max_list_entries: params.max_list_entries,
281            },
282            vec![prep, meta],
283            Shape::new(&[out_elems], dtype),
284            None,
285        )
286    }
287
288    /// Decomposed strict-IR forward: prepare → rasterize.
289    pub fn gaussian_splat_render_decomposed(
290        &mut self,
291        inputs: GaussianSplatInputs,
292        params: GaussianSplatRenderParams,
293    ) -> NodeId {
294        let meta = inputs.meta;
295        let prep = self.gaussian_splat_prepare(inputs, params);
296        self.gaussian_splat_rasterize(prep, meta, params)
297    }
298
299    /// Backward pass for [`Op::GaussianSplatRender`] (packed scene gradients).
300    pub fn gaussian_splat_render_backward(
301        &mut self,
302        inputs: GaussianSplatInputs,
303        d_loss_rgba: NodeId,
304        params: GaussianSplatBackwardParams,
305    ) -> NodeId {
306        let count = self.shape(inputs.positions).num_elements().unwrap_or(0) / 3;
307        let sh_len = self.shape(inputs.sh_coeffs).num_elements().unwrap_or(0);
308        let sh_coeff_count = if count == 0 {
309            1
310        } else {
311            (sh_len / (count * 3)).max(1)
312        };
313        let packed_len = gaussian_splat_packed_grad_len(count, sh_coeff_count);
314        let dtype = self.shape(inputs.positions).dtype();
315        let r = params.render;
316        self.push(
317            Op::GaussianSplatRenderBackward {
318                width: r.width,
319                height: r.height,
320                tile_size: r.tile_size,
321                radius_scale: r.radius_scale,
322                alpha_cutoff: r.alpha_cutoff,
323                max_splat_steps: r.max_splat_steps,
324                transmittance_threshold: r.transmittance_threshold,
325                max_list_entries: r.max_list_entries,
326                loss_grad_clip: params.loss_grad_clip,
327                sh_band: params.sh_band,
328                max_anisotropy: params.max_anisotropy,
329            },
330            vec![
331                inputs.positions,
332                inputs.scales,
333                inputs.rotations,
334                inputs.opacities,
335                inputs.colors,
336                inputs.sh_coeffs,
337                inputs.meta,
338                d_loss_rgba,
339            ],
340            Shape::new(&[packed_len], dtype),
341            None,
342        )
343    }
344}
rlx_ir/ops/splat.rs

rlx_ir/ops/
splat.rs
