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.5 KiB52 linesraw
1#pragma once
2#include "LargeBuffer.h"
3
4namespace CpuCompute
5{
6	// Callback interface for the parallel `for`
7	__interface iComputeRange
8	{
9		// The implementation calls this method on multiple thread pool threads in parallel, and aggregates status codes.
10		HRESULT __stdcall compute( size_t begin, size_t end ) const;
11	};
12
13	// Similar to ThreadPoolWork in parallelFor.h, optimized to be used as a direct replacement of OpenMP pool.
14	class alignas( 64 ) ParallelForRunner
15	{
16	public:
17		ParallelForRunner( int threads );
18		~ParallelForRunner();
19
20		HRESULT setThreadsCount( int threads );
21
22		HRESULT parallelFor( iComputeRange& compute, size_t length, size_t minBatch = 1 );
23
24		// Allocate a temporary buffer for the calling thread.
25		// The pointer is guaranteed to be aligned by page size = 4kb
26		void* threadLocalBuffer( size_t cb );
27
28	private:
29
30		int maxThreads;
31		PTP_WORK work = nullptr;
32		iComputeRange* computeRange = nullptr;
33		size_t countItems = 0;
34		size_t countThreads = 0;
35
36		// Aligning by cache lines.
37		// Avoiding cache line sharing between CPU cores improves performance, despite wasting a few bytes of memory.
38		struct alignas( 64 ) ThreadBuffer
39		{
40			LargeBuffer memory;
41			size_t cb = 0;
42		};
43		std::vector<ThreadBuffer> threadBuffers;
44
45		alignas( 64 ) volatile long threadIndex = 0;
46		volatile HRESULT status = S_OK;
47
48		void runBatch( size_t ith ) noexcept;
49
50		static void __stdcall workCallbackStatic( PTP_CALLBACK_INSTANCE Instance, void* pv, PTP_WORK Work ) noexcept;
51	};
52}