@meimu/mutracer 0.1.5 → 0.3.0

package/lib/renderer/WebGLRayTracePass.js

@@ -1,3 +1,11 @@
+import { researchTransport } from './researchTransport.js';
+import { EmissiveTriangleCache } from './lightSampling.js';
+import { prepareWebGLLightSampling } from './webglLightSampling.js';
+import { AovPrimitiveCache } from './AovPrimitiveCache.js';
+import { aovTraceFragment } from './glsl/aovTrace.frag.js';
+import { packAovScene } from './aovScene.js';
+import { aovMask } from './aov.js';
+import { WebGLSkySource } from './WebGLSkySource.js';
 import { normalizeBounces } from './pathLimits.js';
 import { updateMeshInstanceColors } from './prepareMeshBVH.js';
 import { updateTLASInstanceColors } from './bvhAccel.js';
@@ -18,7 +26,7 @@ import { constants, Texture2D as WebGLTexture2D } from './webgl.js';
 import { WebGLFullQuadPass } from './WebGLFullQuadPass.js';
 import { mat4 } from '../math.js';
 import { buildBLAS, buildTLAS, concatBLASBuffers, refitBLAS, refitTLAS, mergeRefitRanges } from './bvhAccel.js';
-import { generateBackgroundMapFromSceneBackground, environmentRotation, backgroundMapTransform } from './envMapCreation.js';
+import { generateBackgroundMapFromSceneBackground, environmentRotation, backgroundMapTransform, backgroundMapRotation } from './envMapCreation.js';
 import { envMapDistribution, environmentMapSourcesMatch } from './envMapDistribution.js';
 import createRayTraceFragmentShader from './glsl/rayTrace.frag.js';
 import shadowCatcherFrag from './glsl/shadowCatcher.frag.js';
