WebGPU: Sobol random number generator

example/viewerTest.js

@@ -18,7 +18,7 @@ import { HDRLoader } from 'three/examples/jsm/loaders/HDRLoader.js';
 import { GLTFLoader } from 'three/examples/jsm/loaders/GLTFLoader.js';
 import { GUI } from 'three/examples/jsm/libs/lil-gui.module.min.js';
 import { WebGLPathTracer } from 'three-gpu-pathtracer';
-import { WebGPUPathTracer } from 'three-gpu-pathtracer/webgpu';
+import { RNG_PCG, RNG_SOBOL, RNG_SOBOL_TEXTURE, WebGPUPathTracer } from 'three-gpu-pathtracer/webgpu';
 import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
 import { ParallelMeshBVHWorker } from 'three-mesh-bvh/worker';
 import { LoaderElement } from './utils/LoaderElement.js';
@@ -37,6 +37,7 @@ const params = {
 
 	isWebGPU,
 	useMegakernel: true,
+	webgpuPRNG: RNG_SOBOL,
 
 	enable: true,
 	bounces: 10,
@@ -328,14 +329,28 @@ function buildGui() {
 
 	} );
 
-	webgpuOptions = pathTracingFolder.add( params, 'useMegakernel' );
-	webgpuOptions.onChange( () => {
+	webgpuOptions = pathTracingFolder.addFolder( 'WebGPU Options' );
+
+	webgpuOptions.add( params, 'useMegakernel' ).onChange( () => {
 
 		pathTracer.useMegakernel( params.useMegakernel );
 		pathTracer.reset();
 		detailedSampleCount = null;
 
 	} );
+
+	webgpuOptions.add( params, 'webgpuPRNG', {
+		PCG: RNG_PCG,
+		SOBOL: RNG_SOBOL,
+		SOBOL_TEXTURE: RNG_SOBOL_TEXTURE
+	} ).onChange( () => {
+
+		pathTracer.setPRNGType( params.webgpuPRNG );
+		pathTracer.reset();
+
+	} );
+
+
 	webgpuOptions.show( params.isWebGPU );
 
 	pathTracingFolder.add( params, 'enable' );

src/webgpu/MegaKernelPathTracer.js

@@ -5,18 +5,17 @@ import { PathTracerBackend } from './PathTracerBackend.js';
 
 export class MegaKernelPathTracer extends PathTracerBackend {
 
-	constructor( renderer ) {
+	constructor( renderer, rngData ) {
 
-		super( renderer );
+		super( renderer, rngData );
 
 		// options
 		this.tiles = new Vector2( 2, 2 );
 		this.envInfo = new EquirectHdrInfoUniform();
 		this.samples = 0;
 
 		// kernels
-		this.kernel = new PathTracerMegaKernel().setWorkgroupSize( 8, 8, 1 );
-		this.setMaterial( this.material );
+		this.kernel = new PathTracerMegaKernel( this.rngData ).setWorkgroupSize( 8, 8, 1 );
 
 	}
 
@@ -32,12 +31,7 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 
 		this.kernel.textures = texture;
 		this.kernel.kernel.computeNode.parameters.textureSampler.node.value = texture;
-
-	}
-
-	getMaterial() {
-
-		return this.material;
+		this.reset();
 
 	}
 

src/webgpu/PathTracerBackend.js

@@ -1,11 +1,10 @@
 import { ColorManagement, FloatType, RGBAFormat } from 'three';
 import { RedIntegerFormat, StorageTexture, UnsignedIntType } from 'three/webgpu';
 import { ZeroOutKernel } from './compute/ZeroOutKernel.js';
-import { GltfCompliantMaterial } from './materials/GltfCompliantMaterial.js';
 
 export class PathTracerBackend {
 
-	constructor( renderer ) {
+	constructor( renderer, rngData ) {
 
 		this.renderer = renderer;
 		this.camera = null;
@@ -38,7 +37,7 @@ export class PathTracerBackend {
 		this.sampleCountClearKernel = new ZeroOutKernel().setWorkgroupSize( 8, 8, 1 );
 		this.outputTargetClearKernel = new ZeroOutKernel().setWorkgroupSize( 8, 8, 1 );
 
-		this.material = new GltfCompliantMaterial();
+		this.rngData = rngData;
 
 	}
 
@@ -122,13 +121,6 @@ export class PathTracerBackend {
 
 		}
 
-		if ( ! this.material.initialized ) {
-
-			this.material.init( renderer );
-			this.material.initialized = true;
-
-		}
-
 		if ( ! this._renderTask ) {
 
 			this._renderTask = this.createRenderTask();

src/webgpu/RNGData.js

@@ -0,0 +1,22 @@
+import { proxyFn } from './lib/nodes/NodeProxy';
+
+export class RNGData {
+
+	constructor() {
+
+		this.init = proxyFn( 'fns.init', this );
+		this.nextBounce = proxyFn( 'fns.nextBounce', this );
+		this.f32 = proxyFn( 'fns.f32', this );
+		this.vec2f = proxyFn( 'fns.vec2f', this );
+		this.vec3f = proxyFn( 'fns.vec3f', this );
+		this.vec4f = proxyFn( 'fns.vec4f', this );
+
+	}
+
+	setFunctions( fns ) {
+
+		this.fns = fns;
+
+	}
+
+}

src/webgpu/SobolNumberMapGenerator.js

@@ -0,0 +1,51 @@
+import { FloatType, MeshBasicNodeMaterial, NearestFilter, RenderTarget, RGBAFormat } from 'three/webgpu';
+import { FullScreenQuad } from 'three/examples/jsm/Addons.js';
+import { generateSobolPointFunc } from './nodes/random.wgsl';
+import { uv } from 'three/tsl';
+
+const _quad = new FullScreenQuad( new MeshBasicNodeMaterial() );
+export class SobolNumberMapGenerator {
+
+	constructor( renderer, dimensions ) {
+
+		this.target = new RenderTarget( dimensions, dimensions, {
+
+			type: FloatType,
+			format: RGBAFormat,
+			minFilter: NearestFilter,
+			maxFilter: NearestFilter,
+			generateMipmaps: false,
+
+		} );
+
+		this.renderer = renderer;
+		this.dimensions = dimensions;
+		this.isGenerated = false;
+
+	}
+
+	get texture() {
+
+		return this.target.texture;
+
+	}
+
+	generate() {
+
+		const { renderer, dimensions, target } = this;
+
+		const ogTarget = renderer.getRenderTarget();
+		renderer.setRenderTarget( target );
+
+		_quad.material.colorNode = generateSobolPointFunc(
+			uv().x.mul( dimensions ).toUint().add( uv().y.mul( dimensions ).toUint().mul( dimensions ) )
+		);
+		_quad.render( renderer );
+
+		renderer.setRenderTarget( ogTarget );
+
+		this.isGenerated = true;
+
+	}
+
+}

src/webgpu/WaveFrontPathTracer.js

@@ -18,9 +18,9 @@ const MAX_HIT_COUNT = Math.floor( MAX_BUFFER_SIZE / ( queuedHitStruct.getLength(
 
 export class WaveFrontPathTracer extends PathTracerBackend {
 
-	constructor( renderer ) {
+	constructor( renderer, rngData ) {
 
-		super( renderer );
+		super( renderer, rngData );
 
 		// options
 		this.tiles = new Vector2( 3, 3 );
@@ -44,10 +44,10 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 
 		// kernels
 		this.primeRayGenerationDispatchKernel = new PrimeRayGenerationDispatchKernel().setWorkgroupSize( 1, 1, 1 );
-		this.enqueueRaysKernel = new RayGenerationKernel().setWorkgroupSize( 8, 8, 1 );
-		this.rayIntersectionKernel = new RayIntersectionKernel().setWorkgroupSize( 64, 1, 1 );
+		this.enqueueRaysKernel = new RayGenerationKernel( this.rngData ).setWorkgroupSize( 8, 8, 1 );
+		this.rayIntersectionKernel = new RayIntersectionKernel( this.rngData ).setWorkgroupSize( 64, 1, 1 );
 		this.updateRayQueueParamsKernel = new UpdateRayQueueParamsKernel().setWorkgroupSize( 1, 1, 1 );
-		this.hitProcessKernel = new ProcessHitsKernel( this.material ).setWorkgroupSize( 64, 1, 1 );
+		this.hitProcessKernel = new ProcessHitsKernel( this.rngData ).setWorkgroupSize( 64, 1, 1 );
 
 		// clear kernels
 		this.zeroDispatchKernel = new ZeroOutBufferKernel().setWorkgroupSize( 1, 1, 1 );
@@ -58,6 +58,19 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 
 	}
 
+	setRandomFunctions( randomFunctions ) {
+
+		this.enqueueRaysKernel.random = randomFunctions;
+		this.enqueueRaysKernel.needsUpdate = true;
+
+		this.rayIntersectionKernel.random = randomFunctions;
+		this.rayIntersectionKernel.needsUpdate = true;
+
+		this.hitProcessKernel.random = randomFunctions;
+		this.hitProcessKernel.needsUpdate = true;
+
+	}
+
 	setBVHData( bvhData ) {
 
 		this.rayIntersectionKernel.bvhData = bvhData;
@@ -79,8 +92,8 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 
 	setMaterial( material ) {
 
-		this.material = material;
-		this.hitProcessKernel = new ProcessHitsKernel( this.material ).setWorkgroupSize( 64, 1, 1 );
+		this.hitProcessKernel.material = material.getData();
+		this.hitProcessKernel.needsUpdate = true;
 		this.reset();
 
 	}
@@ -172,6 +185,8 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 
 		super.reset();
 
+		this.enqueueRaysKernel.seed = 0;
+
 		const { width, height } = sampleCountTarget;
 		const dispatchSize = sampleCountClearKernel.getDispatchSize( width, height );
 

src/webgpu/WebGPUPathTracer.js

@@ -8,9 +8,18 @@ import { CubeToEquirectGenerator } from '../utils/CubeToEquirectGenerator.js';
 import { PathtracerBVHComputeData } from './nodes/PathtracerBVHComputeData.js';
 import { RenderTarget2DArray } from './RenderTarget2DArray.js';
 import { setCommonAttributes } from '../core/utils/GeometryPreparationUtils.js';
+import { pcgFunctions, sobolFunctions, sobolTextureFunctions } from './nodes/random.wgsl.js';
+import { GltfCompliantMaterial } from './materials/GltfCompliantMaterial.js';
+import { SobolNumberMapGenerator } from './SobolNumberMapGenerator.js';
+import { RNGData } from './RNGData.js';
 
 const _resolution = new Vector2();
 const _color = new Color();
+
+export const RNG_PCG = 0;
+export const RNG_SOBOL = 1;
+export const RNG_SOBOL_TEXTURE = 2;
+
 export class WebGPUPathTracer {
 
 	get bounces() {
@@ -41,9 +50,10 @@ export class WebGPUPathTracer {
 	useMegakernel( value ) {
 
 		this._pathTracer.dispose();
-		this._pathTracer = value ? new MegaKernelPathTracer( this._renderer ) : new WaveFrontPathTracer( this._renderer );
+		this._pathTracer = value ? new MegaKernelPathTracer( this._renderer, this.rngData ) : new WaveFrontPathTracer( this._renderer, this.rngData );
 		this._pathTracer.setBVHData( this._bvhData );
 		this._pathTracer.setTextures( this.textureArray.texture );
+		this._pathTracer.setMaterial( this.material );
 		this.setCamera( this.camera );
 		this.updateEnvironment();
 
@@ -53,7 +63,6 @@ export class WebGPUPathTracer {
 
 		// members
 		this._renderer = renderer;
-		this._pathTracer = new MegaKernelPathTracer( renderer );
 		this._timer = new Timer();
 
 		this._envColorTexture = new DataTexture( );
@@ -91,6 +100,14 @@ export class WebGPUPathTracer {
 
 		this.textureArray = new RenderTarget2DArray( 1024, 1024 );
 
+		this._sobolMap = new SobolNumberMapGenerator( renderer, 256 );
+		this.rngData = new RNGData();
+		this.setPRNGType( RNG_SOBOL_TEXTURE );
+
+		this.material = new GltfCompliantMaterial();
+		this.material.setRNGData( this.rngData );
+		this._pathTracer = new MegaKernelPathTracer( renderer, this.rngData );
+
 		// initialize the scene so it doesn't fail
 		this.setScene( new Scene(), new PerspectiveCamera() );
 
@@ -135,7 +152,7 @@ export class WebGPUPathTracer {
 		} );
 
 		// Build TLAS and compute functions
-		const bvhData = new PathtracerBVHComputeData( scene );
+		const bvhData = new PathtracerBVHComputeData( scene, this.rngData );
 		bvhData.update();
 		bvhData.useTransparencyRaycastFn( this.textureArray.texture );
 
@@ -152,12 +169,34 @@ export class WebGPUPathTracer {
 
 	getMaterial() {
 
-		return this._pathTracer.getMaterial();
+		return this.material;
+
+	}
+
+	setPRNGType( type ) {
+
+		switch ( type ) {
+
+			case RNG_PCG:
+				this.rngData.setFunctions( pcgFunctions );
+				break;
+
+			case RNG_SOBOL_TEXTURE:
+				this.rngData.setFunctions( sobolTextureFunctions( this._sobolMap.texture ) );
+				break;
+
+			case RNG_SOBOL:
+			default:
+				this.rngData.setFunctions( sobolFunctions );
+
+		}
 
 	}
 
 	setMaterial( material ) {
 
+		this.material = material;
+		this.material.setRNGData( this.rngData );
 		this._pathTracer.setMaterial( material );
 
 	}
@@ -256,6 +295,19 @@ export class WebGPUPathTracer {
 
 		}
 
+		if ( ! this.material.initialized ) {
+
+			this.material.init( renderer );
+			this.material.initialized = true;
+
+		}
+
+		if ( ! this._sobolMap.isGenerated ) {
+
+			this._sobolMap.generate();
+
+		}
+
 		const delta = 1000 * timer.getDelta();
 		this._resetTime += delta;