init branch

Author: cmhhelgeson

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

@@ -80,6 +80,7 @@ export const batch = TSL.batch;
 export const bentNormalView = TSL.bentNormalView;
 export const billboarding = TSL.billboarding;
 export const bitAnd = TSL.bitAnd;
+export const bitCount = TSL.bitCount;
 export const bitNot = TSL.bitNot;
 export const bitOr = TSL.bitOr;
 export const bitXor = TSL.bitXor;
@@ -136,6 +137,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;
@@ -180,6 +184,8 @@ export const expression = TSL.expression;
 export const faceDirection = TSL.faceDirection;
 export const faceForward = TSL.faceForward;
 export const faceforward = TSL.faceforward;
+export const findLSB = TSL.findLSB;
+export const findMSB = TSL.findMSB;
 export const float = TSL.float;
 export const floatBitsToInt = TSL.floatBitsToInt;
 export const floatBitsToUint = TSL.floatBitsToUint;

src/nodes/math/MathNode.js

@@ -364,6 +364,9 @@ MathNode.FWIDTH = 'fwidth';
 MathNode.TRANSPOSE = 'transpose';
 MathNode.DETERMINANT = 'determinant';
 MathNode.INVERSE = 'inverse';
+MathNode.COUNT_TRAILING_ZEROS = 'countTrailingZeros';
+MathNode.COUNT_LEADING_ZEROS = 'countLeadingZeros';
+MathNode.COUNT_ONE_BITS = 'countOneBits';
 
 // 2 inputs
 
@@ -1099,10 +1102,50 @@ export const atan2 = ( y, x ) => { // @deprecated, r172
 
 };
 
+
+/**
+ * Finds the number of consecutive 0 bits from the least significant bit of the input value,
+ * which is also the index of the least significant bit of the input value.
+ *
+ * Can only be used with {@link WebGPURenderer} and a WebGPU backend.
+ *
+ * @tsl
+ * @function
+ * @param {Node | number} x - The input value.
+ * @returns {Node}
+ */
+export const countTrailingZeros = /*@__PURE__*/ nodeProxyIntent( MathNode, MathNode.COUNT_TRAILING_ZEROS ).setParameterLength( 1 );
+
+/**
+ * Finds the number of consecutive 0 bits starting from the most significant bit of the input value.
+ *
+ * Can only be used with {@link WebGPURenderer} and a WebGPU backend.
+ *
+ * @tsl
+ * @function
+ * @param {Node | number} x - The input value.
+ * @returns {Node}
+ */
+export const countLeadingZeros = /*@__PURE__*/ nodeProxyIntent( MathNode, MathNode.COUNT_LEADING_ZEROS ).setParameterLength( 1 );
+
+/**
+ * Finds the number of '1' bits set in the input value
+ *
+ * Can only be used with {@link WebGPURenderer} and a WebGPU backend.
+ *
+ * @tsl
+ * @function
+ * @returns {Node}
+ */
+export const countOneBits = /*@__PURE__*/ nodeProxyIntent( MathNode, MathNode.COUNT_ONE_BITS ).setParameterLength( 1 );
+
 // GLSL alias function
 
 export const faceforward = faceForward;
 export const inversesqrt = inverseSqrt;
+export const findLSB = countTrailingZeros;
+export const findMSB = countLeadingZeros;
+export const bitCount = countOneBits;
 
 // Method chaining
 
@@ -1165,3 +1208,6 @@ addMethodChaining( 'transpose', transpose );
 addMethodChaining( 'determinant', determinant );
 addMethodChaining( 'inverse', inverse );
 addMethodChaining( 'rand', rand );
+addMethodChaining( 'countTrailingZeros', countTrailingZeros );
+addMethodChaining( 'countLeadingZeros', countLeadingZeros );
+addMethodChaining( 'countOneBits', countOneBits );