From 5291351790dc129bae0bdd319b157756777ddc26 Mon Sep 17 00:00:00 2001
From: Andreas Hehn <ahehn@nvidia.com>
Date: Tue, 24 Mar 2020 17:00:24 +0100
Subject: [PATCH] [common] Added a device_memory_allocation_exception

---
 .../include/claragenomics/utils/allocator.hpp | 26 +++++++++--
 .../claragenomics/utils/exceptions.hpp        | 43 +++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)
 create mode 100644 common/utils/include/claragenomics/utils/exceptions.hpp

diff --git a/common/utils/include/claragenomics/utils/allocator.hpp b/common/utils/include/claragenomics/utils/allocator.hpp
index 054de2614..d5cc18014 100644
--- a/common/utils/include/claragenomics/utils/allocator.hpp
+++ b/common/utils/include/claragenomics/utils/allocator.hpp
@@ -18,6 +18,7 @@
 #include <claragenomics/utils/device_preallocated_allocator.cuh>
 
 #include <claragenomics/utils/cudautils.hpp>
+#include <claragenomics/utils/exceptions.hpp>
 
 namespace claragenomics
 {
@@ -105,8 +106,20 @@ class CudaMallocAllocator
     pointer allocate(std::size_t n, cudaStream_t stream = 0)
     {
         static_cast<void>(stream);
-        void* ptr = 0;
-        CGA_CU_CHECK_ERR(cudaMalloc(&ptr, n * sizeof(T)));
+        void* ptr       = nullptr;
+        cudaError_t err = cudaMalloc(&ptr, n * sizeof(T));
+        if (err == cudaErrorMemoryAllocation)
+        {
+            // Clear the error from the runtime...
+            err = cudaGetLastError();
+            // Did a different (async) error happen in the meantime?
+            if (err != cudaErrorMemoryAllocation)
+            {
+                CGA_CU_CHECK_ERR(err);
+            }
+            throw device_memory_allocation_exception();
+        }
+        CGA_CU_CHECK_ERR(err);
         return static_cast<pointer>(ptr);
     }
 
@@ -213,8 +226,13 @@ class CachingDeviceAllocator
     /// @return pointer to allocated array
     pointer allocate(std::size_t n, cudaStream_t stream = 0)
     {
-        void* ptr = 0;
-        CGA_CU_CHECK_ERR(memory_resource_->DeviceAllocate(&ptr, n * sizeof(T), stream));
+        void* ptr       = nullptr;
+        cudaError_t err = memory_resource_->DeviceAllocate(&ptr, n * sizeof(T), stream);
+        if (err == cudaErrorMemoryAllocation)
+        {
+            throw device_memory_allocation_exception();
+        }
+        CGA_CU_CHECK_ERR(err);
         return static_cast<pointer>(ptr);
     }
 
diff --git a/common/utils/include/claragenomics/utils/exceptions.hpp b/common/utils/include/claragenomics/utils/exceptions.hpp
new file mode 100644
index 000000000..099022b32
--- /dev/null
+++ b/common/utils/include/claragenomics/utils/exceptions.hpp
@@ -0,0 +1,43 @@
+/*
+* Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+*
+* NVIDIA CORPORATION and its licensors retain all intellectual property
+* and proprietary rights in and to this software, related documentation
+* and any modifications thereto.  Any use, reproduction, disclosure or
+* distribution of this software and related documentation without an express
+* license agreement from NVIDIA CORPORATION is strictly prohibited.
+*/
+
+#pragma once
+
+#include <exception>
+
+namespace claragenomics
+{
+
+/// @brief Exception class for out-of-(device-)memory errors.
+///
+/// Exceptions of this class are thrown if a memory allocation fails on the device.
+class device_memory_allocation_exception : public std::exception
+{
+public:
+    device_memory_allocation_exception() = default;
+    /// Copy constructor
+    device_memory_allocation_exception(device_memory_allocation_exception const&) = default;
+    /// Move constructor
+    device_memory_allocation_exception(device_memory_allocation_exception&&) = default;
+    /// Assignment
+    device_memory_allocation_exception& operator=(device_memory_allocation_exception const&) = default;
+    /// Move-Assignment
+    device_memory_allocation_exception& operator=(device_memory_allocation_exception&&) = default;
+    /// Destructor
+    virtual ~device_memory_allocation_exception() = default;
+
+    /// Returns the error message of the exception
+    virtual const char* what() const noexcept
+    {
+        return "Could not allocate device memory!";
+    }
+};
+
+} // namespace claragenomics