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.9 KiB118 linesraw
1// Dispatch with [ neq1*neq2*neq3, 1, 1 ] thread groups
2#include "flashAttentionCommon.hlsli"
3#include "groupReduce.hlsli"
4#include "miscUtils.hlsli"
5
6inline void roundTempVector( uint i, const uint len, const uint thread )
7{
8	const uint iEnd = i + len;
9	for( i += thread; i < iEnd; i += 32 )
10	{
11		float f = temp[ i ];
12		f = roundToFp16( f );
13		temp[ i ] = f;
14	}
15}
16
17inline void computeDotProduct( Buffer<float> buff0, RWBuffer<float> buff1, uint s0, uint s1, const uint len, const uint thread, inout float acc )
18{
19	acc = 0;
20/*	const uint s0End = s0 + len;
21	s0 += thread;
22	s1 += thread;
23	for( ; s0 < s0End; s0 += 32, s1 += 32 )
24		acc = mad( buff0[ s0 ], buff1[ s1 ], acc );
25
26	horizontalSumCompatNew( thread, acc ); */
27	const uint completeVectors = len / 32;
28	uint i;
29	for( i = 0; i < completeVectors; i++, s0 += 32, s1 += 32 )
30		acc = mad( buff0[ s0 + thread ], buff1[ s1 + thread ], acc );
31
32	horizontalSumCompatNew( thread, acc );
33
34	if( 0 == thread )
35	{
36		const uint rem = len % 32;
37		if( 0 != rem )
38		{
39			double f64 = acc;
40			for( i = 0; i < rem; i++ )
41			{
42				precise float a = buff0[ s0 + i ];
43				precise float b = buff1[ s1 + i ];
44				precise float prod = a * b;
45				f64 += prod;
46			}
47			acc = (float)f64;
48		}
49	}
50}
51
52[ numthreads( 32, 1, 1 ) ]
53void main( uint3 group: SV_GroupID, uint thread : SV_GroupIndex )
54{
55	const uint neq0 = q_elements[ 0 ];
56	const uint neq1 = q_elements[ 1 ];
57	const uint neq2 = q_elements[ 2 ];
58	const uint neq3 = q_elements[ 3 ];
59
60	const uint nek0 = k_elements[ 0 ];
61	const uint nek1 = k_elements[ 1 ];
62
63	const uint nev1 = v_elements[ 1 ];
64
65	const uint ne0 = res_elements[ 0 ];
66	const uint ne1 = res_elements[ 1 ];
67
68	const uint nbk0 = k_strides[ 0 ];
69	const uint nbk1 = k_strides[ 1 ];
70	const uint nbk2 = k_strides[ 2 ];
71	const uint nbk3 = k_strides[ 3 ];
72
73	const uint nbq0 = q_strides[ 0 ];
74	const uint nbq1 = q_strides[ 1 ];
75	const uint nbq2 = q_strides[ 2 ];
76	const uint nbq3 = q_strides[ 3 ];
77
78	const uint nbv0 = v_strides[ 0 ];
79	const uint nbv1 = v_strides[ 1 ];
80	const uint nbv2 = v_strides[ 2 ];
81	const uint nbv3 = v_strides[ 3 ];
82
83	const uint nb0 = res_strides[ 0 ];
84	const uint nb1 = res_strides[ 1 ];
85	const uint nb2 = res_strides[ 2 ];
86	const uint nb3 = res_strides[ 3 ];
87
88	const uint D = neq0;
89	const uint N = neq1;
90	const uint P = nek1 - N;
91	// const uint M = P + N;
92	const uint M = nek1;
93
94	const uint ir = group.x;
95	const uint iq3 = ir / ( neq2 * neq1 );
96	const uint iq2 = ( ir - iq3 * neq2 * neq1 ) / neq1;
97	const uint iq1 = ( ir - iq3 * neq2 * neq1 - iq2 * neq1 );
98
99	const uint tempIndex = ir * tempBufferStride;
100
101	roundTempVector( tempIndex, nek1, thread );
102	AllMemoryBarrierWithGroupSync();
103
104	uint rdi = iq1 * nb1 + iq2 * nb2 + iq3 * nb3;
105	for( uint ic = 0; ic < nev1; ic++, rdi += nb0 )
106	{
107		// dst indices
108		const uint i1 = iq1;
109		const uint i2 = iq2;
110		const uint i3 = iq3;
111
112		const uint s0 = ic * nbv1 + i2 * nbv2 + i3 * nbv3;
113		float dp;
114		computeDotProduct( v, temp, s0, tempIndex, nek1, thread, dp );
115		if( 0 == thread )
116			result[ rdi ] = dp;
117	}
118}