span.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <cudf/types.hpp>
#include <cudf/utilities/export.hpp>
 
#include <rmm/device_buffer.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/device_vector.hpp>
 
#include <cuda/std/span>
#include <thrust/detail/raw_pointer_cast.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/memory.h>
 
#include <cstddef>
#include <span>
#include <type_traits>
#include <utility>
 
namespace CUDF_EXPORT cudf {
constexpr std::size_t dynamic_extent = cuda::std::dynamic_extent;
 
// ===== host_span =================================================================================
 
template <typename T>
struct is_host_span_supported_container : std::false_type {};
 
template <typename T, typename Alloc>
struct is_host_span_supported_container<  //
  std::vector<T, Alloc>> : std::true_type {};
 
template <typename T, typename Alloc>
struct is_host_span_supported_container<  //
  thrust::host_vector<T, Alloc>> : std::true_type {};
 
template <typename T, typename Alloc>
struct is_host_span_supported_container<  //
  std::basic_string<T, std::char_traits<T>, Alloc>> : std::true_type {};
 
template <typename T, std::size_t Extent = cudf::dynamic_extent>
struct host_span {
 private:
  using span_type = cuda::std::span<T, Extent>;  
 
 public:
  using element_type    = typename span_type::element_type;     
  using value_type      = typename span_type::value_type;       
  using size_type       = typename span_type::size_type;        
  using difference_type = typename span_type::difference_type;  
  using pointer         = typename span_type::pointer;          
  using const_pointer   = typename span_type::const_pointer;  
  using reference       = typename span_type::reference;      
  using const_reference = typename span_type::const_reference;  
  using iterator        = pointer;  
 
  static constexpr std::size_t extent = span_type::extent;  
 
  constexpr host_span() noexcept {}  // required to compile on centos
 
  CUDF_HOST_DEVICE constexpr host_span(T* data, std::size_t size) : _span{data, size} {}
 
  CUDF_HOST_DEVICE constexpr host_span(T* data, std::size_t size, bool is_device_accessible)
    : _span{data, size}, _is_device_accessible{is_device_accessible}
  {
  }
 
  template <typename C,
            // Only supported containers of types convertible to T
            std::enable_if_t<is_host_span_supported_container<C>::value &&
                             std::is_convertible_v<
                               std::remove_pointer_t<decltype(thrust::raw_pointer_cast(  // NOLINT
                                 std::declval<C&>().data()))> (*)[],
                               T (*)[]>>* = nullptr>  // NOLINT
  constexpr host_span(C& in) : _span{thrust::raw_pointer_cast(in.data()), in.size()}
  {
  }
 
  template <typename C,
            // Only supported containers of types convertible to T
            std::enable_if_t<is_host_span_supported_container<C>::value &&
                             std::is_convertible_v<
                               std::remove_pointer_t<decltype(thrust::raw_pointer_cast(  // NOLINT
                                 std::declval<C&>().data()))> (*)[],
                               T (*)[]>>* = nullptr>  // NOLINT
  constexpr host_span(C const& in) : _span{thrust::raw_pointer_cast(in.data()), in.size()}
  {
  }
 
  // Copy construction to support const conversion
  template <typename OtherT,
            std::size_t OtherExtent,
            std::enable_if_t<(Extent == OtherExtent || Extent == dynamic_extent) &&
                               std::is_convertible_v<OtherT (*)[], T (*)[]>,  // NOLINT
                             void>* = nullptr>
  constexpr host_span(host_span<OtherT, OtherExtent> const& other) noexcept
    : _span{other.data(), other.size()}, _is_device_accessible{other.is_device_accessible()}
  {
  }
 
  // not noexcept due to undefined behavior when idx < 0 || idx >= size
  constexpr reference operator[](size_type idx) const
  {
    static_assert(sizeof(idx) >= sizeof(size_t), "index type must not be smaller than size_t");
    return _span[idx];
  }
 
