three-gpu-pathtracer 0.0.23 → 0.0.25

package/src/webgpu/compute/wavefront/structs.js

@@ -0,0 +1,155 @@
+import { StructTypeNode } from 'three/webgpu';
+
+// StructTypeNode that force-emits dependency structs. TSL can't discover a struct referenced only via
+// a type string, so we build each dependency first, which registers its definition ahead of this struct.
+// TODO: remove this once TSL resolves struct dependencies from string member types itself.
+class DependentStructTypeNode extends StructTypeNode {
+
+	constructor( members, name, dependencies = [] ) {
+
+		super( members, name );
+		this.dependencies = dependencies;
+
+	}
+
+	setup( builder ) {
+
+		for ( const dep of this.dependencies ) {
+
+			dep.build( builder );
+
+		}
+
+		super.setup( builder );
+
+	}
+
+}
+
+// Persistent per-path-slot state. One slot per in-flight path, holding everything needed to resolve
+// the previous frame's trace results and stage the next bounce.
+export const rayDataStruct = new StructTypeNode( {
+
+	throughputColor: 'vec3f',
+	currentBounce: 'uint',
+
+	resultColor: 'vec4f',
+
+	emission: 'vec3f',
+	scatterPdf: 'float',
+
+	scatterColor: 'vec3f',
+
+	// the smallest scatter pdf seen along the path, for the glossy filter
+	minPdf: 'float',
+
+	origin: 'vec3f',
+
+	// whether every scatter so far has been transmissive, for the background modes
+	isFullyTransmissive: 'uint',
+
+	direction: 'vec3f',
+	side: 'float',
+
+	normal: 'vec3f',
+	objectIndex: 'int',
+
+	barycoord: 'vec3f',
+	pixelIndex: 'uint',
+
+	indices: 'vec3u',
+	seed: 'uint',
+
+	lightDirection: 'vec3f',
+	lightPdf: 'float',
+
+	lightEmission: 'vec3f',
+	lightDist: 'float',
+
+	lightBsdf: 'vec3f',
+	lightBsdfPdf: 'float',
+
+	rayIntersectionIndex: 'int',
+	shadowRayIntersectionIndex: 'int',
+	lightType: 'int',
+
+	// the traced segment length, for the transmission attenuation applied by MaterialKernel
+	dist: 'float',
+
+	// alpha test pass throughs, counted separately from the bounce count
+	alphaDepth: 'uint',
+
+	// the camera segment's maximum trace distance, carried across alpha pass throughs
+	maxDist: 'float',
+
+	// the path's hero wavelength, negative until its reconstruction weight is applied at the
+	// first dispersive interaction
+	dispersionWavelength: 'float',
+
+}, 'RayData' );
+
+// A ray queued for BVH traversal by the trace kernels. A "maxDist" of zero traces unbounded.
+export const traceQueuedRayStruct = new StructTypeNode( {
+
+	origin: 'vec3f',
+	pixelIndex: 'uint',
+
+	direction: 'vec3f',
+	currentBounce: 'uint',
+
+	seed: 'uint',
+	alphaDepth: 'uint',
+	maxDist: 'float',
+	_alignment0: 'uint',
+
+}, 'TraceQueuedRay' );
+
+// Compact trace result, written by the trace kernels at the ray's queue index and consumed by
+// LogicKernel the following frame. objectIndex < 0 encodes a miss.
+export const intersectionResultStruct = new StructTypeNode( {
+
+	barycoord: 'vec3f',
+	objectIndex: 'int',
+
+	position: 'vec3f',
+	dist: 'float',
+
+	normal: 'vec3f',
+	side: 'float',
+
+	indices: 'vec3u',
+	_alignment0: 'uint',
+
+}, 'TraceResult' );
+
+// Queue wrappers that keep an append-only length counter in a header ahead of the elements. The
+// atomic variant is bound where rays are pushed with atomicAdd; the plain one where the length is
+// only read or reset. getLength returns just the header since the trailing array is runtime-sized.
+export const rayQueueStruct = new DependentStructTypeNode( {
+	length: 'uint',
+	elements: `array<${ traceQueuedRayStruct.name }>`,
+}, 'RayQueue', [ traceQueuedRayStruct ] );
+rayQueueStruct.getLength = () => 4;
+
+export const rayQueueAtomicStruct = new DependentStructTypeNode( {
+	length: { type: 'uint', atomic: true },
+	elements: `array<${ traceQueuedRayStruct.name }>`,
+}, 'RayQueue', [ traceQueuedRayStruct ] );
+rayQueueAtomicStruct.getLength = () => 4;
+
+// Round-robin queue of pixel indices waiting for a free path slot when the output resolution exceeds
+// the ray data pool. Sized to the pixel count so retiring slots can always return their pixel. The
+// non-atomic variant is used for contention-free initialization.
+export const pixelQueueStruct = new DependentStructTypeNode( {
+	current: { type: 'uint', atomic: true },
+	elementCount: { type: 'uint', atomic: true },
+	elements: 'array<atomic<u32>>',
+}, 'PixelQueue' );
+pixelQueueStruct.getLength = () => 2;
+
+export const pixelQueueNonAtomicStruct = new DependentStructTypeNode( {
+	current: 'uint',
+	elementCount: 'uint',
+	elements: 'array<u32>',
+}, 'PixelQueue' );
+pixelQueueNonAtomicStruct.getLength = () => 2;

