tang_expr/
split.rs

1//! One graph as several kernels, for graphs too large for one shader
2//! (a traced contact step's adjoint is hundreds of thousands of nodes, which
3//! shader compilers take minutes and tens of GB to optimize, or reject).
4
5use std::collections::HashMap;
6use std::fmt::Write;
7
8use crate::codegen::{expr_str, ref_str};
9use crate::graph::ExprGraph;
10use crate::node::{ExprId, Node};
11
12/// WGSL kernels that together compute a graph's outputs, run in order over
13/// the same environments: each reads `inputs` (`n_inputs` per environment)
14/// and a per-environment `scratch` (`n_scratch` floats), and the last writes
15/// `outputs`. Bindings: 0 `inputs` (read), 1 `scratch` (read_write),
16/// 2 `outputs` (read_write), 3 the uniform `count`; entry point `main`.
17pub struct SplitKernels {
18    /// Kernel sources, in launch order.
19    pub parts: Vec<String>,
20    /// Scratch floats per environment.
21    pub n_scratch: usize,
22    /// Workgroup size of every part.
23    pub workgroup_size: u32,
24}
25
26fn children(node: Node) -> [Option<ExprId>; 3] {
27    match node {
28        Node::Var(_) | Node::Lit(_) => [None; 3],
29        Node::Add(a, b) | Node::Mul(a, b) | Node::Atan2(a, b) => [Some(a), Some(b), None],
30        Node::Neg(a)
31        | Node::Recip(a)
32        | Node::Sqrt(a)
33        | Node::Sin(a)
34        | Node::Exp2(a)
35        | Node::Log2(a) => [Some(a), None, None],
36        Node::Select(c, a, b) => [Some(c), Some(a), Some(b)],
37    }
38}
39
40impl ExprGraph {
41    /// [`to_kernel`](Self::to_kernel) for WGSL, cut into kernels of at most
42    /// `max_nodes` evaluated nodes each (topological order); a value used
43    /// outside the kernel that computes it goes through `scratch`.
44    pub fn to_wgsl_split(
45        &self,
46        outputs: &[ExprId],
47        n_inputs: usize,
48        max_nodes: usize,
49    ) -> SplitKernels {
50        self.split_wgsl(outputs, n_inputs, max_nodes, false)
51    }
52
53    /// [`to_wgsl_split`](Self::to_wgsl_split) in `f64` (buffers, scratch and
54    /// arithmetic), for a device with wgpu's `SHADER_F64`. Vulkan has no f64
55    /// `sin`/`exp2`/`log2`/`atan2`, so those come from WGSL helpers accurate
56    /// to a few ulp. A buffer of `n` values takes `2n` f32 slots.
57    pub fn to_wgsl_split_f64(
58        &self,
59        outputs: &[ExprId],
60        n_inputs: usize,
61        max_nodes: usize,
62    ) -> SplitKernels {
63        self.split_wgsl(outputs, n_inputs, max_nodes, true)
64    }
65
66    fn split_wgsl(
67        &self,
68        outputs: &[ExprId],
69        n_inputs: usize,
70        max_nodes: usize,
71        f64: bool,
72    ) -> SplitKernels {
73        let wg = 256u32;
74        let (ty, suffix) = if f64 { ("f64", "lf") } else { ("f32", "") };
75        let live = self.live_set(outputs);
76        let mut order: Vec<usize> = live
77            .iter()
78            .copied()
79            .filter(|&i| !matches!(self.node(ExprId(i as u32)), Node::Var(_) | Node::Lit(_)))
80            .collect();
81        order.sort_unstable();
82        let chunks: Vec<&[usize]> = order.chunks(max_nodes.max(1)).collect();
83        let nparts = chunks.len().max(1);
84        let mut part_of: HashMap<usize, usize> = HashMap::new();
85        for (p, c) in chunks.iter().enumerate() {
86            for &i in *c {
87                part_of.insert(i, p);
88            }
89        }
90        // scratch slots: values read by a later part (or an output, read by
91        // the last part) outside the part that computes them
92        let mut slot: HashMap<usize, usize> = HashMap::new();
93        let need = |i: usize, reader: usize, slot: &mut HashMap<usize, usize>| {
94            if let Some(&p) = part_of.get(&i) {
95                if p != reader && !slot.contains_key(&i) {
96                    let n = slot.len();
97                    slot.insert(i, n);
98                }
99            }
100        };
101        for (p, c) in chunks.iter().enumerate() {
102            for &i in *c {
103                for ch in children(self.node(ExprId(i as u32))).into_iter().flatten() {
104                    need(ch.0 as usize, p, &mut slot);
105                }
106            }
107        }
108        for o in outputs {
109            need(o.0 as usize, nparts - 1, &mut slot);
110        }
111        let ns = slot.len().max(1);
112        let helpers = if f64 { F64_HELPERS } else { "" };
113        let header = format!(
114            "{helpers}struct Params {{ count: u32, _p1: u32, _p2: u32, _p3: u32 }}\n\
115             @group(0) @binding(0) var<storage, read> inputs: array<{ty}>;\n\
116             @group(0) @binding(1) var<storage, read_write> scratch: array<{ty}>;\n\
117             @group(0) @binding(2) var<storage, read_write> outputs: array<{ty}>;\n\
118             @group(0) @binding(3) var<uniform> params: Params;\n\
119             @compute @workgroup_size({wg})\n\
120             fn main(@builtin(global_invocation_id) gid: vec3<u32>) {{\n    let idx = gid.x;\n    if (idx >= params.count) {{ return; }}\n"
121        );
122        let mut parts = Vec::with_capacity(nparts);
123        for p in 0..nparts {
124            let chunk: &[usize] = chunks.get(p).copied().unwrap_or(&[]);
125            let mut src = header.clone();
126            let mut loaded: std::collections::HashSet<usize> = std::collections::HashSet::new();
127            let mut body = String::new();
128            let load = |id: ExprId,
129                        pre: &mut String,
130                        loaded: &mut std::collections::HashSet<usize>| {
131                let i = id.0 as usize;
132                if !loaded.insert(i) {
133                    return;
134                }
135                match self.node(id) {
136                    Node::Var(n) => {
137                        writeln!(pre, "    let x{n} = inputs[idx * {n_inputs}u + {n}u];").unwrap()
138                    }
139                    Node::Lit(_) => {}
140                    _ => {
141                        if part_of.get(&i) != Some(&p) {
142                            writeln!(pre, "    let t{i} = scratch[idx * {ns}u + {}u];", slot[&i])
143                                .unwrap();
144                        }
145                    }
146                }
147            };
148            let mut pre = String::new();
149            for &i in chunk {
150                let node = self.node(ExprId(i as u32));
151                for ch in children(node).into_iter().flatten() {
152                    load(ch, &mut pre, &mut loaded);
153                }
154                loaded.insert(i);
155                let mut rhs = expr_str(self, node, suffix, crate::codegen::Dialect::Wgsl);
156                if f64 {
157                    for f in ["sin(", "exp2(", "log2(", "atan2("] {
158                        if let Some(args) = rhs.strip_prefix(f) {
159                            rhs = format!("{}_64({args}", &f[..f.len() - 1]);
160                            break;
161                        }
162                    }
163                }
164                writeln!(body, "    let t{i} = {rhs};").unwrap();
165                if let Some(s) = slot.get(&i) {
166                    writeln!(body, "    scratch[idx * {ns}u + {s}u] = t{i};").unwrap();
167                }
168            }
169            if p == nparts - 1 {
170                for o in outputs {
171                    load(*o, &mut pre, &mut loaded);
172                }
173                let no = outputs.len();
174                for (k, o) in outputs.iter().enumerate() {
175                    writeln!(
176                        body,
177                        "    outputs[idx * {no}u + {k}u] = {};",
178                        ref_str(self, *o, suffix)
179                    )
180                    .unwrap();
181                }
182            }
183            src += &pre;
184            src += &body;
185            src += "}\n";
186            parts.push(src);
187        }
188        SplitKernels {
189            parts,
190            n_scratch: ns,
191            workgroup_size: wg,
192        }
193    }
194}
195
196/// f64 `sin`, `exp2`, `log2`, `atan2` for WGSL on Vulkan (GLSL.std.450 has
197/// them for 16- and 32-bit floats only): Cody-Waite reduction and
198/// polynomials to a few ulp; `atan2` refines the f32 one by Newton steps.
199const F64_HELPERS: &str = "fn sin_64(x: f64) -> f64 {
200    let k = round(x * 0.6366197723675814lf);
201    let r = (x - k * 1.5707963267948966lf) - k * 6.123233995736766e-17lf;
202    let r2 = r * r;
203    let s = r * (1.0lf + r2 * (-0.16666666666666666lf + r2 * (0.008333333333333333lf + r2 * (-0.0001984126984126984lf
204        + r2 * (2.755731922398589e-6lf + r2 * (-2.505210838544172e-8lf + r2 * (1.6059043836821613e-10lf
205        + r2 * (-7.647163731819816e-13lf))))))));
206    let c = 1.0lf + r2 * (-0.5lf + r2 * (0.041666666666666664lf + r2 * (-0.001388888888888889lf
207        + r2 * (2.48015873015873e-5lf + r2 * (-2.755731922398589e-7lf + r2 * (2.08767569878681e-9lf
208        + r2 * (-1.1470745597729725e-11lf)))))));
209    let q = i32(k - 4.0lf * floor(k * 0.25lf));
210    if (q == 0) { return s; }
211    if (q == 1) { return c; }
212    if (q == 2) { return -s; }
213    return -c;
214}
215fn pow2i64(n: i32) -> f64 {
216    var m = n;
217    var b = 2.0lf;
218    if (m < 0) { b = 0.5lf; m = -m; }
219    var p = 1.0lf;
220    for (var i = 0; i < 11; i++) {
221        if ((m & 1) == 1) { p = p * b; }
222        b = b * b;
223        m = m >> 1u;
224    }
225    return p;
226}
227fn exp2_64(x: f64) -> f64 {
228    let n = floor(clamp(x, -1070.0lf, 1023.0lf) + 0.5lf);
229    let y = (x - n) * 0.6931471805599453lf;
230    var e = 1.0lf;
231    var term = 1.0lf;
232    for (var i = 1; i < 18; i++) { term = term * y / f64(i); e = e + term; }
233    return e * pow2i64(i32(n));
234}
235fn log2_64(x: f64) -> f64 {
236    var m = x;
237    var e = 0.0lf;
238    for (var i = 0; i < 20 && m >= 1.8446744073709552e19lf; i++) { m = m * 5.421010862427522e-20lf; e = e + 64.0lf; }
239    for (var i = 0; i < 20 && m < 5.421010862427522e-20lf; i++) { m = m * 1.8446744073709552e19lf; e = e - 64.0lf; }
240    for (var i = 0; i < 70 && m >= 1.4142135623730951lf; i++) { m = m * 0.5lf; e = e + 1.0lf; }
241    for (var i = 0; i < 70 && m < 0.7071067811865476lf; i++) { m = m * 2.0lf; e = e - 1.0lf; }
242    let t = (m - 1.0lf) / (m + 1.0lf);
243    let t2 = t * t;
244    var acc = 0.0lf;
245    for (var k = 12; k >= 0; k--) { acc = acc * t2 + 1.0lf / f64(2 * k + 1); }
246    return e + 2.0lf * t * acc * 1.4426950408889634lf;
247}
248fn atan2_64(y: f64, x: f64) -> f64 {
249    var th = f64(atan2(f32(y), f32(x)));
250    for (var i = 0; i < 2; i++) {
251        let s = sin_64(th);
252        let c = sin_64(th + 1.5707963267948966lf);
253        th = th + (y * c - x * s) / (x * c + y * s);
254    }
255    return th;
256}
257";
tang_expr/split.rs

tang_expr/
split.rs