  // not noexcept due to undefined behavior when size = 0
  [[nodiscard]] constexpr reference front() const { return _span.front(); }
  // not noexcept due to undefined behavior when size = 0
  [[nodiscard]] constexpr reference back() const { return _span.back(); }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr iterator begin() const noexcept { return _span.data(); }
  [[nodiscard]] CUDF_HOST_DEVICE constexpr iterator end() const noexcept
  {
    return _span.data() + _span.size();
  }
  [[nodiscard]] CUDF_HOST_DEVICE constexpr pointer data() const noexcept { return _span.data(); }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr size_type size() const noexcept { return _span.size(); }
  [[nodiscard]] CUDF_HOST_DEVICE constexpr size_type size_bytes() const noexcept
  {
    return _span.size_bytes();
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr bool empty() const noexcept { return _span.empty(); }
 
  [[nodiscard]] constexpr host_span first(size_type count) const noexcept
  {
    return host_span{_span.data(), count, _is_device_accessible};
  }
 
  [[nodiscard]] constexpr host_span last(size_type count) const noexcept
  {
    return host_span{_span.data() + _span.size() - count, count, _is_device_accessible};
  }
 
  [[nodiscard]] bool is_device_accessible() const { return _is_device_accessible; }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr host_span subspan(size_type offset,
                                                             size_type count) const noexcept
  {
    return host_span{_span.data() + offset, count, _is_device_accessible};
  }
 
  [[nodiscard]] constexpr operator std::span<T>() const noexcept
  {
    return std::span<T>(_span.data(), _span.size());
  }
 
 private:
  // TODO: could be std::span once base_2dspan moves to cuda::std::mdspan and host_span no longer
  // needs to be device-usable.
  span_type _span;
  bool _is_device_accessible{false};
};
 
// ===== device_span ===============================================================================
 
template <typename T, std::size_t Extent = cuda::std::dynamic_extent>
using device_span = cuda::std::span<T, Extent>;  // end of group
 
namespace detail {
 
template <typename T, template <typename, std::size_t> typename RowType>
class base_2dspan {
 public:
  using size_type =
    std::pair<size_t, size_t>;  
 
  constexpr base_2dspan() noexcept = default;
  constexpr base_2dspan(RowType<T, dynamic_extent> flat_view, size_t columns)
    : _flat{flat_view}, _size{columns == 0 ? 0 : flat_view.size() / columns, columns}
  {
#ifndef __CUDA_ARCH__
    CUDF_EXPECTS(_size.first * _size.second == flat_view.size(), "Invalid 2D span size");
#endif
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr auto data() const noexcept { return _flat.data(); }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr auto size() const noexcept { return _size; }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr auto count() const noexcept { return _flat.size(); }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr bool is_empty() const noexcept { return count() == 0; }
 
  CUDF_HOST_DEVICE constexpr RowType<T, dynamic_extent> operator[](std::size_t row) const
  {
    return _flat.subspan(row * _size.second, _size.second);
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE constexpr RowType<T, dynamic_extent> flat_view() const
  {
    return _flat;
  }
 
  template <typename OtherT,
            template <typename, size_t> typename OtherRowType,
            std::enable_if_t<std::is_convertible_v<OtherRowType<OtherT, dynamic_extent>,
                                                   RowType<T, dynamic_extent>>,
                             void>* = nullptr>
  constexpr base_2dspan(base_2dspan<OtherT, OtherRowType> const& other) noexcept
    : _flat{other.flat_view()}, _size{other.size()}
  {
  }
 
 protected:
  RowType<T, dynamic_extent> _flat;  
  size_type _size{0, 0};             
};
 
template <class T>
using host_2dspan = base_2dspan<T, host_span>;
 
template <class T>
using device_2dspan = base_2dspan<T, device_span>;
 
}  // namespace detail
}  // namespace CUDF_EXPORT cudf