diff --git a/include/rmm/device_buffer.hpp b/include/rmm/device_buffer.hpp index 48e1fc6b3..78f8ba184 100644 --- a/include/rmm/device_buffer.hpp +++ b/include/rmm/device_buffer.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2021, NVIDIA CORPORATION. + * Copyright (c) 2019-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -285,7 +285,7 @@ class device_buffer { * Reallocates and copies on stream `stream` the contents of the device memory * allocation to reduce `capacity()` to `size()`. * - * If `size() == capacity()`, no allocations nor copies occur. + * If `size() == capacity()`, no allocations or copies occur. * * @throws rmm::bad_alloc If creating the new allocation fails * @throws rmm::cuda_error If the copy from the old to new allocation fails @@ -315,13 +315,21 @@ class device_buffer { void* data() noexcept { return _data; } /** - * @brief Returns size in bytes that was requested for the device memory - * allocation + * @brief Returns the number of bytes. */ [[nodiscard]] std::size_t size() const noexcept { return _size; } /** - * @brief Returns whether the size in bytes of the `device_buffer` is zero. + * @brief Returns the signed number of bytes. + */ + [[nodiscard]] std::int64_t ssize() const noexcept + { + assert(size() < std::numeric_limits::max() && "Size overflows signed integer"); + return static_cast(size()); + } + + /** + * @brief returns the number of bytes that can be held in currently allocated storage. * * If `is_empty() == true`, the `device_buffer` may still hold an allocation * if `capacity() > 0`. diff --git a/include/rmm/device_uvector.hpp b/include/rmm/device_uvector.hpp index ff1220ec3..57ef1149a 100644 --- a/include/rmm/device_uvector.hpp +++ b/include/rmm/device_uvector.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2021, NVIDIA CORPORATION. + * Copyright (c) 2020-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -480,12 +480,19 @@ class device_uvector { [[nodiscard]] const_iterator end() const noexcept { return cend(); } /** - * @brief Returns the number of elements in the vector. - * - * @return The number of elements. + * @brief Returns the number of elements. */ [[nodiscard]] std::size_t size() const noexcept { return bytes_to_elements(_storage.size()); } + /** + * @brief Returns the signed number of elements. + */ + [[nodiscard]] std::int64_t ssize() const noexcept + { + assert(size() < std::numeric_limits::max() && "Size overflows signed integer"); + return static_cast(size()); + } + /** * @brief Returns true if the vector contains no elements, i.e., `size() == 0`. * diff --git a/tests/device_buffer_tests.cu b/tests/device_buffer_tests.cu index fd772cfea..9cab76323 100644 --- a/tests/device_buffer_tests.cu +++ b/tests/device_buffer_tests.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2021, NVIDIA CORPORATION. + * Copyright (c) 2019-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -67,6 +67,7 @@ TYPED_TEST(DeviceBufferTest, DefaultMemoryResource) rmm::device_buffer buff(this->size, rmm::cuda_stream_view{}); EXPECT_NE(nullptr, buff.data()); EXPECT_EQ(this->size, buff.size()); + EXPECT_EQ(this->size, buff.ssize()); EXPECT_EQ(this->size, buff.capacity()); EXPECT_EQ(rmm::mr::get_current_device_resource(), buff.memory_resource()); EXPECT_EQ(rmm::cuda_stream_view{}, buff.stream()); diff --git a/tests/device_uvector_tests.cpp b/tests/device_uvector_tests.cpp index 539dd2c26..c1505bc0a 100644 --- a/tests/device_uvector_tests.cpp +++ b/tests/device_uvector_tests.cpp @@ -1,6 +1,6 @@ /* - * Copyright (c) 2020-2021, NVIDIA CORPORATION. + * Copyright (c) 2020-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -52,7 +52,8 @@ TYPED_TEST(TypedUVectorTest, NonZeroSizeConstructor) { auto const size{12345}; rmm::device_uvector vec(size, this->stream()); - EXPECT_EQ(vec.size(), 12345); + EXPECT_EQ(vec.size(), size); + EXPECT_EQ(vec.ssize(), size); EXPECT_NE(vec.data(), nullptr); EXPECT_EQ(vec.end(), vec.begin() + vec.size()); EXPECT_FALSE(vec.is_empty());