@@ -56,8 +64,13 @@ function packLightsToTexture(texture, lights, texturesIndices, selectionPmfs, se
 // TODO use one method to set all required properties.
 // Ensure necessary properties are set.
 class WebGLRayTracePass {
+    _aovPrimitiveCache = new AovPrimitiveCache();
     _renderer;
     _meshes = [];
+    _emissiveActive = false;
+    _directLightLegacy;
+    _emissiveCache = new EmissiveTriangleCache();
+    _emissiveTex = new WebGLTexture2D(commonNearestFloatProps);
     _sceneEnvironment = new SceneEnvironmentCache();
     _rectLightTextures = new RectLightTextureTracker();
     _analyticLights = [];
@@ -100,12 +113,14 @@ class WebGLRayTracePass {
     _envMap = new WebGLTexture2D(commonLinearFloatProps);
     _envMapDistribution = new WebGLTexture2D(commonLinearFloatProps);
     _envImageSource;
+    _skySource = new WebGLSkySource();
     _backgroundTex = new WebGLTexture2D(commonLinearFloatProps);
     _hasEnvMap = false;
     _hasBackground = false;
     _groundedBackgroundParams = new Float32Array([0, 0, 0, 0]);
     _groundedBackgroundCenter = new Float32Array([0, 0]);
     _backgroundMapTransform = backgroundMapTransform();
+    _backgroundMapRotation = new Float32Array([1, 0, 0, 0, 1, 0, 0, 0, 1]);
     _bounces = 2;
     _bevelSamples = DEFAULT_BEVEL_SAMPLES;
     _clamping;
@@ -134,6 +149,8 @@ class WebGLRayTracePass {
     _viewProjInv = mat4.create();
     _apertureRadius = 0;
     _focusDistance = 10;
+    _aovPass;
+    _aovCenter;
     _tracePasses;
     _depthPass;
     _depthAbort = new AbortController();
@@ -176,13 +193,22 @@ class WebGLRayTracePass {
             shader: depthQueryFrag, outputs: ['depth'],
             outputParams: { depth: commonNearestFloatProps }
         });
-        this._tracePasses = [this._mainPass, this._shadowCatcherPass, this._depthPass];
+        this._aovPass = new WebGLFullQuadPass('AOV accumulation', renderer, { shader: aovTraceFragment(false), outputs: ['geometry', 'position', 'diffuse', 'specular'], outputParams: Object.fromEntries(['geometry', 'position', 'diffuse', 'specular'].map(name => [name, commonNearestFloatProps])) });
+        this._aovCenter = new WebGLFullQuadPass('AOV center', renderer, { shader: aovTraceFragment(true), outputs: ['id', 'motion'], outputParams: { id: { ...commonNearestFloatProps, type: 5125 }, motion: commonNearestFloatProps } });
+        this._tracePasses = [this._aovPass, this._aovCenter, this._mainPass, this._shadowCatcherPass, this._depthPass];
         this._tracePasses.forEach((pass) => {
             pass.setDefines({
                 OES_texture_float_linear: !!renderer.gl.getExtension('OES_texture_float_linear')
             });
         });
         this._mainPass.setUniform('openPbrEnergyLut', this._openPbrEnergyLut);
+        this._aovPass.setUniform('openPbrEnergyLut', this._openPbrEnergyLut);
+    }
+    preparePrograms() {
+        this.prepare();
+        const main = this._mainPass.isProgramReady();
+        const shadow = !this._hasShadowCatcher || this._shadowCatcherPass.isProgramReady();
+        return main && shadow;
     }
     createPass(clear, tile, disableShadowCatcher = false) {
         this.prepare();
@@ -239,7 +265,9 @@ class WebGLRayTracePass {
     }
     setBackground(background) {
         this._backgroundMapTransform.set(backgroundMapTransform(background));
-        const backgroundImage = background && generateBackgroundMapFromSceneBackground(background);
+        const rotation = backgroundMapRotation(background);
+        this._backgroundMapRotation.set([rotation[0], rotation[1], rotation[2], rotation[4], rotation[5], rotation[6], rotation[8], rotation[9], rotation[10]]);
+        const backgroundImage = background ? this._skySource.resolve(generateBackgroundMapFromSceneBackground(background)) : undefined;
         if (backgroundImage) {
             this._backgroundTex.source = backgroundImage;
         }
@@ -391,10 +419,14 @@ class WebGLRayTracePass {
         }
         return this._debugBvhData;
     }
+    /** @internal */
+    _loadDiagnosticAccumulation(inputs, catcher) {
+        this._mainPass.loadDiagnosticOutputs({ color: inputs.color, albedo: inputs.albedo, normal: inputs.normal, moments: inputs.moments });
+        this._shadowCatcherPass.loadDiagnosticOutputs({ shadowCatcher: inputs.shadowCatcher, catcherBackground: inputs.catcherBackground, catcherMoments: inputs.catcherMoments });
+        this._hasShadowCatcher = catcher;
+    }
     getDiagnostics() {
-        return {
-            textureAtlas: this._materialData?.textureAtlasStats
-        };
+        return { textureAtlas: this._materialData?.textureAtlasStats };
     }
     _updateMeshData() {
         this._pendingGeometryUpdates.clear();
@@ -635,7 +667,7 @@ class WebGLRayTracePass {
             light.type === 'SpotLight' ||
             light.type === 'RectLight');
         const numAnalyticLights = analyticLights.length;
-        const envImage = this._sceneEnvironment.update(lights);
+        const envImage = this._skySource.resolve(this._sceneEnvironment.update(lights));
         const rotation = environmentRotation(lights);
         const rotation3 = [rotation[0], rotation[1], rotation[2], rotation[4], rotation[5], rotation[6], rotation[8], rotation[9], rotation[10]];
         this._tracePasses.forEach(pass => pass.setUniform('environmentRotation', rotation3));
@@ -654,6 +686,7 @@ class WebGLRayTracePass {
             ? (1 - POWER_WEIGHTED_LIGHT_MIX) * fallbackPmf +
                 POWER_WEIGHTED_LIGHT_MIX * (weight / totalWeight)
             : fallbackPmf);
+        this._directLightLegacy = { env: this._envLightSelectionPmf, pmfs: this._analyticLightSelectionPmfs };
         let cumulativePmf = this._envLightSelectionPmf;
         this._analyticLightSelectionCdfs = Float32Array.from(this._analyticLightSelectionPmfs, (pmf) => {
             cumulativePmf += pmf;
@@ -726,15 +759,16 @@ class WebGLRayTracePass {
         this._shadowCatcherPass.useTarget(target);
     }
     setQuality(options) {
+        researchTransport()?.configure(this, 'webgl');
         this._tracePasses.forEach(pass => {
             pass.setUniform('bounceLimit', normalizeBounces(options.bounces));
             pass.setUniform('bevelEnabled', options.bevel);
             pass.setUniform('dispersionWavelengths', options.dispersionWavelengths);
         });
     }
-    setWavelength(wavelength) {
+    setSpectralChannel(channel) {
         this._tracePasses.forEach((pass) => {
-            pass.setUniform('wavelength', wavelength);
+            pass.setUniform('spectralChannel', channel);
         });
     }
     setSampleCount(sampleCount) {
@@ -766,6 +800,29 @@ class WebGLRayTracePass {
             this._updateMeshColors();
         if (this._pendingMaterialUpdate)
             this._updateMaterialData();
+        researchTransport()?.prepare(this, 'webgl');
+        const res = this._preparedMeshData;
+        if (res) {
+            const triangles = this._emissiveCache.collect(res.emissiveBindings, res.materials);
+            this._mainPass.setDefines({ HAS_EMISSIVE_LIGHTS: triangles.length > 0 });
+            if (!triangles.length && this._emissiveActive && this._directLightLegacy) {
+                this._envLightSelectionPmf = this._directLightLegacy.env;
+                this._analyticLightSelectionPmfs = Float32Array.from(this._directLightLegacy.pmfs);
+                this._tracePasses.forEach(pass => pass.setUniform('envLightSelectionPmf', this._envLightSelectionPmf));
+                this._needsUpdateAnalyticLightsBuffer = true;
+            }
+            this._emissiveActive = triangles.length > 0;
+            if (triangles.length) {
+                const lighting = prepareWebGLLightSampling(this._analyticLights, triangles, estimateEnvironmentPower(this._envImageSource));
+                this._envLightSelectionPmf = lighting.environmentPmf;
+                this._analyticLightSelectionPmfs = Float32Array.from(lighting.analyticPmfs);
+                const columns = setDataTexture(this._emissiveTex, lighting.data, 4).columnsLog;
+                this._mainPass.setUniforms({ emissiveLightData: this._emissiveTex, emissiveLightColumns: columns,
+                    emissiveChoiceOffset: lighting.emissiveTexels, emissiveChoiceCount: lighting.choiceCount,
+                    emissivePmfOffset: lighting.pmfOffset, emissiveWeightTotal: lighting.totalWeight, envLightSelectionPmf: lighting.environmentPmf });
+                this._needsUpdateAnalyticLightsBuffer = true;
+            }
+        }
         if (this._needsUpdateAnalyticLightsBuffer) {
             const analyticLights = this._analyticLights;
             analyticLights.length > 0 &&
@@ -785,6 +842,7 @@ class WebGLRayTracePass {
             pass.setUniform('groundedBackgroundParams', this._groundedBackgroundParams);
             pass.setUniform('groundedBackgroundCenter', this._groundedBackgroundCenter);
             pass.setUniform('backgroundMapTransform', this._backgroundMapTransform);
+            pass.setUniform('backgroundMapRotation', this._backgroundMapRotation);
             pass.setDefines({
                 // Double bounce fix on intel mac
                 // Limit the bounce to avoid long compilation time.
@@ -816,6 +874,57 @@ class WebGLRayTracePass {
     }
     // WebGL has uniform per-pixel counts; setSampleCount already includes offset.
     setSampleSequenceOffset(_offset) { }
+    prepareAovs(state, camera, width, height, kinds) {
+        if (!kinds.length)
+            return true;
+        // An AOV channel may be the first thing to render a new scene.
+        this.prepare();
+        if (!this._preparedMeshData || !this._tlas)
+            return true;
+        const gl = this._renderer.gl;
+        if (!gl.getExtension('EXT_color_buffer_float'))
+            throw new Error('AOV requires EXT_color_buffer_float');
+        state.resize(width, height);
+        this._aovPass.setDefines({ AOV_REFLECTANCE: (aovMask(kinds) & 24) !== 0 });
+        if (!state.frame) {
+            const packed = packAovScene(this._preparedMeshData, this._tlas, camera, width, height, state.history, this._aovPrimitiveCache.get(this._preparedMeshData, this._tlas));
+            state.frame = packed.frame;
+            state.setScene(packed.words);
+        }
+        this._aovCenter.setUniform('aovScene', state.scene);
+        this._aovCenter.setUniform('previousId', state.previousCenter.id);
+        this._aovCenter.setUniform('previousMotion', state.previousCenter.motion);
+        this._aovPass.setUniform('aovScene', state.scene);
+        this._aovCenter.setSize(width, height);
+        this._aovPass.setSize(width, height);
+        const previous = state.accumulation[(state.samples + 1) % 2];
+        const center = this._aovCenter.isProgramReady();
+        const accumulation = this._aovPass.isProgramReady({ previousGeometry: previous.geometry, previousPosition: previous.position, previousDiffuse: previous.diffuse, previousSpecular: previous.specular });
+        return center && accumulation;
+    }
+    renderAovs(state, camera, width, height, samples, cap, kinds, sequenceOffset = 0) {
+        if (!this.prepareAovs(state, camera, width, height, kinds))
+            return false;
+        if (!kinds.length || !this._preparedMeshData || !this._tlas)
+            return true;
+        if (!state.centerReady) {
+            this._aovCenter.createPass({}, state.center);
+            state.centerReady = true;
+            state.history.accept(state.frame);
+        }
+        const target = Math.min(samples, cap);
+        if ((aovMask(kinds) & 31) === 0)
+            state.samples = target;
+        while (state.samples < target) {
+            const output = state.accumulation[state.samples % 2], previous = state.accumulation[(state.samples + 1) % 2];
+            this._aovPass.setUniform('aovMask', aovMask(kinds));
+            this._aovPass.setUniform('aovSample', state.samples);
+            this._aovPass.setUniform('sampleCount', state.samples + sequenceOffset);
+            this._aovPass.createPass({ previousGeometry: previous.geometry, previousPosition: previous.position, previousDiffuse: previous.diffuse, previousSpecular: previous.specular }, output);
+            state.samples++;
+        }
+        return true;
+    }
     async readDepthSamples(area, aspect) {
         if (this._depthAbort.signal.aborted || this._invalid || !this._preparedMeshData ||
             !this._renderer.isAllProgramCompiled())
@@ -827,6 +936,8 @@ class WebGLRayTracePass {
         return unpackDepthSamples(rgba);
     }
     dispose() {
+        researchTransport()?.dispose(this, 'webgl');
+        this._skySource.dispose();
         this._depthAbort.abort();
         this._materialData?.dispose(this._renderer.gl);
         this._materialData = undefined;
@@ -834,6 +945,7 @@ class WebGLRayTracePass {
         const gl = this._renderer.gl;
         [
             this._lightTex,
+            this._emissiveTex,
             this._openPbrEnergyLut,
             this._envMap,
             this._envMapDistribution,