Functionality to async-read texture from GPU on WebGPU (#6370)

Co-authored-by: Martin Valigursky <[email protected]>

src/platform/graphics/texture.js

@@ -999,6 +999,10 @@ class Texture {
  }
  await Promise.all(promises);
  }
+
+ read(x, y, width, height, mipLevel = 0, face = 0, data = null, immediate = false) {
+ return this.impl.read?.(x, y, width, height, mipLevel, face, data, immediate);
+ }
 }
 
 export { Texture };

src/platform/graphics/webgpu/webgpu-graphics-device.js

@@ -920,6 +920,13 @@ class WebgpuGraphicsDevice extends GraphicsDevice {
  this.addCommandBuffer(cb);
  }
 
+ return this.readBuffer(stagingBuffer, size, data, immediate);
+ }
+
+ readBuffer(stagingBuffer, size, data = null, immediate = false) {
+
+ const destBuffer = stagingBuffer.buffer;
+
  // return a promise that resolves with the data
  return new Promise((resolve, reject) => {
 

src/platform/graphics/webgpu/webgpu-texture.js

@@ -8,7 +8,9 @@ import {
  PIXELFORMAT_RGBA16F, PIXELFORMAT_RGBA32F, PIXELFORMAT_DEPTHSTENCIL,
  SAMPLETYPE_UNFILTERABLE_FLOAT, SAMPLETYPE_DEPTH,
  FILTER_NEAREST, FILTER_LINEAR, FILTER_NEAREST_MIPMAP_NEAREST, FILTER_NEAREST_MIPMAP_LINEAR,
- FILTER_LINEAR_MIPMAP_NEAREST, FILTER_LINEAR_MIPMAP_LINEAR, isIntegerPixelFormat, SAMPLETYPE_INT, SAMPLETYPE_UINT
+ FILTER_LINEAR_MIPMAP_NEAREST, FILTER_LINEAR_MIPMAP_LINEAR, isIntegerPixelFormat, SAMPLETYPE_INT, SAMPLETYPE_UINT,
+ BUFFERUSAGE_READ,
+ BUFFERUSAGE_COPY_DST
 } from '../constants.js';
 import { TextureUtils } from '../texture-utils.js';
 import { WebgpuDebug } from './webgpu-debug.js';
@@ -498,6 +500,73 @@ class WebgpuTexture {
  Debug.trace(TRACEID_RENDER_QUEUE, `WRITE-TEX: mip:${mipLevel} index:${index} ${this.texture.name}`);
  wgpu.queue.writeTexture(dest, data, dataLayout, size);
  }
+
+ read(x, y, width, height, mipLevel, face, data, immediate) {
+
+ const texture = this.texture;
+ const formatInfo = pixelFormatInfo.get(texture.format);
+ Debug.assert(formatInfo);
+
+ const bytesPerRow = width * formatInfo.size;
+ Debug.assert(bytesPerRow);
+
+ // bytesPerRow must be a multiple of 256
+ const paddedBytesPerRow = math.roundUp(bytesPerRow, 256);
+ const size = paddedBytesPerRow * height;
+
+ // create a temporary staging buffer
+ /** @type {import('./webgpu-graphics-device.js').WebgpuGraphicsDevice} */
+ const device = texture.device;
+ const stagingBuffer = device.createBufferImpl(BUFFERUSAGE_READ | BUFFERUSAGE_COPY_DST);
+ stagingBuffer.allocate(device, size);
+
+ // use existing or create new encoder
+ const commandEncoder = device.commandEncoder ?? device.wgpu.createCommandEncoder();
+
+ const src = {
+ texture: this.gpuTexture,
+ mipLevel: mipLevel,
+ origin: [x, y, face]
+ };
+
+ const dst = {
+ buffer: stagingBuffer.buffer,
+ offset: 0,
+ bytesPerRow: paddedBytesPerRow
+ };
+
+ const copySize = {
+ width,
+ height,
+ depthOrArrayLayers: 1 // single layer
+ };
+
+ // copy the GPU texture to the staging buffer
+ commandEncoder.copyTextureToBuffer(src, dst, copySize);
+
+ // if we created new encoder
+ if (!device.commandEncoder) {
+ DebugHelper.setLabel(commandEncoder, 'copyTextureToBuffer-Encoder');
+ const cb = commandEncoder.finish();
+ DebugHelper.setLabel(cb, 'copyTextureToBuffer-CommandBuffer');
+ device.addCommandBuffer(cb);
+ }
+
+ // async read data from the staging buffer to a temporary array
+ return device.readBuffer(stagingBuffer, size, null, immediate).then((temp) => {
+
+ // remove the 256 alignment padding from the end of each row
+ data ??= new Uint8Array(height * bytesPerRow);
+ for (let i = 0; i < height; i++) {
+ const srcOffset = i * paddedBytesPerRow;
+ const dstOffset = i * bytesPerRow;
+ const sub = temp.subarray(srcOffset, srcOffset + bytesPerRow);
+ data.set(sub, dstOffset);
+ }
+
+ return data;
+ });
+ }
 }
 
 export { WebgpuTexture };