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 →
3.3 KiB145 linesraw
1#include <stdafx.h>
2#include "BufferAllocator.h"
3#include <immintrin.h>
4#include <ammintrin.h>
5using namespace CpuCompute;
6
7HRESULT BufferAllocator::create( size_t cb )
8{
9	CHECK( buffer.allocate( cb ) );
10	head = 0;
11	size = cb;
12	dbgMarkUninitializedMemory( buffer.pointer(), cb );
13	return S_OK;
14}
15
16namespace
17{
18	// Round up the integer by 32 bytes
19	__forceinline size_t roundUpAlloc( size_t cb )
20	{
21		const size_t mask = 31;
22		cb += mask;
23		// We require AVX1+FMA3 support, might as well use BMI1
24		return _andn_u64( mask, cb );
25	}
26}
27
28void* BufferAllocator::allocate( size_t cb, size_t align ) noexcept
29{
30	assert( align <= 32 );
31	cb = roundUpAlloc( cb );
32
33	uint8_t* pointer = buffer.pointer();
34	if( head + cb > size || nullptr == pointer )
35	{
36		logError( u8"BufferAllocator.allocate, not enough capacity" );
37		return nullptr;
38	}
39
40	void* const res = pointer + head;
41	head += cb;
42	assert( head <= size );
43	dbgMarkUninitializedMemory( res, cb );
44	return res;
45}
46
47namespace
48{
49	// 2 MB of memory, we hope the OS kernel will then be smart enough to give us large pages.
50	constexpr size_t virtualAllocGranularityExp2 = 21;
51
52	constexpr size_t virtualAllocGranularityMask = ( ( (size_t)1 ) << virtualAllocGranularityExp2 ) - 1;
53
54	// Round up the integer by 2 megabytes
55	__forceinline size_t roundUpVirtualAlloc( size_t cb )
56	{
57		const size_t mask = virtualAllocGranularityMask;
58		cb += mask;
59		return _andn_u64( mask, cb );
60	}
61}
62
63HRESULT VirtualAllocator::create( size_t cb )
64{
65	if( nullptr != pointer )
66		return HRESULT_FROM_WIN32( ERROR_ALREADY_INITIALIZED );
67	cb = roundUpVirtualAlloc( cb );
68	pointer = (uint8_t*)VirtualAlloc( NULL, cb, MEM_RESERVE, PAGE_READWRITE );
69	if( nullptr != pointer )
70	{
71		head = 0;
72		sizeAllocated = 0;
73		sizeVirtual = cb;
74		return S_OK;
75	}
76
77	const HRESULT hr = getLastHr();
78	logErrorHr( hr, u8"VirtualAlloc failed" );
79	return hr;
80}
81
82void* VirtualAllocator::allocate( size_t cb, size_t align ) noexcept
83{
84	assert( align <= 32 );
85	cb = roundUpAlloc( cb );
86
87	const size_t newHead = head + cb;
88	if( newHead <= sizeAllocated )
89	{
90		void* const res = pointer + head;
91		head = newHead;
92		dbgMarkUninitializedMemory( res, cb );
93		return res;
94	}
95
96	if( newHead <= sizeVirtual )
97	{
98		uint8_t* const ptrCommit = pointer + sizeAllocated;
99		const size_t cbCommit = roundUpVirtualAlloc( newHead ) - sizeAllocated;
100		void* const res = VirtualAlloc( ptrCommit, cbCommit, MEM_COMMIT, PAGE_READWRITE );
101		if( nullptr != res )
102		{
103			sizeAllocated += cbCommit;
104			assert( sizeAllocated <= sizeVirtual );
105			void* const res = pointer + head;
106			head = newHead;
107			dbgMarkUninitializedMemory( res, cb );
108			return res;
109		}
110
111		const HRESULT hr = getLastHr();
112		logErrorHr( hr, u8"VirtualAllocator.allocate, VirtualAlloc failed" );
113		return nullptr;
114	}
115
116	logError( u8"VirtualAllocator.allocate, not enough arena capacity" );
117	return nullptr;
118}
119
120VirtualAllocator::~VirtualAllocator()
121{
122	if( nullptr == pointer )
123		return;
124
125	if( VirtualFree( pointer, 0, MEM_RELEASE ) )
126	{
127		pointer = nullptr;
128		return;
129	}
130
131	const HRESULT hr = getLastHr();
132	logErrorHr( hr, u8"VirtualFree failed" );
133}
134
135#ifndef NDEBUG
136// Reusing Microsoft's magic numbers: https://asawicki.info/news_1292_magic_numbers_in_visual_c
137void CpuCompute::dbgMarkUninitializedMemory( void* pv, size_t cb )
138{
139	__stosb( (uint8_t*)pv, 0xCD, cb );
140}
141void CpuCompute::dbgMarkFreedMemory( void* pv, size_t cb )
142{
143	__stosd( (DWORD*)pv, 0xFEEEFEEEu, cb / 4 );
144}
145#endif