package/src/webgpu/constants.js

@@ -0,0 +1,62 @@
+import * as pcg from './nodes/rand/pcg.wgsl.js';
+import * as sobol from './nodes/rand/sobol.wgsl.js';
+import * as bluedither from './nodes/rand/bluedither.wgsl.js';
+
+// Sampling strategies the kernels draw random numbers from. Pass one to
+// "WebGPUPathTracer.setRandom".
+
+/**
+ * Uncorrelated pseudo random samples. Converges slowest but never shows structured artifacts.
+ * @section Random Strategies
+ * @type {Object}
+ */
+export const RANDOM_PCG = pcg;
+
+/**
+ * Sobol low discrepancy samples, scrambled per pixel. Converges faster than `RANDOM_PCG`.
+ * @section Random Strategies
+ * @type {Object}
+ */
+export const RANDOM_SOBOL = sobol;
+
+/**
+ * The same Sobol sequence for every pixel, offset by a blue noise dither so early frames read as
+ * smooth grain rather than clumped noise. The default.
+ * @section Random Strategies
+ * @type {Object}
+ */
+export const RANDOM_BLUE_DITHER = bluedither;
+
+// How rays that have only passed through transmissive surfaces treat the background on a miss.
+// Assign one to "WebGPUPathTracer.transmissiveBackground".
+
+/**
+ * Sample the environment map, so glass shows what is behind it.
+ * @section Transmissive Background Modes
+ * @type {number}
+ */
+export const TRANSMISSIVE_BACKGROUND_ENVIRONMENT = 0;
+
+/**
+ * Set the opacity from the transmitted light intensity, tinted by the environment lighting.
+ * @section Transmissive Background Modes
+ * @type {number}
+ */
+export const TRANSMISSIVE_BACKGROUND_OVERLAY = 1;
+
+/**
+ * Attenuate the background by the transmitted light, so glass renders against transparency.
+ * @section Transmissive Background Modes
+ * @type {number}
+ */
+export const TRANSMISSIVE_BACKGROUND_TRANSPARENT = 2;
+
+// Layout of the r32uint "sample count" target: two flag bits then a 30 bit count.
+
+// this pixel already has a ray on the queue
+export const SAMPLE_ACTIVE_FLAG = 0x80000000;
+
+// this pixel has had a camera ray at least once
+export const SAMPLE_DISPATCHED_FLAG = 0x40000000;
+
+export const SAMPLE_COUNT_MASK = 0x3FFFFFFF;

package/src/webgpu/denoise/OIDNDenoiser.js

