WebGPURenderer: Introduce ReadbackBuffer and reuse getArrayBufferAsync() buffers (#33300)

Author: sunag

examples/jsm/inspector/tabs/Memory.js

@@ -58,6 +58,9 @@ class Memory extends Tab {
 		this.programs = new Item( 'Programs', createValueSpan(), createValueSpan() );
 		this.memoryStats.add( this.programs );
 
+		this.readbackBuffers = new Item( 'Readback Buffers', createValueSpan(), createValueSpan() );
+		this.memoryStats.add( this.readbackBuffers );
+
 		this.renderTargets = new Item( 'Render Targets', createValueSpan(), 'N/A' );
 		this.memoryStats.add( this.renderTargets );
 
@@ -108,6 +111,9 @@ class Memory extends Tab {
 		setText( this.programs.data[ 1 ], memory.programs.toString() );
 		setText( this.programs.data[ 2 ], formatBytes( memory.programsSize ) );
 
+		setText( this.readbackBuffers.data[ 1 ], memory.readbackBuffers.toString() );
+		setText( this.readbackBuffers.data[ 2 ], formatBytes( memory.readbackBuffersSize ) );
+
 		setText( this.renderTargets.data[ 1 ], memory.renderTargets.toString() );
 
 		setText( this.storageAttributes.data[ 1 ], memory.storageAttributes.toString() );

examples/webgpu_compute_reduce.html

@@ -963,7 +963,15 @@ <h3 id="panel-title" style="flex: 0 0 auto;">Subgroup Reduction Explanation</h3>
 				functionObj[ logFunctionName ] = async() => {
 
 					const selectedBuffer = buffers[ unifiedEffectController.loggedBuffer ];
-					console.log( new Uint32Array( await renderer.getArrayBufferAsync( selectedBuffer.value ) ) );
+					const readbackBuffer = new THREE.ReadbackBuffer( selectedBuffer.value );
+
+					const result = new Uint32Array( await renderer.getArrayBufferAsync( readbackBuffer ) );
+
+					console.log( result );
+
+					// Remove GPU/CPU readback buffer from memory
+
+					readbackBuffer.dispose();
 
 				};
 

src/Three.WebGPU.Nodes.js

@@ -10,6 +10,7 @@ export { default as QuadMesh } from './renderers/common/QuadMesh.js';
 export { default as PMREMGenerator } from './renderers/common/extras/PMREMGenerator.js';
 export { default as RenderPipeline } from './renderers/common/RenderPipeline.js';
 export { default as PostProcessing } from './renderers/common/PostProcessing.js';
+export { default as ReadbackBuffer } from './renderers/common/ReadbackBuffer.js';
 import * as RendererUtils from './renderers/common/RendererUtils.js';
 export { RendererUtils };
 export { default as StorageTexture } from './renderers/common/StorageTexture.js';

src/Three.WebGPU.js

@@ -11,6 +11,7 @@ export { default as QuadMesh } from './renderers/common/QuadMesh.js';
 export { default as PMREMGenerator } from './renderers/common/extras/PMREMGenerator.js';
 export { default as RenderPipeline } from './renderers/common/RenderPipeline.js';
 export { default as PostProcessing } from './renderers/common/PostProcessing.js';
+export { default as ReadbackBuffer } from './renderers/common/ReadbackBuffer.js';
 import * as RendererUtils from './renderers/common/RendererUtils.js';
 export { RendererUtils };
 export { default as CubeRenderTarget } from './renderers/common/CubeRenderTarget.js';

src/core/BufferAttribute.js

@@ -3,6 +3,7 @@ import { Vector2 } from '../math/Vector2.js';
 import { denormalize, normalize } from '../math/MathUtils.js';
 import { StaticDrawUsage, FloatType } from '../constants.js';
 import { fromHalfFloat, toHalfFloat } from '../extras/DataUtils.js';
+import { EventDispatcher } from './EventDispatcher.js';
 
 const _vector = /*@__PURE__*/ new Vector3();
 const _vector2 = /*@__PURE__*/ new Vector2();
@@ -17,7 +18,7 @@ let _id = 0;
  * When working with vector-like data, the `fromBufferAttribute( attribute, index )`
  * helper methods on vector and color class might be helpful. E.g. {@link Vector3#fromBufferAttribute}.
  */
-class BufferAttribute {
+class BufferAttribute extends EventDispatcher {
 
 	/**
 	 * Constructs a new buffer attribute.
@@ -28,6 +29,8 @@ class BufferAttribute {
 	 */
 	constructor( array, itemSize, normalized = false ) {
 
+		super();
+
 		if ( Array.isArray( array ) ) {
 
 			throw new TypeError( 'THREE.BufferAttribute: array should be a Typed Array.' );
@@ -674,6 +677,15 @@ class BufferAttribute {
 
 	}
 
+	/**
+	 * Disposes of the buffer attribute. Available only in {@link WebGPURenderer}.
+	 */
+	dispose() {
+
+		this.dispatchEvent( { type: 'dispose' } );
+
+	}
+
 }
 
 /**

src/renderers/common/Info.js

@@ -100,6 +100,7 @@ class Info {
 		 * @property {number} indexAttributes - The number of active index attributes.
 		 * @property {number} storageAttributes - The number of active storage attributes.
 		 * @property {number} indirectStorageAttributes - The number of active indirect storage attributes.
+		 * @property {number} readbackBuffers - The number of active readback buffers.
 		 * @property {number} programs - The number of active programs.
 		 * @property {number} renderTargets - The number of active renderTargets.
 		 * @property {number} total - The total memory size in bytes.
@@ -108,6 +109,7 @@ class Info {
 		 * @property {number} indexAttributesSize - The memory size of active index attributes in bytes.
 		 * @property {number} storageAttributesSize - The memory size of active storage attributes in bytes.
 		 * @property {number} indirectStorageAttributesSize - The memory size of active indirect storage attributes in bytes.
+		 * @property {number} readbackBuffersSize - The memory size of active readback buffers in bytes.
 		 * @property {number} programsSize - The memory size of active programs in bytes.
 		 */
 		this.memory = {
@@ -117,6 +119,7 @@ class Info {
 			indexAttributes: 0,
 			storageAttributes: 0,
 			indirectStorageAttributes: 0,
+			readbackBuffers: 0,
 			programs: 0,
 			renderTargets: 0,
 			total: 0,
@@ -125,6 +128,7 @@ class Info {
 			indexAttributesSize: 0,
 			storageAttributesSize: 0,
 			indirectStorageAttributesSize: 0,
+			readbackBuffersSize: 0,
 			programsSize: 0
 		};
 
@@ -329,6 +333,33 @@ class Info {
 
 	}
 
+	/**
+	 * Tracks a readback buffer memory explicitly.
+	 *
+	 * @param {ReadbackBuffer} readbackBuffer - The readback buffer to track.
+	 */
+	createReadbackBuffer( readbackBuffer ) {
+
+		const size = this._getAttributeMemorySize( readbackBuffer.attribute );
+		this.memoryMap.set( readbackBuffer, { size, type: 'readbackBuffers' } );
+
+		this.memory.readbackBuffers ++;
+		this.memory.total += size;
+		this.memory.readbackBuffersSize += size;
+
+	}
+
+	/**
+	 * Tracks a readback buffer memory explicitly.
+	 *
+	 * @param {ReadbackBuffer} readbackBuffer - The readback buffer to track.
+	 */
+	destroyReadbackBuffer( readbackBuffer ) {
+
+		this.destroyAttribute( readbackBuffer );
+
+	}
+
 	/**
 	 * Tracks program memory explicitly, updating counts and byte tracking.
 	 *

src/renderers/common/ReadbackBuffer.js

@@ -0,0 +1,62 @@
+import { EventDispatcher } from '../../core/EventDispatcher.js';
+
+/**
+ * A readback buffer is used to transfer data from the GPU to the CPU.
+ * It is primarily used to read back compute shader results.
+ *
+ * @augments EventDispatcher
+ */
+class ReadbackBuffer extends EventDispatcher {
+
+	/**
+	 * Constructs a new readback buffer.
+	 *
+	 * @param {BufferAttribute} attribute - The buffer attribute.
+	 */
+	constructor( attribute ) {
+
+		super();
+
+		/**
+		 * The buffer attribute.
+		 *
+		 * @type {BufferAttribute}
+		 */
+		this.attribute = attribute;
+
+		/**
+		 * This flag can be used for type testing.
+		 *
+		 * @type {boolean}
+		 * @readonly
+		 * @default true
+		 */
+		this.isReadbackBuffer = true;
+
+	}
+
+	/**
+	 * Releases the mapped buffer data so the GPU buffer can be
+	 * used by the GPU again.
+	 *
+	 * Note: Any `ArrayBuffer` data associated with this readback buffer
+	 * are removed and no longer accessible after calling this method.
+	 */
+	release() {
+
+		this.dispatchEvent( { type: 'release' } );
+
+	}
+
+	/**
+	 * Frees internal resources.
+	 */
+	dispose() {
+
+		this.dispatchEvent( { type: 'dispose' } );
+
+	}
+
+}
+
+export default ReadbackBuffer;

src/renderers/common/Renderer.js

@@ -19,6 +19,7 @@ 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';
 
@@ -1918,12 +1919,61 @@ class Renderer {
 	 * from the GPU to the CPU in context of compute shaders.
 	 *
 	 * @async
-	 * @param {StorageBufferAttribute} attribute - The storage buffer attribute.
+	 * @param {StorageBufferAttribute|ReadbackBuffer} buffer - The storage buffer attribute.
 	 * @return {Promise<ArrayBuffer>} A promise that resolves with the buffer data when the data are ready.
 	 */
-	async getArrayBufferAsync( attribute ) {
+	async getArrayBufferAsync( buffer ) {
 
-		return await this.backend.getArrayBufferAsync( attribute );
+		let readbackBuffer = buffer;
+
+		if ( readbackBuffer.isReadbackBuffer !== true ) {
+
+			const attribute = buffer;
+			const attributeData = this.backend.get( attribute );
+
+			readbackBuffer = attributeData.readbackBuffer;
+
+			if ( readbackBuffer === undefined ) {
+
+				readbackBuffer = new ReadbackBuffer( attribute );
+
+				const dispose = () => {
+
+					attribute.removeEventListener( 'dispose', dispose );
+
+					readbackBuffer.dispose();
+
+					delete attributeData.readbackBuffer;
+
+				};
+
+				attribute.addEventListener( 'dispose', dispose );
+
+				attributeData.readbackBuffer = readbackBuffer;
+
+			}
+
+		}
+
+		if ( this.info.memoryMap.has( readbackBuffer ) === false ) {
+
+			this.info.createReadbackBuffer( readbackBuffer );
+
+			const disposeInfo = () => {
+
+				readbackBuffer.removeEventListener( 'dispose', disposeInfo );
+
+				this.info.destroyReadbackBuffer( readbackBuffer );
+
+			};
+
+			readbackBuffer.addEventListener( 'dispose', disposeInfo );
+
+		}
+
+		readbackBuffer.release();
+
+		return await this.backend.getArrayBufferAsync( readbackBuffer );
 
 	}
 

src/renderers/webgl-fallback/WebGLBackend.js

@@ -312,12 +312,12 @@ class WebGLBackend extends Backend {
 	 * a storage buffer attribute from the GPU to the CPU.
 	 *
 	 * @async
-	 * @param {StorageBufferAttribute} attribute - The storage buffer attribute.
+	 * @param {ReadbackBuffer} readbackBuffer - The readback buffer.
 	 * @return {Promise<ArrayBuffer>} A promise that resolves with the buffer data when the data are ready.
 	 */
-	async getArrayBufferAsync( attribute ) {
+	async getArrayBufferAsync( readbackBuffer ) {
 
-		return await this.attributeUtils.getArrayBufferAsync( attribute );
+		return await this.attributeUtils.getArrayBufferAsync( readbackBuffer );
 
 	}
 

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

@@ -257,14 +257,15 @@ class WebGLAttributeUtils {
 	 * a storage buffer attribute from the GPU to the CPU.
 	 *
 	 * @async
-	 * @param {StorageBufferAttribute} attribute - The storage buffer attribute.
+	 * @param {ReadbackBuffer} readbackBuffer - The readback buffer.
 	 * @return {Promise<ArrayBuffer>} A promise that resolves with the buffer data when the data are ready.
 	 */
-	async getArrayBufferAsync( attribute ) {
+	async getArrayBufferAsync( readbackBuffer ) {
 
 		const backend = this.backend;
 		const { gl } = backend;
 
+		const attribute = readbackBuffer.attribute;
 		const bufferAttribute = attribute.isInterleavedBufferAttribute ? attribute.data : attribute;
 		const { bufferGPU } = backend.get( bufferAttribute );
 
@@ -273,10 +274,40 @@ class WebGLAttributeUtils {
 
 		gl.bindBuffer( gl.COPY_READ_BUFFER, bufferGPU );
 
-		const writeBuffer = gl.createBuffer();
+		const readbackBufferData = backend.get( readbackBuffer );
 
-		gl.bindBuffer( gl.COPY_WRITE_BUFFER, writeBuffer );
-		gl.bufferData( gl.COPY_WRITE_BUFFER, byteLength, gl.STREAM_READ );
+		let { writeBuffer } = readbackBufferData;
+
+		if ( writeBuffer === undefined ) {
+
+			writeBuffer = gl.createBuffer();
+
+			gl.bindBuffer( gl.COPY_WRITE_BUFFER, writeBuffer );
+			gl.bufferData( gl.COPY_WRITE_BUFFER, byteLength, gl.STREAM_READ );
+
+			// dispose
+
+			const dispose = () => {
+
+				gl.deleteBuffer( writeBuffer );
+
+				backend.delete( readbackBuffer );
+
+				readbackBuffer.removeEventListener( 'dispose', dispose );
+
+			};
+
+			readbackBuffer.addEventListener( 'dispose', dispose );
+
+			// register
+
+			readbackBufferData.writeBuffer = writeBuffer;
+
+		} else {
+
+			gl.bindBuffer( gl.COPY_WRITE_BUFFER, writeBuffer );
+
+		}
 
 		gl.copyBufferSubData( gl.COPY_READ_BUFFER, gl.COPY_WRITE_BUFFER, 0, 0, byteLength );
 
@@ -289,12 +320,10 @@ class WebGLAttributeUtils {
 
 		gl.getBufferSubData( gl.COPY_WRITE_BUFFER, 0, dstBuffer );
 
-		gl.deleteBuffer( writeBuffer );
-
 		gl.bindBuffer( gl.COPY_READ_BUFFER, null );
 		gl.bindBuffer( gl.COPY_WRITE_BUFFER, null );
 
-		return dstBuffer.buffer;
+		return dstBuffer;
 
 	}