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

KonstantinMinor, micro-optimization15dbcac

master
branchesmasterall refs →
3.4 KiB120 linesraw
1#pragma once
2#include <stdint.h>
3#include <array>
4#include <smmintrin.h>
5
6struct ggml_tensor;
7using HRESULT = long;
8
9namespace DirectCompute
10{
11	// This POD structure describes the shape of a tensor.
12	// It’s used for both GPU tensors in VRAM, and tensors in system memory used by the Hybrid model.
13	struct TensorShape
14	{
15		// Count of elements, up to 4 coordinates
16		// The unused coordinates are set to 1
17		std::array<uint32_t, 4> ne;
18
19		// Strides of the tensor
20		// For a dense row-major tensor, these numbers are [ 1, ne[0], ne[0]*ne[1], ne[0]*ne[1]*ne[2] ]
21		// Note that unlike GGML code, these numbers are expressed in elements not bytes, but the meaning is the same
22		// GPU matrices reshaped into panels are keeping different values here: [ 0, panelSize, panelSize * panelsCount, panelSize * panelsCount * ne[ 2 ] ]
23		std::array<uint32_t, 4> nb;
24
25		TensorShape();
26		TensorShape( const TensorShape& that );
27		void operator=( const TensorShape& that );
28		HRESULT create( const ggml_tensor& ggml );
29		TensorShape( const ggml_tensor& ggml );
30
31		__m128i __vectorcall sizeVec() const
32		{
33			return load( ne );
34		}
35		__m128i __vectorcall stridesVec() const
36		{
37			return load( nb );
38		}
39
40		uint32_t countRows() const
41		{
42			return ne[ 1 ] * ne[ 2 ] * ne[ 3 ];
43		}
44
45		uint32_t countElements() const
46		{
47			// return ne[ 0 ] * countRows();
48			const __m128i a = sizeVec();
49			const __m128i b = _mm_srli_si128( a, 4 );
50			const __m128i p2 = _mm_mul_epu32( a, b );
51			uint64_t res = (uint64_t)_mm_extract_epi64( p2, 1 );
52			res *= (uint64_t)_mm_cvtsi128_si64( p2 );
53			assert( 0 == ( res >> 32 ) );
54			return (uint32_t)res;
55		}
56
57		// Compute strides from sizes, assuming dense row-major memory layout of the tensor
58		void setDenseStrides();
59
60		bool isMatrix() const
61		{
62			// return ne[ 2 ] == 1 && ne[ 3 ] == 1;
63			const uint64_t num = *(const uint64_t*)&ne[ 2 ];
64			return num == 0x100000001ull;
65		}
66		bool isVector() const
67		{
68			return 1 == ne[ 1 ] && isMatrix();
69		}
70
71		// True of this tensor is dense and row-major
72		bool isContinuous() const
73		{
74			/* return 1 == nb[ 0 ] &&
75				nb[ 1 ] == nb[ 0 ] * ne[ 0 ] &&
76				nb[ 2 ] == nb[ 1 ] * ne[ 1 ] &&
77				nb[ 3 ] == nb[ 2 ] * ne[ 2 ]; */
78
79			const __m128i nbv = stridesVec();
80			const __m128i nev = sizeVec();
81			__m128i tmp = _mm_mullo_epi32( nbv, nev );	// Vertical product of int32 lanes
82			tmp = _mm_shuffle_epi32( tmp, _MM_SHUFFLE( 2, 1, 0, 0 ) );	// Shift left by 1 int32 lane
83			tmp = _mm_insert_epi32( tmp, 1, 0 );	// Reset X lane to 1
84			return vectorEqual( tmp, nbv );
85		}
86
87		// Reset all fields to zero
88		void setZero()
89		{
90			const __m128i z = _mm_setzero_si128();
91			_mm_storeu_si128( ( __m128i* )ne.data(), z );
92			_mm_storeu_si128( ( __m128i* )nb.data(), z );
93		}
94	};
95
96	// True when two tensors have equal count of elements
97	inline bool isSameShape( const TensorShape& t0, const TensorShape& t1 )
98	{
99		__m128i a = t0.sizeVec();
100		__m128i b = t1.sizeVec();
101		return vectorEqual( a, b );
102	}
103
104	// True when two tensors have equal count of elements, and equal VRAM layout too
105	inline bool isSameShapeAndLayout( const TensorShape& t0, const TensorShape& t1 )
106	{
107		__m128i a, b, x;
108		a = t0.sizeVec();
109		b = t1.sizeVec();
110		x = _mm_xor_si128( a, b );
111
112		a = t0.stridesVec();
113		b = t1.stridesVec();
114		x = _mm_or_si128( x, _mm_xor_si128( a, b ) );
115		return (bool)_mm_testz_si128( x, x );
116	}
117
118	// True when we can multiply two tensors of the provided shapes
119	bool canMulMat( const TensorShape& t0, const TensorShape& t1 );
120}