diff --git a/src/webgpu/MegaKernelPathTracer.js b/src/webgpu/MegaKernelPathTracer.js index 82b1c6e0..64ad762b 100644 --- a/src/webgpu/MegaKernelPathTracer.js +++ b/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; } diff --git a/src/webgpu/PathTracerCore.js b/src/webgpu/PathTracerCore.js index 5a3973aa..428d58a7 100644 --- a/src/webgpu/PathTracerCore.js +++ b/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; diff --git a/src/webgpu/WaveFrontPathTracer.js b/src/webgpu/WaveFrontPathTracer.js index 5aad3a6a..1723f442 100644 --- a/src/webgpu/WaveFrontPathTracer.js +++ b/src/webgpu/WaveFrontPathTracer.js @@ -1,4 +1,4 @@ -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'; @@ -6,6 +6,7 @@ import { RayIntersectionKernel } from './compute/wavefront/RayIntersectionKernel 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; } diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index f592a513..de621968 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/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; + +} diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index cea76dd4..3f7d4505 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/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, + envMapSampler: sampler, + envMapRotation: mat3x3f, + envMapIntensity: f32, + + background: texture_2d, + 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 ) ); diff --git a/src/webgpu/compute/wavefront/RayGenerationKernel.js b/src/webgpu/compute/wavefront/RayGenerationKernel.js index 696103f0..a50dba5f 100644 --- a/src/webgpu/compute/wavefront/RayGenerationKernel.js +++ b/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; diff --git a/src/webgpu/compute/wavefront/RayIntersectionKernel.js b/src/webgpu/compute/wavefront/RayIntersectionKernel.js index 1c0e2898..ea6fbeb1 100644 --- a/src/webgpu/compute/wavefront/RayIntersectionKernel.js +++ b/src/webgpu/compute/wavefront/RayIntersectionKernel.js @@ -1,9 +1,10 @@ -import { IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; +import { DataTexture, Matrix3, IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; import { ComputeKernel } from '../ComputeKernel.js'; -import { storage, wgslFn, textureStore, globalId } from 'three/tsl'; -import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; +import { uniform, texture, sampler, storage, wgslFn, textureStore, globalId } from 'three/tsl'; +import { pcgRand2, pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; import { proxy } from '../../lib/nodes/NodeProxy.js'; +import { sampleEnvironmentFn } from '../../nodes/sampling.wgsl.js'; export class RayIntersectionKernel extends ComputeKernel { @@ -23,6 +24,18 @@ export class RayIntersectionKernel extends ComputeKernel { hitQueue: storage( new IndirectStorageBufferAttribute( 1, QUEUED_HIT_SIZE ), 'QueuedHit' ), hitQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ).toAtomic(), + // 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 ), + globalId: globalId, }; @@ -42,9 +55,33 @@ export class RayIntersectionKernel extends ComputeKernel { hitQueue: ptr, read_write>, hitQueueSize: ptr>, read_write>, + // environment + envMap: texture_2d, + envMapSampler: sampler, + envMapRotation: mat3x3f, + envMapIntensity: f32, + + background: texture_2d, + backgroundSampler: sampler, + backgroundRotation: mat3x3f, + backgroundIntensity: f32, + backgroundBlurriness: f32, + globalId: vec3u ) -> void { + let envInfo = EnvironmentInfo( + envMapRotation, + envMapIntensity, + 0.0 // blur, + ); + + let backgroundInfo = EnvironmentInfo( + backgroundRotation, + backgroundIntensity, + backgroundBlurriness, + ); + // skip any rays invocations beyond the ray count let queueCapacity = arrayLength( rayQueue ); let rayIndex = ( globalId.x + rayQueueSize[ 0 ] ); @@ -75,12 +112,21 @@ export class RayIntersectionKernel extends ComputeKernel { hitQueue[ index ].pixel_y = input.pixel.y; hitQueue[ index ].objectIndex = hitResult.objectIndex; hitQueue[ index ].throughputColor = input.throughputColor; - hitQueue[ index ].currentBounce = input.currentBounce;; + hitQueue[ index ].currentBounce = input.currentBounce; } else { - let background = vec3f( 0.5 ); - let newColor = background * input.throughputColor; + var light: vec3f; + if ( input.currentBounce > 0u ) { + + light = sampleEnvironment( envMap, envMapSampler, envInfo, input.ray.direction, pcgRand2() ); + + } else { + + light = sampleEnvironment( background, backgroundSampler, backgroundInfo, input.ray.direction, pcgRand2() ); + + } + let newColor = light * input.throughputColor; let sampleCount = ( textureLoad( sampleCountTarget, indexUV ).r & ( ~ ACTIVE_FLAG ) ) + 1; var color = textureLoad( prevOutputTarget, indexUV ).xyz; @@ -95,7 +141,9 @@ export class RayIntersectionKernel extends ComputeKernel { `, [ proxy( 'bvhData.value.fns.raycastFirstHit', parameters ), proxy( 'bvhData.value.structs.material', parameters ), - queuedRayStruct, pcgRand3, pcgInit, queuedHitStruct ] ); + queuedRayStruct, pcgRand2, pcgRand3, pcgInit, queuedHitStruct, + sampleEnvironmentFn + ] ); super( fn( parameters ) ); diff --git a/src/webgpu/nodes/megakernel.wgsl.js b/src/webgpu/nodes/megakernel.wgsl.js deleted file mode 100644 index 9ef2e7f8..00000000 --- a/src/webgpu/nodes/megakernel.wgsl.js +++ /dev/null @@ -1,120 +0,0 @@ -import { wgslFn } from 'three/tsl'; -import { ndcToCameraRay, bvhIntersectFirstHit, constants, getVertexAttribute } from 'three-mesh-bvh/webgpu'; -import { pcgRand3, pcgInit } from './random.wgsl.js'; -import { lambertBsdfFunc } from './sampling.wgsl.js'; -import { materialStruct, surfaceRecordStruct } from './structs.wgsl.js'; - -export const megakernelShader = wgslFn( /* wgsl */` - - fn compute( - - // indices and target - globalId: vec3u, - prevOutputTarget: texture_storage_2d, - outputTarget: texture_storage_2d, - sampleCountTarget: texture_storage_2d, - - // tiles - offset: vec2u, - tileSize: vec2u, - - // settings - smoothNormals: u32, - inverseProjectionMatrix: mat4x4f, - cameraToModelMatrix: mat4x4f, - seed: u32, - bounces: u32, - - // scene - geom_position: ptr, read>, - geom_index: ptr, read>, - geom_normals: ptr, read>, - geom_material_index: ptr, read>, - bvh: ptr, read>, - - materials: ptr, read>, - - ) -> void { - - // make sure we don't bleed over the edge of our tile - if ( globalId.x >= tileSize.x || globalId.y >= tileSize.y ) { - - return; - - } - - // to screen coordinates - let indexUV = offset + globalId.xy; - let targetDimensions = textureDimensions( outputTarget ); - if ( indexUV.x >= targetDimensions.x || indexUV.y >= targetDimensions.y ) { - - return; - - } - - let uv = vec2f( indexUV ) / vec2f( targetDimensions ); - let ndc = uv * 2.0 - vec2f( 1.0 ); - - pcgInitialize( indexUV, seed ); - - // scene ray - // TODO: jittering the ray by [-1, 1] seems to look better but is larger than a pixel? - var jitter = 2.0 * ( pcgRand2() - vec2( 0.5 ) ) / vec2f( targetDimensions.xy ); - var ray = ndcToCameraRay( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix ); - - var resultColor = vec3f( 0.0 ); - var throughputColor = vec3f( 1.0 ); - - // TODO: fix shadow acne? RTIOW says we could just ignore ray hits that are too close - for ( var bounce = 0u; bounce < bounces; bounce ++ ) { - - let hitResult = bvhIntersectFirstHit( geom_index, geom_position, bvh, ray ); - - // write result - if ( hitResult.didHit ) { - - let material = materials[ geom_material_index[ hitResult.indices.x ] ]; - // var surfaceRecord: SurfaceRecord; - // surfaceRecord.normal = hitResult.normal; - // surfaceRecord.albedo = material.albedo; - // surfaceRecord.roughness = material.roughness; - // surfaceRecord.metalness = material.metalness; - - let hitPosition = getVertexAttribute( hitResult.barycoord, hitResult.indices.xyz, geom_position ); - let hitNormal = getVertexAttribute( hitResult.barycoord, hitResult.indices.xyz, geom_normals ); - - let scatterRec = bsdfEval( hitNormal, - ray.direction ); - // let scatterRec = bsdfEval(hitResult.normal, - ray.direction); - // TODO: fix shadow acne - // if (bounce == 1) { - // resultColor = vec3f( 0.0, 1.0, 0.0 ); // dot( scatterRec.direction, hitNormal ) ); // ( vec3f( 1.0 ) + scatterRec.direction ) * 0.5; - // sampleCount = 1; - // break; - // } - - throughputColor *= material.albedo * scatterRec.value / scatterRec.pdf; - - ray.origin = hitPosition; - ray.direction = scatterRec.direction; - - } else { - - let background = vec3f( 0.5 ); - resultColor += background * throughputColor; - break; - - } - - } - - let sampleCount = textureLoad( sampleCountTarget, indexUV ).r + 1; - var color = textureLoad( prevOutputTarget, indexUV ).xyz; - color += ( resultColor - color.xyz ) / f32( sampleCount ); - - textureStore( sampleCountTarget, indexUV, vec4( sampleCount ) ); - textureStore( outputTarget, indexUV, vec4( color, 1.0 ) ); - - } -`, [ ndcToCameraRay, bvhIntersectFirstHit, constants, getVertexAttribute, materialStruct, surfaceRecordStruct, pcgRand3, pcgInit, lambertBsdfFunc ] ); - -export default megakernelShader; diff --git a/src/webgpu/nodes/sampling.wgsl.js b/src/webgpu/nodes/sampling.wgsl.js index f1e51681..23604282 100644 --- a/src/webgpu/nodes/sampling.wgsl.js +++ b/src/webgpu/nodes/sampling.wgsl.js @@ -1,6 +1,6 @@ import { wgslFn } from 'three/tsl'; import { pcgRand2 } from './random.wgsl.js'; -import { scatterRecordStruct, constants } from './structs.wgsl.js'; +import { scatterRecordStruct, environmentInfoStruct, constants } from './structs.wgsl.js'; // TODO: Move to a local (s, t, n) coordinate system // From RayTracingGems v1.9 chapter 16.6.2 -- Its shit! @@ -37,25 +37,67 @@ export const lambertBsdfFunc = wgslFn( /* wgsl */` } `, [ scatterRecordStruct, sampleSphereCosineFn, pcgRand2, constants ] ); -// const equirectDirectionToUvFn = wgslFn( /* wgsl */` -// fn equirectDirectionToUv(direction: vec3f) -> vec2f { -// -// // from Spherical.setFromCartesianCoords -// vec2 uv = vec2f( atan2( direction.z, direction.x ), acos( direction.y ) ); -// uv /= vec2f( 2.0 * PI, PI ); -// -// // apply adjustments to get values in range [0, 1] and y right side up -// uv.x += 0.5; -// uv.y = 1.0 - uv.y; -// return uv; -// -// } -// ` ); - -// const sampleEquirectColorFn = wgslFn( /* wgsl */ ` -// fn sampleEquirectColor( envMap: texture_2d, envMapSampler: sampler, direction: vec3f ) -> vec3f { - -// return texture2D( envMap, equirectDirectionToUv( direction ) ).rgb; - -// } -// `, [ equirectDirectionToUvFn ] ); +const equirectDirectionToUvFn = wgslFn( /* wgsl */` + fn equirectDirectionToUv(direction: vec3f) -> vec2f { + + // from Spherical.setFromCartesianCoords + var uv = vec2f( atan2( direction.z, direction.x ), acos( direction.y ) ); + uv /= vec2f( 2.0 * PI, PI ); + + // apply adjustments to get values in range [0, 1] and y right side up + uv.x += 0.5; + uv.y = 1.0 - uv.y; + return uv; + + } +` ); + +const sampleEquirectColorFn = wgslFn( /* wgsl */ ` + fn sampleEquirectColor( envMap: texture_2d, envMapSampler: sampler, direction: vec3f ) -> vec3f { + + return textureSampleLevel( envMap, envMapSampler, equirectDirectionToUv( direction ), 0 ).rgb; + + } +`, [ equirectDirectionToUvFn ] ); + +const sampleHemisphereFn = wgslFn( /* wgsl */ ` + + fn sampleHemisphere( n: vec3f, uv: vec2f ) -> vec3f { + + // https://www.rorydriscoll.com/2009/01/07/better-sampling/ + // https://graphics.pixar.com/library/OrthonormalB/paper.pdf + let sign = select( sign( n.z ), 1.0, n.z == 0.0 ); + let a = - 1.0 / ( sign + n.z ); + let b = n.x * n.y * a; + let b1 = vec3( 1.0 + sign * n.x * n.x * a, sign * b, - sign * n.x ); + let b2 = vec3( b, sign + n.y * n.y * a, - n.y ); + + let r = sqrt( uv.x ); + let theta = 2.0 * PI * uv.y; + let x = r * cos( theta ); + let y = r * sin( theta ); + return x * b1 + y * b2 + sqrt( 1.0 - uv.x ) * n; + + } + +`, [ constants ] ); + +export const sampleEnvironmentFn = wgslFn( /* wgsl */ ` + + fn sampleEnvironment( + envMap: texture_2d, + envMapSampler: sampler, + env: EnvironmentInfo, + direction: vec3f, + uv: vec2f, + ) -> vec3f { + + let offsetDir = sampleHemisphere( direction, uv ) * 0.5 * env.blur; + + let sampleDir = normalize( env.rotation * direction + offsetDir ); + + return env.intensity * sampleEquirectColor( envMap, envMapSampler, sampleDir ); + + } + +`, [ sampleEquirectColorFn, sampleHemisphereFn, environmentInfoStruct ] ); diff --git a/src/webgpu/nodes/structs.wgsl.js b/src/webgpu/nodes/structs.wgsl.js index 8d872b77..a0f84cbe 100644 --- a/src/webgpu/nodes/structs.wgsl.js +++ b/src/webgpu/nodes/structs.wgsl.js @@ -64,3 +64,14 @@ export const hitResultQueueElementStruct = wgsl( /* wgsl */` }; ` ); +export const environmentInfoStruct = wgsl( /* wgsl */ ` + + struct EnvironmentInfo { + + rotation: mat3x3f, + intensity: f32, + blur: f32, + + }; + +` );