Since last few weeks, I have been working on a memory manager for Kokkos Views (see KMem). Although, originally, I wanted the manager to have a strict max_pool_memory limit for the pool, having a sensible default value that is also portable across different architectures, is not trivial. Also, a consensus among the users of KMem was that the default should be “unlimited”.
However, such a feature means that if the pool uses up all the available memory, new allocations will fail. A proper memory manager should handle such situations and clean up oldest used memory in the pool to make space for new allocations. This means, we need to catch the exception thrown by Kokkos when an allocation fails, and clear the pool, and then retry the allocation. But Kokkos throws a simple std::runtime_error exception when an allocation fails with the following message:
void Kokkos::Impl::throw_bad_alloc(std::string_view memory_space_name,
std::size_t size, std::string_view label) {
std::stringstream ss;
ss << "Kokkos ERROR: " << memory_space_name
<< " memory space failed to allocate " << human_memory_size(size)
<< " (label=\"" << label << "\").";
throw std::runtime_error(ss.str());
}But this means, that we need to parse the error message to check if it is an out-of-memory error via string matching, which is not a good idea. So, I decided to add a custom BadAlloc exception to Kokkos that is thrown when an allocation fails due to insufficient memory. This way, we can catch the BadAlloc exception and handle it properly in KMem. BadAlloc is derived from std::runtime_error and constructs the base with the same error message as before. This ensures that the code is completely backward compatible with existing code that catches std::runtime_error exceptions.
namespace Kokkos::Experimental {
/// Exception thrown when a Kokkos memory allocation fails.
///
/// Inherits from std::runtime_error for backward compatibility — existing
/// catch blocks that catch std::runtime_error still work. Provides structured
/// access to the memory space, requested allocation size, and label of the
/// allocation that failed.
class BadAlloc : public std::runtime_error {
std::string m_memory_space_name;
std::size_t m_allocation_size;
std::string m_label;
static std::string make_message(std::string_view memory_space_name,
std::size_t size, std::string_view label) {
std::stringstream ss;
ss << "Kokkos ERROR: " << memory_space_name
<< " memory space failed to allocate "
<< Kokkos::Impl::human_memory_size(size) << " (label=\"" << label
<< "\").";
return ss.str();
}
public:
BadAlloc(std::string_view memory_space_name, std::size_t size,
std::string_view label)
: std::runtime_error(make_message(memory_space_name, size, label)),
m_memory_space_name(memory_space_name),
m_allocation_size(size),
m_label(label) {}
const std::string& memory_space_name() const noexcept {
return m_memory_space_name;
}
std::size_t allocation_size() const noexcept { return m_allocation_size; }
const std::string& label() const noexcept { return m_label; }
};
} // namespace Kokkos::Experimentaland the corresponding implementation in Kokkos_Error.cpp:
void Kokkos::Impl::throw_bad_alloc(std::string_view memory_space_name,
std::size_t size, std::string_view label) {
throw Kokkos::Experimental::BadAlloc(memory_space_name, size, label);
}This got added into Kokkos in PR #9260