WebGPU: Sample environment

src/webgpu/MegaKernelPathTracer.js

@@ -1,6 +1,7 @@
-import { StorageTexture, Vector2, FloatType, RGBAFormat, LinearFilter, RedIntegerFormat, UnsignedIntType, ColorManagement } from 'three/webgpu';
+import { Matrix4, StorageTexture, Vector2, FloatType, RGBAFormat, LinearFilter, RedIntegerFormat, UnsignedIntType, ColorManagement } from 'three/webgpu';
 import { PathTracerMegaKernel } from './compute/PathTracerMegaKernel.js';
 import { ZeroOutKernel } from './compute/ZeroOutKernel.js';
+import { EquirectHdrInfoUniform } from '../uniforms/EquirectHdrInfoUniform.js';
 
 function* renderTask() {
 
@@ -69,6 +70,8 @@ export class MegaKernelPathTracer {
 		this.bounces = 7;
 		this.tiles = new Vector2( 2, 2 );
 
+		this.envInfo = new EquirectHdrInfoUniform();
+
 		// targets
 		this.outputTarget = new StorageTexture( 1, 1, );
 		this.outputTarget.format = RGBAFormat;
@@ -107,6 +110,38 @@ export class MegaKernelPathTracer {
 
 	}
 
+	setEnvironment(
+		envMap,
+		envMapIntensity,
+		envMapRotation,
+
+		background,
+		backgroundIntensity,
+		backgroundRotation,
+		backgroundBlurriness,
+	) {
+
+		if ( envMap !== null ) {
+
+			this.envInfo.updateFrom( envMap );
+			this.kernel.envMap = this.envInfo.map;
+			this.kernel.kernel.computeNode.parameters.envMapSampler.node.value = this.envInfo.map;
+
+		}
+
+		const rotationMatrix = new Matrix4().makeRotationFromEuler( envMapRotation ).invert();
+		this.kernel.envMapRotation.setFromMatrix4( rotationMatrix );
+		this.kernel.envMapIntensity = envMapIntensity;
+
+		this.kernel.background = background;
+		this.kernel.kernel.computeNode.parameters.backgroundSampler.node.value = background;
+		rotationMatrix.makeRotationFromEuler( backgroundRotation ).invert();
+		this.kernel.backgroundRotation.setFromMatrix4( rotationMatrix );
+		this.kernel.backgroundIntensity = backgroundIntensity;
+		this.kernel.backgroundBlurriness = backgroundBlurriness;
+
+	}
+
 	setCamera( camera ) {
 
 		this.camera = camera;
@@ -168,6 +203,7 @@ export class MegaKernelPathTracer {
 	dispose() {
 
 		// TODO: dispose of all buffers
+		this.envInfo.dispose();
 		this._task = null;
 
 	}

src/webgpu/PathTracerCore.js

@@ -6,6 +6,7 @@ import {
 	writeTraceRayDispatchSize, writeBsdfDispatchSize, writeEscapedRayDispatchSize,
 } from './nodes/wavefront.wgsl.js';
 import { PathTracerMegaKernel } from './compute/PathTracerMegaKernel.js';
+import { EquirectHdrInfoUniform } from '../uniforms/EquirectHdrInfoUniform.js';
 
 function* renderTask() {
 
@@ -143,6 +144,7 @@ export class PathTracerCore {
 
 		};
 
+		this.envInfo = new EquirectHdrInfoUniform();
 		this.geometry = {
 			bvh: new StorageBufferAttribute(),
 			index: new StorageBufferAttribute(),
@@ -363,6 +365,23 @@ export class PathTracerCore {
 
 	}
 
+	setEnvironment( environment, intensity, rotation, blur ) {
+
+		if ( environment !== null ) {
+
+			this.envInfo.updateFrom( environment );
+			this.logicKernel.envMap = this.envInfo.map;
+			this.logicKernel.kernel.computeNode.parameters.envMapSampler.node.value = this.envInfo.map;
+
+		}
+
+		const rotationMatrix = new Matrix4().makeRotationFromEuler( rotation ).invert();
+		this.logicKernel.envMapRotation.setFromMatrix4( rotationMatrix );
+		this.logicKernel.envMapIntensity = intensity;
+		this.logicKernel.envMapBlur = blur;
+
+	}
+
 	setCamera( camera ) {
 
 		this.camera = camera;

src/webgpu/WaveFrontPathTracer.js

@@ -1,11 +1,12 @@
-import { IndirectStorageBufferAttribute, StorageTexture, Vector2, FloatType, RGBAFormat, LinearFilter, RedIntegerFormat, UnsignedIntType, ColorManagement } from 'three/webgpu';
+import { Matrix4, IndirectStorageBufferAttribute, StorageTexture, Vector2, FloatType, RGBAFormat, LinearFilter, RedIntegerFormat, UnsignedIntType, ColorManagement } from 'three/webgpu';
 import { ZeroOutKernel } from './compute/ZeroOutKernel.js';
 import { PrimeRayGenerationDispatchKernel } from './compute/wavefront/PrimeRayGenerationDispatchKernel.js';
 import { RayGenerationKernel } from './compute/wavefront/RayGenerationKernel.js';
 import { RayIntersectionKernel } from './compute/wavefront/RayIntersectionKernel.js';
 import { UpdateRayQueueParamsKernel } from './compute/wavefront/UpdateRayQueueParamsKernel.js';
 import { ZeroOutBufferKernel } from './compute/ZeroOutBufferKernel.js';
 import { ProcessHitsKernel } from './compute/wavefront/ProcessHitsKernel.js';
+import { EquirectHdrInfoUniform } from '../uniforms/EquirectHdrInfoUniform.js';
 
 // set the buffers to the max possible size supported by default (128MB)
 // TODO: this can be increased based on platform.
@@ -168,6 +169,8 @@ export class WaveFrontPathTracer {
 			materials: null,
 		};
 
+		this.envInfo = new EquirectHdrInfoUniform();
+
 		// targets
 		this.outputTarget = new StorageTexture( 1, 1, );
 		this.outputTarget.format = RGBAFormat;
@@ -237,6 +240,40 @@ export class WaveFrontPathTracer {
 
 	}
 
+	setEnvironment(
+		envMap,
+		envMapIntensity,
+		envMapRotation,
+
+		background,
+		backgroundIntensity,
+		backgroundRotation,
+		backgroundBlurriness,
+	) {
+
+		const kernel = this.rayIntersectionKernel;
+
+		if ( envMap !== null ) {
+
+			this.envInfo.updateFrom( envMap );
+			kernel.envMap = this.envInfo.map;
+			kernel.kernel.computeNode.parameters.envMapSampler.node.value = this.envInfo.map;
+
+		}
+
+		const rotationMatrix = new Matrix4().makeRotationFromEuler( envMapRotation ).invert();
+		kernel.envMapRotation.setFromMatrix4( rotationMatrix );
+		kernel.envMapIntensity = envMapIntensity;
+
+		kernel.background = background;
+		kernel.kernel.computeNode.parameters.backgroundSampler.node.value = background;
+		rotationMatrix.makeRotationFromEuler( backgroundRotation ).invert();
+		kernel.backgroundRotation.setFromMatrix4( rotationMatrix );
+		kernel.backgroundIntensity = backgroundIntensity;
+		kernel.backgroundBlurriness = backgroundBlurriness;
+
+	}
+
 	setCamera( camera ) {
 
 		this.camera = camera;
@@ -298,6 +335,7 @@ export class WaveFrontPathTracer {
 	dispose() {
 
 		// TODO: dispose of all buffers
+		this.envInfo.dispose();
 		this._task = null;
 
 	}

src/webgpu/WebGPUPathTracer.js

@@ -1,9 +1,10 @@
-import { Vector2, Scene, PerspectiveCamera } from 'three/webgpu';
+import { DataTexture, LinearFilter, Vector2, Scene, PerspectiveCamera } from 'three/webgpu';
 import { MeshBVH, SAH } from 'three-mesh-bvh';
 import { FullScreenQuad } from 'three/examples/jsm/postprocessing/Pass.js';
 import { RenderToScreenNodeMaterial } from './materials/RenderToScreenMaterial.js';
 import { MegaKernelPathTracer } from './MegaKernelPathTracer.js';
 import { WaveFrontPathTracer } from './WaveFrontPathTracer.js';
+import { CubeToEquirectGenerator } from '../utils/CubeToEquirectGenerator.js';
 import { ObjectBVH } from './lib/ObjectBVH.js';
 import { PathtracerBVHComputeData } from './nodes/PathtracerBVHComputeData.js';
 
@@ -37,6 +38,18 @@ export class WebGPUPathTracer {
 		this._renderer = renderer;
 		this._pathTracer = new MegaKernelPathTracer( renderer );
 
+		this._envColorTexture = new DataTexture( );
+		this._envColorTexture.image.data = new Uint8Array( [ 255, 255, 255, 255 ] );
+		this._envColorTexture.needsUpdate = true;
+		this._envColorTexture.minFilter = LinearFilter;
+		this._envColorTexture.magFilter = LinearFilter;
+
+		this._backgroundColorTexture = new DataTexture( );
+		this._backgroundColorTexture.image.data = new Uint8Array( [ 255, 255, 255, 255 ] );
+		this._backgroundColorTexture.needsUpdate = true;
+		this._backgroundColorTexture.minFilter = LinearFilter;
+		this._backgroundColorTexture.magFilter = LinearFilter;
+
 		// options
 		this.renderScale = 1;
 		this.synchronizeRenderSize = true;
@@ -73,6 +86,20 @@ export class WebGPUPathTracer {
 		this._pathTracer.setBVHData( bvhData );
 		this.setCamera( camera );
 
+		const environment = convertToTexture( this._renderer, scene.environment, this._envColorTexture );
+		const background = convertToTexture( this._renderer, scene.background, this._backgroundColorTexture );
+
+		this._pathTracer.setEnvironment(
+			environment,
+			scene.environmentIntensity,
+			scene.environmentRotation,
+
+			background,
+			scene.backgroundIntensity,
+			scene.backgroundRotation,
+			scene.backgroundBlurriness,
+		);
+
 	}
 
 	setCamera( camera ) {
@@ -155,3 +182,34 @@ export class WebGPUPathTracer {
 	}
 
 }
+
+function convertToTexture( renderer, value, colorTexture ) {
+
+	if ( ! value ) {
+
+		for ( let i = 0; i < 3; i ++ ) {
+
+			colorTexture.image.data[ i ] = 0;
+
+		}
+
+		colorTexture.needsUpdate = true;
+		value = colorTexture;
+
+	} else if ( value.isColor ) {
+
+		colorTexture.image.data[ 0 ] = value.r * 255;
+		colorTexture.image.data[ 1 ] = value.g * 255;
+		colorTexture.image.data[ 2 ] = value.b * 255;
+		colorTexture.needsUpdate = true;
+		value = colorTexture;
+
+	} else if ( value?.isCubeTexture ) {
+
+		value = new CubeToEquirectGenerator( renderer ).generate( value );
+
+	}
+
+	return value;
+
+}

src/webgpu/compute/PathTracerMegaKernel.js

@@ -1,10 +1,11 @@
-import { Matrix4, Vector2, StorageTexture } from 'three/webgpu';
+import { DataTexture, Matrix3, Matrix4, Vector2, StorageTexture } from 'three/webgpu';
 import { ndcToCameraRay } from '../lib/wgsl/common.wgsl.js';
 import { ComputeKernel } from './ComputeKernel.js';
-import { uniform, globalId, textureStore, wgslFn } from 'three/tsl';
-import { pcgRand3, pcgInit } from '../nodes/random.wgsl.js';
+import { texture, sampler, uniform, globalId, textureStore, wgslFn } from 'three/tsl';
+import { pcgRand2, pcgRand3, pcgInit } from '../nodes/random.wgsl.js';
 import { lambertBsdfFunc } from '../nodes/sampling.wgsl.js';
 import { proxy } from '../lib/nodes/NodeProxy.js';
+import { sampleEnvironmentFn } from '../nodes/sampling.wgsl.js';
 
 export class PathTracerMegaKernel extends ComputeKernel {
 
@@ -26,6 +27,18 @@ export class PathTracerMegaKernel extends ComputeKernel {
 			inverseProjectionMatrix: uniform( new Matrix4() ),
 			cameraToModelMatrix: uniform( new Matrix4() ),
 
+			// environment
+			envMap: texture( new DataTexture() ),
+			envMapSampler: sampler( new DataTexture() ),
+			envMapRotation: uniform( new Matrix3() ),
+			envMapIntensity: uniform( 1 ),
+
+			background: texture( new DataTexture() ),
+			backgroundSampler: sampler( new DataTexture() ),
+			backgroundRotation: uniform( new Matrix3() ),
+			backgroundIntensity: uniform( 1 ),
+			backgroundBlurriness: uniform( 0 ),
+
 			// compute variables
 			globalId: globalId,
 		};
@@ -50,8 +63,32 @@ export class PathTracerMegaKernel extends ComputeKernel {
 				seed: u32,
 				bounces: u32,
 
+				// environment
+				envMap: texture_2d<f32>,
+				envMapSampler: sampler,
+				envMapRotation: mat3x3f,
+				envMapIntensity: f32,
+
+				background: texture_2d<f32>,
+				backgroundSampler: sampler,
+				backgroundRotation: mat3x3f,
+				backgroundIntensity: f32,
+				backgroundBlurriness: f32,
+
 			) -> void {
 
+				let envInfo = EnvironmentInfo(
+					envMapRotation,
+					envMapIntensity,
+					0.0 // blur,
+				);
+
+				let backgroundInfo = EnvironmentInfo(
+					backgroundRotation,
+					backgroundIntensity,
+					backgroundBlurriness,
+				);
+
 				// make sure we don't bleed over the edge of our tile
 				if ( globalId.x >= tileSize.x || globalId.y >= tileSize.y ) {
 
@@ -76,6 +113,7 @@ export class PathTracerMegaKernel extends ComputeKernel {
 				// scene ray
 				var jitter = 2.0 * ( pcgRand2() - vec2( 0.5 ) ) / vec2f( targetDimensions.xy );
 				var ray = ndcToCameraRay( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix );
+				ray.direction = normalize( ray.direction );
 
 				var resultColor = vec3f( 0.0 );
 				var throughputColor = vec3f( 1.0 );
@@ -100,8 +138,18 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 					} else {
 
-						let background = vec3f( 0.5 );
-						resultColor += background * throughputColor;
+						var light: vec3f;
+						if ( bounce > 0u ) {
+
+							light = sampleEnvironment( envMap, envMapSampler, envInfo, ray.direction, pcgRand2() );
+
+						} else {
+
+							light = sampleEnvironment( background, backgroundSampler, backgroundInfo, ray.direction, pcgRand2() );
+
+						}
+
+						resultColor += light * throughputColor;
 						break;
 
 					}
@@ -123,7 +171,8 @@ export class PathTracerMegaKernel extends ComputeKernel {
 			proxy( 'bvhData.value.structs.transform', parameters ),
 			proxy( 'bvhData.value.fns.raycastFirstHit', parameters ),
 			proxy( 'bvhData.value.fns.sampleTrianglePoint', parameters ),
-			ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc,
+			ndcToCameraRay, pcgRand2, pcgRand3, pcgInit, lambertBsdfFunc,
+			sampleEnvironmentFn,
 		] );
 
 		super( shader( parameters ) );

src/webgpu/compute/wavefront/RayGenerationKernel.js

@@ -86,7 +86,10 @@ export class RayGenerationKernel extends ComputeKernel {
 
 				// write the ray data
 				var jitter = 2.0 * ( pcgRand2() - vec2( 0.5 ) ) / vec2f( targetDimensions.xy );
-				rayQueue[ index ].ray = ndcToCameraRay( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix );;
+				var ray = ndcToCameraRay( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix );
+				ray.direction = normalize( ray.direction );
+
+				rayQueue[ index ].ray = ray;
 				rayQueue[ index ].pixel = indexUV;
 				rayQueue[ index ].throughputColor = vec3f( 1.0 );
 				rayQueue[ index ].currentBounce = 0;