op: read-only storage textures in vertex shader (#3420)

* op: read-only storage textures in vertex shader

* Small fix in shader

* Remove TODO about vertex shader

* Improve comments

* Address reviewer's comments

src/webgpu/api/operation/storage_texture/read_only.spec.ts

@@ -2,7 +2,6 @@ export const description = `
 Tests for the behavior of read-only storage textures.
 
 TODO:
-- Test the use of read-only storage textures in vertex shader
 - Test 1D and 3D textures
 - Test mipmap level > 0
 - Test resource usage transitions with read-only storage textures
@@ -17,7 +16,7 @@ import {
   kTextureFormatInfo,
 } from '../../../format_info.js';
 import { GPUTest } from '../../../gpu_test.js';
-import { TValidShaderStage } from '../../../util/shader.js';
+import { kValidShaderStages, TValidShaderStage } from '../../../util/shader.js';
 
 function ComponentCount(format: ColorTextureFormat): number {
   switch (format) {
@@ -386,10 +385,126 @@ class F extends GPUTest {
         renderPassEncoder.end();
         break;
       }
-      case 'vertex':
-        // Not implemented yet.
-        unreachable();
+      case 'vertex': {
+        // For each texel location (coordX, coordY), draw one point at (coordX + 0.5, coordY + 0.5)
+        // in the storageTexture.width * storageTexture.height grid, and save all the texel values
+        // at (coordX, coordY, z) (z >= 0 && z < storageTexture.depthOrArrayLayers) into the
+        // corresponding vertex shader outputs.
+        let vertexOutputs = '';
+        for (let layer = 0; layer < storageTexture.depthOrArrayLayers; ++layer) {
+          vertexOutputs = vertexOutputs.concat(
+            `
+            @location(${layer + 1}) @interpolate(flat)
+            vertex_out${layer}: ${this.GetOutputBufferWGSLType(format)},`
+          );
+        }
+
+        let loadFromTextureWGSL = '';
+        if (storageTexture.depthOrArrayLayers === 1) {
+          loadFromTextureWGSL = `
+            output.vertex_out0 = textureLoad(readOnlyTexture, vec2u(coordX, coordY));`;
+        } else {
+          for (let z = 0; z < storageTexture.depthOrArrayLayers; ++z) {
+            loadFromTextureWGSL = loadFromTextureWGSL.concat(
+              `
+              output.vertex_out${z} = textureLoad(readOnlyTexture, vec2u(coordX, coordY), ${z});`
+            );
+          }
+        }
+
+        let outputToBufferWGSL = '';
+        for (let layer = 0; layer < storageTexture.depthOrArrayLayers; ++layer) {
+          outputToBufferWGSL = outputToBufferWGSL.concat(
+            `
+            let outputIndex${layer} =
+              storageTextureTexelCountPerImage * ${layer}u +
+              fragmentInput.tex_coord.y * ${storageTexture.width}u + fragmentInput.tex_coord.x;
+            outputBuffer[outputIndex${layer}] = fragmentInput.vertex_out${layer};`
+          );
+        }
+
+        const shader = `
+        ${bindingResourceDeclaration}
+        struct VertexOutput {
+          @builtin(position) my_pos: vec4f,
+          @location(0) @interpolate(flat) tex_coord: vec2u,
+          ${vertexOutputs}
+        }
+        @vertex
+        fn vs_main(@builtin(vertex_index) vertexIndex : u32) -> VertexOutput {
+            var output : VertexOutput;
+            let coordX = vertexIndex % ${storageTexture.width}u;
+            let coordY = vertexIndex / ${storageTexture.width}u;
+            // Each vertex in the mesh take an even step along X axis from -1.0 to 1.0.
+            let posXStep = f32(${2.0 / storageTexture.width});
+            // As well as along Y axis.
+            let posYStep = f32(${2.0 / storageTexture.height});
+            // And the vertex located in the middle of the step, i.e. with a bias of 0.5 step.
+            let outputPosX = -1.0 + posXStep * 0.5 + posXStep * f32(coordX);
+            let outputPosY = -1.0 + posYStep * 0.5 + posYStep * f32(coordY);
+            output.my_pos = vec4f(outputPosX, outputPosY, 0.0, 1.0);
+            output.tex_coord = vec2u(coordX, coordY);
+            ${loadFromTextureWGSL}
+            return output;
+        }
+        @fragment
+        fn fs_main(fragmentInput : VertexOutput) -> @location(0) vec4f {
+          let storageTextureTexelCountPerImage = ${storageTexture.width * storageTexture.height}u;
+          ${outputToBufferWGSL}
+          return vec4f(0.0, 1.0, 0.0, 1.0);
+        }
+        `;
+
+        const renderPipeline = this.device.createRenderPipeline({
+          layout: 'auto',
+          vertex: {
+            module: this.device.createShaderModule({
+              code: shader,
+            }),
+          },
+          fragment: {
+            module: this.device.createShaderModule({
+              code: shader,
+            }),
+            targets: [
+              {
+                format: 'rgba8unorm',
+              },
+            ],
+          },
+          primitive: {
+            topology: 'point-list',
+          },
+        });
+
+        const bindGroup = this.device.createBindGroup({
+          layout: renderPipeline.getBindGroupLayout(0),
+          entries: bindGroupEntries,
+        });
+
+        const placeholderColorTexture = this.device.createTexture({
+          size: [storageTexture.width, storageTexture.height, 1],
+          usage: GPUTextureUsage.RENDER_ATTACHMENT,
+          format: 'rgba8unorm',
+        });
+        this.trackForCleanup(placeholderColorTexture);
+
+        const renderPassEncoder = commandEncoder.beginRenderPass({
+          colorAttachments: [
+            {
+              view: placeholderColorTexture.createView(),
+              loadOp: 'clear',
+              clearValue: { r: 0, g: 0, b: 0, a: 0 },
+              storeOp: 'store',
+            },
+          ],
+        });
+        renderPassEncoder.setPipeline(renderPipeline);
+        renderPassEncoder.setBindGroup(0, bindGroup);
+        renderPassEncoder.draw(storageTexture.width * storageTexture.height);
+        renderPassEncoder.end();
         break;
+      }
     }
 
     this.queue.submit([commandEncoder.finish()]);
@@ -410,7 +525,7 @@ g.test('basic')
       .filter(
         p => p.format === 'bgra8unorm' || kTextureFormatInfo[p.format].color?.storage === true
       )
-      .combine('shaderStage', ['fragment', 'compute'] as const)
+      .combine('shaderStage', kValidShaderStages)
       .combine('depthOrArrayLayers', [1, 2] as const)
   )
   .beforeAllSubcases(t => {