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 KiB78 linesraw
1#pragma once
2#include "TensorShape.h"
3#include "TensorGpuViews.h"
4#include "../D3D/enums.h"
5
6namespace DirectCompute
7{
8	// A minimal tensor object sufficient to compute things on GPU, with compute shaders
9	// This class only takes 48 bytes in system memory, and is very cheap to make copies 'coz GPU objects are reference counted.
10	class Tensor : public TensorShape, public TensorGpuViews
11	{
12		CComPtr<ID3D11Buffer> getBuffer() const;
13
14		struct TensorType
15		{
16			eDataType type;
17			eBufferUse usage;
18			bool hasInitialData;
19		};
20#ifdef _DEBUG
21		// In debug builds, we include a few pieces of data to this class.
22		TensorType dbgType;
23#endif
24	protected:
25		HRESULT create( eDataType type, std::initializer_list<uint32_t> sizeElements, eBufferUse usage, CComPtr<ID3D11Buffer>& buffer, const void* rsi, ID3D11Buffer** ppStagingBuffer );
26
27		static uint32_t dxgiSizeof( DXGI_FORMAT format );
28
29		void downloadImpl( const D3D11_SHADER_RESOURCE_VIEW_DESC& viewDesc, uint32_t countElements, size_t cbElement, void* rdi ) const;
30
31	public:
32		Tensor() = default;
33
34		// These copy operators don't copy any data, they merely increment ref.counter of the GPU resources
35		Tensor( const Tensor& );
36		Tensor( Tensor&& that ) noexcept;
37		Tensor& operator=( const Tensor& that );
38		Tensor& operator=( Tensor&& that ) noexcept;
39
40		// Move the provided buffer views into this newly created tensor, and assign the shape
41		// This destroys old values in the smart pointers
42		Tensor( const TensorShape& shape, CComPtr<ID3D11ShaderResourceView>& srv, CComPtr<ID3D11UnorderedAccessView>& uav ) noexcept;
43
44		Tensor( const TensorShape& shape, const TensorGpuViews& views );
45
46		// Create a tensor from the GGML's one
47		HRESULT create( const ggml_tensor& ggml, eBufferUse usage, bool uploadData );
48
49		// Create a new dense tensor of the specified size in elements, without initial data
50		HRESULT create( eDataType type, std::initializer_list<uint32_t> sizeElements );
51		HRESULT create( eDataType type, const std::array<uint32_t, 4>& sizeElements );
52		HRESULT createImmutable( eDataType type, const std::array<int, 4>& size, const void* rsi );
53
54		eDataType getType() const;
55
56		// This method should probably only be used to test things
57		// TensorEx is better for production usage, because it creates staging buffer in advance.
58		void download( std::vector<float>& vec ) const;
59		void download( std::vector<uint16_t>& vec ) const;
60
61		// ggml_reshape_3d
62		Tensor reshape3d( uint32_t ne0, uint32_t ne1, uint32_t ne2 ) const;
63
64		inline void dbgSetType( eDataType dt, bool hasData = false, eBufferUse use = eBufferUse::ReadWrite )
65		{
66#ifdef _DEBUG
67			dbgType.type = dt;
68			dbgType.hasInitialData = hasData;
69			dbgType.usage = use;
70#endif
71		}
72
73		__m128i getMemoryUse() const
74		{
75			return resourceMemoryUsage( srv );
76		}
77	};
78}