Backend-agnostic GPGPU on Nodes

examples/files.json

@@ -413,6 +413,9 @@
 		"misc_exporter_usdz",
 		"misc_lookat"
 	],
+	"misc / nodes": [
+		"misc_nodes_gpgpu"
+	],
 	"css2d": [
 		"css2d_label"
 	],

examples/jsm/gpgpu/ComputationRenderer.js

@@ -0,0 +1,30 @@
+export default class ComputationRenderer {
+
+	constructor( renderer ) {
+
+		this.renderer = renderer;
+
+		this._buffers = [];
+
+	}
+
+	createBuffer( /* ...params */ ) {
+
+		console.warn( 'Abstract function.' );
+		// this._buffers.push( buffer );
+
+	}
+
+	async compute( /* shaderNode, outBuffer, populateTypedArray = true */ ) {
+
+		console.warn( 'Abstract function.' );
+
+	}
+
+	disposeBuffers() {
+
+		this._buffers.forEach( buffer => buffer.dispose() );
+
+	}
+
+}

examples/jsm/gpgpu/TypedBuffer.js

@@ -0,0 +1,136 @@
+import { float, int, uint, vec2, ivec2, uvec2, vec3, ivec3, uvec3, vec4, ivec4, uvec4 } from 'three/nodes';
+
+export default class TypedBuffer {
+
+	constructor( bufferAttribute ) {
+
+		if ( ( bufferAttribute.array instanceof Float64Array ) || ( bufferAttribute.array instanceof BigInt64Array ) || ( bufferAttribute.array instanceof BigUint64Array ) ) {
+
+			throw new Error( 'Float64Array, BigInt64Array, and BigUint64Array are not supported, because float64 and int64 are not supported either in WebGL or WebGPU' );
+
+		}
+
+		this._attribute = bufferAttribute;
+		this._buffer = null;
+
+	}
+
+	get typedArray() {
+
+		return this._attribute.array;
+
+	}
+
+	set typedArray( typedArray ) {
+
+		this._attribute.array.set( new this._attribute.array.constructor( typedArray ) );
+
+	}
+
+	get buffer() {
+
+		return this._buffer;
+
+	}
+
+	set buffer( value ) {
+
+		throw new Error( 'GPU buffer of a TypedBuffer cannot be changed' );
+
+	}
+
+	get arrayLength() {
+
+		return this._attribute.array.length;
+
+	}
+
+	get length() {
+
+		return this._attribute.count;
+
+	}
+
+	getBufferElement( /* i */ ) {
+
+		console.warn( 'Abstract function.' );
+
+	}
+
+	setBufferElement( /* i, value */ ) {
+
+		console.warn( 'Abstract function.' );
+
+	}
+
+	set needsUpdate( value ) {
+
+		this._buffer.needsUpdate = value;
+
+	}
+
+	dispose() {
+
+		this._buffer.dispose();
+
+	}
+
+}
+
+export function getFunction( bufferAttribute ) {
+
+	const array = bufferAttribute.array;
+
+	let functionType;
+
+	if ( ( array instanceof Int8Array ) || ( array instanceof Int16Array ) || ( array instanceof Int32Array ) ) {
+
+		functionType = 'int';
+
+	} else if ( ( array instanceof Uint8Array ) || ( array instanceof Uint8ClampedArray ) || ( array instanceof Uint16Array ) || ( array instanceof Uint32Array ) ) {
+
+		functionType = 'uint';
+
+	} else if ( array instanceof Float32Array ) {
+
+		functionType = 'float';
+
+	}
+
+	switch ( bufferAttribute.itemSize ) {
+
+		case 1:
+			return ( functionType === 'uint' ) ? uint : ( functionType === 'int' ) ? int : float;
+
+		case 2:
+			return ( functionType === 'uint' ) ? uvec2 : ( functionType === 'int' ) ? ivec2 : vec2;
+
+		case 3:
+			return ( functionType === 'uint' ) ? uvec3 : ( functionType === 'int' ) ? ivec3 : vec3;
+
+		case 4:
+			return ( functionType === 'uint' ) ? uvec4 : ( functionType === 'int' ) ? ivec4 : vec4;
+
+	}
+
+}
+
+export function getFloatFunction( bufferAttribute ) {
+
+	switch ( bufferAttribute.itemSize ) {
+
+		case 1:
+			return float;
+
+		case 2:
+			return vec2;
+
+		case 3:
+			return vec3;
+
+		case 4:
+			return vec4;
+
+	}
+
+}

examples/jsm/gpgpu/WebGLComputationRenderer.js

