Improved VRAM memory management, both speed and memory usage

3 files changed, 38 insertions, 25 deletions

diff --git a/Whisper/ML/TensorsArena.cpp b/Whisper/ML/TensorsArena.cpp
index f1a4bed..8e5d449 100644
--- a/Whisper/ML/TensorsArena.cpp
+++ b/Whisper/ML/TensorsArena.cpp
@@ -1,23 +1,41 @@
 #include "stdafx.h"
 #include "TensorsArena.h"
 #include "../D3D/createBuffer.h"
-#include <bit>
+
+static inline uint32_t roundUpPower2( uint32_t x )
+{
+	// std::bit_ceil from C++/20 standard library implements runtime dispatch, uses LZCNT when AVX2 is available, otherwise BSR
+	// That's not what we want.
+	// BSR is only slightly slower than LZCNT: same speed on Intel, on AMD it's 3 versus 1 cycles.
+	// defaultNewCapacity function is only called occasionally, that branch is therefore unpredictable.
+	assert( x > 1 );
+	unsigned long idx;
+	_BitScanReverse( &idx, x - 1 );
+	return 2u << idx;
+}
 
 uint32_t DirectCompute::defaultNewCapacity( uint32_t current, uint32_t requested )
 {
-	if( 0 == current )
-	{
-		// When the current capacity is 0 this means it's the first resize for the pooled tensor
-		// Create tensor of the exact requested size, as most tensors on these pools are never actually resized.
-		return requested;
-	}
-	else
+	// Implement some reasonable tactics to compute capacity of these buffers
+
+	constexpr uint32_t minAlloc = 1024;	// 1k elements = 4kb of VRAM for FP32 tensors
+	constexpr uint32_t allocGranularity = 1u << 14;	// 16k elements = 64kb of VRAM for FP32 tensors
+
+	if( requested > minAlloc )
 	{
-		// Implement some reasonable tactics to grow an old tensor
-		const uint32_t res = std::max( 1024u, std::bit_ceil( requested ) );
+		const uint32_t roundedUpPowerOf2 = roundUpPower2( requested );
+
+		constexpr uint32_t lowMask = allocGranularity - 1;
+		constexpr uint32_t highMask = ~lowMask;
+		const uint32_t roundedUpGranularity = ( requested + lowMask ) & highMask;
+
+		const uint32_t res = std::min( roundedUpPowerOf2, roundedUpGranularity );
+
 		assert( res >= requested );
 		return res;
 	}
+
+	return minAlloc;
 }
 
 using namespace DirectCompute;
diff --git a/Whisper/Whisper/WhisperContext.cpp b/Whisper/Whisper/WhisperContext.cpp
index 9e62766..6b541f4 100644
--- a/Whisper/Whisper/WhisperContext.cpp
+++ b/Whisper/Whisper/WhisperContext.cpp
@@ -25,7 +25,7 @@ namespace
 }
 
 WhisperContext::Arenas::Arenas() :
-	enc( defaultArenaConfigs() ), encLayer( defaultArenaConfigs() ), dec( defaultArenaConfigs() ), decLayer( defaultArenaConfigs() )
+	outer( defaultArenaConfigs() ), layer( defaultArenaConfigs() )
 { }
 
 Tensor WhisperContext::DecoderLayerPool::tensor( eDataType type, const std::array<uint32_t, 4>& ne )
@@ -158,7 +158,7 @@ Tensor WhisperContext::convolutionAndGelu( const Tensor& mel, uint32_t n_ctx )
 Tensor WhisperContext::encodeLayer( const Tensor& source, size_t index, uint32_t n_state, uint32_t n_head, uint32_t n_ctx )
 {
 	auto prof = profiler.block( eProfilerBlock::EncodeLayer );
-	ArenaRaii arenaRaii{ *this, arenas.encLayer };
+	ArenaRaii arenaRaii{ *this, arenas.layer };
 
 	const LayerEncoder& layer = gpuModel.enc.layers[ index ];
 	// norm
@@ -318,8 +318,8 @@ Tensor WhisperContext::encode( Whisper::iSpectrogram& spectrogram, const sEncode
 	check( melInput.create( spectrogram, encParams ) );
 	Tracing::tensor( "enc.input", melInput );
 
-	arenas.enc.clear();
-	ArenaRaii arenaRaii{ *this, arenas.enc };
+	arenas.outer.clear();
+	ArenaRaii arenaRaii{ *this, arenas.outer };
 
 	// Initial few steps
 	Tensor cur = convolutionAndGelu( melInput, encParams.n_ctx );
@@ -409,7 +409,7 @@ Tensor WhisperContext::decodeLayer( const Tensor& inpL, size_t il, const sLayerD
 {
 	auto prof = profiler.block( eProfilerBlock::DecodeLayer );
 	const auto& layer = gpuModel.dec.layers[ il ];
-	std::optional<ArenaRaii> arenaRaii{ std::in_place, *this, arenas.decLayer };
+	std::optional<ArenaRaii> arenaRaii{ std::in_place, *this, arenas.layer };
 	if( 0 == il ) Tracing::tensor( "dec-inpL", inpL );
 
 	// norm
@@ -594,7 +594,7 @@ void WhisperContext::decode( const int* tokens, const int n_tokens, const sDecod
 	auto prof = profiler.block( eProfilerBlock::DecodeStep );
 	CaptureRaii renderdocCapture;
 	profiler.profileShaders = profileDecodeShaders;
-	ArenaRaii arenaRaii{ *this, arenas.dec };
+	ArenaRaii arenaRaii{ *this, arenas.outer };
 
 	assert( n_tokens > 0 );
 	const uint32_t N = (uint32_t)n_tokens;
@@ -641,10 +641,8 @@ void WhisperContext::decode( const int* tokens, const int n_tokens, const sDecod
 
 __m128i WhisperContext::Arenas::getMemoryUse() const
 {
-	__m128i res = enc.getMemoryUse();
-	res = _mm_add_epi64( res, encLayer.getMemoryUse() );
-	res = _mm_add_epi64( res, dec.getMemoryUse() );
-	res = _mm_add_epi64( res, decLayer.getMemoryUse() );
+	__m128i res = outer.getMemoryUse();
+	res = _mm_add_epi64( res, layer.getMemoryUse() );
 	return res;
 }
 
diff --git a/Whisper/Whisper/WhisperContext.h b/Whisper/Whisper/WhisperContext.h
index 227e6b6..0e4319f 100644
--- a/Whisper/Whisper/WhisperContext.h
+++ b/Whisper/Whisper/WhisperContext.h
@@ -22,11 +22,8 @@ namespace DirectCompute
 	{
 		struct Arenas
 		{
-			TensorsArena enc;
-			TensorsArena encLayer;
-			TensorsArena dec;
-			TensorsArena decLayer;
-
+			TensorsArena outer;
+			TensorsArena layer;
 			Arenas();
 			__m128i getMemoryUse() const;
 		};