Merge pull request #750 from gkjohnson/webgpu/mobile-2

WebGPUPathTracer: Use custom float texture filter for mobile

Author: gkjohnson

package-lock.json

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

package.json

@@ -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/textures/ProceduralEquirectTexture.js

@@ -1,4 +1,4 @@
-import { ColorManagement, FloatType, LinearFilter, RGBAFormat } from 'three';
+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';
@@ -18,15 +18,13 @@ export class PathTracerBackend {
 		this.outputTarget = new StorageTexture( 1, 1, );
 		this.outputTarget.format = RGBAFormat;
 		this.outputTarget.type = FloatType;
-		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.magFilter = LinearFilter;
 		this.prevOutputTarget.colorSpace = ColorManagement.workingColorSpace;
 		this.prevOutputTarget.name = 'Output #1';
 		this.prevOutputTarget.generateMipmaps = false;

src/webgpu/PathTracerBackend.js

@@ -1,4 +1,5 @@
 import {
+	MeshBasicMaterial,
 	RenderTarget,
 	RGBAFormat,
 	UnsignedByteType,
@@ -7,7 +8,6 @@ import {
 	NoToneMapping,
 	QuadMesh,
 	NoBlending,
-	MeshBasicMaterial,
 } from 'three/webgpu';
 
 function getTextureHash( texture ) {

src/webgpu/RenderTarget2DArray.js

@@ -70,9 +70,12 @@ export class WebGPUPathTracer {
 		this._resetTime = - 1;
 		this._fadeState = 0;
 		this._size = new Vector2();
+		this._blitQuad = new FullScreenQuad( new RenderToScreenNodeMaterial() );
+
+		// avoid mipmap gen on copy, not always supported with float type
 		this._lowResTarget = new StorageTexture( 1, 1 );
 		this._lowResTarget.type = FloatType;
-		this._blitQuad = new FullScreenQuad( new RenderToScreenNodeMaterial() );
+		this._lowResTarget.generateMipmaps = false;
 
 		// options
 		this.minSamples = 1;

src/webgpu/WebGPUPathTracer.js

@@ -1,5 +1,5 @@
 import { wgslFn, texture, textureStore, globalId } from 'three/tsl';
-import { StorageTexture, RedFormat, LinearFilter, FloatType, TextureLoader } from 'three/webgpu';
+import { StorageTexture, RedFormat, LinearFilter, TextureLoader, HalfFloatType } from 'three/webgpu';
 import { wgslTagFn } from '../lib/nodes/WGSLTagFnNode';
 import { PathtracingMaterial } from './PathtracingMaterial';
 import { specularBrdfFunc, diffuseBrdfFunc, fresnelMixFunc, conductorFresnelFunc, albedoIntegralMetallic } from '../nodes/material.wgsl';
@@ -29,6 +29,7 @@ export class GltfCompliantMaterial extends PathtracingMaterial {
 		if ( calculateTurquinTexture ) {
 
 			this.turquinTexture = new StorageTexture( 32, 32 );
+			this.turquinTexture.type = HalfFloatType;
 
 		} else {
 
@@ -38,7 +39,6 @@ export class GltfCompliantMaterial extends PathtracingMaterial {
 		}
 
 		this.turquinTexture.format = RedFormat;
-		this.turquinTexture.type = FloatType;
 		this.turquinTexture.minFilter = LinearFilter;
 		this.turquinTexture.magFilter = LinearFilter;
 

src/webgpu/materials/GltfCompliantMaterial.js

@@ -1,5 +1,35 @@
 import { MeshBasicNodeMaterial, NoToneMapping, StorageTexture } from 'three/webgpu';
 import { uv, varying, texture, vec4, toneMapping, uniform, wgslFn } from 'three/tsl';
+import { wgslTagFn } from '../lib/nodes/WGSLTagFnNode.js';
+
+// TODO: we could fall back to hardware-based filtering if available but it has to be specifically
+// requested and available on the renderer which we don't have access to immediately. It's possible this
+// can be detected during build time via a custom node to adjust the sample approach used?
+const sampleTexelFn = wgslFn( /* wgsl */`
+	fn sampleTexel( tex: texture_2d<f32>, coord: vec2f ) -> vec4f {
+
+		// Manual bilinear filtering using textureLoad to support filterable float32 textures
+		// on all devices
+		let size = vec2f( textureDimensions( tex, 0 ) );
+		let pxCoord = coord * size - 0.5;
+		let px = vec2i( floor( pxCoord ) );
+		let fr = fract( pxCoord );
+
+		// get the four sibling samples
+		let s00 = textureLoad( tex, clamp( px,                 vec2i( 0 ), vec2i( size ) - 1 ), 0 );
+		let s10 = textureLoad( tex, clamp( px + vec2i( 1, 0 ), vec2i( 0 ), vec2i( size ) - 1 ), 0 );
+		let s01 = textureLoad( tex, clamp( px + vec2i( 0, 1 ), vec2i( 0 ), vec2i( size ) - 1 ), 0 );
+		let s11 = textureLoad( tex, clamp( px + vec2i( 1, 1 ), vec2i( 0 ), vec2i( size ) - 1 ), 0 );
+
+		// interpolate on the x axis
+		let y0 = mix( s00, s10, fr.x );
+		let y1 = mix( s01, s11, fr.x );
+
+		// then y
+		return mix( y0, y1, fr.y );
+
+	}
+` );
 
 // Material to apply tone mapping _before_ applying alpha blending
 export class RenderToScreenNodeMaterial extends MeshBasicNodeMaterial {
@@ -68,28 +98,41 @@ export class RenderToScreenNodeMaterial extends MeshBasicNodeMaterial {
 
 		super();
 
-		const texNode = texture( new StorageTexture(), varying( uv() ) );
+		const texNode = texture( new StorageTexture() );
 		this._texNode = texNode;
 
-		const fromTexNode = texture( new StorageTexture(), varying( uv() ) );
+		const fromTexNode = texture( new StorageTexture() );
 		this._fromTexNode = fromTexNode;
 
+		const texUV = varying( uv() );
+
 		const transitionUniform = uniform( 1.0 );
 		this._transitionUniform = transitionUniform;
 
-		const fadedColor = wgslFn( /* wgsl */`
-			fn fade( col0: vec4f, col1: vec4f, transition: f32 ) -> vec4f {
+		// NOTE: varyings cannot be referenced directly and must be passed as arguments
+		const getFadedColorFn = wgslTagFn/* wgsl */`
+			fn fade( uv: vec2f ) -> vec4f {
+
+				if ( ${ transitionUniform } <= 0.0 ) {
+
+					return ${ sampleTexelFn }( ${ fromTexNode }, uv );
+
+				} else if ( ${ transitionUniform } >= 1.0 ) {
+
+					return ${ sampleTexelFn }( ${ texNode }, uv );
+
+				} else {
+
+					let col0 = ${ sampleTexelFn }( ${ fromTexNode }, uv );
+					let col1 = ${ sampleTexelFn }( ${ texNode }, uv );
+					return mix( col0, col1, ${ transitionUniform } );
 
-				return mix( col0, col1, transition );
+				}
 
 			}
-		` )( {
-			col0: fromTexNode,
-			col1: texNode,
-			transition: transitionUniform,
-		} );
+		`;
 
-		const toneMappingNode = toneMapping( NoToneMapping, 1.0, fadedColor );
+		const toneMappingNode = toneMapping( NoToneMapping, 1.0, getFadedColorFn( texUV ) );
 		this._toneMapping = toneMappingNode;
 
 		// apply alpha _after_ applying tone mapping

src/webgpu/materials/RenderToScreenMaterial.js

@@ -358,7 +358,7 @@ export const conductorFresnelFunc = ( turquinTexture ) => wgslFn( /* wgsl */ `
 	  let ss = bsdf * schlickFresnelVec( abs( VdotH ), f0, vec3f( 1 ) );
 
 		let uv = vec2( NdotV, sqrt( alpha ) );
-		let energySs = max( textureSampleLevel( turquinTexture, turquinTexture_sampler, uv, 0 ).r, 1e-5);
+		let energySs = max( textureSampleLevel( turquinTexture, turquinTexture_sampler, uv, 0 ).r, 1e-5 );
 
 		return ss * ( 1.0 + f0 * ( 1.0 - energySs ) / energySs );
 
@@ -371,7 +371,7 @@ export const conductorFresnelFunc = ( turquinTexture ) => wgslFn( /* wgsl */ `
 export const albedoIntegralMetallic = wgslFn( /* wgsl */ `
 
 	fn albedo(
-		texture: texture_storage_2d<r32float, write>,
+		texture: texture_storage_2d<r16float, write>,
 
 		globalId: vec3u,
 	) -> void {
@@ -413,7 +413,7 @@ export const albedoIntegralMetallic = wgslFn( /* wgsl */ `
 
 		result /= f32( INTEGRATION_SAMPLES );
 
-		textureStore(texture, globalId.xy, vec4( result ));
+		textureStore( texture, globalId.xy, vec4( result ) );
 
 	}
 

src/webgpu/nodes/material.wgsl.js

@@ -1,7 +1,10 @@
 import { searchForWorkspaceRoot } from 'vite';
+import basicSsl from '@vitejs/plugin-basic-ssl';
 import fs from 'fs';
 
-export default {
+export default ( { mode } ) => ( {
+
+	plugins: mode === 'ssl' ? [ basicSsl() ] : [],
 
 	root: './example/',
 	base: '',
@@ -26,4 +29,4 @@ export default {
 	optimizeDeps: {
     	exclude: [ 'three-mesh-bvh' ],
   	},
-};
+} );