Expand the sample_texture_combo test to include storage textures. (#4294)

Author: greggman

src/webgpu/api/operation/sampling/sampler_texture.spec.ts

@@ -16,26 +16,41 @@ g.test('sample_texture_combos')
   .desc(
     `
 Test that you can use the maximum number of textures with the maximum number of samplers.
+and the maximum number of storage textures.
 
-The test works by making the maximum number of texture+sampler combos.
-Each texture is [maxSamplersPerShaderStage, 1] in size where each texel is [textureId, samplerId]
-A function "useCombo<StageNum>(comboId)" is made that returns stage[stageNum].combo[comboId].texel[id, 0]
-or to put it another way, it returns the nth texel from the nth combo for that stage.
+The test works by making the maximum number of texture+sampler combos and the max storage
+textures per stage. Each texture is [maxSamplersPerShaderStage + maxStorageTexturesInStage, 1]
+in size and each texel is [textureId, samplerId]. A function "useCombo<StageNum>(comboId)" is
+made that returns stage[stageNum].combo[comboId].texel[id, 0] or to put it another way, it
+returns the nth texel from the nth combo for that stage.
 
-These are read in both the vertex shader and fragment shader and written to a [maxSamplerPerShaderStage, 2]
-texture where the top row is the values from the vertex shader and the bottom row from the fragment shader.
+These are read in both the vertex shader and fragment shader and written to a
+[maxSamplerPerShaderStage + maxStorageTexturesInStage, 2] texture where the top row is the
+values from the vertex shader and the bottom row from the fragment shader.
 
 The result should be a texture that has a value in each texel unique to a particular combo
+or storage texture.
 `
   )
   .fn(t => {
     const { device } = t;
-    const { maxSampledTexturesPerShaderStage, maxSamplersPerShaderStage, maxBindingsPerBindGroup } =
-      device.limits;
+    const {
+      maxSampledTexturesPerShaderStage,
+      maxSamplersPerShaderStage,
+      maxBindingsPerBindGroup,
+      maxStorageTexturesInVertexStage,
+      maxStorageTexturesInFragmentStage,
+      maxStorageTexturesPerShaderStage,
+    } = device.limits;
 
     assert(maxSampledTexturesPerShaderStage < 0xfffe);
     assert(maxSamplersPerShaderStage < 0xfffe);
 
+    const numStorageTexturesInVertexStage =
+      maxStorageTexturesInVertexStage ?? maxStorageTexturesPerShaderStage;
+    const numStorageTexturesInFragmentStage =
+      maxStorageTexturesInFragmentStage ?? maxStorageTexturesPerShaderStage;
+
     const maxTestableCombosPerStage = t.isCompatibility
       ? Math.min(maxSampledTexturesPerShaderStage, maxSamplersPerShaderStage)
       : maxSampledTexturesPerShaderStage * maxSamplersPerShaderStage;
@@ -44,12 +59,17 @@ The result should be a texture that has a value in each texel unique to a partic
     const declarationLines: string[] = [];
     const groups: GPUBindGroupEntry[][] = [[]];
     const layouts: GPUBindGroupLayoutEntry[][] = [[]];
-    const textureIds = new Set<string>();
+    const textureIdToTexelValue = new Map<string, number>();
     const samplerIds = new Set<string>();
     // per stage, per texel, each texel has 2 numbers, the texture id, and sampler id
     const expected: number[][][] = [[], []];
 
-    function addResource(stage: number, resourceId: string, resource: GPUTextureView | GPUSampler) {
+    function addResource(
+      stage: number,
+      resourceId: string,
+      resource: GPUTextureView | GPUSampler,
+      storageTexture?: boolean
+    ) {
       let bindGroupEntries = groups[groups.length - 1];
       let bindGroupLayoutEntries = layouts[groups.length - 1];
       if (bindGroupEntries.length === maxBindingsPerBindGroup) {
@@ -58,7 +78,12 @@ The result should be a texture that has a value in each texel unique to a partic
         groups.push(bindGroupEntries);
         layouts.push(bindGroupLayoutEntries);
       }
-      const resourceType = resource instanceof GPUSampler ? 'sampler' : 'texture_2d<f32>';
+      const resourceType =
+        resource instanceof GPUSampler
+          ? 'sampler'
+          : storageTexture
+          ? 'texture_storage_2d<rgba8unorm, read>'
+          : 'texture_2d<f32>';
       const binding = bindGroupEntries.length;
       declarationLines.push(
         `    @group(${groups.length - 1}) @binding(${binding}) var ${resourceId}: ${resourceType};`
@@ -74,37 +99,54 @@ The result should be a texture that has a value in each texel unique to a partic
           ? {
               sampler: {},
             }
+          : storageTexture
+          ? {
+              storageTexture: {
+                access: 'read-only',
+                format: 'rgba8unorm',
+              },
+            }
           : {
               texture: {},
             }),
       });
     }
 
-    function addTexture(stage: number, textureNum: number) {
+    const width =
+      maxSamplersPerShaderStage +
+      Math.max(numStorageTexturesInVertexStage, numStorageTexturesInFragmentStage);
+    t.debug(`width: ${width}`);
+
+    function addTexture(stage: number, textureNum: number, storageTexture: boolean) {
       const textureId = `tex${stage}_${textureNum}`;
-      if (!textureIds.has(textureId)) {
-        textureIds.add(textureId);
+      let texelValue = textureIdToTexelValue.get(textureId);
+      if (texelValue === undefined) {
+        texelValue = textures.length + 1;
+        textureIdToTexelValue.set(textureId, texelValue);
         const texture = t.createTextureTracked({
           format: 'rgba8unorm',
-          size: [maxSamplersPerShaderStage, 1],
-          usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST,
+          size: [width, 1],
+          usage:
+            GPUTextureUsage.STORAGE_BINDING |
+            GPUTextureUsage.TEXTURE_BINDING |
+            GPUTextureUsage.COPY_DST,
         });
         textures.push(texture);
         // Encode an rgba8unorm texture with rg16uint data where each texel is
-        // [(textureId + 1) | (stage << 15), {samplerId + 1}]
+        // [texelValue | (stage << 15), {samplerId + 1}]
         // The +1 is to avoid 0.
-        const data = new Uint16Array(maxSamplersPerShaderStage * 2);
-        const rg = (textureNum + 1) | (stage << 15);
-        for (let x = 0; x < maxSamplersPerShaderStage; ++x) {
+        const data = new Uint16Array(width * 2);
+        const rg = texelValue | (stage << 15);
+        for (let x = 0; x < width; ++x) {
           const offset = x * 2;
-          const samplerNum = x + 1;
+          const samplerNum = (x % maxSamplersPerShaderStage) + 1;
           data[offset + 0] = rg;
-          data[offset + 1] = samplerNum;
+          data[offset + 1] = storageTexture ? 0 : samplerNum;
         }
-        device.queue.writeTexture({ texture }, data, {}, [maxSamplersPerShaderStage]);
-        addResource(stage, textureId, texture.createView());
+        device.queue.writeTexture({ texture }, data, {}, [width]);
+        addResource(stage, textureId, texture.createView(), storageTexture);
       }
-      return textureId;
+      return { textureId, texelValue };
     }
 
     const kAddressModes = ['repeat', 'clamp-to-edge', 'mirror-repeat'] as const;
@@ -135,32 +177,52 @@ The result should be a texture that has a value in each texel unique to a partic
       return samplerId;
     }
 
+    const numStorageTexturesInStage = [
+      numStorageTexturesInVertexStage,
+      numStorageTexturesInFragmentStage,
+    ];
+
     // Note: We are storing textureId, samplerId in the texture. That suggests we could use rgba32uint
     // texture but we can't do that because we want to be able to set the samplers to linear.
     // Similarly we can't use rgba32float since they're not filterable by default.
     // So, we encode via rgba8unorm where rg is a 16bit textureId and ba is a 16bit samplerId
     const code = `
     // maxTestableCombosPerStage: ${maxTestableCombosPerStage}
+    // numStorageTexturesPerVertexStage: ${numStorageTexturesInVertexStage}
+    // numStorageTexturesPerFragmentStage: ${numStorageTexturesInFragmentStage}
 
     fn sample(t: texture_2d<f32>, s: sampler, validId: u32, currentId: u32, c: vec4f) -> vec4f {
       let size = textureDimensions(t, 0);
-      let uv = vec2f((f32(currentId) + 0.5) / f32(size.x), 0.5);
+      let uv = vec2f((f32(currentId % ${maxSamplersPerShaderStage}) + 0.5) / f32(size.x), 0.5);
       let v = textureSampleLevel(t, s, uv, 0);
       return select(c, v, currentId == validId);
     }
 
+    fn load(t: texture_storage_2d<rgba8unorm, read>, validId: u32, currentId: u32, c: vec4f) -> vec4f {
+      let size = textureDimensions(t);
+      let uv = vec2u(currentId % size.x, 0);
+      let v = textureLoad(t, uv);
+      return select(c, v, currentId == validId);
+    }
+
     ${range(
       2,
       stage => `
       fn useCombos${stage}(id: u32) -> vec4f {
         var c: vec4f;
 ${range(maxTestableCombosPerStage, i => {
   const texNum = (i / maxSamplersPerShaderStage) | 0;
-  const textureId = addTexture(stage, texNum);
+  const { textureId, texelValue } = addTexture(stage, texNum, false);
   const smpNum = i % maxSamplersPerShaderStage;
   const samplerId = addSampler(stage, smpNum);
-  expected[stage].push([(texNum + 1) | (stage << 15), smpNum + 1]);
+  expected[stage].push([texelValue | (stage << 15), smpNum + 1]);
   return `        c = sample(${textureId}, ${samplerId}, ${i}, id, c);`;
+}).join('\n')}
+${range(numStorageTexturesInStage[stage], i => {
+  const texNum = textures.length;
+  const { textureId, texelValue } = addTexture(stage, texNum, true);
+  expected[stage].push([texelValue | (stage << 15), 0]);
+  return `        c = load(${textureId}, ${i + maxTestableCombosPerStage}, id, c);`;
 }).join('\n')}
         return c;
       }
@@ -192,6 +254,8 @@ ${declarationLines.join('\n')}
     }
     `;
 
+    t.debug(code);
+
     const module = device.createShaderModule({ code });
     const bindGroupLayouts = layouts.map(entries => device.createBindGroupLayout({ entries }));
 
@@ -214,9 +278,14 @@ ${declarationLines.join('\n')}
       })
     );
 
+    const numAcross =
+      maxTestableCombosPerStage +
+      numStorageTexturesInVertexStage +
+      numStorageTexturesInFragmentStage;
+
     const renderTarget = t.createTextureTracked({
       format: 'rg16uint',
-      size: [maxTestableCombosPerStage, 2],
+      size: [numAcross, 2],
       usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_SRC,
     });
     textures.push(renderTarget);
@@ -234,7 +303,7 @@ ${declarationLines.join('\n')}
     pass.setPipeline(pipeline);
     bindGroups.forEach((bindGroup, i) => pass.setBindGroup(i, bindGroup));
     for (let y = 0; y < 2; ++y) {
-      for (let x = 0; x < maxTestableCombosPerStage; ++x) {
+      for (let x = 0; x < numAcross; ++x) {
         pass.setViewport(x, y, 1, 1, 0, 1);
         pass.draw(1, 1, 0, x);
       }
@@ -243,10 +312,10 @@ ${declarationLines.join('\n')}
 
     device.queue.submit([encoder.finish()]);
 
-    const expectedData = new Uint16Array(maxTestableCombosPerStage * 2 * 2);
+    const expectedData = new Uint16Array(numAcross * 2 * 2);
     for (let stage = 0; stage < 2; ++stage) {
       expected[stage].forEach(([tid, sid], i) => {
-        const offset = (maxTestableCombosPerStage * stage + i) * 2;
+        const offset = (numAcross * stage + i) * 2;
         expectedData[offset + 0] = tid;
         expectedData[offset + 1] = sid;
       });
@@ -256,14 +325,14 @@ ${declarationLines.join('\n')}
       'rg16uint',
       new Uint8Array(expectedData.buffer),
       {
-        bytesPerRow: maxTestableCombosPerStage * 4,
+        bytesPerRow: numAcross * 4,
         rowsPerImage: 2,
         subrectOrigin: [0, 0, 0],
-        subrectSize: [maxTestableCombosPerStage, 2],
+        subrectSize: [numAcross, 2],
       }
     );
 
-    const size = [maxSamplersPerShaderStage, 2];
+    const size = [numAcross, 2];
     t.expectTexelViewComparisonIsOkInTexture({ texture: renderTarget }, expTexelView, size);
 
     textures.forEach(texture => texture.destroy());