TSL: Add bitCount functions (#31990)

Author: cmhhelgeson

examples/jsm/tsl/display/SSGINode.js

@@ -1,5 +1,5 @@
 import { RenderTarget, Vector2, TempNode, QuadMesh, NodeMaterial, RendererUtils, MathUtils } from 'three/webgpu';
-import { clamp, normalize, reference, nodeObject, Fn, NodeUpdateType, uniform, vec4, passTexture, uv, logarithmicDepthToViewZ, viewZToPerspectiveDepth, getViewPosition, screenCoordinate, float, sub, fract, dot, vec2, rand, vec3, Loop, mul, PI, cos, sin, uint, cross, acos, sign, pow, luminance, If, max, abs, Break, sqrt, HALF_PI, div, ceil, shiftRight, convertToTexture, bool, getNormalFromDepth, interleavedGradientNoise } from 'three/tsl';
+import { clamp, normalize, reference, nodeObject, Fn, NodeUpdateType, uniform, vec4, passTexture, uv, logarithmicDepthToViewZ, viewZToPerspectiveDepth, getViewPosition, screenCoordinate, float, sub, fract, dot, vec2, rand, vec3, Loop, mul, PI, cos, sin, uint, cross, acos, sign, pow, luminance, If, max, abs, Break, sqrt, HALF_PI, div, ceil, shiftRight, convertToTexture, bool, getNormalFromDepth, countOneBits, interleavedGradientNoise } from 'three/tsl';
 
 const _quadMesh = /*@__PURE__*/ new QuadMesh();
 const _size = /*@__PURE__*/ new Vector2();
@@ -435,22 +435,6 @@ class SSGINode extends TempNode {
 			]
 		} );
 
-		const bitCount = Fn( ( [ value ] ) => {
-
-			const v = uint( value );
-			v.assign( v.sub( v.shiftRight( uint( 1 ) ).bitAnd( uint( 0x55555555 ) ) ) );
-			v.assign( v.bitAnd( uint( 0x33333333 ) ).add( v.shiftRight( uint( 2 ) ).bitAnd( uint( 0x33333333 ) ) ) );
-
-			return v.add( v.shiftRight( uint( 4 ) ) ).bitAnd( uint( 0xF0F0F0F ) ).mul( uint( 0x1010101 ) ).shiftRight( uint( 24 ) );
-
-		} ).setLayout( {
-			name: 'bitCount',
-			type: 'uint',
-			inputs: [
-				{ name: 'value', type: 'uint' }
-			]
-		} );
-
 		const horizonSampling = Fn( ( [ directionIsRight, RADIUS, viewPosition, slideDirTexelSize, initialRayStep, uvNode, viewDir, viewNormal, n ] ) => {
 
 			const STEP_COUNT = this.stepCount.toConst();
@@ -513,7 +497,7 @@ class SSGINode extends TempNode {
 				currentOccludedBitfield = currentOccludedBitfield.bitAnd( globalOccludedBitfield.bitNot() );
 
 				globalOccludedBitfield.assign( globalOccludedBitfield.bitOr( currentOccludedBitfield ) );
-				const numOccludedZones = bitCount( currentOccludedBitfield );
+				const numOccludedZones = countOneBits( currentOccludedBitfield );
 
 				//
 
@@ -597,7 +581,7 @@ class SSGINode extends TempNode {
 				color.addAssign( horizonSampling( bool( true ), RADIUS, viewPosition, slideDirTexelSize, initialRayStep, uvNode, viewDir, viewNormal, n ) );
 				color.addAssign( horizonSampling( bool( false ), RADIUS, viewPosition, slideDirTexelSize, initialRayStep, uvNode, viewDir, viewNormal, n ) );
 
-				ao.addAssign( float( bitCount( globalOccludedBitfield ) ).div( float( MAX_RAY ) ) );
+				ao.addAssign( float( countOneBits( globalOccludedBitfield ) ).div( float( MAX_RAY ) ) );
 
 			} );
 

examples/webgpu_compute_reduce.html

@@ -190,7 +190,7 @@ <h3 id="panel-title" style="flex: 0 0 auto;">Subgroup Reduction Explanation</h3>
 		<script type="module">
 
 			import * as THREE from 'three/webgpu';
-			import { instancedArray, Loop, If, vec3, dot, clamp, storage, uvec4, subgroupAdd, uniform, uv, uint, float, Fn, vec2, invocationLocalIndex, invocationSubgroupIndex, uvec2, floor, instanceIndex, workgroupId, workgroupBarrier, workgroupArray, subgroupSize, select, log2 } from 'three/tsl';
+			import { instancedArray, Loop, If, vec3, dot, clamp, storage, uvec4, subgroupAdd, uniform, uv, uint, float, Fn, vec2, invocationLocalIndex, invocationSubgroupIndex, uvec2, floor, instanceIndex, workgroupId, workgroupBarrier, workgroupArray, subgroupSize, select, countTrailingZeros } from 'three/tsl';
 
 			import WebGPU from 'three/addons/capabilities/WebGPU.js';
 
@@ -831,12 +831,12 @@ <h3 id="panel-title" style="flex: 0 0 auto;">Subgroup Reduction Explanation</h3>
 
 					// Multiple approaches here
 					// log2(subgroupSize) -> TSL log2 function
-					// countTrailingZeros/findLSB(subgroupSize) -> Currently unsupported function in TSL that counts trailing zeros in number bit representation
+					// countTrailingZeros/findLSB(subgroupSize) -> TSL function that counts trailing zeros in number bit representation
 					// Can technically petition GPU for subgroupSize in shader and calculate logs on CPU at cost of shader being generalizable across devices
 					// May also break if subgroupSize changes when device is lost or if program is rerun on lower power device
-					const subgroupSizeLog = uint( log2( float( subgroupSize ) ) ).toVar( 'subgroupSizeLog' );
+					const subgroupSizeLog = countTrailingZeros( subgroupSize ).toVar( 'subgroupSizeLog' );
 					const spineSize = uint( workgroupSize ).shiftRight( subgroupSizeLog );
-					const spineSizeLog = uint( log2( float( spineSize ) ) ).toVar( 'spineSizeLog' );
+					const spineSizeLog = countTrailingZeros( spineSize ).toVar( 'spineSizeLog' );
 
 
 					// Align size to powers of subgroupSize

src/Three.TSL.js

@@ -136,6 +136,9 @@ export const context = TSL.context;
 export const convert = TSL.convert;
 export const convertColorSpace = TSL.convertColorSpace;
 export const convertToTexture = TSL.convertToTexture;
+export const countLeadingZeros = TSL.countLeadingZeros;
+export const countOneBits = TSL.countOneBits;
+export const countTrailingZeros = TSL.countTrailingZeros;
 export const cos = TSL.cos;
 export const cross = TSL.cross;
 export const cubeTexture = TSL.cubeTexture;

src/nodes/TSL.js

@@ -20,6 +20,7 @@ export * from './core/MRTNode.js';
 
 // math
 export * from './math/BitcastNode.js';
+export * from './math/BitcountNode.js';
 export * from './math/Hash.js';
 export * from './math/MathUtils.js';
 export * from './math/TriNoise3D.js';