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
branchesmasterall refs →
2.7 KiB82 linesraw
1#pragma once
2#include <immintrin.h>
3
4void addF16to32( float* rdi, const uint16_t* a, const uint16_t* b, size_t length );
5void addF16to32( float* rdi, const uint16_t* a, const float* b, size_t length );
6
7class AlignedSpan
8{
9	float* pointer;
10
11public:
12	AlignedSpan( void* data )
13	{
14		size_t i = (size_t)data;
15		constexpr size_t mask32 = ~(size_t)31;
16		i = ( i + 31 ) & mask32;
17		pointer = (float*)i;
18	}
19
20	operator float* ( ) { return pointer; }
21};
22
23inline size_t tempBufferForFloats( size_t count )
24{
25	// Round up by 8 to be able to use full-vector loads and stores
26	constexpr size_t mask8 = ~(size_t)7;
27	count = ( count + 7 ) & mask8;
28
29	// Add 32 more bytes to align the temporary buffer
30	return ( count * 4 ) + 32;
31}
32
33#define ALIGNED_SPAN( name, countFloats ) AlignedSpan name{ _alloca( tempBufferForFloats( countFloats ) ) }
34
35void norm( float* rdi, float* temp, const float* rsi, size_t length );
36
37void fmaRepeatRow( float* rdi, size_t len, const float* w, const float* b, size_t lenPattern );
38void __vectorcall addRepeatScaleRow( float* rdi, size_t len, const float* b, size_t lenPattern, const __m256 scale );
39void addRepeatRow( float* rdi, size_t len, const float* b, size_t lenPattern );
40void __vectorcall scaleRow( float* rdi, size_t len, const __m256 scale );
41
42namespace DirectCompute
43{
44	struct LookupTablesData;
45}
46const DirectCompute::LookupTablesData& getLookupTables();
47void addRepeatGeluRow( float* rdi, size_t len, const float* b, size_t lenPattern, const DirectCompute::LookupTablesData& lookup );
48
49void softMax( float* rdi, size_t length, const float inputScale );
50
51// A cache line-aligned array where first 8 elements have all bits set, last 8 elements are zeros
52extern const std::array<int, 16> s_zeroTailMask;
53
54// Load a tail mask as FP32 vector, for use with _mm256_and_ps or _mm256_blendv_ps instructions
55__forceinline __m256 loadTailMaskFloats( size_t remainder )
56{
57	assert( remainder > 0 && remainder < 8 );
58	const float* rsi = (const float*)&s_zeroTailMask;
59	rsi += 8;
60	return _mm256_loadu_ps( rsi - remainder );
61}
62
63// Load a tail mask as int32 vector, for use with _mm256_maskstore_ps instruction
64template<bool assertIncomplete = true>
65__forceinline __m256i loadTailMaskInt( size_t remainder )
66{
67	if constexpr( assertIncomplete )
68		assert( remainder > 0 && remainder < 8 );
69	else
70		assert( remainder >= 0 && remainder <= 8 );
71
72	const int* rsi = (const int*)&s_zeroTailMask;
73	rsi += 8;
74	return _mm256_loadu_si256( ( const __m256i* )( rsi - remainder ) );
75}
76
77void floatsUpcast( float* rdi, const uint16_t* rsi, size_t length );
78
79void floatsDowncast( uint16_t* rdi, const float* rsi, size_t length );
80
81void addRowInPlace( float* rdi, const float* rsi, size_t length );
82void addRow( float* rdi, const float* a, const float* b, size_t length );