Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions cpp/include/rmm/detail/format.hpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

Expand All @@ -20,8 +20,8 @@ inline std::string format_bytes(std::size_t value)
{
static std::array units{"B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"};

int index = 0;
auto size = static_cast<double>(value);
std::size_t index = 0;
auto size = static_cast<double>(value);
while (size > 1024) {
size /= 1024;
index++;
Expand Down
32 changes: 32 additions & 0 deletions cpp/include/rmm/detail/safe_cast.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

#pragma once

#include <rmm/detail/error.hpp>

#include <type_traits>

namespace rmm::detail {

/**
* @brief Checked narrowing/sign-converting cast.
*
* Casts `value` to type `To`, asserting at runtime that the value is
* representable in `To`. In release builds the assertion compiles away when
* the compiler can prove it is always true.
*/
template <typename To, typename From>
[[nodiscard]] constexpr To safe_cast(From value)
{
static_assert(std::is_integral_v<From> && std::is_integral_v<To>,
"safe_cast is only defined for integral types");
if constexpr (std::is_signed_v<From> && std::is_unsigned_v<To>) {
RMM_EXPECTS(value >= From{0}, "safe_cast: negative value cannot be represented as unsigned");
}
return static_cast<To>(value);
}

} // namespace rmm::detail
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/format.hpp>
#include <rmm/detail/safe_cast.hpp>
#include <rmm/logger.hpp>
#include <rmm/mr/device_memory_resource.hpp>

Expand Down Expand Up @@ -267,9 +268,10 @@ class stream_ordered_memory_resource : public crtp<PoolResource>, public device_
// the CUDA runtime and thread_local destructors (can) run below
// main: it is undefined behaviour to call into the CUDA
// runtime below main.
thread_local std::vector<cudaEvent_t> events_tls(rmm::get_num_cuda_devices());
thread_local std::vector<cudaEvent_t> events_tls(
rmm::detail::safe_cast<std::size_t>(rmm::get_num_cuda_devices()));
auto event = [device_id = this->device_id_]() {
auto& e = events_tls[device_id.value()];
auto& e = events_tls[rmm::detail::safe_cast<std::size_t>(device_id.value())];
Comment thread
Matt711 marked this conversation as resolved.
Outdated
if (!e) {
// These events are deliberately not destructed and therefore live until
// program exit.
Expand Down
Loading