From 29f70b5410eb6c256c6c9b69257d7a63b95b23d2 Mon Sep 17 00:00:00 2001
From: Daniel McNab <36049421+DJMcNab@users.noreply.github.com>
Date: Wed, 28 Aug 2024 09:43:40 +0100
Subject: [PATCH] Properly write-protect segments

Extracted from #606
---
 vello/src/render.rs                         |  2 +-
 vello/src/shaders.rs                        |  2 +-
 vello_shaders/shader/path_tiling_setup.wgsl | 21 ++++++++++++++++-----
 vello_shaders/src/cpu/path_tiling_setup.rs  | 11 ++++++-----
 4 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/vello/src/render.rs b/vello/src/render.rs
index 2cca8e6b2..b035d2539 100644
--- a/vello/src/render.rs
+++ b/vello/src/render.rs
@@ -429,7 +429,7 @@ impl Render {
         recording.dispatch(
             shaders.path_tiling_setup,
             wg_counts.path_tiling_setup,
-            [bump_buf, indirect_count_buf.into(), ptcl_buf],
+            [config_buf, bump_buf, indirect_count_buf.into(), ptcl_buf],
         );
         recording.dispatch_indirect(
             shaders.path_tiling,
diff --git a/vello/src/shaders.rs b/vello/src/shaders.rs
index a58e0edfe..2aa793a81 100644
--- a/vello/src/shaders.rs
+++ b/vello/src/shaders.rs
@@ -194,7 +194,7 @@ pub(crate) fn full_shaders(
             Buffer,
         ]
     );
-    let path_tiling_setup = add_shader!(path_tiling_setup, [Buffer, Buffer, Buffer]);
+    let path_tiling_setup = add_shader!(path_tiling_setup, [Uniform, Buffer, Buffer, Buffer]);
     let path_tiling = add_shader!(
         path_tiling,
         [
diff --git a/vello_shaders/shader/path_tiling_setup.wgsl b/vello_shaders/shader/path_tiling_setup.wgsl
index 4d5bf2e30..ebb6d62d1 100644
--- a/vello_shaders/shader/path_tiling_setup.wgsl
+++ b/vello_shaders/shader/path_tiling_setup.wgsl
@@ -3,15 +3,19 @@
 
 // Set up dispatch size for path tiling stage.
 
+#import config
 #import bump
 
 @group(0) @binding(0)
-var<storage, read_write> bump: BumpAllocators;
+var<uniform> config: Config;
 
 @group(0) @binding(1)
-var<storage, read_write> indirect: IndirectCount;
+var<storage, read_write> bump: BumpAllocators;
 
 @group(0) @binding(2)
+var<storage, read_write> indirect: IndirectCount;
+
+@group(0) @binding(3)
 var<storage, read_write> ptcl: array<u32>;
 
 // Partition size for path tiling stage
@@ -19,7 +23,16 @@ let WG_SIZE = 256u;
 
 @compute @workgroup_size(1)
 fn main() {
-    if atomicLoad(&bump.failed) != 0u {
+    indirect.count_y = 1u;
+    indirect.count_z = 1u;
+    let segments = atomicLoad(&bump.seg_counts);
+    let overflowed = segments > config.segments_size;
+    if atomicLoad(&bump.failed) != 0u  || overflowed {
+        if overflowed {
+            // Report the failure so that the CPU can know we have failed.
+            atomicOr(&bump.failed, STAGE_COARSE);
+        }
+        // Cancel path_tiling
         indirect.count_x = 0u;
         // signal fine rasterizer that failure happened (it doesn't bind bump)
         ptcl[0] = ~0u;
@@ -27,6 +40,4 @@ fn main() {
         let segments = atomicLoad(&bump.seg_counts);
         indirect.count_x = (segments + (WG_SIZE - 1u)) / WG_SIZE;
     }
-    indirect.count_y = 1u;
-    indirect.count_z = 1u;
 }
diff --git a/vello_shaders/src/cpu/path_tiling_setup.rs b/vello_shaders/src/cpu/path_tiling_setup.rs
index 9b6303691..cdae3a86d 100644
--- a/vello_shaders/src/cpu/path_tiling_setup.rs
+++ b/vello_shaders/src/cpu/path_tiling_setup.rs
@@ -1,13 +1,13 @@
 // Copyright 2023 the Vello Authors
 // SPDX-License-Identifier: Apache-2.0 OR MIT OR Unlicense
 
-use vello_encoding::{BumpAllocators, IndirectCount};
+use vello_encoding::{BumpAllocators, ConfigUniform, IndirectCount};
 
 use super::CpuBinding;
 
 const WG_SIZE: usize = 256;
 
-fn path_tiling_setup_main(bump: &BumpAllocators, indirect: &mut IndirectCount) {
+fn path_tiling_setup_main(_: &ConfigUniform, bump: &BumpAllocators, indirect: &mut IndirectCount) {
     let segments = bump.seg_counts;
     indirect.count_x = (segments + (WG_SIZE as u32 - 1)) / WG_SIZE as u32;
     indirect.count_y = 1;
@@ -15,8 +15,9 @@ fn path_tiling_setup_main(bump: &BumpAllocators, indirect: &mut IndirectCount) {
 }
 
 pub fn path_tiling_setup(_n_wg: u32, resources: &[CpuBinding]) {
-    let bump = resources[0].as_typed();
-    let mut indirect = resources[1].as_typed_mut();
+    let config = resources[0].as_typed();
+    let bump = resources[1].as_typed();
+    let mut indirect = resources[2].as_typed_mut();
     // binding 2 is ptcl, which we would need if we propagate failure
-    path_tiling_setup_main(&bump, &mut indirect);
+    path_tiling_setup_main(&config, &bump, &mut indirect);
 }