Added various functions to waveforms and buffers to enable freeing of GPU memory in OOM conditions

Author: azonenberg

scopehal/AcceleratorBuffer.h

@@ -344,7 +344,7 @@ class AcceleratorBuffer
 	~AcceleratorBuffer()
 	{
 		FreeCpuBuffer();
-		FreeGpuBuffer();
+		FreeGpuBuffer(true);
 	}
 
 	////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1073,15 +1073,25 @@ class AcceleratorBuffer
 		}
 	}
 
+public:
 	/**
 		@brief Free the GPU-side buffer and underlying physical memory
+
+		@param dataLossOK		True if we do not intend to use the contents of this buffer again
+								(and thus it's OK to remove the only copy of the data)
 	 */
-	void FreeGpuBuffer()
+	void FreeGpuBuffer(bool dataLossOK = false)
 	{
 		//Early out if buffer is already null
 		if(m_gpuPhysMem == nullptr)
 			return;
 
+		//If we do NOT have a CPU-side buffer, we're deleting all of our data! Warn for now
+		if( (m_cpuMemoryType == MEM_TYPE_NULL) && m_gpuPhysMemIsStale && !empty() && !dataLossOK)
+		{
+			LogWarning("Freeing a GPU buffer without any CPU backing, may cause data loss\n");
+		}
+
 		//If we have a CPU-side buffer, and it's stale, move our about-to-be-deleted content over before we free it
 		if( (m_cpuMemoryType != MEM_TYPE_NULL) && m_cpuPhysMemIsStale && !empty() )
 			CopyToCpu();
@@ -1325,12 +1335,34 @@ class AcceleratorBuffer
 				}
 			}
 
+			//Retry one more time.
+			//If we OOM simultaneously in two threads, it's possible to have the second OnMemoryPressure call
+			//return false because the first one already freed all it could. But we might have enough free to continue.
+			if(!ok)
+			{
+				LogDebug("Final retry\n");
+				try
+				{
+					m_gpuPhysMem = std::make_unique<vk::raii::DeviceMemory>(*g_vkComputeDevice, info);
+					ok = true;
+				}
+				catch(vk::OutOfDeviceMemoryError& ex2)
+				{
+					LogDebug("Allocation failed again\n");
+				}
+			}
+
 			//If we get here, we couldn't allocate no matter what
-			LogError(
-				"Failed to allocate %s of GPU memory despite our best efforts to reclaim space\n"
-				"This is unrecoverable (for now).\n",
-				Unit(Unit::UNIT_BYTES).PrettyPrint(req.size, 4).c_str());
-			exit(1);
+			//TODO: Fall back to a CPU-side allocation
+			if(!ok)
+			{
+				LogError(
+					"Failed to allocate %s of GPU memory despite our best efforts to reclaim space\n"
+					"This is unrecoverable (for now).\n",
+					Unit(Unit::UNIT_BYTES).PrettyPrint(req.size, 4).c_str());
+
+				std::abort();
+			}
 		}
 		m_gpuMemoryType = MEM_TYPE_GPU_ONLY;
 

scopehal/ConstellationWaveform.h

@@ -85,6 +85,12 @@ class ConstellationWaveform : public DensityFunctionWaveform
 	 */
 	float m_saturationLevel;
 
+	virtual void FreeGpuMemory() override
+	{}
+
+	virtual bool HasGpuBuffer() override
+	{ return false; }
+
 protected:
 
 	/**

scopehal/EyeWaveform.h

@@ -2,7 +2,7 @@
 *                                                                                                                      *
 * libscopehal                                                                                                          *
 *                                                                                                                      *
-* Copyright (c) 2012-2023 Andrew D. Zonenberg and contributors                                                         *
+* Copyright (c) 2012-2024 Andrew D. Zonenberg and contributors                                                         *
 * All rights reserved.                                                                                                 *
 *                                                                                                                      *
 * Redistribution and use in source and binary forms, with or without modification, are permitted provided that the     *
@@ -86,6 +86,12 @@ class EyeWaveform : public DensityFunctionWaveform
 	EyeType GetType()
 	{ return m_type; }
 
+	virtual void FreeGpuMemory() override
+	{}
+
+	virtual bool HasGpuBuffer() override
+	{ return false; }
+
 protected:
 	int64_t* m_accumdata;
 

scopehal/Oscilloscope.h

@@ -959,6 +959,15 @@ class Oscilloscope	: public virtual Instrument
 	}
 
 public:
+
+	/**
+		@brief Free all waveforms in our pool to reclaim memory
+
+		@return True if memory was freed, false if pools were already empty
+	 */
+	bool FreeWaveformPools()
+	{ return m_analogWaveformPool.clear() || m_digitalWaveformPool.clear();	}
+
 	void AddWaveformToAnalogPool(WaveformBase* w)
 	{ m_analogWaveformPool.Add(w); }
 

