three-gpu-pathtracer 0.0.24 → 0.0.25

package/src/webgpu/materials/GraphMaterial.js

@@ -0,0 +1,331 @@
+import { Color, SRGBColorSpace, Vector2 } from 'three';
+import { MeshBasicNodeMaterial, StorageBufferAttribute } from 'three/webgpu';
+import { globalId, storage, uniform, uv, varying, wgslFn } from 'three/tsl';
+import { wgslTagFn, wgslTagCode } from 'three-mesh-bvh/webgpu';
+
+// pastel palette evenly spaced along the color wheel, looped when there are more graphs
+const GRAPH_COLOR_COUNT = 10;
+
+function getGraphColor( i ) {
+
+	const hue = ( i % GRAPH_COLOR_COUNT ) / GRAPH_COLOR_COUNT;
+	return new Color().setHSL( hue, 0.65, 0.7, SRGBColorSpace );
+
+}
+
+// anti-aliased background with x / y axes and 1-unit markers
+const getBackgroundFunc = wgslFn( /* wgsl */`
+	fn getBackground( point: vec2f, steepness: f32 ) -> vec3f {
+
+		let pw = fwidth( point );
+		let halfWidth = pw * 0.5;
+
+		// x, y axes
+		var distToZero = smoothstep(
+			- halfWidth * 0.5,
+			halfWidth * 0.5,
+			abs( point ) - pw
+		);
+
+		// 1 unit markers
+		let modAxis = fract( abs( point + vec2f( 0.5 ) ) ) - 0.5;
+		var distToAxis = smoothstep(
+			- halfWidth,
+			halfWidth,
+			abs( modAxis ) - pw * 0.5
+		);
+
+		// if we're at a chart boundary then remove the artifacts
+		if ( abs( pw.y ) > steepness * 0.5 ) {
+
+			distToZero.y = 1.0;
+			distToAxis.y = 1.0;
+
+		}
+
+		// mix colors into a background color
+		let axisIntensity = 1.0 - min( distToZero.x, distToZero.y );
+		let markerIntensity = 1.0 - min( distToAxis.x, distToAxis.y );
+
+		let markerColor = mix( vec3f( 0.005 ), vec3f( 0.05 ), markerIntensity );
+		return mix( markerColor, vec3f( 0.2 ), axisIntensity );
+
+	}
+` );
+
+// Fullscreen material that plots an arbitrary number of WGSL functions of the form
+// "fn( x: f32 ) -> f32" provided via "setGraphs", either overlaid or in stacked sections.
+export class GraphMaterial extends MeshBasicNodeMaterial {
+
+	get dim() {
+
+		return this._dim.value > 0.5;
+
+	}
+
+	set dim( v ) {
+
+		this._dim.value = v ? 1 : 0;
+
+	}
+
+	get overlay() {
+
+		return this._overlay.value > 0.5;
+
+	}
+
+	set overlay( v ) {
+
+		this._overlay.value = v ? 1 : 0;
+
+	}
+
+	get thickness() {
+
+		return this._thickness.value;
+
+	}
+
+	set thickness( v ) {
+
+		this._thickness.value = v;
+
+	}
+
+	get xRange() {
+
+		return this._xRange.value;
+
+	}
+
+	get yRange() {
+
+		return this._yRange.value;
+
+	}
+
+	get graphCount() {
+
+		return this._graphs.length;
+
+	}
+
+	// the display names of the plotted graphs, in slot order
+	get graphNames() {
+
+		return this._graphs.map( entry => {
+
+			return entry.name || entry.code.match( /fn ([^(]+)/ )[ 1 ];
+
+		} );
+
+	}
+
+	// the mouse position in graph space - set to a distant point to hide the markers
+	get mousePoint() {
+
+		return this._mousePoint.value;
+
+	}
+
+	constructor( options = {} ) {
+
+		const { graphs = [], ...params } = options;
+
+		super();
+
+		this.depthWrite = false;
+		this.depthTest = false;
+
+		// uniforms - created once and shared across shader rebuilds
+		this._xRange = uniform( new Vector2( - 2, 2 ) );
+		this._yRange = uniform( new Vector2( - 2, 2 ) );
+		this._thickness = uniform( window.devicePixelRatio );
+		this._dim = uniform( 1 );
+		this._overlay = uniform( 1 );
+		this._displayMask = uniform( 0xffffffff, 'uint' );
+		this._mousePoint = uniform( new Vector2( 1e10, 1e10 ) );
+		this._graphs = [];
+
+		this.setGraphs( graphs );
+		this.setValues( params );
+
+	}
+
+	// toggles the visibility of the graph at the given index
+	setGraphVisible( index, visible ) {
+
+		const mask = this._displayMask;
+		if ( visible ) {
+
+			mask.value = ( mask.value | ( 1 << index ) ) >>> 0;
+
+		} else {
+
+			mask.value = ( mask.value & ~ ( 1 << index ) ) >>> 0;
+
+		}
+
+	}
+
+	getGraphVisible( index ) {
+
+		return ( ( this._displayMask.value >> index ) & 1 ) === 1;
+
+	}
+
+	getGraphColor( index ) {
+
+		return getGraphColor( index );
+
+	}
+
+	// Replaces the set of plotted functions - each an "fn( x: f32 ) -> f32" wgsl function node -
+	// and rebuilds the shader to display them.
+	setGraphs( graphs ) {
+
+		const { _xRange, _yRange, _thickness, _dim, _overlay, _displayMask, _mousePoint } = this;
+
+		this._graphs = [ ...graphs ];
+
+		// unroll one evaluation block per graph, accumulated into a single code chunk
+		let graphBody = '';
+		graphs.forEach( ( graph, i ) => {
+
+			const color = getGraphColor( i );
+			graphBody = wgslTagCode/* wgsl */`
+				${ graphBody }
+
+				{
+					let visible = ( ( displayMask >> ${ i }u ) & 1u ) == 1u;
+					if ( visible && ( sectionCount < 1.5 || section == ${ i } ) ) {
+
+						let delta = ${ graph }( point.x ) - point.y;
+						var halfDdf = fwidth( delta ) * 0.5;
+						if ( pointYWidth > yWidth * 0.5 ) {
+
+							halfDdf = 0.0;
+
+						}
+
+						var graph = smoothstep( - halfDdf, halfDdf, abs( delta ) - ${ _thickness } * halfDdf );
+						if ( dimmed ) {
+
+							graph = mix( 1.0, graph, 0.1 );
+
+						}
+
+						color = mix( vec3f( ${ color.r }, ${ color.g }, ${ color.b } ), color, graph );
+
+						// circle marker at the mouse line's intersection with the curve
+						let intersection = vec2f( mousePoint.x, ${ graph }( mousePoint.x ) );
+						let circleDist = length( ( point - intersection ) / pw );
+						let circle = smoothstep( 4.0, 5.0, circleDist );
+						color = mix( vec3f( ${ color.r }, ${ color.g }, ${ color.b } ), color, circle );
+
+					}
+
+				}
+			`;
+
+		} );
+
+		// the section math assumes at least one graph
+		const count = Math.max( graphs.length, 1 );
+		const graphShader = wgslTagFn/* wgsl */`
+			fn graphShader( vUv: vec2f ) -> vec4f {
+
+				var sectionCount = ${ count }.0;
+				if ( ${ _overlay } > 0.5 ) {
+
+					sectionCount = 1.0;
+
+				}
+
+				let yWidth = abs( ${ _yRange }.y - ${ _yRange }.x );
+
+				// separate into sections
+				let scaledY = sectionCount * vUv.y;
+				let sectionY = fract( scaledY );
+				let section = clamp( i32( sectionCount - floor( scaledY ) - 1.0 ), 0, ${ count - 1 } );
+
+				// the graph-space point for this pixel
+				let point = vec2f(
+					mix( ${ _xRange }.x, ${ _xRange }.y, vUv.x ),
+					mix( ${ _yRange }.x, ${ _yRange }.y, sectionY )
+				);
+
+				let pw = fwidth( point );
+				let pointYWidth = pw.y;
+				let dimmed = ${ _dim } > 0.5 && ( point.x < 0.0 || point.y < 0.0 );
+				let displayMask = ${ _displayMask };
+				let mousePoint = ${ _mousePoint };
+
+				var color = ${ getBackgroundFunc }( point, yWidth );
+
+				// vertical marker line at the mouse position
+				let mouseLine = smoothstep( - pw.x * 0.5, pw.x * 0.5, abs( point.x - mousePoint.x ) - pw.x );
+				color = mix( vec3f( 0.35 ), color, mouseLine );
+
+				${ graphBody }
+
+				return vec4f( color, 1.0 );
+
+			}
+		`;
+
+		// NOTE: varyings cannot be referenced directly and must be passed as arguments
+		this.colorNode = graphShader( varying( uv() ) );
+		this.needsUpdate = true;
+
+		// compute kernel that evaluates every graph at the mouse x so the values can be read back.
+		// The buffer is padded to a minimum length of 2 to avoid TSL treating it as a scalar.
+		const valuesAttribute = new StorageBufferAttribute( new Float32Array( Math.max( count, 2 ) ), 1 );
+		const valuesStorage = storage( valuesAttribute, 'float' );
+
+		let valuesBody = '';
+		graphs.forEach( ( graph, i ) => {
+
+			valuesBody = wgslTagCode/* wgsl */`
+				${ valuesBody }
+				${ valuesStorage }[ ${ i } ] = ${ graph }( x );
+			`;
+
+		} );
+
+		const valuesFn = wgslTagFn/* wgsl */`
+			fn computeGraphValues( globalId: vec3u ) -> void {
+
+				let x = ${ _mousePoint }.x;
+				${ valuesBody }
+
+			}
+		`;
+
+		this._valuesAttribute?.dispose();
+		this._valuesAttribute = valuesAttribute;
+		this._valuesKernel = valuesFn( { globalId } ).computeKernel( [ 1, 1, 1 ] );
+
+	}
+
+	dispose() {
+
+		super.dispose();
+		this._valuesAttribute.dispose();
+
+	}
+
+	// Evaluates every graph function at the current "mousePoint.x" on the gpu and resolves to a
+	// Float32Array of the resulting values.
+	async readGraphValues( renderer ) {
+
+		const count = this._graphs.length;
+		renderer.compute( this._valuesKernel, [ 1, 1, 1 ] );
+
+		const buffer = await renderer.getArrayBufferAsync( this._valuesAttribute );
+		return new Float32Array( buffer ).slice( 0, count );
+
+	}
+
+}

package/src/webgpu/materials/PathtracingMaterial.js

@@ -0,0 +1,89 @@
+import { wgslTagFn } from 'three-mesh-bvh/webgpu';
+import { rand2, RNG_INDEX_SCATTER_DIRECTION } from '../nodes/random.wgsl.js';
+import { diffuseDirectionFunc } from '../nodes/sampling.wgsl.js';
+
+/**
+ * Defines a material sampled by the pathtracer
+ */
+export class PathtracingMaterial {
+
+	/**
+	 *
+	 * Called once per material
+	 * Adds ability to initialize state
+	 *
+	 */
+	init( /* renderer */ ) {
+
+	}
+
+	/**
+	 *
+	 * Must return a bsdf sampling function node with signature
+	 * ( worldView: vec3f, surface: Surface ) -> ScatterRecord
+	 *
+	 */
+	getBsdfNode() {
+
+		return wgslTagFn/* wgsl */`
+
+			fn bsdfSample( worldWo: vec3f, surf: SurfaceRecord ) -> ScatterRecord {
+
+				var record: ScatterRecord;
+
+				let wo = normalize( surf.normalInvBasis * worldWo );
+				let wi = ${ diffuseDirectionFunc }( wo, ${ rand2 }( ${ RNG_INDEX_SCATTER_DIRECTION } ) );
+				record.color = surf.color * max( wi.z, 0.0 );
+				record.pdf = max( wi.z, 0.0 ) / PI;
+				record.direction = normalize( surf.normalBasis * wi );
+
+				return record;
+
+			}
+
+		`;
+
+	}
+
+	/**
+	 *
+	 * Must return a bsdf evaluation function node with signature
+	 * ( worldView: vec3f, worldLight: vec3f, surface: Surface ) -> ScatterRecord
+	 * returning the bsdf value and sampling pdf for a given light direction.
+	 * Used by next event estimation to weight a chosen light direction.
+	 *
+	 */
+	getBsdfEvalPdfNode() {
+
+		return wgslTagFn/* wgsl */`
+
+			fn bsdfEvalPdf( worldWo: vec3f, worldWi: vec3f, surf: SurfaceRecord ) -> ScatterRecord {
+
+				var record: ScatterRecord;
+				record.direction = worldWi;
+
+				let wo = normalize( surf.normalInvBasis * worldWo );
+				let wi = normalize( surf.normalInvBasis * worldWi );
+				record.color = surf.color * max( wi.z, 0.0 );
+				record.pdf = max( wi.z, 0.0 ) / PI;
+
+				return record;
+
+			}
+
+		`;
+
+	}
+
+	getData() {
+
+		return {
+
+			bsdfSample: this.getBsdfNode(),
+			bsdfEvalPdf: this.getBsdfEvalPdfNode(),
+
+		};
+
+	}
+
+}

package/src/webgpu/materials/RenderToScreenMaterial.js

@@ -0,0 +1,150 @@
+import { MeshBasicNodeMaterial, NoToneMapping, StorageTexture } from 'three/webgpu';
+import { uv, varying, texture, vec4, toneMapping, uniform, wgslFn } from 'three/tsl';
+import { wgslTagFn } from 'three-mesh-bvh/webgpu';
+
+// TODO: we could fall back to hardware-based filtering if available but it has to be specifically
+// requested and available on the renderer which we don't have access to immediately. It's possible this
+// can be detected during build time via a custom node to adjust the sample approach used?
+const sampleTexelFn = wgslFn( /* wgsl */`
+	fn sampleTexel( tex: texture_2d<f32>, coord: vec2f ) -> vec4f {
+
+		// Manual bilinear filtering using textureLoad to support filterable float32 textures
+		// on all devices
+		let size = vec2f( textureDimensions( tex, 0 ) );
+
+		// the path tracer stores its color rows top down
+		let pxCoord = vec2f( coord.x, 1.0 - coord.y ) * size - 0.5;
+		let px = vec2i( floor( pxCoord ) );
+		let fr = fract( pxCoord );
+
+		// get the four sibling samples
+		let s00 = textureLoad( tex, clamp( px,                 vec2i( 0 ), vec2i( size ) - 1 ), 0 );
+		let s10 = textureLoad( tex, clamp( px + vec2i( 1, 0 ), vec2i( 0 ), vec2i( size ) - 1 ), 0 );
+		let s01 = textureLoad( tex, clamp( px + vec2i( 0, 1 ), vec2i( 0 ), vec2i( size ) - 1 ), 0 );
+		let s11 = textureLoad( tex, clamp( px + vec2i( 1, 1 ), vec2i( 0 ), vec2i( size ) - 1 ), 0 );
+
+		// interpolate on the x axis
+		let y0 = mix( s00, s10, fr.x );
+		let y1 = mix( s01, s11, fr.x );
+
+		// then y
+		return mix( y0, y1, fr.y );
+
+	}
+` );
+
+// Material to apply tone mapping _before_ applying alpha blending
+export class RenderToScreenNodeMaterial extends MeshBasicNodeMaterial {
+
+	get texture() {
+
+		return this._texNode.value;
+
+	}
+
+	set texture( v ) {
+
+		this._texNode.value = v;
+
+	}
+
+	get fromTexture() {
+
+		return this._fromTexNode.value;
+
+	}
+
+	set fromTexture( v ) {
+
+		this._fromTexNode.value = v;
+
+	}
+
+	get toneMapping() {
+
+		return this._toneMapping.toneMapping;
+
+	}
+
+	set toneMapping( v ) {
+
+		this._toneMapping.setToneMapping( v );
+
+	}
+
+	get transition() {
+
+		return this._transitionUniform.value;
+
+	}
+
+	set transition( v ) {
+
+		this._transitionUniform.value = v;
+
+	}
+
+	get exposure() {
+
+		return this._toneMapping.exposureNode.value;
+
+	}
+
+	set exposure( v ) {
+
+		this._toneMapping.exposureNode.value = v;
+
+	}
+
+	constructor( params ) {
+
+		super();
+
+		const texNode = texture( new StorageTexture() );
+		this._texNode = texNode;
+
+		const fromTexNode = texture( new StorageTexture() );
+		this._fromTexNode = fromTexNode;
+
+		const texUV = varying( uv() );
+
+		const transitionUniform = uniform( 1.0 );
+		this._transitionUniform = transitionUniform;
+
+		// NOTE: varyings cannot be referenced directly and must be passed as arguments
+		const getFadedColorFn = wgslTagFn/* wgsl */`
+			fn fade( uv: vec2f ) -> vec4f {
+
+				if ( ${ transitionUniform } <= 0.0 ) {
+
+					return ${ sampleTexelFn }( ${ fromTexNode }, uv );
+
+				} else if ( ${ transitionUniform } >= 1.0 ) {
+
+					return ${ sampleTexelFn }( ${ texNode }, uv );
+
+				} else {
+
+					let col0 = ${ sampleTexelFn }( ${ fromTexNode }, uv );
+					let col1 = ${ sampleTexelFn }( ${ texNode }, uv );
+					return mix( col0, col1, ${ transitionUniform } );
+
+				}
+
+			}
+		`;
+
+		const toneMappingNode = toneMapping( NoToneMapping, 1.0, getFadedColorFn( texUV ) );
+		this._toneMapping = toneMappingNode;
+
+		// apply alpha _after_ applying tone mapping
+		// NOTE: alpha is being multiplied twice here to accommodate some odd blending in three.js
+		// See mrdoob/three.js#33104. It's possible this should be removed or rethought once fixed.
+		this.transparent = true;
+		this.colorNode = vec4( toneMappingNode.rgb, toneMappingNode.a );
+
+		this.setValues( params );
+
+	}
+
+}

package/src/webgpu/materials/debug/AtlasDebugMaterial.js

@@ -0,0 +1,41 @@
+import { MeshBasicNodeMaterial, DataTexture, NoBlending } from 'three/webgpu';
+import { uv, varying, texture, uniform, wgslFn } from 'three/tsl';
+
+export class AtlasDebugMaterial extends MeshBasicNodeMaterial {
+
+	set texture( v ) {
+
+		this._texNode.value = v;
+
+	}
+
+	set layer( v ) {
+
+		this._layerUniform.value = v;
+
+	}
+
+	constructor() {
+
+		super();
+
+		this.blending = NoBlending;
+
+		this._texNode = texture( new DataTexture() );
+		this._layerUniform = uniform( 0 );
+
+		// Shows one layer of the atlas array texture.
+		this.colorNode = wgslFn( /* wgsl */ `
+			fn atlasDebug( tex: texture_2d_array<f32>, coord: vec2f, layer: f32 ) -> vec4f {
+
+				// flip y so texel origin shows at the top-left
+				let dims = vec2f( textureDimensions( tex, 0 ).xy );
+				let texel = vec2i( vec2f( coord.x, 1.0 - coord.y ) * dims );
+				return textureLoad( tex, texel, i32( layer ), 0 );
+
+			}
+		` )( this._texNode, varying( uv() ), this._layerUniform );
+
+	}
+
+}

package/src/webgpu/materials/debug/SampleDensityMaterial.js

@@ -0,0 +1,61 @@
+import { MeshBasicNodeMaterial, NoBlending, StorageTexture } from 'three/webgpu';
+import { uv, varying, texture, uniform } from 'three/tsl';
+import { wgslTagFn } from 'three-mesh-bvh/webgpu';
+import { heatColorFn } from '../../nodes/debugBounds.wgsl.js';
+import { SAMPLE_COUNT_MASK } from '../../constants.js';
+
+// Shows how many samples each pixel has accumulated
+export class SampleDensityMaterial extends MeshBasicNodeMaterial {
+
+	set texture( v ) {
+
+		this._texNode.value = v;
+
+	}
+
+	set minCount( v ) {
+
+		this._minUniform.value = v;
+
+	}
+
+	set maxCount( v ) {
+
+		this._maxUniform.value = v;
+
+	}
+
+	constructor() {
+
+		super();
+
+		this.blending = NoBlending;
+
+		this._texNode = texture( new StorageTexture( 1, 1 ) );
+		this._minUniform = uniform( 0 );
+		this._maxUniform = uniform( 1 );
+
+		const sampleDensity = wgslTagFn/* wgsl */`
+			fn sampleDensity( tex: texture_2d<u32>, coord: vec2f, minCount: f32, maxCount: f32 ) -> vec4f {
+
+				let dims = vec2f( textureDimensions( tex, 0 ) );
+				let texel = vec2i( coord * dims );
+
+				// the high bits hold the ray state flags rather than part of the count
+				let samples = textureLoad( tex, texel, 0 ).r & ${ SAMPLE_COUNT_MASK }u;
+
+				// stretched across the recorded range rather than from zero, otherwise a converged
+				// image with counts bunched near the max reads as uniformly hot
+				let range = max( maxCount - minCount, 1.0 );
+				let t = clamp( ( f32( samples ) - minCount ) / range, 0.0, 1.0 );
+
+				return vec4f( ${ heatColorFn }( t ), 1.0 );
+
+			}
+		`;
+
+		this.colorNode = sampleDensity( this._texNode, varying( uv() ), this._minUniform, this._maxUniform );
+
+	}
+
+}