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.1 KiB149 linesraw
1#include "stdafx.h"
2#include "ParallelForRunner.h"
3using namespace CpuCompute;
4
5ParallelForRunner::ParallelForRunner( int threads ) :
6	maxThreads( threads )
7{
8	if( maxThreads <= 1 )
9	{
10		threadBuffers.resize( 1 );
11		return;
12	}
13
14	work = CreateThreadpoolWork( &workCallbackStatic, this, nullptr );
15	if( nullptr == work )
16		throw getLastHr();
17	threadBuffers.resize( maxThreads );
18}
19
20HRESULT ParallelForRunner::setThreadsCount( int threads )
21{
22	maxThreads = threads;
23	if( threads <= 1 )
24	{
25		threadBuffers.resize( 1 );
26		return S_OK;
27	}
28
29	threadBuffers.resize( maxThreads );
30	if( nullptr == work )
31	{
32		work = CreateThreadpoolWork( &workCallbackStatic, this, nullptr );
33		if( nullptr == work )
34			return getLastHr();
35	}
36	return S_OK;
37}
38
39ParallelForRunner::~ParallelForRunner()
40{
41	if( nullptr != work )
42	{
43		if( S_FALSE == status )
44			WaitForThreadpoolWorkCallbacks( work, FALSE );
45		CloseThreadpoolWork( work );
46	}
47}
48
49namespace
50{
51	thread_local uint32_t currentThreadIndex = UINT_MAX;
52}
53
54void ParallelForRunner::runBatch( size_t ith ) noexcept
55{
56	currentThreadIndex = (uint32_t)ith;
57	const size_t begin = ( ith * countItems ) / countThreads;
58	const size_t end = ( ( ith + 1 ) * countItems ) / countThreads;
59
60	HRESULT hr = E_UNEXPECTED;
61	try
62	{
63		hr = computeRange->compute( begin, end );
64	}
65	catch( HRESULT code )
66	{
67		hr = code;
68	}
69	catch( const std::bad_alloc& )
70	{
71		hr = E_OUTOFMEMORY;
72	}
73	catch( const std::exception& )
74	{
75		hr = E_FAIL;
76	}
77	currentThreadIndex = UINT_MAX;
78	if( SUCCEEDED( hr ) )
79		return;
80	InterlockedCompareExchange( &status, hr, S_FALSE );
81}
82
83void* ParallelForRunner::threadLocalBuffer( size_t cb )
84{
85	const uint32_t idx = currentThreadIndex;
86	if( idx < threadBuffers.size() )
87	{
88		ThreadBuffer& tb = threadBuffers[ idx ];
89		if( tb.cb >= cb )
90		{
91			// We already have large enough buffer for the current thread
92			return tb.memory.pointer();
93		}
94		tb.memory.deallocate();
95		check( tb.memory.allocate( cb ) );
96		tb.cb = cb;
97		return tb.memory.pointer();
98	}
99	if( idx != UINT_MAX )
100		throw E_BOUNDS;
101	else
102	{
103		logError( u8"threadLocalBuffer() method only works from inside a pool callback" );
104		throw E_UNEXPECTED;
105	}
106}
107
108void __stdcall ParallelForRunner::workCallbackStatic( PTP_CALLBACK_INSTANCE Instance, void* pv, PTP_WORK Work ) noexcept
109{
110	ParallelForRunner& context = *(ParallelForRunner*)pv;
111	const size_t ith = (uint32_t)( InterlockedIncrement( &context.threadIndex ) );
112	context.runBatch( ith );
113}
114
115HRESULT ParallelForRunner::parallelFor( iComputeRange& compute, size_t length, size_t minBatch )
116{
117	if( maxThreads <= 1 )
118	{
119		currentThreadIndex = 0;
120		const HRESULT hr1 = compute.compute( 0, length );
121		currentThreadIndex = UINT_MAX;
122		return hr1;
123	}
124	assert( minBatch > 0 );
125
126	size_t nth = length / minBatch;
127	nth = std::min( nth, (size_t)(uint32_t)maxThreads );
128
129	computeRange = &compute;
130	countItems = length;
131	countThreads = nth;
132	threadIndex = 0;
133	status = S_FALSE;
134
135	for( size_t i = 1; i < nth; i++ )
136		SubmitThreadpoolWork( work );
137	runBatch( 0 );
138
139	if( nth > 1 )
140		WaitForThreadpoolWorkCallbacks( work, FALSE );
141
142	computeRange = nullptr;
143	const HRESULT hr = status;
144	status = S_OK;
145	if( SUCCEEDED( hr ) )
146		return S_OK;
147
148	return hr;
149}