graphic design tool
git clone https://git.lucas.co/cce-designer.git
src/springs.rs (9.4K)
1 //! The edge-length spring solve behind Relax's Springs mode, on the CPU and
2 //! on the GPU — the first operator of Phase 7 step 4, and the pattern for
3 //! the ones that follow: one algorithm, one data layout, two backends held
4 //! to each other by a cross-check.
5 //!
6 //! It is a JACOBI solve. Every point gathers the corrections of all its
7 //! incident edges from the positions as they stood at the start of the
8 //! pass, averages them, and moves once; the pass is one dispatch, one
9 //! invocation per point. Until 2026-09-24 the CPU solver was Gauss–Seidel
10 //! over the edge list in sequence, each correction visible to the next
11 //! edge, which no per-point kernel can reproduce — so the parallel form is
12 //! the one both backends share, and it is what the CPU runs too, rather
13 //! than letting the two drift apart. Jacobi converges more slowly per
14 //! iteration (roughly half the rate), which Iterations already controls.
15 //!
16 //! The layout is the GPU's: positions as a flat `xyz` float array (a
17 //! `vec3<f32>` in a WGSL storage array is padded to 16 bytes), the rest
18 //! topology as CSR — `offsets[p]..offsets[p + 1]` index the incident edges
19 //! of point `p`, each with its neighbour and its rest length — and pins as
20 //! one `u32` per point. The CPU solver walks exactly the same arrays in
21 //! exactly the same order, so the two agree to floating-point noise.
22
23 use crate::detail::Detail;
24 use cce_ui::vk::{Binding, ComputeDevice, Kernel};
25 use glam::Vec3;
26
27 /// The rest topology in CSR form, plus the pins, ready for either backend.
28 pub struct SpringSystem {
29 pub n: usize,
30 pub offsets: Vec<u32>,
31 pub neighbour: Vec<u32>,
32 pub rest: Vec<f32>,
33 pub pinned: Vec<u32>,
34 }
35
36 impl SpringSystem {
37 /// From the REST geometry's unique edges and a per-point pin flag.
38 pub fn build(rest: &Detail, pinned: &[bool]) -> Self {
39 let n = rest.num_points();
40 let mut counts = vec![0u32; n];
41 for e in rest.edges() {
42 counts[e[0] as usize] += 1;
43 counts[e[1] as usize] += 1;
44 }
45 let mut offsets = vec![0u32; n + 1];
46 for p in 0..n {
47 offsets[p + 1] = offsets[p] + counts[p];
48 }
49 let total = offsets[n] as usize;
50 let mut neighbour = vec![0u32; total];
51 let mut rest_len = vec![0f32; total];
52 let mut fill = offsets[..n].to_vec();
53 for e in rest.edges() {
54 let (a, b) = (e[0] as usize, e[1] as usize);
55 let len = (rest.pos(b) - rest.pos(a)).length();
56 let ia = fill[a] as usize;
57 neighbour[ia] = b as u32;
58 rest_len[ia] = len;
59 fill[a] += 1;
60 let ib = fill[b] as usize;
61 neighbour[ib] = a as u32;
62 rest_len[ib] = len;
63 fill[b] += 1;
64 }
65 SpringSystem {
66 n,
67 offsets,
68 neighbour,
69 rest: rest_len,
70 pinned: (0..n).map(|p| u32::from(pinned.get(p).copied().unwrap_or(false))).collect(),
71 }
72 }
73
74 pub fn num_edges(&self) -> usize {
75 self.neighbour.len() / 2
76 }
77 }
78
79 /// One Jacobi pass on the CPU: `out` from `pos`.
80 fn pass_cpu(sys: &SpringSystem, stiffness: f32, pos: &[f32], out: &mut [f32]) {
81 for p in 0..sys.n {
82 let base = p * 3;
83 let x = Vec3::new(pos[base], pos[base + 1], pos[base + 2]);
84 let mut x_new = x;
85 if sys.pinned[p] == 0 {
86 let (start, end) = (sys.offsets[p] as usize, sys.offsets[p + 1] as usize);
87 let mut sum = Vec3::ZERO;
88 for e in start..end {
89 let q = sys.neighbour[e] as usize;
90 let y = Vec3::new(pos[q * 3], pos[q * 3 + 1], pos[q * 3 + 2]);
91 let d = y - x;
92 let len = d.length();
93 if len < 1e-6 {
94 continue;
95 }
96 // Half the error toward a free neighbour (it moves the other
97 // half); all of it toward a pinned one, which does not move.
98 let w = if sys.pinned[q] != 0 { 1.0 } else { 0.5 };
99 sum += d * ((len - sys.rest[e]) / len * w * stiffness);
100 }
101 let count = end - start;
102 if count > 0 {
103 x_new = x + sum / count as f32;
104 }
105 }
106 out[base] = x_new.x;
107 out[base + 1] = x_new.y;
108 out[base + 2] = x_new.z;
109 }
110 }
111
112 /// `iterations` Jacobi passes on the CPU, in place.
113 pub fn solve_cpu(sys: &SpringSystem, stiffness: f32, iterations: usize, pos: &mut Vec<f32>) {
114 let mut out = vec![0f32; pos.len()];
115 for _ in 0..iterations {
116 pass_cpu(sys, stiffness, pos, &mut out);
117 std::mem::swap(pos, &mut out);
118 }
119 }
120
121 /// The same pass as a WGSL kernel: one invocation per point.
122 pub const SPRINGS_WGSL: &str = r#"
123 struct Params { stiffness: f32, n: u32, pad0: u32, pad1: u32 }
124 @group(0) @binding(0) var<storage, read> pos_in: array<f32>;
125 @group(0) @binding(1) var<storage, read_write> pos_out: array<f32>;
126 @group(0) @binding(2) var<storage, read> offsets: array<u32>;
127 @group(0) @binding(3) var<storage, read> neighbour: array<u32>;
128 @group(0) @binding(4) var<storage, read> rest: array<f32>;
129 @group(0) @binding(5) var<storage, read> pinned: array<u32>;
130 @group(0) @binding(6) var<uniform> params: Params;
131
132 @compute @workgroup_size(64)
133 fn springs(@builtin(global_invocation_id) id: vec3<u32>) {
134 let p = id.x;
135 if (p >= params.n) { return; }
136 let base = p * 3u;
137 let x = vec3<f32>(pos_in[base], pos_in[base + 1u], pos_in[base + 2u]);
138 var x_new = x;
139 if (pinned[p] == 0u) {
140 let start = offsets[p];
141 let end = offsets[p + 1u];
142 var sum = vec3<f32>(0.0, 0.0, 0.0);
143 for (var e = start; e < end; e = e + 1u) {
144 let q = neighbour[e];
145 let y = vec3<f32>(pos_in[q * 3u], pos_in[q * 3u + 1u], pos_in[q * 3u + 2u]);
146 let d = y - x;
147 let len = length(d);
148 if (len < 1e-6) { continue; }
149 var w = 0.5;
150 if (pinned[q] != 0u) { w = 1.0; }
151 sum = sum + d * ((len - rest[e]) / len * w * params.stiffness);
152 }
153 let count = end - start;
154 if (count > 0u) { x_new = x + sum / f32(count); }
155 }
156 pos_out[base] = x_new.x;
157 pos_out[base + 1u] = x_new.y;
158 pos_out[base + 2u] = x_new.z;
159 }"#;
160
161 #[repr(C)]
162 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
163 struct Params {
164 stiffness: f32,
165 n: u32,
166 pad0: u32,
167 pad1: u32,
168 }
169
170 /// `iterations` Jacobi passes on the GPU, in place: ONE submission, the
171 /// passes chained by memory barriers and the positions ping-ponging between
172 /// two device buffers, so the topology goes up once and the result comes
173 /// back once. The first cut submitted a pass at a time — upload, dispatch,
174 /// wait, read back, sixteen times — and lost to the CPU at every mesh size
175 /// measured, 134k points included: a submission's round trip is about half
176 /// a millisecond on an integrated GPU whatever the dispatch inside it, and
177 /// the solve itself is far cheaper than that. `springs_timing` has the
178 /// numbers for both shapes.
179 pub fn solve_gpu(
180 dev: &mut ComputeDevice,
181 sys: &SpringSystem,
182 stiffness: f32,
183 iterations: usize,
184 pos: &mut Vec<f32>,
185 ) -> Result<(), String> {
186 if sys.n == 0 || iterations == 0 {
187 return Ok(());
188 }
189 let kernel = Kernel::new(SPRINGS_WGSL, "springs");
190 let params = Params { stiffness, n: sys.n as u32, pad0: 0, pad1: 0 };
191 let mut out = vec![0f32; pos.len()];
192 dev.run_passes_over(
193 &kernel,
194 &mut [
195 Binding::input(pos.as_slice()),
196 Binding::rw(out.as_mut_slice()),
197 Binding::input(&sys.offsets),
198 Binding::input(&sys.neighbour),
199 Binding::input(&sys.rest),
200 Binding::input(&sys.pinned),
201 Binding::uniform(¶ms),
202 ],
203 sys.n as u32,
204 iterations as u32,
205 Some((0, 1)),
206 )?;
207 *pos = out;
208 Ok(())
209 }
210
211 /// Below this many points the CPU is faster; in auto mode it takes them.
212 /// From `springs_timing` in release on an Intel Iris Xe, sixteen passes,
213 /// one submission: 1.5k points cpu 0.25 ms / gpu 1.5 ms; 15k points cpu
214 /// 2.6 ms / gpu 3.7 ms; 135k points cpu 25 ms / gpu 14 ms. Break-even is
215 /// in the tens of thousands, and a first run adds ~15 ms of pipeline
216 /// compile on top; 32k is where the GPU is clearly ahead on every run
217 /// after the first.
218 pub const GPU_MIN_POINTS: usize = 32_768;
219
220 /// The solve as Relax runs it: the backend `CCE_COMPUTE` and the size
221 /// choose, a GPU failure in auto mode falling back to the CPU with a note,
222 /// and in forced-GPU mode reported back for the node-error slot.
223 pub fn solve(sys: &SpringSystem, stiffness: f32, iterations: usize, pos: &mut Vec<f32>) -> Result<(), String> {
224 if crate::gpu::use_gpu(sys.n, GPU_MIN_POINTS) {
225 let attempt = crate::gpu::with_any_device(|dev| solve_gpu(dev, sys, stiffness, iterations, pos));
226 match attempt {
227 Ok(Ok(())) => return Ok(()),
228 Ok(Err(e)) | Err(e) => {
229 if crate::gpu::choice() == crate::gpu::Choice::Gpu {
230 solve_cpu(sys, stiffness, iterations, pos);
231 return Err(format!("CCE_COMPUTE=gpu but the springs solve could not run there ({e}); solved on the CPU"));
232 }
233 note_fallback_once(&e);
234 }
235 }
236 }
237 solve_cpu(sys, stiffness, iterations, pos);
238 Ok(())
239 }
240
241 fn note_fallback_once(e: &str) {
242 static ONCE: std::sync::Once = std::sync::Once::new();
243 ONCE.call_once(|| eprintln!("cce-designer: springs solve fell back to the CPU: {e}"));
244 }