diff --git a/ComputeShaders/ComputeShaders.vcxproj b/ComputeShaders/ComputeShaders.vcxproj
index b79f550..46f5ac7 100644
--- a/ComputeShaders/ComputeShaders.vcxproj
+++ b/ComputeShaders/ComputeShaders.vcxproj
@@ -172,6 +172,7 @@
+
diff --git a/ComputeShaders/ComputeShaders.vcxproj.filters b/ComputeShaders/ComputeShaders.vcxproj.filters
index fbf30d6..2acfcca 100644
--- a/ComputeShaders/ComputeShaders.vcxproj.filters
+++ b/ComputeShaders/ComputeShaders.vcxproj.filters
@@ -50,6 +50,7 @@
+
diff --git a/ComputeShaders/dbgFindNaN.hlsl b/ComputeShaders/dbgFindNaN.hlsl
new file mode 100644
index 0000000..206d8a0
--- /dev/null
+++ b/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 tensor: register( t0 );
+// A buffer with a single element for the output boolean. Zero means there were no NAN values in the tensor.
+RWBuffer 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 );
+}
\ No newline at end of file
diff --git a/ComputeShaders/zeroMemory.hlsl b/ComputeShaders/zeroMemory.hlsl
index c486636..5c4774b 100644
--- a/ComputeShaders/zeroMemory.hlsl
+++ b/ComputeShaders/zeroMemory.hlsl
@@ -3,6 +3,7 @@ RWBuffer 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;
}
\ No newline at end of file
diff --git a/Whisper/D3D/shaderNames.cpp b/Whisper/D3D/shaderNames.cpp
index a631f08..97d1589 100644
--- a/Whisper/D3D/shaderNames.cpp
+++ b/Whisper/D3D/shaderNames.cpp
@@ -2,7 +2,7 @@
#include "stdafx.h"
#include "shaderNames.h"
-static const std::array s_shaderNames =
+static const std::array s_shaderNames =
{
"add",
"addInPlace",
@@ -18,6 +18,7 @@ static const std::array s_shaderNames =
"convolutionPrep2",
"copyConvert",
"copyTranspose",
+ "dbgFindNaN",
"diagMaskInf",
"flashAttention",
"flashAttentionCompat1",
diff --git a/Whisper/D3D/shaderNames.h b/Whisper/D3D/shaderNames.h
index 80fc4fa..75654ec 100644
--- a/Whisper/D3D/shaderNames.h
+++ b/Whisper/D3D/shaderNames.h
@@ -20,32 +20,33 @@ namespace DirectCompute
convolutionPrep2 = 11,
copyConvert = 12,
copyTranspose = 13,
- diagMaskInf = 14,
- flashAttention = 15,
- flashAttentionCompat1 = 16,
- flashAttentionCompat2 = 17,
- flashAttentionCompat3 = 18,
- fmaRepeat1 = 19,
- fmaRepeat2 = 20,
- matReshapePanels = 21,
- mulMatByRow = 22,
- mulMatByRowTiled = 23,
- mulMatByRowTiledEx = 24,
- mulMatByScalar = 25,
- mulMatDotMain = 26,
- mulMatDotReshape = 27,
- mulMatMadMain = 28,
- mulMatTiled = 29,
- mulMatTiledEx = 30,
- norm = 31,
- normCompat = 32,
- normFixed = 33,
- scaleInPlace = 34,
- softMax = 35,
- softMaxCompat = 36,
- softMaxFixed = 37,
- softMaxLong = 38,
- zeroMemory = 39,
+ dbgFindNaN = 14,
+ diagMaskInf = 15,
+ flashAttention = 16,
+ flashAttentionCompat1 = 17,
+ flashAttentionCompat2 = 18,
+ flashAttentionCompat3 = 19,
+ fmaRepeat1 = 20,
+ fmaRepeat2 = 21,
+ matReshapePanels = 22,
+ mulMatByRow = 23,
+ mulMatByRowTiled = 24,
+ mulMatByRowTiledEx = 25,
+ mulMatByScalar = 26,
+ mulMatDotMain = 27,
+ mulMatDotReshape = 28,
+ mulMatMadMain = 29,
+ mulMatTiled = 30,
+ mulMatTiledEx = 31,
+ norm = 32,
+ normCompat = 33,
+ normFixed = 34,
+ scaleInPlace = 35,
+ softMax = 36,
+ softMaxCompat = 37,
+ softMaxFixed = 38,
+ softMaxLong = 39,
+ zeroMemory = 40,
};
const char* computeShaderName( eComputeShader cs );