@@ -0,0 +1,63 @@
+import { Camera, Mesh, PlaneGeometry, Scene, WebGLRenderTarget } from 'three';
+import { int, viewportCoordinate, MeshBasicNodeMaterial, ShaderNode } from 'three/nodes';
+import { nodeFrame } from 'three/addons/renderers/webgl/nodes/WebGLNodes.js';
+import ComputationRenderer from './ComputationRenderer.js';
+import WebGLTypedBuffer from './WebGLTypedBuffer.js';
+
+export default class WebGLComputationRenderer extends ComputationRenderer {
+
+	constructor( renderer ) {
+
+		super( renderer );
+
+		this._material = new MeshBasicNodeMaterial();
+		this._scene = new Scene().add( new Mesh( new PlaneGeometry( 2, 2 ), this._material ) );
+		this._camera = new Camera();
+
+	}
+
+	createBuffer( ...params ) {
+
+		const buffer = new WebGLTypedBuffer( ...params );
+		this._buffers.push( buffer );
+		return buffer;
+
+	}
+
+	async compute( shaderNode, outBuffer, populateTypedArray = true ) {
+
+		nodeFrame.update();
+
+		const outGPUBuffer = outBuffer.buffer;
+		outGPUBuffer.isRenderTargetTexture = true;
+		const outTypedArray = outBuffer.typedArray;
+
+		const renderTarget = new WebGLRenderTarget( outGPUBuffer.image.width, outGPUBuffer.image.height, { depthBuffer: false } );
+		renderTarget.texture = outGPUBuffer;
+
+		const index = int( viewportCoordinate.y ).mul( outGPUBuffer.image.width ).add( int( viewportCoordinate.x ) );
+		this._material.colorNode = new ShaderNode( ( inputs, builder ) => {
+
+			return outBuffer.setBufferElement( index, shaderNode.call( { index }, builder ) );
+
+		} );
+		this._material.needsUpdate = true;
+
+		const currentRenderTarget = this.renderer.getRenderTarget();
+		this.renderer.setRenderTarget( renderTarget );
+		this.renderer.render( this._scene, this._camera );
+		if ( populateTypedArray ) {
+
+			this.renderer.readRenderTargetPixels( renderTarget, 0, 0, renderTarget.width, renderTarget.height, outTypedArray );
+							// The .render call populates the GPU buffer, the .readRenderTargetPixels call populates the typed array
+
+		} else {
+
+			outBuffer.typedArray = outTypedArray.length; // null array
+
+		}
+		this.renderer.setRenderTarget( currentRenderTarget );
+
+	}
+
+}

examples/jsm/gpgpu/WebGLTypedBuffer.js

@@ -0,0 +1,136 @@
+import {
+	DataTexture,
+
+	UnsignedByteType,
+	ByteType,
+	ShortType,
+	UnsignedShortType,
+	IntType,
+	UnsignedIntType,
+	FloatType,
+
+	RedFormat,
+	RGFormat,
+	RGBAFormat
+} from 'three';
+import { texture, add, div, remainder, floor, vec2, mul, float, int } from 'three/nodes';
+import TypedBuffer, { getFunction, getFloatFunction } from './TypedBuffer.js';
+
+function getTextureElement( dataTexture, i, width, height ) {
+
+	const x = div( add( 0.5, remainder( i, width ) ), width );
+	const y = div( add( 0.5,       div( i, width ) ), height );
+
+	return texture( dataTexture, vec2( x, y ) );
+
+}
+
+function getTextureType( bufferAttribute ) {
+
+	const array = bufferAttribute.array;
+
+	let textureType;
+
+	if ( array instanceof Int8Array ) {
+
+		textureType = ByteType;
+
+	} else if ( ( array instanceof Uint8Array ) || ( array instanceof Uint8ClampedArray ) ) {
+
+		textureType = UnsignedByteType;
+
+	} else if ( array instanceof Int16Array ) {
+
+		textureType = ShortType;
+
+	} else if ( array instanceof Uint16Array ) {
+
+		textureType = UnsignedShortType;
+
+	} else if ( array instanceof Int32Array ) {
+
+		textureType = IntType;
+
+	} else if ( array instanceof Uint32Array ) {
+
+		textureType = UnsignedIntType;
+
+	} else if ( array instanceof Float32Array ) {
+
+		textureType = FloatType;
+
+	}
+
+	return textureType;
+
+}
+
+export function getTextureFormat( bufferAttribute ) { // @TODO: possibly support impossible type-format combinations by using padding
+
+	switch ( bufferAttribute.itemSize ) {
+
+		case 1:
+			return RedFormat;
+
+		case 2:
+			return RGFormat;
+
+		case 3:
+			return null;
+
+		case 4:
+			return RGBAFormat;
+
+	}
+
+}
+
+export function calculateWidthHeight( length ) {
+
+	let width; // @TODO: maybe just set width and height to Math.sqrt( length ) and pad the texture with zeroes?
+	for ( width = Math.floor( Math.sqrt( length ) ); width > 0; width ++ ) {
+
+		if ( length % width === 0 ) break;
+
+	}
+	const height = length / width;
+
+	return { width, height };
+
+}
+
+export default class WebGLTypedBuffer extends TypedBuffer {
+
+	constructor( bufferAttribute ) {
+
+		// @TODO: add support for UBOs
+
+		super( bufferAttribute );
+
+		const { width, height } = calculateWidthHeight( this.length );
+		this.width = width;
+		this.height = height;
+
+		const buffer = this._buffer = new DataTexture( this.typedArray, width, height, getTextureFormat( bufferAttribute ), getTextureType( bufferAttribute ) );
+		buffer.needsUpdate = true;
+
+		this._function = getFunction( bufferAttribute );
+		this._floatFunction = getFloatFunction( bufferAttribute );
+
+	}
+
+	getBufferElement( i ) {
+
+		const gpuBuffer = this._buffer;
+		const textureElement = getTextureElement( gpuBuffer, int( i ), gpuBuffer.image.width, gpuBuffer.image.height );
+		return this._function( mul( textureElement, 255 ) );
+
+	}
+
+	setBufferElement( i, value ) {
+
+		return div( this._floatFunction( value ), 255 );
+
+	}
+
+}