rlx_cpu/
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//! CPU dispatch hooks for [`rlx_ir::Op::GaussianSplatRender`].
5//!
6//! The renderer itself is not in this workspace — `rlx-splat` is a separate
7//! repo that depends on rlx, so the dependency points inward and this module is
8//! the seam it plugs into. Everything here is the *contract*: rlx-ir owns the
9//! opcodes and the packed layouts, `rlx-cpu` owns dispatch, and the six bodies
10//! arrive at runtime through [`register_splat_executors`].
11//!
12//! The executor aliases are public because they are half of that contract. A
13//! downstream renderer needs to name them to hold one in a struct or hand one
14//! back from a builder; leaving them private made the registration signature
15//! uncallable to read even though the function itself was `pub`.
16
17use std::sync::OnceLock;
18
19/// Arena-resident forward render. Registered by the downstream renderer.
20pub type RenderExec = Box<dyn Fn(ArenaRenderArgs) + Send + Sync>;
21/// Arena-resident render backward.
22pub type RenderBwdExec = Box<dyn Fn(ArenaRenderBwdArgs) + Send + Sync>;
23/// Arena-resident `GaussianSplatPrepare`.
24pub type PrepareExec = Box<dyn Fn(ArenaPrepareArgs) + Send + Sync>;
25/// Arena-resident `GaussianSplatRasterize`.
26pub type RasterizeExec = Box<dyn Fn(ArenaRasterizeArgs) + Send + Sync>;
27/// Host-slice forward render, returning the framebuffer.
28pub type HostRenderExec = Box<dyn Fn(HostRenderArgs) -> Vec<f32> + Send + Sync>;
29/// Host-slice backward, returning packed gradients.
30pub type HostBackwardExec = Box<dyn Fn(HostBackwardArgs) -> Vec<f32> + Send + Sync>;
31
32static RENDER: OnceLock<RenderExec> = OnceLock::new();
33static RENDER_BWD: OnceLock<RenderBwdExec> = OnceLock::new();
34static PREPARE: OnceLock<PrepareExec> = OnceLock::new();
35static RASTERIZE: OnceLock<RasterizeExec> = OnceLock::new();
36static HOST_RENDER: OnceLock<HostRenderExec> = OnceLock::new();
37static HOST_BACKWARD: OnceLock<HostBackwardExec> = OnceLock::new();
38
39/// Arena arguments for forward splat.
40#[allow(clippy::struct_excessive_bools)]
41pub struct ArenaRenderArgs {
42    pub positions_off: usize,
43    pub positions_len: usize,
44    pub scales_off: usize,
45    pub scales_len: usize,
46    pub rotations_off: usize,
47    pub rotations_len: usize,
48    pub opacities_off: usize,
49    pub opacities_len: usize,
50    pub colors_off: usize,
51    pub colors_len: usize,
52    pub sh_coeffs_off: usize,
53    pub sh_coeffs_len: usize,
54    pub meta_off: usize,
55    pub dst_off: usize,
56    pub dst_len: usize,
57    pub width: u32,
58    pub height: u32,
59    pub tile_size: u32,
60    pub radius_scale: f32,
61    pub alpha_cutoff: f32,
62    pub max_splat_steps: u32,
63    pub transmittance_threshold: f32,
64    pub max_list_entries: u32,
65    pub base: *mut u8,
66}
67
68/// Arena arguments for [`Op::GaussianSplatPrepare`](rlx_ir::Op::GaussianSplatPrepare).
69pub struct ArenaPrepareArgs {
70    pub positions_off: usize,
71    pub positions_len: usize,
72    pub scales_off: usize,
73    pub scales_len: usize,
74    pub rotations_off: usize,
75    pub rotations_len: usize,
76    pub opacities_off: usize,
77    pub opacities_len: usize,
78    pub colors_off: usize,
79    pub colors_len: usize,
80    pub sh_coeffs_off: usize,
81    pub sh_coeffs_len: usize,
82    pub meta_off: usize,
83    pub meta_len: usize,
84    pub prep_off: usize,
85    pub prep_len: usize,
86    pub width: u32,
87    pub height: u32,
88    pub tile_size: u32,
89    pub radius_scale: f32,
90    pub alpha_cutoff: f32,
91    pub max_splat_steps: u32,
92    pub transmittance_threshold: f32,
93    pub max_list_entries: u32,
94    pub base: *mut u8,
95}
96
97/// Arena arguments for [`Op::GaussianSplatRasterize`](rlx_ir::Op::GaussianSplatRasterize).
98pub struct ArenaRasterizeArgs {
99    pub prep_off: usize,
100    pub prep_len: usize,
101    pub meta_off: usize,
102    pub meta_len: usize,
103    pub dst_off: usize,
104    pub dst_len: usize,
105    pub count: usize,
106    pub width: u32,
107    pub height: u32,
108    pub tile_size: u32,
109    pub alpha_cutoff: f32,
110    pub max_splat_steps: u32,
111    pub transmittance_threshold: f32,
112    pub max_list_entries: u32,
113    pub base: *mut u8,
114}
115
116/// Arena arguments for backward splat.
117pub struct ArenaRenderBwdArgs {
118    pub positions_off: usize,
119    pub positions_len: usize,
120    pub scales_off: usize,
121    pub scales_len: usize,
122    pub rotations_off: usize,
123    pub rotations_len: usize,
124    pub opacities_off: usize,
125    pub opacities_len: usize,
126    pub colors_off: usize,
127    pub colors_len: usize,
128    pub sh_coeffs_off: usize,
129    pub sh_coeffs_len: usize,
130    pub meta_off: usize,
131    pub d_loss_off: usize,
132    pub d_loss_len: usize,
133    pub packed_off: usize,
134    pub packed_len: usize,
135    pub width: u32,
136    pub height: u32,
137    pub tile_size: u32,
138    pub radius_scale: f32,
139    pub alpha_cutoff: f32,
140    pub max_splat_steps: u32,
141    pub transmittance_threshold: f32,
142    pub max_list_entries: u32,
143    pub loss_grad_clip: f32,
144    pub sh_band: u32,
145    pub max_anisotropy: f32,
146    pub base: *mut u8,
147}
148
149/// Host-buffer forward splat.
150pub struct HostRenderArgs {
151    pub positions: Vec<f32>,
152    pub scales: Vec<f32>,
153    pub rotations: Vec<f32>,
154    pub opacities: Vec<f32>,
155    pub colors: Vec<f32>,
156    pub sh_coeffs: Vec<f32>,
157    pub meta: Vec<f32>,
158    pub width: u32,
159    pub height: u32,
160    pub tile_size: u32,
161    pub radius_scale: f32,
162    pub alpha_cutoff: f32,
163    pub max_splat_steps: u32,
164    pub transmittance_threshold: f32,
165    pub max_list_entries: u32,
166}
167
168/// Host-buffer backward splat.
169pub struct HostBackwardArgs {
170    pub positions: Vec<f32>,
171    pub scales: Vec<f32>,
172    pub rotations: Vec<f32>,
173    pub opacities: Vec<f32>,
174    pub colors: Vec<f32>,
175    pub sh_coeffs: Vec<f32>,
176    pub meta: Vec<f32>,
177    pub d_loss_rgba: Vec<f32>,
178    pub width: u32,
179    pub height: u32,
180    pub tile_size: u32,
181    pub radius_scale: f32,
182    pub alpha_cutoff: f32,
183    pub max_splat_steps: u32,
184    pub transmittance_threshold: f32,
185    pub max_list_entries: u32,
186    pub loss_grad_clip: f32,
187    pub sh_band: u32,
188    pub max_anisotropy: f32,
189}
190
191/// Register the six splat executor bodies. Call once per process, before
192/// compiling any graph that references `Op::GaussianSplat*` — the dispatch
193/// hooks below panic with a pointed message if a body is missing.
194///
195/// This is what a downstream renderer (the out-of-tree `rlx-splat`) calls
196/// from its own `register()`.
197pub fn register_splat_executors(
198    render: RenderExec,
199    backward: RenderBwdExec,
200    prepare: PrepareExec,
201    rasterize: RasterizeExec,
202    host_render: HostRenderExec,
203    host_backward: HostBackwardExec,
204) {
205    let _ = RENDER.set(render);
206    let _ = RENDER_BWD.set(backward);
207    let _ = PREPARE.set(prepare);
208    let _ = RASTERIZE.set(rasterize);
209    let _ = HOST_RENDER.set(host_render);
210    let _ = HOST_BACKWARD.set(host_backward);
211}
212
213#[allow(clippy::too_many_arguments)]
214pub fn render_host_slices(
215    positions: &[f32],
216    scales: &[f32],
217    rotations: &[f32],
218    opacities: &[f32],
219    colors: &[f32],
220    sh_coeffs: &[f32],
221    meta: &[f32],
222    width: u32,
223    height: u32,
224    tile_size: u32,
225    radius_scale: f32,
226    alpha_cutoff: f32,
227    max_splat_steps: u32,
228    transmittance_threshold: f32,
229    max_list_entries: u32,
230) -> Vec<f32> {
231    HOST_RENDER
232        .get()
233        .expect("call `rlx_splat::register()` before host splat render")(HostRenderArgs {
234        positions: positions.to_vec(),
235        scales: scales.to_vec(),
236        rotations: rotations.to_vec(),
237        opacities: opacities.to_vec(),
238        colors: colors.to_vec(),
239        sh_coeffs: sh_coeffs.to_vec(),
240        meta: meta.to_vec(),
241        width,
242        height,
243        tile_size,
244        radius_scale,
245        alpha_cutoff,
246        max_splat_steps,
247        transmittance_threshold,
248        max_list_entries,
249    })
250}
251
252#[allow(clippy::too_many_arguments)]
253pub fn backward_host_slices(
254    positions: &[f32],
255    scales: &[f32],
256    rotations: &[f32],
257    opacities: &[f32],
258    colors: &[f32],
259    sh_coeffs: &[f32],
260    meta: &[f32],
261    d_loss_rgba: &[f32],
262    width: u32,
263    height: u32,
264    tile_size: u32,
265    radius_scale: f32,
266    alpha_cutoff: f32,
267    max_splat_steps: u32,
268    transmittance_threshold: f32,
269    max_list_entries: u32,
270    loss_grad_clip: f32,
271    sh_band: u32,
272    max_anisotropy: f32,
273) -> Vec<f32> {
274    HOST_BACKWARD
275        .get()
276        .expect("call `rlx_splat::register()` before host splat backward")(HostBackwardArgs {
277        positions: positions.to_vec(),
278        scales: scales.to_vec(),
279        rotations: rotations.to_vec(),
280        opacities: opacities.to_vec(),
281        colors: colors.to_vec(),
282        sh_coeffs: sh_coeffs.to_vec(),
283        meta: meta.to_vec(),
284        d_loss_rgba: d_loss_rgba.to_vec(),
285        width,
286        height,
287        tile_size,
288        radius_scale,
289        alpha_cutoff,
290        max_splat_steps,
291        transmittance_threshold,
292        max_list_entries,
293        loss_grad_clip,
294        sh_band,
295        max_anisotropy,
296    })
297}
298
299/// Execute [`Op::GaussianSplatPrepare`](rlx_ir::Op::GaussianSplatPrepare).
300#[allow(unsafe_op_in_unsafe_fn, clippy::too_many_arguments)]
301pub unsafe fn execute_gaussian_splat_prepare(
302    positions_off: usize,
303    positions_len: usize,
304    scales_off: usize,
305    scales_len: usize,
306    rotations_off: usize,
307    rotations_len: usize,
308    opacities_off: usize,
309    opacities_len: usize,
310    colors_off: usize,
311    colors_len: usize,
312    sh_coeffs_off: usize,
313    sh_coeffs_len: usize,
314    meta_off: usize,
315    meta_len: usize,
316    prep_off: usize,
317    prep_len: usize,
318    width: u32,
319    height: u32,
320    tile_size: u32,
321    radius_scale: f32,
322    alpha_cutoff: f32,
323    max_splat_steps: u32,
324    transmittance_threshold: f32,
325    max_list_entries: u32,
326    base: *mut u8,
327) {
328    PREPARE
329        .get()
330        .expect("call `rlx_splat::register()` before GaussianSplatPrepare")(ArenaPrepareArgs {
331        positions_off,
332        positions_len,
333        scales_off,
334        scales_len,
335        rotations_off,
336        rotations_len,
337        opacities_off,
338        opacities_len,
339        colors_off,
340        colors_len,
341        sh_coeffs_off,
342        sh_coeffs_len,
343        meta_off,
344        meta_len,
345        prep_off,
346        prep_len,
347        width,
348        height,
349        tile_size,
350        radius_scale,
351        alpha_cutoff,
352        max_splat_steps,
353        transmittance_threshold,
354        max_list_entries,
355        base,
356    });
357}
358
359/// Execute [`Op::GaussianSplatRasterize`](rlx_ir::Op::GaussianSplatRasterize).
360#[allow(unsafe_op_in_unsafe_fn, clippy::too_many_arguments)]
361pub unsafe fn execute_gaussian_splat_rasterize(
362    prep_off: usize,
363    prep_len: usize,
364    meta_off: usize,
365    meta_len: usize,
366    dst_off: usize,
367    dst_len: usize,
368    count: usize,
369    width: u32,
370    height: u32,
371    tile_size: u32,
372    alpha_cutoff: f32,
373    max_splat_steps: u32,
374    transmittance_threshold: f32,
375    max_list_entries: u32,
376    base: *mut u8,
377) {
378    RASTERIZE
379        .get()
380        .expect("call `rlx_splat::register()` before GaussianSplatRasterize")(
381        ArenaRasterizeArgs {
382            prep_off,
383            prep_len,
384            meta_off,
385            meta_len,
386            dst_off,
387            dst_len,
388            count,
389            width,
390            height,
391            tile_size,
392            alpha_cutoff,
393            max_splat_steps,
394            transmittance_threshold,
395            max_list_entries,
396            base,
397        },
398    );
399}
400
401/// Execute [`Op::GaussianSplatRender`](rlx_ir::Op::GaussianSplatRender) against the arena `base` pointer.
402#[allow(unsafe_op_in_unsafe_fn, clippy::too_many_arguments)]
403pub unsafe fn execute_gaussian_splat_render(
404    positions_off: usize,
405    positions_len: usize,
406    scales_off: usize,
407    scales_len: usize,
408    rotations_off: usize,
409    rotations_len: usize,
410    opacities_off: usize,
411    opacities_len: usize,
412    colors_off: usize,
413    colors_len: usize,
414    sh_coeffs_off: usize,
415    sh_coeffs_len: usize,
416    meta_off: usize,
417    dst_off: usize,
418    dst_len: usize,
419    width: u32,
420    height: u32,
421    tile_size: u32,
422    radius_scale: f32,
423    alpha_cutoff: f32,
424    max_splat_steps: u32,
425    transmittance_threshold: f32,
426    max_list_entries: u32,
427    base: *mut u8,
428) {
429    RENDER
430        .get()
431        .expect("call `rlx_splat::register()` before GaussianSplatRender")(ArenaRenderArgs {
432        positions_off,
433        positions_len,
434        scales_off,
435        scales_len,
436        rotations_off,
437        rotations_len,
438        opacities_off,
439        opacities_len,
440        colors_off,
441        colors_len,
442        sh_coeffs_off,
443        sh_coeffs_len,
444        meta_off,
445        dst_off,
446        dst_len,
447        width,
448        height,
449        tile_size,
450        radius_scale,
451        alpha_cutoff,
452        max_splat_steps,
453        transmittance_threshold,
454        max_list_entries,
455        base,
456    });
457}
458
459/// Execute [`Op::GaussianSplatRenderBackward`](rlx_ir::Op::GaussianSplatRenderBackward).
460#[allow(unsafe_op_in_unsafe_fn, clippy::too_many_arguments)]
461pub unsafe fn execute_gaussian_splat_render_backward(
462    positions_off: usize,
463    positions_len: usize,
464    scales_off: usize,
465    scales_len: usize,
466    rotations_off: usize,
467    rotations_len: usize,
468    opacities_off: usize,
469    opacities_len: usize,
470    colors_off: usize,
471    colors_len: usize,
472    sh_coeffs_off: usize,
473    sh_coeffs_len: usize,
474    meta_off: usize,
475    d_loss_off: usize,
476    d_loss_len: usize,
477    packed_off: usize,
478    packed_len: usize,
479    width: u32,
480    height: u32,
481    tile_size: u32,
482    radius_scale: f32,
483    alpha_cutoff: f32,
484    max_splat_steps: u32,
485    transmittance_threshold: f32,
486    max_list_entries: u32,
487    loss_grad_clip: f32,
488    sh_band: u32,
489    max_anisotropy: f32,
490    base: *mut u8,
491) {
492    RENDER_BWD
493        .get()
494        .expect("call `rlx_splat::register()` before GaussianSplatRenderBackward")(
495        ArenaRenderBwdArgs {
496            positions_off,
497            positions_len,
498            scales_off,
499            scales_len,
500            rotations_off,
501            rotations_len,
502            opacities_off,
503            opacities_len,
504            colors_off,
505            colors_len,
506            sh_coeffs_off,
507            sh_coeffs_len,
508            meta_off,
509            d_loss_off,
510            d_loss_len,
511            packed_off,
512            packed_len,
513            width,
514            height,
515            tile_size,
516            radius_scale,
517            alpha_cutoff,
518            max_splat_steps,
519            transmittance_threshold,
520            max_list_entries,
521            loss_grad_clip,
522            sh_band,
523            max_anisotropy,
524            base,
525        },
526    );
527}
rlx_cpu/splat.rs

rlx_cpu/
splat.rs
