From 096f9fd5f088b0b49dff3e595abade79d70aea53 Mon Sep 17 00:00:00 2001
From: alan-baker <alanbaker@google.com>
Date: Wed, 4 Dec 2024 10:26:17 -0500
Subject: [PATCH] Improve fragment shader checks in subgroupBitwise (#4070)

* Modifies the fragment tests in subgroupBitwise to avoid any subgroup
  in the last row or column
  * this fixes failures seen on some devices performing subgroupXor when
    helpers are in bounds, but not on the edge of the framebuffer
* Tests skip cases that could not possibly yield useful results, but
  checking is also terminated early if no useful results can be found
  * Enough framebuffer variants are swept to ensure some tests will
    cover useful subgroups, but the sizes are only known after a pipeline
    compile
---
 .../call/builtin/subgroupBitwise.spec.ts      | 59 ++++++++++++++++---
 1 file changed, 52 insertions(+), 7 deletions(-)

diff --git a/src/webgpu/shader/execution/expression/call/builtin/subgroupBitwise.spec.ts b/src/webgpu/shader/execution/expression/call/builtin/subgroupBitwise.spec.ts
index dc46310937a0..a8f423e1f729 100644
--- a/src/webgpu/shader/execution/expression/call/builtin/subgroupBitwise.spec.ts
+++ b/src/webgpu/shader/execution/expression/call/builtin/subgroupBitwise.spec.ts
@@ -433,7 +433,7 @@ fn main(
 /**
  * Checks bitwise ops results from a fragment shader.
  *
- * Avoids the last row and column to skip potential helper invocations.
+ * Avoids subgroups in last row or column to skip potential helper invocations.
  * @param data Framebuffer output
  *             * component 0 is result
  *             * component 1 is generated subgroup id
@@ -453,11 +453,38 @@ function checkBitwiseFragment(
 ): Error | undefined {
   const { uintsPerRow, uintsPerTexel } = getUintsPerFramebuffer(format, width, height);
 
-  // Iteration skips last row and column to avoid helper invocations because it is not
-  // guaranteed whether or not they participate in the subgroup operation.
+  // Determine if the subgroup should be included in the checks.
+  const inBounds = new Map<number, boolean>();
+  for (let row = 0; row < height; row++) {
+    for (let col = 0; col < width; col++) {
+      const offset = uintsPerRow * row + col * uintsPerTexel;
+      const subgroup_id = data[offset + 1];
+      if (subgroup_id === 0) {
+        return new Error(`Internal error: helper invocation at (${col}, ${row})`);
+      }
+
+      let ok = inBounds.get(subgroup_id) ?? true;
+      ok = ok && row !== height - 1 && col !== width - 1;
+      inBounds.set(subgroup_id, ok);
+    }
+  }
+
+  let anyInBounds = false;
+  for (const [_, value] of inBounds) {
+    const ok = Boolean(value);
+    anyInBounds = anyInBounds || ok;
+  }
+  if (!anyInBounds) {
+    // This variant would not reliably test behavior.
+    return undefined;
+  }
+
+  // Iteration skips subgroups in the last row or column to avoid helper
+  // invocations because it is not guaranteed whether or not they participate
+  // in the subgroup operation.
   const expected = new Map<number, number>();
-  for (let row = 0; row < height - 1; row++) {
-    for (let col = 0; col < width - 1; col++) {
+  for (let row = 0; row < height; row++) {
+    for (let col = 0; col < width; col++) {
       const offset = uintsPerRow * row + col * uintsPerTexel;
       const subgroup_id = data[offset + 1];
 
@@ -465,14 +492,19 @@ function checkBitwiseFragment(
         return new Error(`Internal error: helper invocation at (${col}, ${row})`);
       }
 
+      const subgroupInBounds = inBounds.get(subgroup_id) ?? true;
+      if (!subgroupInBounds) {
+        continue;
+      }
+
       let v = expected.get(subgroup_id) ?? identity(op);
       v = bitwise(op, v, input[row * width + col]);
       expected.set(subgroup_id, v);
     }
   }
 
-  for (let row = 0; row < height - 1; row++) {
-    for (let col = 0; col < width - 1; col++) {
+  for (let row = 0; row < height; row++) {
+    for (let col = 0; col < width; col++) {
       const offset = uintsPerRow * row + col * uintsPerTexel;
       const res = data[offset];
       const subgroup_id = data[offset + 1];
@@ -482,6 +514,11 @@ function checkBitwiseFragment(
         continue;
       }
 
+      const subgroupInBounds = inBounds.get(subgroup_id) ?? true;
+      if (!subgroupInBounds) {
+        continue;
+      }
+
       const expected_v = expected.get(subgroup_id) ?? 0;
       if (expected_v !== res) {
         return new Error(`Row ${row}, col ${col}: incorrect results:
@@ -509,6 +546,14 @@ g.test('fragment,all_active')
   })
   .fn(async t => {
     const numInputs = t.params.size[0] * t.params.size[1];
+
+    interface SubgroupLimits extends GPUSupportedLimits {
+      minSubgroupSize: number;
+    }
+    const { minSubgroupSize } = t.device.limits as SubgroupLimits;
+    const innerTexels = (t.params.size[0] - 1) * (t.params.size[1] - 1);
+    t.skipIf(innerTexels < minSubgroupSize, 'Too few texels to be reliable');
+
     const inputData = generateInputData(t.params.case, numInputs, identity(t.params.op));
 
     const ident = identity(t.params.op) === 0 ? '0' : '0xffffffff';