Add device_buffer::ssize() and device_uvector::ssize()

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<int64_t>::max() && "Size overflows signed integer");
+    return static_cast<int64_t>(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`.

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<int64_t>::max() && "Size overflows signed integer");
+    return static_cast<int64_t>(size());
+  }
+
   /**
    * @brief Returns true if the vector contains no elements, i.e., `size() == 0`.
    *

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());

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<TypeParam> 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());