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

KonstantinBugfix, incorrect output of command-line examples when launched with multiple input files3ba8e63

master
branchesmasterall refs →
4.1 KiB165 linesraw
1#include "stdafx.h"
2#include "TensorsArena.h"
3#include "../D3D/createBuffer.h"
4#include "TempBuffers.h"
5
6static inline uint32_t roundUpPower2( uint32_t x )
7{
8	// std::bit_ceil from C++/20 standard library implements runtime dispatch, uses LZCNT when AVX2 is available, otherwise BSR
9	// That's not what we want.
10	// BSR is only slightly slower than LZCNT: same speed on Intel, on AMD it's 3 versus 1 cycles.
11	// defaultNewCapacity function is only called occasionally, that branch is therefore unpredictable.
12	assert( x > 1 );
13	unsigned long idx;
14	_BitScanReverse( &idx, x - 1 );
15	return 2u << idx;
16}
17
18uint32_t DirectCompute::defaultNewCapacity( uint32_t current, uint32_t requested )
19{
20	// Implement some reasonable tactics to compute capacity of these buffers
21
22	constexpr uint32_t minAlloc = 1024;	// 1k elements = 4kb of VRAM for FP32 tensors
23	constexpr uint32_t allocGranularity = 1u << 14;	// 16k elements = 64kb of VRAM for FP32 tensors
24
25	if( requested > minAlloc )
26	{
27		const uint32_t roundedUpPowerOf2 = roundUpPower2( requested );
28
29		constexpr uint32_t lowMask = allocGranularity - 1;
30		constexpr uint32_t highMask = ~lowMask;
31		const uint32_t roundedUpGranularity = ( requested + lowMask ) & highMask;
32
33		const uint32_t res = std::min( roundedUpPowerOf2, roundedUpGranularity );
34
35		assert( res >= requested );
36		return res;
37	}
38
39	return minAlloc;
40}
41
42using namespace DirectCompute;
43
44TensorsArena::ArenaImpl::ArenaImpl( eDataType dataType, const sArenaConfig& config ) :
45	type( dataType ),
46	pfnNewCap( nullptr != config.pfnCapInner ? config.pfnCapInner : &defaultNewCapacity )
47{
48	pool.reserve( config.initialCapOuter );
49}
50
51Tensor PooledTensor::tensor( eDataType type, const std::array<uint32_t, 4>& ne, pfnNewCapacity pfnNewCap )
52{
53	const uint32_t p1 = ne[ 0 ] * ne[ 1 ];
54	const uint32_t p2 = ne[ 2 ] * ne[ 3 ];
55	const uint32_t count = p1 * p2;
56
57	if( count > capacity )
58	{
59		views.clear();
60		const uint32_t newCap = pfnNewCap( capacity, count );
61		assert( newCap >= count );
62
63		const size_t cb = elementSize( type ) * newCap;
64		CComPtr<ID3D11Buffer> buffer;
65		check( createBuffer( eBufferUse::ReadWrite, cb, &buffer, nullptr, nullptr ) );
66		check( views.create( buffer, viewFormat( type ), newCap, true ) );
67		capacity = newCap;
68	}
69
70	TensorShape shape;
71	shape.ne = ne;
72	shape.setDenseStrides();
73	Tensor res{ shape, views };
74	res.dbgSetType( type );
75	return res;
76}
77
78Tensor TensorsArena::ArenaImpl::tensor( const std::array<uint32_t, 4>& ne )
79{
80	PooledTensor* res;
81	if( index >= pool.size() )
82	{
83		assert( index == pool.size() );
84		res = &pool.emplace_back();
85	}
86	else
87		res = &pool[ index ];
88
89	index++;
90	return res->tensor( type, ne, pfnNewCap );
91}
92
93TensorsArena::TensorsArena( const sArenaConfigs& configs ) :
94	arenas{ ArenaImpl{ eDataType::FP16, configs.fp16 }, ArenaImpl{ eDataType::FP32, configs.fp32 } }
95{
96	static_assert( 0 == (uint8_t)eDataType::FP16 );
97	static_assert( 1 == (uint8_t)eDataType::FP32 );
98}
99
100Tensor TensorsArena::tensor( eDataType type, const std::array<uint32_t, 4>& ne )
101{
102	ArenaImpl& arena = arenas[ (uint8_t)type ];
103	return arena.tensor( ne );
104}
105
106void TensorsArena::reset()
107{
108	for( ArenaImpl& a : arenas )
109		a.reset();
110}
111
112void TensorsArena::clear()
113{
114	for( ArenaImpl& a : arenas )
115		a.clear();
116}
117
118__m128i TensorsArena::ArenaImpl::getMemoryUse() const
119{
120	const size_t cbElement = elementSize( type );
121	size_t countElts = 0;
122	for( const auto& t : pool )
123		countElts += t.getCapacity();
124
125	const size_t cbVideo = cbElement * countElts;
126	const size_t cbSystem = vectorMemoryUse( pool );
127	return setr_size( cbSystem, cbVideo );
128}
129
130__m128i TensorsArena::getMemoryUse() const
131{
132	__m128i res = _mm_setzero_si128();
133	for( const auto& a : arenas )
134		res = _mm_add_epi64( res, a.getMemoryUse() );
135	return res;
136}
137
138HRESULT PooledTensor::zeroMemory( CComPtr<ID3D11Buffer>& cb )
139{
140	if( 0 == capacity )
141		return S_FALSE;
142	try
143	{
144		TempBuffers::zeroMemory( views, capacity, cb );
145		return S_OK;
146	}
147	catch( HRESULT hr )
148	{
149		return hr;
150	}
151}
152
153HRESULT TensorsArena::ArenaImpl::zeroMemory( CComPtr<ID3D11Buffer>& cb )
154{
155	for( PooledTensor& e : pool )
156		CHECK( e.zeroMemory( cb ) );
157	return S_OK;
158}
159
160HRESULT TensorsArena::zeroMemory( CComPtr<ID3D11Buffer>& cb )
161{
162	for( ArenaImpl& e : arenas )
163		CHECK( e.zeroMemory( cb ) );
164	return S_OK;
165}