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 →
2.0 KiB80 linesraw
1#include "stdafx.h"
2#include "Reshaper.h"
3#include "../D3D/MappedResource.h"
4#include "../D3D/Binder.h"
5#include "../D3D/shaders.h"
6#include "reshapedMultiply.h"
7
8namespace
9{
10	using namespace DirectCompute;
11	struct Constants
12	{
13		// Size and strides of the source tensor
14		TensorShape arg0;
15		uint32_t zzPadding;
16		// Count of elements per panel
17		uint32_t panelSize;
18		// Layer strides of the output matrix
19		std::array<uint32_t, 2> layerStrides;
20	};
21}
22
23HRESULT DirectCompute::Reshaper::createConstants()
24{
25	constexpr uint32_t cb = sizeof( Constants );
26	CD3D11_BUFFER_DESC desc{ cb, D3D11_BIND_CONSTANT_BUFFER, D3D11_USAGE_DYNAMIC, D3D11_CPU_ACCESS_WRITE };
27	return device()->CreateBuffer( &desc, nullptr, &constantBuffer );
28}
29
30HRESULT DirectCompute::Reshaper::makePanels( Tensor& tensor, eDataType dataType )
31{
32	if( !constantBuffer )
33		CHECK( createConstants() );
34
35	constexpr uint32_t TILE_SIZE = ReshapedMultiply::TILE_SIZE;
36
37	// Reshaping into column major horizontal panels, height = TILE_SIZE, width = width of the source matrix
38
39	std::array<uint32_t, 4> ne = tensor.ne;
40	const uint32_t groupsX = ( ne[ 1 ] + TILE_SIZE - 1 ) / TILE_SIZE;
41	ne[ 1 ] = groupsX * TILE_SIZE;;
42	// Each panel has [ size.x, TILE_SIZE ] elements
43	const uint32_t panelSize = ne[ 0 ] * TILE_SIZE;
44	
45	Tensor result;
46	result.create( dataType, ne );
47
48	{
49		MappedResource mapped;
50		CHECK( mapped.map( constantBuffer, false ) );
51		Constants& cb = *(Constants*)mapped.data();
52
53		store( cb.arg0.ne, tensor.sizeVec() );
54		store( cb.arg0.nb, tensor.stridesVec() );
55		cb.panelSize = panelSize;
56		cb.layerStrides[ 0 ] = result.nb[ 2 ];
57		cb.layerStrides[ 1 ] = result.nb[ 3 ];
58	}
59
60	csSetCB( constantBuffer );
61	{
62		Binder bind;
63		bind.bind( tensor, result );
64		bindShader( eComputeShader::matReshapePanels );
65		context()->Dispatch( groupsX, tensor.ne[ 2 ], tensor.ne[ 3 ] );
66	}
67
68	tensor.nb[ 0 ] = 0;
69	tensor.nb[ 1 ] = panelSize;
70	tensor.nb[ 2 ] = result.nb[ 2 ];
71	tensor.nb[ 3 ] = result.nb[ 3 ];
72	tensor.setGpuViews( result );
73	return S_OK;
74}
75
76DirectCompute::Reshaper::~Reshaper()
77{
78	if( constantBuffer )
79		csSetCB( nullptr );
80}