Extend compute pass multisampled copy texture helper to copy multiple…

… pixels (gpuweb#3340)

* Extend compute pass multisampled copy texture helper to copy multiple pixels

* fix review comments

* fix storage buffer size

src/webgpu/api/operation/render_pipeline/sample_mask.spec.ts

@@ -435,7 +435,7 @@ class F extends TextureTestMixin(GPUTest) {
     sampleMask: number,
     fragmentShaderOutputMask: number
   ) {
-    const buffer = this.copySinglePixelTextureToBufferUsingComputePass(
+    const buffer = this.copy2DTextureToBufferUsingComputePass(
       TypeF32, // correspond to 'rgba8unorm' format
       4,
       texture.createView(),
@@ -459,7 +459,7 @@ class F extends TextureTestMixin(GPUTest) {
     sampleMask: number,
     fragmentShaderOutputMask: number
   ) {
-    const buffer = this.copySinglePixelTextureToBufferUsingComputePass(
+    const buffer = this.copy2DTextureToBufferUsingComputePass(
       // Use f32 as the scalar type for depth (depth24plus, depth32float)
       // Use u32 as the scalar type for stencil (stencil8)
       aspect === 'depth-only' ? TypeF32 : TypeU32,
@@ -702,7 +702,7 @@ color' <= color.
         2
       );
 
-      const colorBuffer = t.copySinglePixelTextureToBufferUsingComputePass(
+      const colorBuffer = t.copy2DTextureToBufferUsingComputePass(
         TypeF32, // correspond to 'rgba8unorm' format
         4,
         color.createView(),
@@ -714,7 +714,7 @@ color' <= color.
       });
       colorResultPromises.push(colorResult);
 
-      const depthBuffer = t.copySinglePixelTextureToBufferUsingComputePass(
+      const depthBuffer = t.copy2DTextureToBufferUsingComputePass(
         TypeF32, // correspond to 'depth24plus-stencil8' format
         1,
         depthStencil.createView({ aspect: 'depth-only' }),
@@ -726,7 +726,7 @@ color' <= color.
       });
       depthResultPromises.push(depthResult);
 
-      const stencilBuffer = t.copySinglePixelTextureToBufferUsingComputePass(
+      const stencilBuffer = t.copy2DTextureToBufferUsingComputePass(
         TypeU32, // correspond to 'depth24plus-stencil8' format
         1,
         depthStencil.createView({ aspect: 'stencil-only' }),

src/webgpu/gpu_test.ts

@@ -60,7 +60,7 @@ import {
   textureContentIsOKByT2B,
 } from './util/texture/texture_ok.js';
 import { createTextureFromTexelView, createTextureFromTexelViews } from './util/texture.js';
-import { reifyOrigin3D } from './util/unions.js';
+import { reifyExtent3D, reifyOrigin3D } from './util/unions.js';
 
 const devicePool = new DevicePool();
 
@@ -816,24 +816,32 @@ export class GPUTestBase extends Fixture<GPUTestSubcaseBatchState> {
 
   /**
    * Emulate a texture to buffer copy by using a compute shader
-   * to load texture value of a single pixel and write to a storage buffer.
-   * For sample count == 1, the buffer contains only one value of the sample.
-   * For sample count > 1, the buffer contains (N = sampleCount) values sorted
+   * to load texture values of a subregion of a 2d texture and write to a storage buffer.
+   * For sample count == 1, the buffer contains extent[0] * extent[1] of the sample.
+   * For sample count > 1, the buffer contains extent[0] * extent[1] * (N = sampleCount) values sorted
    * in the order of their sample index [0, sampleCount - 1]
    *
    * This can be useful when the texture to buffer copy is not available to the texture format
    * e.g. (depth24plus), or when the texture is multisampled.
    *
-   * MAINTENANCE_TODO: extend to read multiple pixels with given origin and size.
+   * MAINTENANCE_TODO: extend texture dimension to 1d and 3d.
    *
    * @returns storage buffer containing the copied value from the texture.
    */
-  copySinglePixelTextureToBufferUsingComputePass(
+  copy2DTextureToBufferUsingComputePass(
     type: ScalarType,
     componentCount: number,
     textureView: GPUTextureView,
-    sampleCount: number
+    sampleCount: number = 1,
+    extent_: GPUExtent3D = [1, 1, 1],
+    origin_: GPUOrigin3D = [0, 0, 0]
   ): GPUBuffer {
+    const origin = reifyOrigin3D(origin_);
+    const extent = reifyExtent3D(extent_);
+    const width = extent.width;
+    const height = extent.height;
+    const kWorkgroupSizeX = 8;
+    const kWorkgroupSizeY = 8;
     const textureSrcCode =
       sampleCount === 1
         ? `@group(0) @binding(0) var src: texture_2d<${type}>;`
@@ -846,13 +854,24 @@ export class GPUTestBase extends Fixture<GPUTestSubcaseBatchState> {
       ${textureSrcCode}
       @group(0) @binding(1) var<storage, read_write> dst : Buffer;
 
-      @compute @workgroup_size(1) fn main() {
-        var coord = vec2<i32>(0, 0);
-        for (var sampleIndex = 0; sampleIndex < ${sampleCount};
+      struct Params {
+        origin: vec2u,
+        extent: vec2u,
+      };
+      @group(0) @binding(2) var<uniform> params : Params;
+
+      @compute @workgroup_size(${kWorkgroupSizeX}, ${kWorkgroupSizeY}, 1) fn main(@builtin(global_invocation_id) id : vec3u) {
+        let boundary = params.origin + params.extent;
+        let coord = params.origin + id.xy;
+        if (any(coord >= boundary)) {
+          return;
+        }
+        let offset = (id.x + id.y * params.extent.x) * ${componentCount} * ${sampleCount};
+        for (var sampleIndex = 0u; sampleIndex < ${sampleCount};
           sampleIndex = sampleIndex + 1) {
-          let o = sampleIndex * ${componentCount};
-          let v = textureLoad(src, coord, sampleIndex);
-          for (var component = 0; component < ${componentCount}; component = component + 1) {
+          let o = offset + sampleIndex * ${componentCount};
+          let v = textureLoad(src, coord.xy, sampleIndex);
+          for (var component = 0u; component < ${componentCount}; component = component + 1) {
             dst.data[o + component] = v[component];
           }
         }
@@ -869,11 +888,16 @@ export class GPUTestBase extends Fixture<GPUTestSubcaseBatchState> {
     });
 
     const storageBuffer = this.device.createBuffer({
-      size: sampleCount * type.size * componentCount,
+      size: sampleCount * type.size * componentCount * width * height,
       usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
     });
     this.trackForCleanup(storageBuffer);
 
+    const uniformBuffer = this.makeBufferWithContents(
+      new Uint32Array([origin.x, origin.y, width, height]),
+      GPUBufferUsage.UNIFORM
+    );
+
     const uniformBindGroup = this.device.createBindGroup({
       layout: computePipeline.getBindGroupLayout(0),
       entries: [
@@ -887,14 +911,24 @@ export class GPUTestBase extends Fixture<GPUTestSubcaseBatchState> {
             buffer: storageBuffer,
           },
         },
+        {
+          binding: 2,
+          resource: {
+            buffer: uniformBuffer,
+          },
+        },
       ],
     });
 
     const encoder = this.device.createCommandEncoder();
     const pass = encoder.beginComputePass();
     pass.setPipeline(computePipeline);
     pass.setBindGroup(0, uniformBindGroup);
-    pass.dispatchWorkgroups(1);
+    pass.dispatchWorkgroups(
+      Math.floor((width + kWorkgroupSizeX - 1) / kWorkgroupSizeX),
+      Math.floor((height + kWorkgroupSizeY - 1) / kWorkgroupSizeY),
+      1
+    );
     pass.end();
     this.device.queue.submit([encoder.finish()]);