TSL: Add bitCount functions

[cmhhelgeson]

Description

Add countTrailingZeros, countLeadingZeros, and countOneBits to TSL. Ideally, this functionality would also be accessible through the WebGLBackend, but I'm not sure what the status or likelihood is of the GLSL 3.1 functionality getting added to WebGL anytime in the future ( see KhronosGroup/WebGL#3714 ).

The impetus for implementing these functions is that they're notionally more concise and/or performant than equivalent operations ( log2, directly bit shifting, etc ).

I've also implemented each of these functions GLSL equivalents (findLSB, findMSB, bitCount) as aliases for the GLSL to TSL transpiration process. If it's preferred by the maintainers, we can get rid of the aliases and/or use the NodeBuilders to specify how the bitCount functions get constructed instead.

• Add functionality to MathNode and/or separate BitCountNode
• Add GLSL polyfills for functionality
• Replace bitCount node in SSGI with TSL function and test output on WebGL/WebGPU backends
• Validate that Transpiler produces correct output

[github-actions]

📦 Bundle size

Full ESM build, minified and gzipped.

	Before	After	Diff
WebGL	355.14 86.33	355.14 86.33	+0 B +0 B
WebGPU	609.66 171.01	612.51 171.86	+2.84 kB +847 B
WebGPU Nodes	608.27 170.75	611.12 171.6	+2.84 kB +856 B

🌳 Bundle size after tree-shaking

Minimal build including a renderer, camera, empty scene, and dependencies.

	Before	After	Diff
WebGL	486.77 121.07	486.77 121.07	+0 B +0 B
WebGPU	678.52 186.29	681.84 187.18	+3.32 kB +886 B
WebGPU Nodes	620.51 169.56	623.83 170.48	+3.32 kB +919 B

[Mugen87]

Related:

There is a custom bitCount() TSL function that I needed for the SSGI. The original Unity code used a built-in function but because of GLSL the code uses an implementation from https://graphics.stanford.edu/%7Eseander/bithacks.html.

three.js/examples/jsm/tsl/display/SSGINode.js

Lines 448 to 462 in 2a02884

	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' }
	]
	} );

It's great to see WGSL support! How about providing a built-in bitCount() TSL function that fallbacks to the above code for the WebGL backend?

[cmhhelgeson]

It's great to see WGSL support! How about providing a built-in bitCount() TSL function that fallbacks to the above code for the WebGL backend?

In hindsight, I think it's possible to add polyfills for the other functions as well, so I will get on doing that. The link you provided should be useful in generating the necessary functions.

[cmhhelgeson]

I'm going to fix the extant issues like documentation, but I'm also going to draft this PR so I can implement a polyfill approach. Basically, I'd be adding GLSL polyfills similar to how the WGSLNodeBuilder has polyfills for certain functions. Would this approach be desired, or would a different approach be preferred ( for instance, generating different nodes depending on the renderer's backend ). Personally, I'd prefer polyfilling in the GLSLNodeBuilder to keep the node code simpler for a simple mathematical operation.

[Mugen87]

Personally, I'd prefer polyfilling in the GLSLNodeBuilder to keep the node code simpler for a simple mathematical operation.

Then I suggest you start with that approach so we can see how it ends up in code.

[cmhhelgeson]

I see that the bid twiddling hacks for counting bits mainly concern Uints, but the WGSL countOneBits function (GLSL bitCount equivalent), can also take in a signed integer or vectors. The vectors are trivial to implement, but are there any references for a signed integer equivalent of bitCount?

Nevermind, we can just case the i32 to a u32 using the GLSL bitcast functionality.

[cmhhelgeson]

Test Examples

1. Validate functionality under WebGL backend

https://raw.githack.com/cmhhelgeson/three.js/bit_count_example/examples/webgpu_bitcount.html

2. Validate SSGI functionality (the visuals seem off here but they match the existing representation of SSGI under the WebGL backend)

https://raw.githack.com/cmhhelgeson/three.js/bit_count_example/examples/webgpu_postprocessing_ssgi.html

Existing SSGI Example with forceWebGL set to true as example. Note how the existing SSGI implementation is significantly more noisy than the WebGPU example.

