yum-archive/TaSTT-Whisper

High-performance GPGPU inference of OpenAI's Whisper automatic speech recognition (ASR) model

git clone https://git.yummers.dev/yum-archive/TaSTT-Whisper

KonstantinSource codes8c4603c

master
8.7 KiB334 linesraw
1#include "stdafx.h"
2#include "testUtils.h"
3#include <immintrin.h>
4#include <atlfile.h>
5#include <atlpath.h>
6
7namespace
8{
9	using DirectCompute::sTensorDiff;
10
11	__forceinline __m256 load( const float* rsi )
12	{
13		return _mm256_loadu_ps( rsi );
14	}
15
16	__forceinline __m256 load( const uint16_t* rsi )
17	{
18		const __m128i iv = _mm_load_si128( ( const __m128i* )rsi );
19		return _mm256_cvtph_ps( iv );
20	}
21
22	__forceinline void loadPartial( const uint16_t* x, const uint16_t* y, size_t count, __m256& fx, __m256& fy )
23	{
24		__m128i ix, iy;
25		switch( count )
26		{
27		case 1: // load 2 bytes
28			ix = _mm_cvtsi32_si128( *x );
29			iy = _mm_cvtsi32_si128( *y );
30			break;
31		case 2: // load 4 bytes
32			ix = _mm_cvtsi32_si128( *(const int*)x );
33			iy = _mm_cvtsi32_si128( *(const int*)y );
34			break;
35		case 3: // load 6 bytes
36			ix = _mm_cvtsi32_si128( *(const int*)x );
37			iy = _mm_cvtsi32_si128( *(const int*)y );
38			ix = _mm_insert_epi16( ix, x[ 2 ], 2 );
39			iy = _mm_insert_epi16( iy, y[ 2 ], 2 );
40			break;
41		case 4: // load 8 bytes
42			ix = _mm_cvtsi64_si128( *(const int64_t*)x );
43			iy = _mm_cvtsi64_si128( *(const int64_t*)y );
44			break;
45		case 5: // load 10 bytes
46			ix = _mm_cvtsi64_si128( *(const int64_t*)x );
47			iy = _mm_cvtsi64_si128( *(const int64_t*)y );
48			ix = _mm_insert_epi16( ix, x[ 4 ], 4 );
49			iy = _mm_insert_epi16( iy, y[ 4 ], 4 );
50			break;
51		case 6: // load 12 bytes
52			ix = _mm_cvtsi64_si128( *(const int64_t*)x );
53			iy = _mm_cvtsi64_si128( *(const int64_t*)y );
54			ix = _mm_insert_epi32( ix, *(const int*)( x + 4 ), 2 );
55			iy = _mm_insert_epi32( iy, *(const int*)( y + 4 ), 2 );
56			break;
57		case 7: // load 14 bytes
58			ix = _mm_cvtsi64_si128( *(const int64_t*)x );
59			iy = _mm_cvtsi64_si128( *(const int64_t*)y );
60			ix = _mm_insert_epi32( ix, *(const int*)( x + 4 ), 2 );
61			iy = _mm_insert_epi32( iy, *(const int*)( y + 4 ), 2 );
62			ix = _mm_insert_epi16( ix, x[ 6 ], 6 );
63			iy = _mm_insert_epi16( iy, y[ 6 ], 6 );
64			break;
65		default:
66			fx = fy = _mm256_setzero_ps();
67			return;
68		}
69
70		fx = _mm256_cvtph_ps( ix );
71		fy = _mm256_cvtph_ps( iy );
72	}
73
74	inline __m128 loadFloat2( const float* rsi )
75	{
76		return _mm_castpd_ps( _mm_load_sd( (const double*)rsi ) );
77	}
78	inline __m128 loadFloat3( const float* rsi )
79	{
80		__m128 f = loadFloat2( rsi );
81		f = _mm_insert_ps( f, _mm_load_ss( rsi + 2 ), 0x20 );
82		return f;
83	}
84	__forceinline void loadPartial( const float* x, const float* y, size_t count, __m256& fx, __m256& fy )
85	{
86		__m128 low1, high1;
87		__m128 low2, high2;
88		high1 = high2 = _mm_setzero_ps();
89		switch( count )
90		{
91		case 1:
92			low1 = _mm_load_ss( x );
93			low2 = _mm_load_ss( y );
94			break;
95		case 2:
96			low1 = loadFloat2( x );
97			low2 = loadFloat2( y );
98			break;
99		case 3:
100			low1 = loadFloat3( x );
101			low2 = loadFloat3( y );
102			break;
103		case 4:
104			low1 = _mm_loadu_ps( x );
105			low2 = _mm_loadu_ps( y );
106			break;
107		case 5:
108			low1 = _mm_loadu_ps( x );
109			low2 = _mm_loadu_ps( y );
110			high1 = _mm_load_ss( x + 4 );
111			high2 = _mm_load_ss( y + 4 );
112			break;
113		case 6:
114			low1 = _mm_loadu_ps( x );
115			low2 = _mm_loadu_ps( y );
116			high1 = loadFloat2( x + 4 );
117			high2 = loadFloat2( y + 4 );
118			break;
119		case 7: // load 14 bytes
120			low1 = _mm_loadu_ps( x );
121			low2 = _mm_loadu_ps( y );
122			high1 = loadFloat3( x + 4 );
123			high2 = loadFloat3( y + 4 );
124			break;
125		default:
126			fx = fy = _mm256_setzero_ps();
127			return;
128		}
129
130		fx = _mm256_setr_m128( low1, high1 );
131		fy = _mm256_setr_m128( low2, high2 );
132	}
133
134	__forceinline float horizontalMaximum( __m256 v )
135	{
136		__m128 s = _mm256_extractf128_ps( v, 1 );
137		s = _mm_max_ps( s, _mm256_castps256_ps128( v ) );
138		s = _mm_max_ps( s, _mm_movehl_ps( s, s ) );
139		s = _mm_max_ss( s, _mm_movehdup_ps( s ) );
140		return _mm_cvtss_f32( s );
141	}
142
143	__forceinline double horizontalSum( __m256 v )
144	{
145		__m256d d = _mm256_cvtps_pd( _mm256_extractf128_ps( v, 1 ) );
146		d = _mm256_add_pd( d, _mm256_cvtps_pd( _mm256_castps256_ps128( v ) ) );
147
148		__m128d s = _mm256_extractf128_pd( d, 1 );
149		s = _mm_add_pd( s, _mm256_castpd256_pd128( d ) );
150		s = _mm_add_sd( s, _mm_unpackhi_pd( s, s ) );
151		return _mm_cvtsd_f64( s );
152	}
153
154	__m256 maskInfNan( __m256 diff, __m256 a, __m256 b )
155	{
156		__m256i ai = _mm256_castps_si256( a );
157		__m256i bi = _mm256_castps_si256( b );
158		__m256i eqi = _mm256_cmpeq_epi32( ai, bi );
159		__m256 eq = _mm256_castsi256_ps( eqi );
160		return _mm256_andnot_ps( eq, diff );
161	}
162
163	class DiffAcc
164	{
165		__m256 maxAbs = _mm256_setzero_ps();
166		__m256 sumSquares = _mm256_setzero_ps();
167
168	public:
169
170		__forceinline void add( __m256 a, __m256 b )
171		{
172			const __m256 neg0 = _mm256_set1_ps( -0.0f );
173			__m256 diff = _mm256_sub_ps( b, a );
174			diff = maskInfNan( diff, a, b );
175			sumSquares = _mm256_fmadd_ps( diff, diff, sumSquares );
176			const __m256 absDiff = _mm256_andnot_ps( neg0, diff );
177			maxAbs = _mm256_max_ps( maxAbs, absDiff );
178		}
179
180		__forceinline sTensorDiff reduce( size_t count )
181		{
182			sTensorDiff res;
183			res.maxAbsDiff = horizontalMaximum( maxAbs );
184			res.avgDiffSquared = (float)( horizontalSum( sumSquares ) / (double)(int64_t)count );
185			res.length = count;
186			return res;
187		}
188	};
189
190	template<class E>
191	static sTensorDiff __declspec( noinline ) diffVectors( const E* a, const E* b, size_t length )
192	{
193		// const E* const aEnd = a + length;
194		const E* const aEndAligned = a + ( length / 8 ) * 8;
195		const size_t remainder = length % 8;
196
197		DiffAcc acc;
198		for( ; a < aEndAligned; a += 8, b += 8 )
199			acc.add( load( a ), load( b ) );
200
201		if( remainder != 0 )
202		{
203			__m256 va, vb;
204			loadPartial( a, b, remainder, va, vb );
205			acc.add( va, vb );
206		}
207
208		return acc.reduce( length );
209	}
210}
211
212sTensorDiff DirectCompute::computeDiff( const float* a, const float* b, size_t length )
213{
214	return diffVectors( a, b, length );
215}
216
217sTensorDiff DirectCompute::computeDiff( const uint16_t* a, const uint16_t* b, size_t length )
218{
219	return diffVectors( a, b, length );
220}
221
222void DirectCompute::sTensorDiff::print( const char* what ) const
223{
224	logDebug( u8"%s: length %zu, maxAbsDiff = %g, avgDiffSquared = %g", what, length, maxAbsDiff, avgDiffSquared );
225}
226void DirectCompute::sTensorDiff::print() const
227{
228	logDebug( u8"%zu elements, maxAbsDiff = %g, avgDiffSquared = %g", length, maxAbsDiff, avgDiffSquared );
229}
230
231HRESULT DirectCompute::dbgWriteBinaryFile( LPCTSTR fileName, const void* rsi, size_t cb )
232{
233	CPath path;
234	path.m_strPath = LR"(C:\Temp\2remove\Whisper)";
235	path.Append( fileName );
236
237	CAtlFile file;
238	CHECK( file.Create( path, GENERIC_WRITE, 0, CREATE_ALWAYS ) );
239	CHECK( file.Write( rsi, (DWORD)cb ) );
240	CHECK( file.Flush() );
241	return S_OK;
242}
243
244#include "Tensor.h"
245
246sTensorDiff DirectCompute::computeDiff( const Tensor& a, const Tensor& b )
247{
248	assert( isSameShapeAndLayout( a, b ) );
249	const eDataType dt = a.getType();
250	assert( dt == b.getType() );
251	switch( dt )
252	{
253	case eDataType::FP32:
254	{
255		std::vector<float> v1, v2;
256		a.download( v1 );
257		b.download( v2 );
258		assert( v1.size() == v2.size() );
259#if 0
260		const size_t firstZero = std::find( v2.begin(), v2.end(), 0.0f ) - v2.begin();
261		
262		std::vector<float> delta;
263		delta.resize( v1.size() );
264		for( size_t i = 0; i < v1.size(); i++ )
265			delta[ i ] = std::abs( v1[ i ] - v2[ i ] );
266		const size_t maxIndex = std::max_element( delta.begin(), delta.end() ) - delta.begin();
267#endif
268		return computeDiff( v1.data(), v2.data(), v1.size() );
269	}
270	}
271	throw E_NOTIMPL;
272}
273
274using namespace DirectCompute;
275
276void PrintUniqueTensorSizes::printImpl( const std::array<uint32_t, 8>& a )
277{
278	auto pair = set.emplace( a );
279	if( !pair.second )
280		return;	// was already there
281
282	const __m128i rhs = _mm_loadu_si128( ( const __m128i* ) ( &a[ 4 ] ) );
283
284	if( _mm_testz_si128( rhs, rhs ) )
285	{
286		logDebug( u8"%s: [ %i, %i, %i, %i ]", what,
287			a[ 0 ], a[ 1 ], a[ 2 ], a[ 3 ] );
288	}
289	else
290	{
291		logDebug( u8"%s: [ %i, %i, %i, %i ], [ %i, %i, %i, %i ]", what,
292			a[ 0 ], a[ 1 ], a[ 2 ], a[ 3 ], a[ 4 ], a[ 5 ], a[ 6 ], a[ 7 ] );
293	}
294}
295
296void PrintUniqueTensorSizes::print( const Tensor& lhs, const Tensor& rhs )
297{
298	std::array<uint32_t, 8> arr;
299	__m128i* const rdi = ( __m128i* )arr.data();
300	_mm_storeu_si128( rdi, lhs.sizeVec() );
301	_mm_storeu_si128( rdi + 1, rhs.sizeVec() );
302
303	printImpl( arr );
304}
305
306void PrintUniqueTensorSizes::print( const int* lhs, const int* rhs )
307{
308	std::array<uint32_t, 8> arr;
309	__m128i* const rdi = ( __m128i* )arr.data();
310	_mm_storeu_si128( rdi, load16( lhs ) );
311	_mm_storeu_si128( rdi + 1, load16( rhs ) );
312
313	printImpl( arr );
314}
315
316void PrintUniqueTensorSizes::print( const Tensor& lhs )
317{
318	std::array<uint32_t, 8> arr;
319	__m128i* const rdi = ( __m128i* )arr.data();
320	_mm_storeu_si128( rdi, lhs.sizeVec() );
321	_mm_storeu_si128( rdi + 1, _mm_setzero_si128() );
322
323	printImpl( arr );
324}
325
326#include "testUtilsC.h"
327
328void printUniqueTensorSize( const char* name, const int* lhs, const int* rhs )
329{
330	using TS = DirectCompute::PrintUniqueTensorSizes;
331	static std::unordered_map<std::string, TS> map;
332	TS& ts = map.try_emplace( name, name ).first->second;
333	ts.print( lhs, rhs );
334}