WebGPUPathTracer: Use half float storage target w/ Kahan summation

package.json

@@ -43,6 +43,7 @@
     "@types/node": "^24.3.1",
     "@types/three": "^0.183.0",
     "@typescript-eslint/eslint-plugin": "^8.40.0",
+    "@vitejs/plugin-basic-ssl": "^2.1.4",
     "canvas-capture": "^2.0.5",
     "eslint": "^8.56.0",
     "eslint-config-mdcs": "^5.0.0",

src/textures/ProceduralEquirectTexture.js

@@ -2,13 +2,14 @@ import {
 	ClampToEdgeWrapping,
 	Color,
 	DataTexture,
+	DataUtils,
 	EquirectangularReflectionMapping,
+	HalfFloatType,
 	LinearFilter,
 	RepeatWrapping,
 	RGBAFormat,
 	Spherical,
 	Vector2,
-	FloatType
 } from 'three';
 
 const _uv = new Vector2();
@@ -20,8 +21,8 @@ export class ProceduralEquirectTexture extends DataTexture {
 	constructor( width = 512, height = 512 ) {
 
 		super(
-			new Float32Array( width * height * 4 ),
-			width, height, RGBAFormat, FloatType, EquirectangularReflectionMapping,
+			new Uint16Array( width * height * 4 ),
+			width, height, RGBAFormat, HalfFloatType, EquirectangularReflectionMapping,
 			RepeatWrapping, ClampToEdgeWrapping, LinearFilter, LinearFilter,
 		);
 
@@ -53,10 +54,10 @@ export class ProceduralEquirectTexture extends DataTexture {
 
 				const i = y * width + x;
 				const i4 = 4 * i;
-				data[ i4 + 0 ] = ( _color.r );
-				data[ i4 + 1 ] = ( _color.g );
-				data[ i4 + 2 ] = ( _color.b );
-				data[ i4 + 3 ] = ( 1.0 );
+				data[ i4 + 0 ] = DataUtils.toHalfFloat( _color.r );
+				data[ i4 + 1 ] = DataUtils.toHalfFloat( _color.g );
+				data[ i4 + 2 ] = DataUtils.toHalfFloat( _color.b );
+				data[ i4 + 3 ] = DataUtils.toHalfFloat( 1.0 );
 
 			}
 

src/webgpu/MegaKernelPathTracer.js

@@ -90,6 +90,7 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 			tiles,
 			outputTarget,
 			sampleCountTarget,
+			compensationTarget,
 			lowResMode,
 		} = this;
 
@@ -98,6 +99,7 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 		// init parameters
 		kernel.outputTarget = outputTarget;
 		kernel.sampleCountTarget = sampleCountTarget;
+		kernel.compensationTarget = compensationTarget;
 
 		kernel.bounces = bounces;
 		kernel.inverseProjectionMatrix.copy( camera.projectionMatrixInverse );

src/webgpu/PathTracerBackend.js

@@ -1,4 +1,4 @@
-import { ColorManagement, FloatType, LinearFilter, RGBAFormat } from 'three';
+import { ColorManagement, HalfFloatType, LinearFilter, RGBAFormat } from 'three';
 import { RedIntegerFormat, StorageTexture, UnsignedIntType } from 'three/webgpu';
 import { ZeroOutKernel } from './compute/ZeroOutKernel.js';
 
@@ -16,15 +16,15 @@ export class PathTracerBackend {
 
 		this.outputTarget = new StorageTexture( 1, 1, );
 		this.outputTarget.format = RGBAFormat;
-		this.outputTarget.type = FloatType;
+		this.outputTarget.type = HalfFloatType;
 		this.outputTarget.magFilter = LinearFilter;
 		this.outputTarget.colorSpace = ColorManagement.workingColorSpace;
 		this.outputTarget.name = 'Output #0';
 		this.outputTarget.generateMipmaps = false;
 
 		this.prevOutputTarget = new StorageTexture( 1, 1, );
 		this.prevOutputTarget.format = RGBAFormat;
-		this.prevOutputTarget.type = FloatType;
+		this.prevOutputTarget.type = HalfFloatType;
 		this.prevOutputTarget.magFilter = LinearFilter;
 		this.prevOutputTarget.colorSpace = ColorManagement.workingColorSpace;
 		this.prevOutputTarget.name = 'Output #1';
@@ -38,6 +38,11 @@ export class PathTracerBackend {
 
 		this.sampleCountClearKernel = new ZeroOutKernel( { textureType: 'r32uint' } ).setWorkgroupSize( 8, 8, 1 );
 		this.outputTargetClearKernel = new ZeroOutKernel( { textureType: 'rgba32float' } ).setWorkgroupSize( 8, 8, 1 );
+		this.compensationTarget = new StorageTexture( 1, 1 );
+		this.compensationTarget.format = RedIntegerFormat;
+		this.compensationTarget.type = UnsignedIntType;
+		this.compensationTarget.name = 'Accumulation Compensation';
+		this.compensationTarget.generateMipmaps = false;
 
 	}
 
@@ -84,14 +89,17 @@ export class PathTracerBackend {
 		this.outputTarget.dispose();
 		this.prevOutputTarget.dispose();
 		this.sampleCountTarget.dispose();
+		this.compensationTarget.dispose();
 
 		this.outputTarget = this.outputTarget.clone();
 		this.prevOutputTarget = this.outputTarget.clone();
 		this.sampleCountTarget = this.sampleCountTarget.clone();
+		this.compensationTarget = this.compensationTarget.clone();
 
 		this.outputTarget.setSize( w, h );
 		this.prevOutputTarget.setSize( w, h );
 		this.sampleCountTarget.setSize( w, h );
+		this.compensationTarget.setSize( w, h );
 
 		this.reset();
 		return true;
@@ -141,6 +149,7 @@ export class PathTracerBackend {
 			outputTarget,
 			prevOutputTarget,
 			sampleCountTarget,
+			compensationTarget,
 		} = this;
 
 		if ( ! renderer.initialized ) {
@@ -161,6 +170,9 @@ export class PathTracerBackend {
 		sampleCountClearKernel.target = sampleCountTarget;
 		renderer.compute( sampleCountClearKernel.kernel, dispatchSize );
 
+		sampleCountClearKernel.target = compensationTarget;
+		renderer.compute( sampleCountClearKernel.kernel, dispatchSize );
+
 		this.samples = 0;
 		this._renderTask = null;
 
@@ -171,6 +183,7 @@ export class PathTracerBackend {
 		this.outputTarget.dispose();
 		this.prevOutputTarget.dispose();
 		this.sampleCountTarget.dispose();
+		this.compensationTarget.dispose();
 
 		this._renderTask = null;
 

src/webgpu/WaveFrontPathTracer.js

@@ -242,8 +242,10 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 				[ this.outputTarget, this.prevOutputTarget ] = [ this.prevOutputTarget, this.outputTarget ];
 				rayIntersectionKernel.prevOutputTarget = this.prevOutputTarget;
 				rayIntersectionKernel.outputTarget = this.outputTarget;
+				rayIntersectionKernel.compensationTarget = this.compensationTarget;
 				hitProcessKernel.prevOutputTarget = this.prevOutputTarget;
 				hitProcessKernel.outputTarget = this.outputTarget;
+				hitProcessKernel.compensationTarget = this.compensationTarget;
 
 				// Step 1: Top up the ray queue
 				// set up the ray prime kernel

src/webgpu/WebGPUPathTracer.js

@@ -1,4 +1,4 @@
-import { DataTexture, LinearFilter, Vector2, Scene, PerspectiveCamera, Color, NoToneMapping, FloatType, Timer, StorageTexture } from 'three/webgpu';
+import { DataTexture, LinearFilter, Vector2, Scene, PerspectiveCamera, Color, NoToneMapping, Timer, StorageTexture, HalfFloatType } from 'three/webgpu';
 import { SkinnedMeshBVH, MeshBVH, SAH } from 'three-mesh-bvh';
 import { FullScreenQuad } from 'three/examples/jsm/postprocessing/Pass.js';
 import { RenderToScreenNodeMaterial } from './materials/RenderToScreenMaterial.js';
@@ -70,7 +70,7 @@ export class WebGPUPathTracer {
 		this._fadeState = 0;
 		this._size = new Vector2();
 		this._lowResTarget = new StorageTexture( 1, 1 );
-		this._lowResTarget.type = FloatType;
+		this._lowResTarget.type = HalfFloatType;
 		this._blitQuad = new FullScreenQuad( new RenderToScreenNodeMaterial() );
 
 		// options

src/webgpu/compute/PathTracerMegaKernel.js

@@ -2,22 +2,24 @@ import { DataTexture, Matrix3, Matrix4, Vector2, StorageTexture } from 'three/we
 import { ndcToCameraRay } from '../lib/wgsl/common.wgsl.js';
 import { ComputeKernel } from './ComputeKernel.js';
 import { texture, sampler, uniform, globalId, textureStore } from 'three/tsl';
-import { pcgRand2, pcgRand3, pcgInit } from '../nodes/random.wgsl.js';
+import { pcgRand2, pcgInit } from '../nodes/random.wgsl.js';
 import { getSurfaceRecordFunc, lambertBsdfFunc } from '../nodes/material.wgsl.js';
 import { sampleEnvironmentFn, weightedAlphaBlendFn } from '../nodes/sampling.wgsl.js';
-import { proxy } from '../lib/nodes/NodeProxy.js';
+import { proxy, proxyFn } from '../lib/nodes/NodeProxy.js';
 import { wgslTagFn } from '../lib/nodes/WGSLTagFnNode.js';
+import { packCompensationFn, unpackCompensationFn } from '../nodes/f16packing.wgsl.js';
 
 export class PathTracerMegaKernel extends ComputeKernel {
 
 	constructor() {
 
-		const parameters = {
+		const params = {
 			bvhData: { value: null },
 
 			prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(),
 			outputTarget: textureStore( new StorageTexture( 1, 1 ) ).toWriteOnly(),
 			sampleCountTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadWrite(),
+			compensationTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadWrite(),
 
 			offset: uniform( new Vector2() ),
 			tileSize: uniform( new Vector2() ),
@@ -47,6 +49,10 @@ export class PathTracerMegaKernel extends ComputeKernel {
 			globalId: globalId,
 		};
 
+		const raycastFirstHitFn = proxyFn( 'bvhData.value.fns.raycastFirstHit', params );
+		const sampleTrianglePointFn = proxyFn( 'bvhData.value.fns.sampleTrianglePoint', params );
+
+
 		const shader = wgslTagFn/* wgsl */`
 
 			fn compute(
@@ -81,6 +87,9 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 			) -> void {
 
+				let transforms = &${ proxy( 'bvhData.value.storage.transforms', params ) };
+				let materials = &${ proxy( 'bvhData.value.storage.materials', params ) };
+
 				let envInfo = EnvironmentInfo(
 					envMapRotation,
 					envMapIntensity,
@@ -102,7 +111,7 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 				// to screen coordinates
 				let indexUV = offset + globalId.xy;
-				let targetDimensions = textureDimensions( ${ parameters.outputTarget } );
+				let targetDimensions = textureDimensions( ${ params.outputTarget } );
 				if ( indexUV.x >= targetDimensions.x || indexUV.y >= targetDimensions.y ) {
 
 					return;
@@ -112,35 +121,35 @@ export class PathTracerMegaKernel extends ComputeKernel {
 				let uv = vec2f( indexUV ) / vec2f( targetDimensions );
 				let ndc = uv * 2.0 - vec2f( 1.0 );
 
-				pcgInitialize( indexUV, seed );
+				${ pcgInit }( indexUV, seed );
 
 				// scene ray
-				var jitter = 2.0 * pcgRand2() / vec2f( targetDimensions.xy );
-				var ray = ndcToCameraRay( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix );
+				var jitter = 2.0 * ${ pcgRand2 }() / vec2f( targetDimensions.xy );
+				var ray = ${ ndcToCameraRay }( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix );
 				ray.direction = normalize( ray.direction );
 
 				var resultColor = vec4f( 0, 0, 0, 1 );
 				var throughputColor = vec3f( 1.0 );
 
 				for ( var bounce = 0u; bounce < bounces; bounce ++ ) {
 
-					let hitResult = bvh_RaycastFirstHit( ray );
+					let hitResult = ${ raycastFirstHitFn }( ray );
 					if ( hitResult.didHit ) {
 
-						let object = bvh_transforms.value[ hitResult.objectIndex ];
-						var material = bvh_materials.value[ object.materialIndex ];
+						let object = transforms[ hitResult.objectIndex ];
+						var material = materials[ object.materialIndex ];
 
 						// apply per-object colors
 						material.color *= object.color.rgb;
 						material.opacity *= object.color.a;
 
-						var vertexData = bvh_sampleTrianglePoint( hitResult.barycoord, hitResult.indices.xyz );
+						var vertexData = ${ sampleTrianglePointFn }( hitResult.barycoord, hitResult.indices.xyz );
 						vertexData.normal = normalize( transpose( object.inverseMatrixWorld ) * vertexData.normal );
 						vertexData.position = object.matrixWorld * vertexData.position;
 
-						let surface = getSurfaceRecord( material, vertexData, hitResult.side, hitResult.normal, textures, textureSampler );
+						let surface = ${ getSurfaceRecordFunc }( material, vertexData, hitResult.side, hitResult.normal, textures, textureSampler );
 
-						let scatterRec = bsdfSample( - ray.direction, surface );
+						let scatterRec = ${ lambertBsdfFunc }( - ray.direction, surface );
 
 						// white diffuse surface
 						throughputColor *= scatterRec.color / scatterRec.pdf;
@@ -152,11 +161,11 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 						if ( bounce > 0u ) {
 
-							resultColor = sampleEnvironment( envMap, envMapSampler, envInfo, ray.direction, pcgRand2() ) * vec4f( throughputColor, 1.0 );
+							resultColor = ${ sampleEnvironmentFn }( envMap, envMapSampler, envInfo, ray.direction, pcgRand2() ) * vec4f( throughputColor, 1.0 );
 
 						} else {
 
-							resultColor = sampleEnvironment( background, backgroundSampler, backgroundInfo, ray.direction, pcgRand2() );
+							resultColor = ${ sampleEnvironmentFn }( background, backgroundSampler, backgroundInfo, ray.direction, pcgRand2() );
 
 						}
 
@@ -166,27 +175,26 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 				}
 
-				let sampleCount = textureLoad( ${ parameters.sampleCountTarget }, indexUV ).r + 1;
-				let prevColor = textureLoad( ${ parameters.prevOutputTarget }, indexUV );
-				let blendedColor = weightedAlphaBlend( prevColor, resultColor, 1.0 / f32( sampleCount ) );
-				textureStore( ${ parameters.sampleCountTarget }, indexUV, vec4( sampleCount ) );
-				textureStore( ${ parameters.outputTarget }, indexUV, blendedColor );
+				// Kahan-compensated running mean: recover true mean before computing delta
+				let prevColor = textureLoad( ${ params.prevOutputTarget }, indexUV );
+				let packedComp = textureLoad( ${ params.compensationTarget }, indexUV ).r;
+				let compensation = ${ unpackCompensationFn }( packedComp, prevColor );
+				let sampleCount = textureLoad( ${ params.sampleCountTarget }, indexUV ).r + 1;
+				let blendedColor = ${ weightedAlphaBlendFn }( prevColor + compensation, resultColor, 1.0 / f32( sampleCount ) );
+
+				// simulate FP16 rounding via pack/unpack to compute the residual that will be lost at store
+				let storedColor = quantizeToF16( blendedColor );
+				let newPackedComp = ${ packCompensationFn }( blendedColor - storedColor, storedColor );
 
-			}
+				textureStore( ${ params.sampleCountTarget }, indexUV, vec4( sampleCount ) );
+				textureStore( ${ params.outputTarget }, indexUV, storedColor );
+				textureStore( ${ params.compensationTarget }, indexUV, vec4( newPackedComp ) );
 
-	${ [
-		proxy( 'bvhData.value.storage.materials', parameters ),
-		proxy( 'bvhData.value.structs.material', parameters ),
-		proxy( 'bvhData.value.structs.transform', parameters ),
-		proxy( 'bvhData.value.fns.raycastFirstHit', parameters ),
-		proxy( 'bvhData.value.fns.sampleTrianglePoint', parameters ),
-		ndcToCameraRay, pcgRand2, pcgRand3, pcgInit, lambertBsdfFunc,
-		sampleEnvironmentFn, getSurfaceRecordFunc, weightedAlphaBlendFn,
-	] }`;
+			}`;
 
-		super( shader( parameters ) );
+		super( shader( params ) );
 
-		this.defineUniformAccessors( parameters );
+		this.defineUniformAccessors( params );
 
 	}