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.0 KiB71 linesraw
1#pragma once
2#include "Tensor.h"
3#include "ParallelForRunner.h"
4
5namespace CpuCompute
6{
7	class MlContext
8	{
9		ParallelForRunner pfor;
10		iMemoryAllocator* allocator = nullptr;
11
12	public:
13		MlContext( int threads );
14		MlContext( const MlContext& ) = delete;
15		~MlContext() = default;
16
17		HRESULT setThreadsCount( int threads )
18		{
19			return pfor.setThreadsCount( threads );
20		}
21
22		iMemoryAllocator* setAllocator( iMemoryAllocator* alloc )
23		{
24			iMemoryAllocator* const ret = allocator;
25			allocator = alloc;
26			return ret;
27		}
28
29		Tensor createTensor( eDataType type, const std::array<uint32_t, 4>& size );
30		Tensor createTensor( eDataType type, std::initializer_list<uint32_t> size );
31
32		Tensor addRows( const Tensor& d_te, const Tensor& d_pe, const int* tokens, const int n_tokens, const int n_past );
33
34		Tensor norm( const Tensor& arg );
35
36		// cur = add( mul( repeat( w, cur ), cur ), repeat( b, cur ) );
37		void fmaRepeat( Tensor& cur, const Tensor& w, const Tensor& b );
38
39		inline void fmaRepeat( Tensor& cur, const TensorPair wb )
40		{
41			fmaRepeat( cur, wb.w, wb.b );
42		}
43
44		// Multiply two matrices
45		Tensor mulMat( const Tensor& a, const Tensor& b );
46
47		// cur = add( repeat( b, cur ), cur ); cur = scale(cur, scaling)
48		void addRepeatScale( Tensor& cur, const Tensor& b, float scaling );
49
50		void addRepeat( Tensor& cur, const Tensor& b );
51
52		Tensor add( const Tensor& a, const Tensor& b );
53		void addInPlace( Tensor& a, const Tensor& b );
54		void addRepeatGelu( Tensor& cur, const Tensor& b );
55
56		// cur = scale(cur, scaling)
57		void scale( Tensor& cur, float scaling );
58
59		void diagMaskInf( Tensor& cur, uint32_t n_past );
60
61		void softMax( Tensor& cur, float inputScale = 1.0f );
62
63		Tensor copy( const Tensor& a, eDataType type, std::initializer_list<uint32_t> size );
64
65		HRESULT copyImpl( Tensor& result, const Tensor& source );
66
67		Tensor permute( const Tensor& a, uint8_t axis0, uint8_t axis1, uint8_t axis2, uint8_t axis3 );
68
69		void copyInPlace( Tensor& dest, const Tensor& a, eDataType type, std::initializer_list<uint32_t> size );
70	};
71}