scopehal/Waveform.h

@@ -234,6 +234,12 @@ class WaveformBase
 
 	virtual void CacheColors();
 
+	///@brief Free GPU-side memory if we are short on VRAM or do not anticipate using this waveform for a while
+	virtual void FreeGpuMemory() =0;
+
+	///@brief Returns true if we have at least one buffer resident on the GPU
+	virtual bool HasGpuBuffer() =0;
+
 protected:
 
 	///@brief Cache of packed RGBA32 data with colors for each protocol decode event. Empty for non-protocol waveforms.
@@ -403,6 +409,12 @@ class UniformWaveform : public UniformWaveformBase
 	///@brief Sample data
 	AcceleratorBuffer<S> m_samples;
 
+	virtual void FreeGpuMemory() override
+	{ m_samples.FreeGpuBuffer(); }
+
+	virtual bool HasGpuBuffer()
+	{ return m_samples.HasGpuBuffer(); }
+
 	virtual void Resize(size_t size) override
 	{ m_samples.resize(size); }
 
@@ -511,6 +523,16 @@ class SparseWaveform : public SparseWaveformBase
 	///@brief Sample data
 	AcceleratorBuffer<S> m_samples;
 
+	virtual void FreeGpuMemory() override
+	{
+		m_offsets.FreeGpuBuffer();
+		m_durations.FreeGpuBuffer();
+		m_samples.FreeGpuBuffer();
+	}
+
+	virtual bool HasGpuBuffer()
+	{ return m_samples.HasGpuBuffer() || m_offsets.HasGpuBuffer() || m_durations.HasGpuBuffer(); }
+
 	virtual void Resize(size_t size) override
 	{
 		m_offsets.resize(size);

scopehal/WaveformPool.h

@@ -101,6 +101,25 @@ class WaveformPool
 		return ret;
 	}
 
+	/**
+		@brief Free all waveforms in the pool to reclaim memory
+
+		@return True if memory was freed, false if pool was empty to begin with
+	 */
+	bool clear()
+	{
+		std::lock_guard<std::mutex> lock(m_mutex);
+
+		if(m_waveforms.empty())
+			return false;
+
+		for(auto w : m_waveforms)
+			delete w;
+		m_waveforms.clear();
+
+		return true;
+	}
+
 protected:
 
 	///@brief Maximum number of waveforms to store in the pool

scopehal/scopehal.cpp

@@ -1003,6 +1003,11 @@ const char* ScopehalGetVersion()
  */
 bool OnMemoryPressure(MemoryPressureLevel level, MemoryPressureType type, size_t requestedSize)
 {
+	//Only allow one OnMemoryPressure() to execute at a time
+	//Anyone else who simultaneously OOM'd has to wait
+	static mutex memoryPressureMutex;
+	lock_guard<mutex> lock(memoryPressureMutex);
+
 	LogWarning("OnMemoryPressure: %s memory exhaustion on %s (tried to allocate %s)\n",
 		(level == MemoryPressureLevel::Hard) ? "Hard" : "Soft",
 		(type == MemoryPressureType::Host) ? "host" : "device",

scopeprotocols/SpectrogramFilter.h

@@ -66,6 +66,12 @@ class SpectrogramWaveform : public DensityFunctionWaveform
 	double GetBottomEdgeFrequency()
 	{ return m_bottomEdgeFrequency; }
 
+	virtual void FreeGpuMemory() override
+	{}
+
+	virtual bool HasGpuBuffer() override
+	{ return false; }
+
 protected:
 	double m_binsize;
 	double m_bottomEdgeFrequency;

scopeprotocols/Waterfall.h

@@ -2,7 +2,7 @@
 *                                                                                                                      *
 * libscopeprotocols                                                                                                    *
 *                                                                                                                      *
-* Copyright (c) 2012-2023 Andrew D. Zonenberg and contributors                                                         *
+* Copyright (c) 2012-2024 Andrew D. Zonenberg and contributors                                                         *
 * All rights reserved.                                                                                                 *
 *                                                                                                                      *
 * Redistribution and use in source and binary forms, with or without modification, are permitted provided that the     *
@@ -57,6 +57,12 @@ class WaterfallWaveform : public DensityFunctionWaveform
 	WaterfallWaveform(const WaterfallWaveform&) =delete;
 	WaterfallWaveform& operator=(const WaterfallWaveform&) =delete;
 
+	virtual void FreeGpuMemory() override
+	{}
+
+	virtual bool HasGpuBuffer() override
+	{ return false; }
+
 	AcceleratorBuffer<float> m_tempBuf;
 };