Compute shader to test a tensor for NaN values

ComputeShaders/ComputeShaders.vcxproj

@@ -172,6 +172,7 @@
     <FxCompile Include="convolutionPrep2.hlsl" />
     <FxCompile Include="copyConvert.hlsl" />
     <FxCompile Include="copyTranspose.hlsl" />
+    <FxCompile Include="dbgFindNaN.hlsl" />
     <FxCompile Include="diagMaskInf.hlsl" />
     <FxCompile Include="flashAttention.hlsl" />
     <FxCompile Include="flashAttentionCompat1.hlsl" />

ComputeShaders/ComputeShaders.vcxproj.filters

@@ -50,6 +50,7 @@
     <FxCompile Include="mulMatByRowTiledEx.hlsl" />
     <FxCompile Include="addRepeatEx.hlsl" />
     <FxCompile Include="softMaxLong.hlsl" />
+    <FxCompile Include="dbgFindNaN.hlsl" />
   </ItemGroup>
   <ItemGroup>
     <None Include="componentwiseBinaryOp.hlsli" />

ComputeShaders/dbgFindNaN.hlsl

@@ -0,0 +1,76 @@
+// When reset = TRUE, write zero to the output buffer
+// When reset = FALSE, test input tensor for NaN, when found at least 1 NaN element, write 1 to the output buffer
+
+// FP32 or FP16 tensor to test for NAN
+Buffer<float> tensor: register( t0 );
+// A buffer with a single element for the output boolean. Zero means there were no NAN values in the tensor.
+RWBuffer<uint> result: register( u0 );
+
+cbuffer Constants: register( b0 )
+{
+	uint elements: packoffset( c0.x );
+	bool reset : packoffset( c0.y );
+}
+
+// Thread group index is 16 bits per coordinate:
+// https://learn.microsoft.com/en-us/windows/win32/api/d3d11/nf-d3d11-id3d11devicecontext-dispatch
+// We want this shader to support buffers up to 2 GB.
+#ifndef THREADS
+static const uint THREADS = 512;
+#endif
+#ifndef ITERATIONS
+static const uint ITERATIONS = 128;
+#endif
+
+static const uint itemsPerGroup = THREADS * ITERATIONS;
+
+inline bool isNaN( float x )
+{
+	// https://sakibsaikia.github.io/graphics/2022/01/04/Nan-Checks-In-HLSL.html
+	return ( asuint( x ) & 0x7fffffff ) > 0x7f800000;
+}
+
+groupshared uint reductionBuffer;
+
+[numthreads( THREADS, 1, 1 )]
+void main( uint3 group: SV_GroupID, uint thread : SV_GroupIndex )
+{
+	if( reset )
+	{
+		if( 0 == thread )
+			result[ 0 ] = 0;
+		return;
+	}
+
+	uint rsi = group.x * itemsPerGroup;
+	const uint rsiEnd = min( rsi + itemsPerGroup, elements );
+
+	// The main loop updates a local variable. There're THREADS instances of that variable for the group of threads.
+	bool foundNan = false;
+	for( rsi += thread; rsi < rsiEnd; rsi += THREADS )
+	{
+		const float val = tensor[ rsi ];
+		if( !isNaN( val ) )
+			continue;
+		foundNan = true;
+		break;
+	}
+
+	// Reduce THREADS booleans to a single one, using group shared memory atomics
+	if( 0 == thread )
+		reductionBuffer = 0;
+	GroupMemoryBarrierWithGroupSync();
+
+	if( foundNan )
+		InterlockedOr( reductionBuffer, 1u );
+
+	GroupMemoryBarrierWithGroupSync();
+
+	// When found, update output value with global memory atomic
+	if( 0 != thread )
+		return;
+	if( 0 == reductionBuffer )
+		return;
+
+	InterlockedOr( result[ 0 ], 1u );
+}

ComputeShaders/zeroMemory.hlsl

@@ -3,6 +3,7 @@ RWBuffer<float> result: register( u0 );
 cbuffer Constants: register( b0 )
 {
 	uint elements: packoffset( c0.x );
+	bool writeNan: packoffset( c0.y );
 }
 
 // Thread group index is 16 bits per coordinate:
@@ -22,6 +23,8 @@ void main( uint3 group: SV_GroupID, uint thread : SV_GroupIndex )
 {
 	uint rdi = group.x * itemsPerGroup;
 	const uint rdiEnd = min( rdi + itemsPerGroup, elements );
+	// https://www.h-schmidt.net/FloatConverter/IEEE754.html
+	const float pattern = writeNan ? asfloat( 0x7FFFFFFFu ) : 0.0;
 	for( rdi += thread; rdi < rdiEnd; rdi += THREADS )
-		result[ rdi ] = 0.0;
+		result[ rdi ] = pattern;
 }

Whisper/D3D/shaderNames.cpp

@@ -2,7 +2,7 @@
 #include "stdafx.h"
 #include "shaderNames.h"
 
-static const std::array<const char*, 40> s_shaderNames = 
+static const std::array<const char*, 41> s_shaderNames = 
 {
 	"add",
 	"addInPlace",
@@ -18,6 +18,7 @@ static const std::array<const char*, 40> s_shaderNames =
 	"convolutionPrep2",
 	"copyConvert",
 	"copyTranspose",
+	"dbgFindNaN",
 	"diagMaskInf",
 	"flashAttention",
 	"flashAttentionCompat1",

Whisper/D3D/shaderNames.h

@@ -20,32 +20,33 @@ namespace DirectCompute
 		convolutionPrep2 = 11,
 		copyConvert = 12,
 		copyTranspose = 13,
-		diagMaskInf = 14,
+		dbgFindNaN = 14,
-		flashAttention = 15,
+		diagMaskInf = 15,
-		flashAttentionCompat1 = 16,
+		flashAttention = 16,
-		flashAttentionCompat2 = 17,
+		flashAttentionCompat1 = 17,
-		flashAttentionCompat3 = 18,
+		flashAttentionCompat2 = 18,
-		fmaRepeat1 = 19,
+		flashAttentionCompat3 = 19,
-		fmaRepeat2 = 20,
+		fmaRepeat1 = 20,
-		matReshapePanels = 21,
+		fmaRepeat2 = 21,
-		mulMatByRow = 22,
+		matReshapePanels = 22,
-		mulMatByRowTiled = 23,
+		mulMatByRow = 23,
-		mulMatByRowTiledEx = 24,
+		mulMatByRowTiled = 24,
-		mulMatByScalar = 25,
+		mulMatByRowTiledEx = 25,
-		mulMatDotMain = 26,
+		mulMatByScalar = 26,
-		mulMatDotReshape = 27,
+		mulMatDotMain = 27,
-		mulMatMadMain = 28,
+		mulMatDotReshape = 28,
-		mulMatTiled = 29,
+		mulMatMadMain = 29,
-		mulMatTiledEx = 30,
+		mulMatTiled = 30,
-		norm = 31,
+		mulMatTiledEx = 31,
-		normCompat = 32,
+		norm = 32,
-		normFixed = 33,
+		normCompat = 33,
-		scaleInPlace = 34,
+		normFixed = 34,
-		softMax = 35,
+		scaleInPlace = 35,
-		softMaxCompat = 36,
+		softMax = 36,
-		softMaxFixed = 37,
+		softMaxCompat = 37,
-		softMaxLong = 38,
+		softMaxFixed = 38,
-		zeroMemory = 39,
+		softMaxLong = 39,
+		zeroMemory = 40,
 	};
 
 	const char* computeShaderName( eComputeShader cs );