https://raw.githack.com/cmhhelgeson/three.js/ssgi_test/examples/webgpu_postprocessing_ssgi.html

[Mugen87]

the visuals seem off here

What deviations are you seeing on your system?

[cmhhelgeson]

Is it possible that #32091 is affecting the ability of displacementmap to pass its test? I'm not sure if the screenshots were updated before or after the light scattering PR, as that's the only reason I could think as to why materials_displacement is failing despite not activating the SSGINode.

[cmhhelgeson]

I really don't know what I keep on getting strange E2E errors like E2E testing stalling for dozens of minutes or a legacy WebGLRenderer example failing when the changes are in the WebGPURenderer.

[Mugen87]

I really don't know what I keep on getting strange E2E errors like E2E testing stalling for dozens of minutes or a legacy WebGLRenderer example failing when the changes are in the WebGPURenderer.

A rerun solved the issue. It seems there are still some fragile examples like webgl_video_kinect that fail from time to time.

[cmhhelgeson]

I'm afraid they are not.

We can't merge yet since the SSGI example looks different with this PR. The GI badly flickers which indicates an issue with countOneBits() (you can debug this by selecting the "GI" output in the inspector). You must either revert the changes to SSGINode or find out why countOneBits() and the TSL version produce different results. The built-in and the fallback must be a match.

I see the flickering you're talking about when zooming in on the SSGI. I'm a little bit perplexed that the results are so different, especially for a native function that should be executing the same functionality. I'll look into it and revert SSGI if I can't find the issue.

WebGPUBackend Bitcount Test:
https://raw.githack.com/cmhhelgeson/three.js/bit_count_example/examples/webgpu_bitcount.html

WebGLBackend Bitcount Test:
https://raw.githack.com/cmhhelgeson/three.js/bit_count_example_webgl/examples/webgpu_bitcount.html

[Mugen87]

I'll look into it and revert SSGI if I can't find the issue.

After thinking about it I think that's not sufficient. If you can't find the root cause, it's best to not use the native countOneBits() but the TSL version for both WebGPU and WebGL since it guarantees the same result. We can't offer a countOneBits() that produces different result for WebGPU and WebGL. This will be sure reported as a bug.

[cmhhelgeson]

Hasty example made to compare WebGL and WebGPU output for values between 1-2000

(https://raw.githack.com/cmhhelgeson/three.js/bit_count_example/examples/webgpu_bitcount_consecutive.html)

Code for hasty example:
https://github.com/cmhhelgeson/three.js/blob/bit_count_example/examples/webgpu_bitcount_consecutive.html

[cmhhelgeson]

I'll look into it and revert SSGI if I can't find the issue.

After thinking about it I think that's not sufficient. If you can't find the root cause, it's best to not use the native countOneBits() but the TSL version for both WebGPU and WebGL since it guarantees the same result. We can't offer a countOneBits() that produces different result for WebGPU and WebGL. This will be sure reported as a bug.

I checked countOneBits and the bitCount fallback against each other multiple times and could not find a difference in the output, so for now I'm going to specify in BitcountNode that the fallback will be used in both WebGPU and WebGL.

[cmhhelgeson]

@Mugen87 I believe the issue was just that the branch was old and not fully rebased onto dev. The example no longer exhibits flickering. I think this along with the extra validation I did of the bitCount functionality should be sufficient.

https://raw.githack.com/cmhhelgeson/three.js/bit_count_functions_test/examples/webgpu_postprocessing_ssgi.html

[cmhhelgeson]

@Mugen87 Sorry to re-ping but just to reiterate, after testing of the functionality across the WebGL and WebGPU backend, I discovered that the issue was with the branch and not with the code itself. Accordingly, I rebased this branch onto dev to get the most up to date version of the SSGINode, then applied the BitcountNode to the SSGI. As a result the output no longer exhibits the pronounced flickering at the bottom edge of the box on the left.

https://raw.githack.com/cmhhelgeson/three.js/bit_count_functions_test/examples/webgpu_postprocessing_ssgi.html

[Mugen87]

Got it! I'll revisit the PR next week!