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 →
1.6 KiB63 linesraw
1#pragma once
2#include "../Whisper/sModelParams.h"
3#include "../Whisper/KeyValueBuffers.h"
4#include "../D3D/MappedResource.h"
5#include "../CPU/Tensor.h"
6
7class KeyValueDownloader
8{
9	CComPtr<ID3D11Buffer> keys, values;
10	uint32_t length = 0;
11
12	using E = uint16_t;
13	static constexpr DirectCompute::eDataType dataType = DirectCompute::eDataType::FP16;
14
15public:
16	// Create the staging resources to download kvCross tensors produced by the GPGPU encoder
17	HRESULT create( const Whisper::sModelParams& mp );
18
19	// Download these two tensors from VRAM to the staging buffers in system RAM
20	HRESULT download( const DirectCompute::KeyValueBuffers& source );
21
22	class ReadMap
23	{
24		const uint32_t length;
25		DirectCompute::MappedResource mappedKeys, mappedValues;
26
27	public:
28		ReadMap( KeyValueDownloader& owner );
29		~ReadMap() = default;
30		ReadMap( const ReadMap& ) = delete;
31
32		// A slice of model.memory_k tensor
33		CpuCompute::Tensor keysView( uint32_t len, uint32_t off ) const
34		{
35			if( len + off <= length )
36			{
37				E* rsi = (E*)mappedKeys.data();
38				rsi += off;
39				return CpuCompute::Tensor::fromData( rsi, dataType, len );
40			}
41			throw E_BOUNDS;
42		}
43
44		// A slice of model.memory_v tensor
45		CpuCompute::Tensor valuesView( uint32_t len, uint32_t off ) const
46		{
47			if( len + off <= length )
48			{
49				E* rsi = (E*)mappedValues.data();
50				rsi += off;
51				return CpuCompute::Tensor::fromData( rsi, dataType, len );
52			}
53			throw E_BOUNDS;
54		}
55	};
56
57	// Map both staging buffers, return RAII object which unmaps when destroyed,
58	// which can supply the data in the shape of CpuCompute::Tensor vector
59	decltype( auto ) map()
60	{
61		return ReadMap( *this );
62	}
63};