[REVIEW] Implement span storage optimization.

cpp/include/raft/detail/span.hpp

@@ -86,5 +86,30 @@ __host__ __device__ constexpr auto lexicographical_compare(InputIt1 first1,
   }
   return first1 == last1 && first2 != last2;
 }
+
+template <typename T, typename SizeType, size_t Extent>
+struct span_storage {
+ private:
+  T* ptr_{nullptr};
+
+ public:
+  constexpr span_storage() = default;
+  constexpr span_storage(T* ptr, SizeType) : ptr_{ptr} {}
+  [[nodiscard]] constexpr auto size() const -> SizeType { return Extent; }
+  constexpr auto data() const -> T* { return ptr_; }
+};
+
+template <typename T, typename SizeType>
+struct span_storage<T, SizeType, dynamic_extent> {
+ private:
+  T* ptr_{nullptr};
+  SizeType size_{0};
+
+ public:
+  constexpr span_storage() = default;
+  constexpr span_storage(T* ptr, SizeType size) : ptr_{ptr}, size_{size} {}
+  [[nodiscard]] constexpr auto size() const -> SizeType { return size_; }
+  constexpr auto data() const -> T* { return ptr_; }
+};
 }  // namespace detail
 }  // namespace raft

cpp/include/raft/span.hpp

@@ -59,23 +59,24 @@ class span {
   /**
    * @brief Constructs a span that is a view over the range [first, first + count);
    */
-  constexpr span(pointer ptr, size_type count) noexcept : size_(count), data_(ptr)
+  constexpr span(pointer ptr, size_type count) noexcept : storage_{ptr, count}
   {
     assert(!(Extent != dynamic_extent && count != Extent));
     assert(ptr || count == 0);
   }
   /**
    * @brief Constructs a span that is a view over the range [first, last)
    */
-  constexpr span(pointer first, pointer last) noexcept : size_(last - first), data_(first)
+  constexpr span(pointer first, pointer last) noexcept
+    : storage_{first, static_cast<size_type>(thrust::distance(first, last))}
   {
     assert(data_ || size_ == 0);
   }
   /**
    * @brief Constructs a span that is a view over the array arr.
    */
   template <std::size_t N>
-  constexpr span(element_type (&arr)[N]) noexcept : size_(N), data_(&arr[0])
+  constexpr span(element_type (&arr)[N]) noexcept : storage_{&arr[0], N}
   {
   }
 
@@ -89,7 +90,7 @@ class span {
               detail::is_allowed_element_type_conversion_t<U, T>::value &&
               detail::is_allowed_extent_conversion_t<OtherExtent, Extent>::value>>
   constexpr span(const span<U, is_device, OtherExtent>& other) noexcept
-    : size_(other.size()), data_(other.data())
+    : storage_{other.data(), other.size()}
   {
   }
 
@@ -139,10 +140,10 @@ class span {
     return data()[_idx];
   }
 
-  constexpr auto data() const noexcept -> pointer { return data_; }
+  constexpr auto data() const noexcept -> pointer { return storage_.data(); }
 
   // Observers
-  [[nodiscard]] constexpr auto size() const noexcept -> size_type { return size_; }
+  [[nodiscard]] constexpr auto size() const noexcept -> size_type { return storage_.size(); }
   [[nodiscard]] constexpr auto size_bytes() const noexcept -> size_type
   {
     return size() * sizeof(T);
@@ -197,8 +198,7 @@ class span {
   }
 
  private:
-  size_type size_{0};
-  pointer data_{nullptr};
+  detail::span_storage<T, size_type, Extent> storage_;
 };
 
 /**