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.6 KiB95 linesraw
1// GGML_TASK_COMPUTE step for matrix*matrix product, where nb01 >= nb00;
2// Dispatch with [ ne11, ne01*ne02*ne03 ] thread groups
3// Each thread group computes a single dot product
4Buffer<float> arg0: register( t0 );
5Buffer<float> arg1: register( t1 );
6RWBuffer<float> result: register( u0 );
7
8cbuffer Constants: register( b0 )
9{
10	uint4 src0_elements: packoffset( c0 );
11	uint4 src0_strides: packoffset( c1 );
12	uint4 src1_elements: packoffset( c2 );
13	uint4 result_elements: packoffset( c4 );
14	uint4 result_strides: packoffset( c5 );
15}
16
17inline uint product( uint3 vec )
18{
19	return vec.x * vec.y * vec.z;
20}
21
22inline uint product( uint4 vec )
23{
24	uint2 tmp = vec.xy * vec.zw;
25	return tmp.x * tmp.y;
26}
27
28inline float dotProductInner( uint i0, uint i1, uint length, uint thread )
29{
30	float res = 0;
31	for( uint i = thread; i < length; i += 32 )
32		res = mad( arg0[ i0 + i ], arg1[ i1 + i ], res );
33	return res;
34}
35
36#include "groupReduce.hlsli"
37
38[numthreads( 32, 1, 1 )]
39void main( uint3 group: SV_GroupID, uint thread : SV_GroupIndex )
40{
41	const uint ne00 = src0_elements.x;
42	const uint ne01 = src0_elements.y;
43	const uint ne02 = src0_elements.z;
44	const uint ne03 = src0_elements.w;
45
46	const uint ne10 = src1_elements.x;
47	const uint ne11 = src1_elements.y;
48	const uint ne12 = src1_elements.z;
49	const uint ne13 = src1_elements.w;
50
51	const int nb00 = src0_strides.x;
52	const int nb01 = src0_strides.y;
53	const int nb02 = src0_strides.z;
54	const int nb03 = src0_strides.w;
55
56	// total rows in src0
57	// const int nr = ne01*ne02*ne03;
58	const uint nr = product( src0_elements.yzw );
59
60	const uint ir = group.y;
61
62	// src0 indices
63	const uint i03 = ir / ( ne02 * ne01 );
64	const uint i02 = ( ir - i03 * ne02 * ne01 ) / ne01;
65	const uint i01 = ( ir - i03 * ne02 * ne01 - i02 * ne01 );
66
67	const uint i13 = i03;
68	const uint i12 = i02;
69
70	const uint i0 = i01;
71	const uint i2 = i02;
72	const uint i3 = i03;
73
74	// src0_row = (ggml_fp16_t *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03));
75	// src1_col = wdata + ( i13 * ne12 * ne11 + i12 * ne11 + 0 ) * ne00;
76	const uint src0_row = i01 * nb01 + i02 * nb02 + i03 * nb03;
77	const uint src1_col = ( i13 * ne12 * ne11 + i12 * ne11 ) * ne00;
78
79	const uint ic = group.x;
80	float curr = dotProductInner( src0_row, src1_col + ic * ne00, ne00, thread );
81	horizontalSumCompatNew( thread, curr );
82
83	if( 0 != thread )
84		return;
85
86	const uint nb0 = result_strides.x;
87	const uint nb1 = result_strides.y;
88	const uint nb2 = result_strides.z;
89	const uint nb3 = result_strides.w;
90
91	const uint ne0 = result_elements.x;
92	// float * dst_col = (float *) ((char *) dst->data + (i0*nb0 + 0*nb1 + i2*nb2 + i3*nb3));
93	const uint dst_col = i0 * nb0 + i2 * nb2 + i3 * nb3;
94	result[ dst_col + ic * ne0 ] = curr;
95}