@@ -0,0 +1,431 @@
+import { ExternalTexture, NearestFilter, NoToneMapping, RenderTarget, UnsignedByteType } from 'three/webgpu';
+import { diffuseColor, mrt, normalView, vec4 } from 'three/tsl';
+
+/** @import { Camera, Scene, Texture, WebGPURenderer } from 'three/webgpu' */
+/** @import { DynamicTileSetting } from 'oidn-web' */
+
+/**
+ * Runs Open Image Denoise over a path traced image. Pass one to
+ * "WebGPUPathTracer.setDenoiser", or drive it directly with "denoise".
+ *
+ * "initUNetFromURL" and the weights are passed in rather than imported so neither the library
+ * nor the network files become a dependency.
+ *
+ * ```js
+ * import { initUNetFromURL } from 'oidn-web';
+ * pathTracer.setDenoiser( new OIDNDenoiser( { initUNetFromURL, auxWeightsUrl } ) );
+ * ```
+ *
+ * Weights come from the oidn-weights repository, where the "_small" and "_large" variants trade
+ * quality against download size and per tile cost.
+ */
+export class OIDNDenoiser {
+
+	/**
+	 * The denoised result, or null until the first tile has been produced.
+	 *
+	 * @type {ExternalTexture|null}
+	 */
+	get texture() {
+
+		return this._texture;
+
+	}
+
+	/**
+	 * Whether a pass has finished. Stays true until reset.
+	 *
+	 * @type {boolean}
+	 */
+	get complete() {
+
+		return this._complete;
+
+	}
+
+	/**
+	 * Whether a pass is running. The work is spread over several frames.
+	 *
+	 * @type {boolean}
+	 */
+	get running() {
+
+		return this._running;
+
+	}
+
+	/**
+	 * Every field below can also be assigned after construction.
+	 *
+	 * @param {Object} options
+	 * @param {Function} options.initUNetFromURL
+	 * @param {string} options.auxWeightsUrl - Weights for the guided model.
+	 * @param {string} [options.colorWeightsUrl] - Weights for the color only model, needed only
+	 * when `useAuxiliaryBuffers` is `false`.
+	 * @param {boolean} [options.useAuxiliaryBuffers]
+	 * @param {number|null} [options.maxTileSize]
+	 * @param {DynamicTileSetting|null} [options.dynamicTile] - `false` pins every tile to
+	 * `maxTileSize`. An object tunes the adaptive sizing.
+	 */
+	constructor( options = {} ) {
+
+		const {
+			initUNetFromURL,
+			auxWeightsUrl,
+			colorWeightsUrl,
+			useAuxiliaryBuffers = true,
+			maxTileSize = null,
+			dynamicTile = null,
+		} = options;
+
+		if ( ! initUNetFromURL ) {
+
+			throw new Error( 'OIDNDenoiser: "initUNetFromURL" from "oidn-web" must be provided.' );
+
+		}
+
+		this.initUNetFromURL = initUNetFromURL;
+		this.auxWeightsUrl = auxWeightsUrl;
+		this.colorWeightsUrl = colorWeightsUrl;
+
+		// `false` falls back to the color only model, which is blurrier but skips a scene render
+		this.useAuxiliaryBuffers = useAuxiliaryBuffers;
+
+		// tiling, passed through to oidn-web. Null keeps its defaults
+		this.maxTileSize = maxTileSize;
+		this.dynamicTile = dynamicTile;
+
+		this.renderer = null;
+		this.scene = null;
+		this.camera = null;
+
+		this._texture = null;
+		this._running = false;
+		this._complete = false;
+		this._abort = null;
+		this._requestId = 0;
+		this._unets = { aux: null, color: null };
+
+		this._resultPipeline = null;
+		this._rawTexture = null;
+		this._auxTarget = null;
+
+	}
+
+	/**
+	 * @param {WebGPURenderer} renderer
+	 */
+	init( renderer ) {
+
+		this.renderer = renderer;
+
+	}
+
+	/**
+	 * @param {Scene} scene
+	 * @param {Camera} camera
+	 */
+	setScene( scene, camera ) {
+
+		this.scene = scene;
+		this.camera = camera;
+
+	}
+
+	/**
+	 * Renders the auxiliary buffers and starts a pass. Safe to call every frame.
+	 *
+	 * @param {Texture} target - The path traced result, in linear HDR.
+	 * @returns {Texture|null}
+	 */
+	update( target ) {
+
+		if ( this._complete || this._running ) {
+
+			return this._texture;
+
+		}
+
+		let albedo = null;
+		let normal = null;
+		if ( this.useAuxiliaryBuffers ) {
+
+			const auxTarget = this._renderAuxiliaryBuffers( target.width, target.height );
+			albedo = auxTarget.textures[ 0 ];
+			normal = auxTarget.textures[ 1 ];
+
+		}
+
+		this.denoise( target, albedo, normal );
+
+		return this._texture;
+
+	}
+
+	/**
+	 * Runs a pass unless one is already running or finished. The optional albedo and normal
+	 * buffers guide the filter and select the guided model. Both hold [0,1] values, with normals
+	 * mapped so a flat normal is (0.5, 0.5, 1).
+	 *
+	 * @param {Texture} color - The path traced result, in linear HDR.
+	 * @param {Texture|null} albedo
+	 * @param {Texture|null} normal
+	 */
+	async denoise( color, albedo = null, normal = null ) {
+
+		if ( this._complete || this._running ) {
+
+			return;
+
+		}
+
+		this._running = true;
+
+		const useAux = Boolean( albedo && normal );
+		const requestId = ++ this._requestId;
+
+		let unet;
+		try {
+
+			unet = await this._initUNet( useAux );
+
+		} catch ( error ) {
+
+			// leaving "running" set would stall every later pass
+			this._running = false;
+			throw error;
+
+		}
+
+		// bail if a reset or a newer call took over while the weights downloaded
+		if ( requestId !== this._requestId || ! this._running ) {
+
+			return;
+
+		}
+
+		// three only holds a WebGPU texture for a three texture once it has been initialized
+		const { renderer } = this;
+		const backend = renderer.backend;
+		renderer.initTexture( color );
+
+		const { width, height } = color;
+		const inputs = {
+			color: { data: backend.get( color ).texture, width, height },
+
+			// the output is seeded with the raw color, so copying per tile shows a progressive wipe
+			progress: output => this._copyOutputToTexture( output, width, height ),
+			done: output => {
+
+				this._copyOutputToTexture( output, width, height );
+				this._running = false;
+				this._complete = true;
+				this._abort = null;
+
+			},
+		};
+
+		// the guided model needs both buffers, and the color only model must not be given them
+		if ( useAux ) {
+
+			renderer.initTexture( albedo );
+			renderer.initTexture( normal );
+			inputs.albedo = { data: backend.get( albedo ).texture, width, height };
+			inputs.normal = { data: backend.get( normal ).texture, width, height };
+
+		}
+
+		this._abort = unet.tileExecute( inputs );
+
+	}
+
+	/**
+	 * Drops any running pass and the last result. Call it whenever the image changes, such as when
+	 * the camera moves.
+	 */
+	reset() {
+
+		this._abort?.();
+		this._abort = null;
+
+		this._texture?.dispose();
+		this._texture = null;
+
+		this._rawTexture?.destroy();
+		this._rawTexture = null;
+
+		this._running = false;
+		this._complete = false;
+
+	}
+
+	dispose() {
+
+		this.reset();
+
+		this._auxTarget?.dispose();
+		this._auxTarget = null;
+
+		// the networks hold the weights on the GPU, and a load may still be in flight
+		for ( const key in this._unets ) {
+
+			this._unets[ key ]?.then( unet => unet.dispose() );
+			this._unets[ key ] = null;
+
+		}
+
+	}
+
+	// Rasterizes the albedo and normal buffers that guide the filter. oidn-web takes normals
+	// mapped into [0,1], with (0.5, 0.5, 1) as a flat normal.
+	_renderAuxiliaryBuffers( width, height ) {
+
+		const { renderer, scene, camera } = this;
+
+		if ( ! this._auxTarget ) {
+
+			// Eight bits per channel since both buffers hold [0,1]. Multisampled so the rasterized
+			// silhouettes match the jittered path traced ones.
+			this._auxTarget = new RenderTarget( 1, 1, {
+				count: 2,
+				type: UnsignedByteType,
+				minFilter: NearestFilter,
+				magFilter: NearestFilter,
+				samples: 4,
+			} );
+
+			// the MRT keys are matched against these names
+			this._auxTarget.textures[ 0 ].name = 'output';
+			this._auxTarget.textures[ 1 ].name = 'normal';
+
+		}
+
+		const auxTarget = this._auxTarget;
+		auxTarget.setSize( width, height );
+
+		const originalTarget = renderer.getRenderTarget();
+		const originalMRT = renderer.getMRT();
+		const originalToneMapping = renderer.toneMapping;
+
+		// the buffers are data rather than an image, so they must not be tone mapped
+		renderer.toneMapping = NoToneMapping;
+		renderer.setRenderTarget( auxTarget );
+		renderer.setMRT( mrt( {
+			output: diffuseColor,
+			normal: vec4( normalView.mul( 0.5 ).add( 0.5 ), 1.0 ),
+		} ) );
+
+		renderer.render( scene, camera );
+
+		renderer.setMRT( originalMRT );
+		renderer.setRenderTarget( originalTarget );
+		renderer.toneMapping = originalToneMapping;
+
+		return auxTarget;
+
+	}
+
+	// loads each model once, with concurrent calls sharing the same promise
+	_initUNet( aux ) {
+
+		const key = aux ? 'aux' : 'color';
+		if ( ! this._unets[ key ] ) {
+
+			this._unets[ key ] = ( async () => {
+
+				const device = this.renderer.backend.device;
+				const url = aux ? this.auxWeightsUrl : this.colorWeightsUrl;
+
+				if ( ! url ) {
+
+					throw new Error( `OIDNDenoiser: no ${ aux ? 'auxWeightsUrl' : 'colorWeightsUrl' } was provided.` );
+
+				}
+
+				const options = { aux, hdr: true };
+				if ( this.maxTileSize !== null ) options.maxTileSize = this.maxTileSize;
+				if ( this.dynamicTile !== null ) options.dynamicTile = this.dynamicTile;
+
+				return this.initUNetFromURL( url, { device, adapterInfo: device.adapterInfo }, options );
+
+			} )();
+
+		}
+
+		return this._unets[ key ];
+
+	}
+
+	// Copies the denoiser's packed output buffer into a sampleable texture, using a compute pass
+	// since copyBufferToTexture requires 256 byte row alignment.
+	// TODO: remove the copy once oidn-web can write into a caller-provided storage texture.
+	_copyOutputToTexture( output, width, height ) {
+
+		const device = this.renderer.backend.device;
+
+		if ( ! this._resultPipeline ) {
+
+			const module = device.createShaderModule( {
+				code: /* wgsl */`
+					@group( 0 ) @binding( 0 ) var<storage, read> src : array<vec4f>;
+					@group( 0 ) @binding( 1 ) var dst : texture_storage_2d<rgba16float, write>;
+
+					@compute @workgroup_size( 8, 8 )
+					fn main( @builtin( global_invocation_id ) gid : vec3u ) {
+
+						let dims = textureDimensions( dst );
+						if ( gid.x >= dims.x || gid.y >= dims.y ) {
+
+							return;
+
+						}
+
+						textureStore( dst, gid.xy, src[ gid.y * dims.x + gid.x ] );
+
+					}
+				`,
+			} );
+
+			this._resultPipeline = device.createComputePipeline( {
+				layout: 'auto',
+				compute: { module, entryPoint: 'main' },
+			} );
+
+		}
+
+		if ( ! this._rawTexture || this._rawTexture.width !== width || this._rawTexture.height !== height ) {
+
+			this._rawTexture?.destroy();
+			this._rawTexture = device.createTexture( {
+				size: [ width, height ],
+				format: 'rgba16float',
+				usage: GPUTextureUsage.STORAGE_BINDING | GPUTextureUsage.TEXTURE_BINDING,
+			} );
+
+			// ExternalTexture lets three sample a gpu texture it did not create. It has no image
+			// behind it, so its size is set explicitly or "texture.width" reports zero.
+			this._texture?.dispose();
+			this._texture = new ExternalTexture( this._rawTexture );
+			this._texture.image = { width, height };
+
+		}
+
+		const bindGroup = device.createBindGroup( {
+			layout: this._resultPipeline.getBindGroupLayout( 0 ),
+			entries: [
+				{ binding: 0, resource: { buffer: output.data } },
+				{ binding: 1, resource: this._rawTexture.createView() },
+			],
+		} );
+
+		const encoder = device.createCommandEncoder();
+		const pass = encoder.beginComputePass();
+		pass.setPipeline( this._resultPipeline );
+		pass.setBindGroup( 0, bindGroup );
+		pass.dispatchWorkgroups( Math.ceil( width / 8 ), Math.ceil( height / 8 ) );
+		pass.end();
+		device.queue.submit( [ encoder.finish() ] );
+
+	}
+
+}