WebGPURenderer: Adjust `getArrayBufferAsync` behavior, add support for partial readback (#33322)

Co-authored-by: sunag <sunagbrasil@gmail.com>

src/renderers/common/Info.js

@@ -340,12 +340,12 @@ class Info {
 	 */
 	createReadbackBuffer( readbackBuffer ) {
 
-		const size = this._getAttributeMemorySize( readbackBuffer.attribute );
-		this.memoryMap.set( readbackBuffer, { size, type: 'readbackBuffers' } );
+		const maxByteLength = readbackBuffer.maxByteLength;
+		this.memoryMap.set( readbackBuffer, { size: maxByteLength, type: 'readbackBuffers' } );
 
 		this.memory.readbackBuffers ++;
-		this.memory.total += size;
-		this.memory.readbackBuffersSize += size;
+		this.memory.total += maxByteLength;
+		this.memory.readbackBuffersSize += maxByteLength;
 
 	}
 
@@ -356,7 +356,12 @@ class Info {
 	 */
 	destroyReadbackBuffer( readbackBuffer ) {
 
-		this.destroyAttribute( readbackBuffer );
+		const { size } = this.memoryMap.get( readbackBuffer );
+		this.memoryMap.delete( readbackBuffer );
+
+		this.memory.readbackBuffers --;
+		this.memory.total -= size;
+		this.memory.readbackBuffersSize -= size;
 
 	}
 

src/renderers/common/ReadbackBuffer.js

@@ -11,18 +11,32 @@ class ReadbackBuffer extends EventDispatcher {
 	/**
 	 * Constructs a new readback buffer.
 	 *
-	 * @param {BufferAttribute} attribute - The buffer attribute.
+	 * @param {number} maxByteLength - The maximum size of the buffer to be read back.
 	 */
-	constructor( attribute ) {
+	constructor( maxByteLength ) {
 
 		super();
 
 		/**
-		 * The buffer attribute.
+		 * Name used for debugging purposes.
 		 *
-		 * @type {BufferAttribute}
+		 * @type {string}
 		 */
-		this.attribute = attribute;
+		this.name = '';
+
+		/**
+		 * The mapped, read back array buffer.
+		 *
+		 * @type {ArrayBuffer|null}
+		 */
+		this.buffer = null;
+
+		/**
+		 * The maximum size of the buffer to be read back.
+		 *
+		 * @type {number}
+		 */
+		this.maxByteLength = maxByteLength;
 
 		/**
 		 * This flag can be used for type testing.
@@ -33,6 +47,8 @@ class ReadbackBuffer extends EventDispatcher {
 		 */
 		this.isReadbackBuffer = true;
 
+		this._mapped = false;
+
 	}
 
 	/**

src/renderers/common/Renderer.js

@@ -19,7 +19,6 @@ import Lighting from './Lighting.js';
 import XRManager from './XRManager.js';
 import InspectorBase from './InspectorBase.js';
 import CanvasTarget from './CanvasTarget.js';
-import ReadbackBuffer from './ReadbackBuffer.js';
 
 import NodeMaterial from '../../materials/nodes/NodeMaterial.js';
 
@@ -1913,61 +1912,42 @@ class Renderer {
 	 * from the GPU to the CPU in context of compute shaders.
 	 *
 	 * @async
-	 * @param {StorageBufferAttribute|ReadbackBuffer} buffer - The storage buffer attribute.
-	 * @return {Promise<ArrayBuffer>} A promise that resolves with the buffer data when the data are ready.
+	 * @param {BufferAttribute} attribute - The storage buffer attribute to read frm.
+	 * @param {ReadbackBuffer|ArrayBuffer} target - The storage buffer attribute.
+	 * @param {number} offset - The storage buffer attribute.
+	 * @param {number} count - The offset from which to start reading the
+	 * @return {Promise<ArrayBuffer|ReadbackBuffer>} A promise that resolves with the buffer data when the data are ready.
 	 */
-	async getArrayBufferAsync( buffer ) {
+	async getArrayBufferAsync( attribute, target = null, offset = 0, count = - 1 ) {
 
-		let readbackBuffer = buffer;
+		// tally the memory for this readback buffer
+		if ( target !== null && target.isReadbackBuffer ) {
 
-		if ( readbackBuffer.isReadbackBuffer !== true ) {
+			if ( this.info.memoryMap.has( target ) === false ) {
 
-			const attribute = buffer;
-			const attributeData = this.backend.get( attribute );
+				this.info.createReadbackBuffer( target );
 
-			readbackBuffer = attributeData.readbackBuffer;
+				const disposeInfo = () => {
 
-			if ( readbackBuffer === undefined ) {
+					target.removeEventListener( 'dispose', disposeInfo );
 
-				readbackBuffer = new ReadbackBuffer( attribute );
-
-				const dispose = () => {
-
-					attribute.removeEventListener( 'dispose', dispose );
-
-					readbackBuffer.dispose();
-
-					delete attributeData.readbackBuffer;
+					this.info.destroyReadbackBuffer( target );
 
 				};
 
-				attribute.addEventListener( 'dispose', dispose );
-
-				attributeData.readbackBuffer = readbackBuffer;
+				target.addEventListener( 'dispose', disposeInfo );
 
 			}
 
 		}
 
-		if ( this.info.memoryMap.has( readbackBuffer ) === false ) {
-
-			this.info.createReadbackBuffer( readbackBuffer );
+		if ( offset % 4 !== 0 || ( count > 0 && count % 4 !== 0 ) ) {
 
-			const disposeInfo = () => {
-
-				readbackBuffer.removeEventListener( 'dispose', disposeInfo );
-
-				this.info.destroyReadbackBuffer( readbackBuffer );
-
-			};
-
-			readbackBuffer.addEventListener( 'dispose', disposeInfo );
+			throw new Error( 'THREE.Renderer: "getArrayBufferAsync()" offset and count must be a multiple of 4.' );
 
 		}
 
-		readbackBuffer.release();
-
-		return await this.backend.getArrayBufferAsync( readbackBuffer );
+		return await this.backend.getArrayBufferAsync( attribute, target, offset, count );
 
 	}
 

src/renderers/webgl-fallback/WebGLBackend.js

@@ -309,15 +309,20 @@ class WebGLBackend extends Backend {
 
 	/**
 	 * This method performs a readback operation by moving buffer data from
-	 * a storage buffer attribute from the GPU to the CPU.
+	 * a storage buffer attribute from the GPU to the CPU. ReadbackBuffer can
+	 * be used to retain and reuse handles to the intermediate buffers and prevent
+	 * new allocation.
 	 *
 	 * @async
-	 * @param {ReadbackBuffer} readbackBuffer - The readback buffer.
-	 * @return {Promise<ArrayBuffer>} A promise that resolves with the buffer data when the data are ready.
+	 * @param {BufferAttribute} attribute - The storage buffer attribute to read frm.
+	 * @param {ReadbackBuffer|ArrayBuffer} target - The storage buffer attribute.
+	 * @param {number} offset - The storage buffer attribute.
+	 * @param {number} count - The offset from which to start reading the
+	 * @return {Promise<ArrayBuffer|ReadbackBuffer>} A promise that resolves with the buffer data when the data are ready.
 	 */
-	async getArrayBufferAsync( readbackBuffer ) {
+	async getArrayBufferAsync( attribute, target = null, offset = 0, count = - 1 ) {
 
-		return await this.attributeUtils.getArrayBufferAsync( readbackBuffer );
+		return await this.attributeUtils.getArrayBufferAsync( attribute, target, offset, count );
 
 	}
 

src/renderers/webgl-fallback/utils/WebGLAttributeUtils.js

@@ -254,76 +254,81 @@ class WebGLAttributeUtils {
 
 	/**
 	 * This method performs a readback operation by moving buffer data from
-	 * a storage buffer attribute from the GPU to the CPU.
+	 * a storage buffer attribute from the GPU to the CPU. ReadbackBuffer can
+	 * be used to retain and reuse handles to the intermediate buffers and prevent
+	 * new allocation.
 	 *
 	 * @async
-	 * @param {ReadbackBuffer} readbackBuffer - The readback buffer.
-	 * @return {Promise<ArrayBuffer>} A promise that resolves with the buffer data when the data are ready.
+	 * @param {BufferAttribute} attribute - The storage buffer attribute to read frm.
+	 * @param {ReadbackBuffer|ArrayBuffer} target - The storage buffer attribute.
+	 * @param {number} offset - The storage buffer attribute.
+	 * @param {number} count - The offset from which to start reading the
+	 * @return {Promise<ArrayBuffer|ReadbackBuffer>} A promise that resolves with the buffer data when the data are ready.
 	 */
-	async getArrayBufferAsync( readbackBuffer ) {
+	async getArrayBufferAsync( attribute, target = null, offset = 0, count = - 1 ) {
 
 		const backend = this.backend;
 		const { gl } = backend;
 
-		const attribute = readbackBuffer.attribute;
 		const bufferAttribute = attribute.isInterleavedBufferAttribute ? attribute.data : attribute;
-		const { bufferGPU } = backend.get( bufferAttribute );
+		const attributeInfo = backend.get( bufferAttribute );
+		const { bufferGPU } = attributeInfo;
 
-		const array = attribute.array;
-		const byteLength = array.byteLength;
-
-		gl.bindBuffer( gl.COPY_READ_BUFFER, bufferGPU );
-
-		const readbackBufferData = backend.get( readbackBuffer );
+		const byteLength = count === - 1 ? attributeInfo.byteLength - offset : count;
 
-		let { writeBuffer } = readbackBufferData;
+		// read the data back
+		let dstBuffer;
+		if ( target === null ) {
 
-		if ( writeBuffer === undefined ) {
+			dstBuffer = new Uint8Array( new ArrayBuffer( byteLength ) );
 
-			writeBuffer = gl.createBuffer();
+		} else if ( target.isReadbackBuffer ) {
 
-			gl.bindBuffer( gl.COPY_WRITE_BUFFER, writeBuffer );
-			gl.bufferData( gl.COPY_WRITE_BUFFER, byteLength, gl.STREAM_READ );
+			if ( target._mapped === true ) {
 
-			// dispose
+				throw new Error( 'WebGPURenderer: ReadbackBuffer must be released before being used again.' );
 
-			const dispose = () => {
-
-				gl.deleteBuffer( writeBuffer );
+			}
 
-				backend.delete( readbackBuffer );
+			const releaseCallback = () => {
 
-				readbackBuffer.removeEventListener( 'dispose', dispose );
+				target.buffer = null;
+				target._mapped = false;
+				target.removeEventListener( 'release', releaseCallback );
+				target.removeEventListener( 'dispose', releaseCallback );
 
 			};
 
-			readbackBuffer.addEventListener( 'dispose', dispose );
+			target.addEventListener( 'release', releaseCallback );
+			target.addEventListener( 'dispose', releaseCallback );
 
-			// register
-
-			readbackBufferData.writeBuffer = writeBuffer;
+			// WebGL has no concept of a "mapped" data buffer so we create a new buffer, instead.
+			dstBuffer = new Uint8Array( new ArrayBuffer( byteLength ) );
+			target.buffer = dstBuffer.buffer;
 
 		} else {
 
-			gl.bindBuffer( gl.COPY_WRITE_BUFFER, writeBuffer );
+			dstBuffer = new Uint8Array( target );
 
 		}
 
-		gl.copyBufferSubData( gl.COPY_READ_BUFFER, gl.COPY_WRITE_BUFFER, 0, 0, byteLength );
+		// Ensure the buffer is bound before reading
+		gl.bindBuffer( gl.COPY_READ_BUFFER, bufferGPU );
+		gl.getBufferSubData( gl.COPY_READ_BUFFER, offset, dstBuffer );
 
-		await backend.utils._clientWaitAsync();
+		gl.bindBuffer( gl.COPY_READ_BUFFER, null );
+		gl.bindBuffer( gl.COPY_WRITE_BUFFER, null );
 
-		const dstBuffer = new attribute.array.constructor( array.length );
+		// return the appropriate type
+		if ( target && target.isReadbackBuffer ) {
 
-		// Ensure the buffer is bound before reading
-		gl.bindBuffer( gl.COPY_WRITE_BUFFER, writeBuffer );
+			return target;
 
-		gl.getBufferSubData( gl.COPY_WRITE_BUFFER, 0, dstBuffer );
+		} else {
 
-		gl.bindBuffer( gl.COPY_READ_BUFFER, null );
-		gl.bindBuffer( gl.COPY_WRITE_BUFFER, null );
+			return dstBuffer.buffer;
 
-		return dstBuffer;
+		}
 
 	}
 

src/renderers/webgpu/WebGPUBackend.js

@@ -325,15 +325,20 @@ class WebGPUBackend extends Backend {
 
 	/**
 	 * This method performs a readback operation by moving buffer data from
-	 * a storage buffer attribute from the GPU to the CPU.
+	 * a storage buffer attribute from the GPU to the CPU. ReadbackBuffer can
+	 * be used to retain and reuse handles to the intermediate buffers and prevent
+	 * new allocation.
 	 *
 	 * @async
-	 * @param {ReadbackBuffer} readbackBuffer - The readback buffer.
-	 * @return {Promise<ArrayBuffer>} A promise that resolves with the buffer data when the data are ready.
+	 * @param {BufferAttribute} attribute - The storage buffer attribute to read frm.
+	 * @param {number} count - The offset from which to start reading the
+	 * @param {number} offset - The storage buffer attribute.
+	 * @param {ReadbackBuffer|ArrayBuffer} target - The storage buffer attribute.
+	 * @return {Promise<ArrayBuffer|ReadbackBuffer>} A promise that resolves with the buffer data when the data are ready.
 	 */
-	async getArrayBufferAsync( readbackBuffer ) {
+	async getArrayBufferAsync( attribute, target = null, offset = 0, count = - 1 ) {
 
-		return await this.attributeUtils.getArrayBufferAsync( readbackBuffer );
+		return await this.attributeUtils.getArrayBufferAsync( attribute, target, offset, count );
 
 	}
 

src/renderers/webgpu/utils/WebGPUAttributeUtils.js

@@ -315,82 +315,137 @@ class WebGPUAttributeUtils {
 
 	/**
 	 * This method performs a readback operation by moving buffer data from
-	 * a storage buffer attribute from the GPU to the CPU.
+	 * a storage buffer attribute from the GPU to the CPU. ReadbackBuffer can
+	 * be used to retain and reuse handles to the intermediate buffers and prevent
+	 * new allocation.
 	 *
 	 * @async
-	 * @param {ReadbackBuffer} readbackBuffer - The storage buffer attribute.
-	 * @return {Promise<ArrayBuffer>} A promise that resolves with the buffer data when the data are ready.
+	 * @param {BufferAttribute} attribute - The storage buffer attribute to read frm.
+	 * @param {number} count - The offset from which to start reading the
+	 * @param {number} offset - The storage buffer attribute.
+	 * @param {ReadbackBuffer|ArrayBuffer} target - The storage buffer attribute.
+	 * @return {Promise<ArrayBuffer|ReadbackBuffer>} A promise that resolves with the buffer data when the data are ready.
 	 */
-	async getArrayBufferAsync( readbackBuffer ) {
+	async getArrayBufferAsync( attribute, target = null, offset = 0, count = - 1 ) {
 
 		const backend = this.backend;
 		const device = backend.device;
-		const attribute = readbackBuffer.attribute;
 
 		const data = backend.get( this._getBufferAttribute( attribute ) );
 		const bufferGPU = data.buffer;
-		const size = bufferGPU.size;
+		const byteLength = count === - 1 ? bufferGPU.size - offset : count;
 
-		const readbackBufferData = backend.get( readbackBuffer );
+		let readBufferGPU;
+		if ( target !== null && target.isReadbackBuffer ) {
 
-		let { readBufferGPU } = readbackBufferData;
+			const readbackInfo = backend.get( target );
 
-		if ( readBufferGPU === undefined ) {
+			if ( target._mapped === true ) {
 
-			readBufferGPU = device.createBuffer( {
-				label: `${ attribute.name }_readback`,
-				size,
-				usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
-			} );
+				throw new Error( 'WebGPURenderer: ReadbackBuffer must be released before being used again.' );
+
+			}
+
+			target._mapped = true;
+
+			// initialize the GPU-side read copy buffer if it is not present
+			if ( readbackInfo.readBufferGPU === undefined ) {
+
+				readBufferGPU = device.createBuffer( {
+					label: `${ target.name }_readback`,
+					size: target.maxByteLength,
+					usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
+				} );
+
+				// release / dispose
+				const releaseCallback = () => {
+
+					target.buffer = null;
+					target._mapped = false;
+
+					readBufferGPU.unmap();
+
+				};
+
+				const disposeCallback = () => {
 
-			// release / dispose
+					target.buffer = null;
+					target._mapped = false;
 
-			const release = () => {
+					readBufferGPU.destroy();
 
-				readBufferGPU.unmap();
+					backend.delete( target );
 
-			};
+					target.removeEventListener( 'release', releaseCallback );
+					target.removeEventListener( 'dispose', disposeCallback );
 
-			const dispose = () => {
+				};
 
-				readBufferGPU.destroy();
+				target.addEventListener( 'release', releaseCallback );
+				target.addEventListener( 'dispose', disposeCallback );
 
-				backend.delete( readbackBuffer );
+				// register
+				readbackInfo.readBufferGPU = readBufferGPU;
 
-				readbackBuffer.removeEventListener( 'release', release );
-				readbackBuffer.removeEventListener( 'dispose', dispose );
+			} else {
 
-			};
+				readBufferGPU = readbackInfo.readBufferGPU;
 
-			readbackBuffer.addEventListener( 'release', release );
-			readbackBuffer.addEventListener( 'dispose', dispose );
+			}
 
-			// register
+		} else {
 
-			readbackBufferData.readBufferGPU = readBufferGPU;
+			// create a new temp buffer for array buffers otherwise
+			readBufferGPU = device.createBuffer( {
+				label: `${ attribute.name }_readback`,
+				size: byteLength,
+				usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
+			} );
 
 		}
 
+		// copy the data
 		const cmdEncoder = device.createCommandEncoder( {
 			label: `readback_encoder_${ attribute.name }`
 		} );
 
 		cmdEncoder.copyBufferToBuffer(
 			bufferGPU,
-			0,
+			offset,
 			readBufferGPU,
 			0,
-			size
+			byteLength,
 		);
 
 		const gpuCommands = cmdEncoder.finish();
 		device.queue.submit( [ gpuCommands ] );
 
-		await readBufferGPU.mapAsync( GPUMapMode.READ );
+		// map the data to the CPU
+		await readBufferGPU.mapAsync( GPUMapMode.READ, 0, byteLength );
 
-		const arrayBuffer = readBufferGPU.getMappedRange();
+		if ( target === null ) {
 
-		return arrayBuffer;
+			// return a new array buffer and clean up the gpu handles
+			const arrayBuffer = readBufferGPU.getMappedRange( 0, byteLength );
+			const result = arrayBuffer.slice();
+			readBufferGPU.destroy();
+			return result;
+
+		} else if ( target.isReadbackBuffer ) {
+
+			// assign the data to the read back handle
+			target.buffer = readBufferGPU.getMappedRange( 0, byteLength );
+			return target;
+
+		} else {
+
+			// copy the data into the target array buffer
+			const arrayBuffer = readBufferGPU.getMappedRange( 0, byteLength );
+			new Uint8Array( target ).set( new Uint8Array( arrayBuffer ) );
+			readBufferGPU.destroy();
+			return target;
+
+		}
 
 	}