Large STEP imports: no more multi-minute hang, shared STEP text, cancel and progress (#925)

crates/ui-egui/src/gpu.rs

@@ -179,9 +179,13 @@ pub struct AccessMap {
     pub scale: f32,
 }
 
+/// Vertices (or line instances) on the GPU, in as many buffers as the device's buffer size
+/// limit needs: a large import's triangles can exceed one buffer (#883).
 struct Batch {
-    buffer: wgpu::Buffer,
-    count: u32,
+    /// Each buffer and its vertex (instance) count.
+    parts: Vec<(wgpu::Buffer, u32)>,
+    /// Bytes per buffer (all but the last are full).
+    chunk: usize,
 }
 
 /// A chunk of planned bodies on the GPU: their evaluated mesh and a colour per vertex.
@@ -289,8 +293,8 @@ impl Batches {
                 let faces = mesh.vertex_face.get(part.vertices.start as usize..part.vertices.end as usize).unwrap_or_default();
                 bytes.extend(faces.iter().flat_map(|f| b.face_colors.get(*f as usize).copied().unwrap_or(last)));
             }
-            if let Some(colors) = upload(device, "sc_planned_colors", &bytes, 4) {
-                out.push(PlannedDraw { mesh, colors: colors.buffer });
+            if let Some((colors, _)) = upload(device, "sc_planned_colors", &bytes, 4).and_then(|c| c.parts.into_iter().next()) {
+                out.push(PlannedDraw { mesh, colors });
             }
         }
         out
@@ -663,13 +667,27 @@ fn upload(device: &wgpu::Device, label: &str, bytes: &[u8], stride: usize) -> Op
 }
 
 fn upload_with(device: &wgpu::Device, label: &str, bytes: &[u8], stride: usize, usage: wgpu::BufferUsages) -> Option<Batch> {
-    if bytes.len() < stride {
+    if bytes.len() < stride || stride == 0 {
         return None;
     }
-    Some(Batch {
-        buffer: device.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some(label), contents: bytes, usage }),
-        count: u32::try_from(bytes.len() / stride).unwrap_or(0),
-    })
+    let chunk = buffer_chunk(stride, device.limits().max_buffer_size);
+    let parts = bytes
+        .chunks(chunk)
+        .filter(|part| part.len() >= stride)
+        .map(|part| {
+            let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some(label), contents: part, usage });
+            (buffer, u32::try_from(part.len() / stride).unwrap_or(0))
+        })
+        .collect();
+    Some(Batch { parts, chunk })
+}
+
+/// Bytes per GPU buffer for vertices of `stride` bytes: within the device's limit and a
+/// whole number of triangles and line pairs, so no primitive straddles two buffers.
+fn buffer_chunk(stride: usize, max_buffer_size: u64) -> usize {
+    let unit = stride.max(1) * 6;
+    let max = usize::try_from(max_buffer_size).unwrap_or(usize::MAX).min(1 << 30);
+    (max / unit).max(1) * unit
 }
 
 /// Triangles (3 vertices of `TRI_SIZE` bytes each, position first) ordered far to near along
@@ -691,7 +709,10 @@ impl Batches {
         let (Some(g), false) = (&self.glass, self.glass_cpu.is_empty()) else { return };
         let turned = self.glass_sorted_for.is_none_or(|s| s[0] * back[0] + s[1] * back[1] + s[2] * back[2] < 0.9994);
         if turned {
-            queue.write_buffer(&g.buffer, 0, &sort_back_to_front(&self.glass_cpu, back));
+            let sorted = sort_back_to_front(&self.glass_cpu, back);
+            for ((buffer, _), bytes) in g.parts.iter().zip(sorted.chunks(g.chunk.max(1))) {
+                queue.write_buffer(buffer, 0, bytes);
+            }
             self.glass_sorted_for = Some(back);
         }
     }
@@ -796,15 +817,19 @@ impl egui_wgpu::CallbackTrait for ViewportCallback {
         let tris = |pass: &mut wgpu::RenderPass<'static>, p: &wgpu::RenderPipeline, b: &Option<Batch>| {
             if let Some(b) = b {
                 pass.set_pipeline(p);
-                pass.set_vertex_buffer(0, b.buffer.slice(..));
-                pass.draw(0..b.count, 0..1);
+                for (buffer, count) in &b.parts {
+                    pass.set_vertex_buffer(0, buffer.slice(..));
+                    pass.draw(0..*count, 0..1);
+                }
             }
         };
         let lines = |pass: &mut wgpu::RenderPass<'static>, p: &wgpu::RenderPipeline, b: &Option<Batch>| {
             if let Some(b) = b {
                 pass.set_pipeline(p);
-                pass.set_vertex_buffer(0, b.buffer.slice(..));
-                pass.draw(0..6, 0..b.count);
+                for (buffer, count) in &b.parts {
+                    pass.set_vertex_buffer(0, buffer.slice(..));
+                    pass.draw(0..6, 0..*count);
+                }
             }
         };
         let (m, sk, h, pv) = (&res.model, &res.sketch, &res.highlight, &res.preview);
@@ -848,8 +873,10 @@ impl egui_wgpu::CallbackTrait for ViewportCallback {
             if let Some((_, bind)) = res.textures.get(id) {
                 pass.set_pipeline(&res.image);
                 pass.set_bind_group(1, bind, &[]);
-                pass.set_vertex_buffer(0, q.buffer.slice(..));
-                pass.draw(0..q.count, 0..1);
+                for (buffer, count) in &q.parts {
+                    pass.set_vertex_buffer(0, buffer.slice(..));
+                    pass.draw(0..*count, 0..1);
+                }
             }
         }
         tris(pass, &res.ghost, &pv.ghost);
@@ -1079,6 +1106,22 @@ mod tests {
 mod glass_tests {
     use super::*;
 
+    /// A large import's triangles once went into one buffer past the device's limit (524 MB
+    /// against 256 MB) and the app crashed (#883): they are split into whole triangles.
+    #[test]
+    fn large_scenes_are_split_into_buffers_within_the_device_limit() {
+        let limit = 256u64 << 20;
+        for stride in [TRI_SIZE, LINE_SIZE, 4] {
+            let chunk = buffer_chunk(stride, limit);
+            assert!(chunk as u64 <= limit && chunk > limit as usize / 2);
+            assert_eq!(chunk % (stride * 3), 0, "whole triangles");
+            assert_eq!(chunk % (stride * 2), 0, "whole line pairs");
+        }
+        let scene = 524_434_680usize;
+        assert_eq!(scene.div_ceil(buffer_chunk(TRI_SIZE, limit)), 2);
+        assert!(buffer_chunk(TRI_SIZE, 1) >= TRI_SIZE * 6, "a tiny limit still holds one unit");
+    }
+
     #[test]
     fn glass_is_drawn_far_to_near() {
         // Two triangles at z = 5 and z = -5, given near first; looking down from +Z.