Attention

The vector search and clustering algorithms in RAFT are being migrated to a new library dedicated to vector search called cuVS. We will continue to support the vector search algorithms in RAFT during this move, but will no longer update them after the RAPIDS 24.06 (June) release. We plan to complete the migration by RAPIDS 24.08 (August) release.

mdbuffer: Multi-dimensional Maybe-Owning Container#

#include <raft/core/mdbuffer.cuh>

template<raft::memory_type MemType, typename Variant> using alternate_from_mem_type = std::variant_alternative_t<variant_index_from_memory_type(MemType) % std::variant_size_v<Variant>, Variant>#: Retrieve a type from a variant based on a given memory type.

template<typename T> using default_container_policy_variant = std::variant<detail::memory_type_to_default_policy_t<T, memory_type_from_variant_index(0)>, detail::memory_type_to_default_policy_t<T, memory_type_from_variant_index(1)>, detail::memory_type_to_default_policy_t<T, memory_type_from_variant_index(2)>, detail::memory_type_to_default_policy_t<T, memory_type_from_variant_index(3)>>#: A variant of container policies for each memory type which can be used to build the default container policy for a buffer.

template<typename T> using is_mdbuffer_t = is_mdbuffer<std::remove_const_t<T>>#

template<typename T> using is_input_mdbuffer_t = is_input_mdbuffer<T>#

template<typename T> using is_output_mdbuffer_t = is_output_mdbuffer<T>#

template<typename ...Tn> using enable_if_mdbuffer = std::enable_if_t<is_mdbuffer_v<Tn...>>#

template<typename ...Tn> using enable_if_input_mdbuffer = std::enable_if_t<is_input_mdbuffer_v<Tn...>>#

template<typename ...Tn> using enable_if_output_mdbuffer = std::enable_if_t<is_output_mdbuffer_v<Tn...>>#

template<typename ...Tn> constexpr bool is_mdbuffer_v = std::conjunction_v<is_mdbuffer_t<Tn>...>#: \brief Boolean to determine if variadic template types Tn are raft::mdbuffer or derived types

template<typename ...Tn> constexpr bool is_input_mdbuffer_v = std::conjunction_v<is_input_mdbuffer_t<Tn>...>#

template<typename ...Tn> constexpr bool is_output_mdbuffer_v = std::conjunction_v<is_output_mdbuffer_t<Tn>...>#

inline auto constexpr variant_index_from_memory_type(raft::memory_type mem_type)#

Retrieve a canonical index associated with a given memory type.

For variants based on memory type, this index can be used to help keep a consistent ordering of the memory types in the variant.

inline auto constexpr memory_type_from_variant_index(std::underlying_type_t<raft::memory_type> index)#: Retrieve the memory type associated with a canonical index.

template<typename ElementType, typename Extents, typename LayoutPolicy, typename ContainerPolicy> void __takes_an_mdbuffer_ptr(mdbuffer<ElementType, Extents, LayoutPolicy, ContainerPolicy>*)#: \brief Template checks and helpers to determine if type T is an mdbuffer or a derived type

template<typename ElementType, typename ContainerPolicyVariant = default_container_policy_variant<std::remove_cv_t<ElementType>>> struct default_buffer_container_policy#: A template used to translate a variant of underlying mdarray container policies into a container policy that can be used by an mdbuffer.

template<typename ElementType, typename Extents, typename LayoutPolicy = layout_c_contiguous, typename ContainerPolicy = default_buffer_container_policy<ElementType>> struct mdbuffer#

A type representing multi-dimensional data which may or may not own its underlying storage. raft::mdbuffer is used to conveniently perform copies of data only when necessary to ensure that the data are accessible in the desired memory space and format.

When developing functions that interact with the GPU, it is often necessary to ensure that the data are in a particular memory space (e.g. device, host, managed, pinned), but those functions may be called with data that may or may not already be in the desired memory space. For instance, when called in one workflow, the data may have been previously transferred to device, rendering a copy unnecessary. In another, the function may be directly invoked on host data.

Even when working strictly with host memory, it is often necessary to ensure that the data are in a particular layout for efficient access (e.g. column major vs row major) or that the the data are of a particular type (e.g. double) even though we wish to call the function with data of another compatible type (e.g. float).

mdbuffer is a tool for ensuring that the data are represented in exactly the desired format and location while flexibly supporting data which may not already be in that format or location. It does so by providing a non-owning view on data which are already in the required form, but it allocates (owned) memory and performs a copy if and only if it is necessary.

Usage example:

template <typename mdspan_type>
void foo_device(raft::resources const& res, mdspan_type data) {
  auto buf = raft::mdbuffer{res, raft::mdbuffer{data}, raft::memory_type::device};
  // Data in buf is now guaranteed to be accessible from device.
  // If it was already accessible from device, no copy was performed. If it
  // was not, a copy was performed.

  some_kernel<<<...>>>(buf.view<raft::memory_type::device>());

  // It is sometimes useful to know whether or not a copy was performed to
  // e.g. determine whether the transformed data should be copied back to its original
  // location. This can be checked via the `is_owning()` method.
  if (buf.is_owning()) {
    raft::copy(res, data, buf.view<raft::memory_type::device>());
  }
}

Note that in this example, the foo_device template can be correctly instantiated for both host and device mdspans. Similarly we can use mdbuffer to coerce data to a particular memory layout and data-type, as in the following example:

template <typename mdspan_type>
void foo_device(raft::resources const& res, mdspan_type data) {
  auto buf = raft::mdbuffer<float, raft::matrix_extent<int>, raft::row_major>{res,
raft::mdbuffer{data}, raft::memory_type::device};
  // Data in buf is now guaranteed to be accessible from device, and
  // represented by floats in row-major order.

  some_kernel<<<...>>>(buf.view<raft::memory_type::device>());

  // The same check can be used to determine whether or not a copy was
  // required, regardless of the cause. I.e. if the data were already on
  // device but in column-major order, the is_owning() method would still
  // return true because new storage needed to be allocated.
  if (buf.is_owning()) {
    raft::copy(res, data, buf.view<raft::memory_type::device>());
  }
}

Note that in this example, the foo_device template can accept data of any float-convertible type in any layout and of any memory type and coerce it to the desired device-accessible representation.

Because mdspan types can be implicitly converted to mdbuffer, it is even possible to avoid multiple template instantiations by directly accepting an mdbuffer as argument, as in the following example:

void foo_device(raft::resources const& res, raft::mdbuffer<float, raft::matrix_extent<int>>&&
data) { auto buf = raft::mdbuffer{res, data, raft::memory_type::device};
  // Data in buf is now guaranteed to be accessible from device.

  some_kernel<<<...>>>(buf.view<raft::memory_type::device>());
}

In this example, foo_device can now accept any row-major mdspan of floats regardless of memory type without requiring separate template instantiations for each type.

While the view method takes an optional compile-time memory type parameter, omitting this parameter will return a std::variant of mdspan types. This allows for straightforward runtime dispatching based on the memory type using std::visit, as in the following example:

void foo(raft::resources const& res, raft::mdbuffer<float, raft::matrix_extent<int>>&& data) {
  std::visit([](auto&& view) {
    // Do something with the view, including (possibly) dispatching based on
    // whether it is a host, device, managed, or pinned mdspan
  }, data.view());
}

For convenience, runtime memory-type dispatching can also be performed without explicit use of mdbuffer using raft::memory_type_dispatcher, as described in Dispatch functor based on memory type. Please see the full documentation of that function template for more extensive discussion of the many ways it can be used. To illustrate its connection to mdbuffer, however, consider the following example, which performs a similar task to the above std::visit call:

void foo_device(raft::resources const& res, raft::device_matrix_view<float> data) {
  // Implement foo solely for device data
};

// Call foo with data of any memory type:
template <typename mdspan_type>
void foo(raft::resources const& res, mdspan_type data) {
  raft::memory_type_dispatcher(res,
    [&res](raft::device_matrix_view<float> dev_data) {foo_device(res, dev_data);},
    data
  );
}

Here, the memory_type_dispatcher implicitly constructs an mdbuffer from the input and performs any necessary conversions before passing the input to foo_device. While mdbuffer does not require the use of memory_type_dispatcher, there are many common use cases in which explicit invocations of mdbuffer can be elided with memory_type_dispatcher.

Finally, we should note that mdbuffer should almost never be passed as a const reference. To indicate const-ness of the underlying data, the mdbuffer should be constructed with a const memory type, but the mdbuffer itself should generally be passed as an rvalue reference in function arguments. Using an mdbuffer that is itself const is not strictly incorrect, but it indicates a likely misuse of the type.

Template Parameters:

ElementType – element type stored in the buffer
Extents – specifies the number of dimensions and their sizes
LayoutPolicy – specifies how data should be laid out in memory
ContainerPolicy – specifies how data should be allocated if necessary and how it should be accessed. This should very rarely need to be customized. For those cases where it must be customized, it is recommended to instantiate default_buffer_container_policy with a std::variant of container policies for each memory type. Note that the accessor policy of each container policy variant is used as the accessor policy for the mdspan view of the buffer for the corresponding memory type.

Public Functions

constexpr mdbuffer() = default#: Construct an empty, uninitialized buffer.

template<typename OtherAccessorPolicy, std::enable_if_t<is_type_in_variant_v<OtherAccessorPolicy, accessor_policy_variant>>* = nullptr> inline mdbuffer(mdspan<ElementType, Extents, LayoutPolicy, OtherAccessorPolicy> other)#: Construct an mdbuffer wrapping an existing mdspan. The resulting mdbuffer will be non-owning and match the memory type, layout, and element type of the mdspan.

template<typename OtherElementType, typename OtherAccessorPolicy, std::enable_if_t<!std::is_same_v<OtherElementType, ElementType> && std::is_same_v<OtherElementType const, ElementType> && is_type_in_variant_v<OtherAccessorPolicy, accessor_policy_variant>>* = nullptr> inline mdbuffer(mdspan<OtherElementType, Extents, LayoutPolicy, OtherAccessorPolicy> other)#: Construct an mdbuffer of const elements wrapping an existing mdspan with non-const elements. The resulting mdbuffer will be non-owning and match the memory type, layout, and element type of the mdspan.

template<typename OtherContainerPolicy, std::enable_if_t<is_type_in_variant_v<host_device_accessor<typename OtherContainerPolicy::accessor_type, OtherContainerPolicy::mem_type>, typename container_policy_type::container_policy_variant>>* = nullptr> inline mdbuffer(mdarray<ElementType, Extents, LayoutPolicy, OtherContainerPolicy> &&other)#: Construct an mdbuffer to hold an existing mdarray rvalue. The mdarray will be moved into the mdbuffer, and the mdbuffer will be owning.

template<typename OtherContainerPolicy, std::enable_if_t<is_type_in_variant_v<host_device_accessor<typename OtherContainerPolicy::accessor_type, OtherContainerPolicy::mem_type>, typename container_policy_type::container_policy_variant>>* = nullptr> inline mdbuffer(mdarray<ElementType, Extents, LayoutPolicy, OtherContainerPolicy> &other)#: Construct an mdbuffer from an existing mdarray lvalue. An mdspan view will be taken from the mdarray in order to construct the mdbuffer, and the mdbuffer will be non-owning.

inline mdbuffer(raft::resources const &res, mdbuffer<ElementType, Extents, LayoutPolicy, ContainerPolicy> &&other, std::optional<memory_type> specified_mem_type = std::nullopt)#

Construct one mdbuffer from another mdbuffer rvalue with matching element type, extents, layout, and container policy.

If the existing mdbuffer is owning and of the correct memory type, the new mdbuffer will take ownership of the underlying memory (preventing a view on memory owned by a moved-from object). The memory type of the new mdbuffer may be specified explicitly, in which case a copy will be performed if and only if it is necessary to do so.

inline mdbuffer(raft::resources const &res, mdbuffer<ElementType, Extents, LayoutPolicy, ContainerPolicy> &other, std::optional<memory_type> specified_mem_type = std::nullopt)#

Construct one mdbuffer from another mdbuffer lvalue with matching element type, extents, layout, and container policy.

Unlike when constructing from an rvalue, the new mdbuffer will take a non-owning view whenever possible, since it is assumed that the caller will manage the lifetime of the lvalue input. Note that the mdbuffer passed here must itself be non-const in order to allow this constructor to provide an equivalent view of the underlying data. To indicate const-ness of the underlying data, mdbuffers should be constructed with a const ElementType.

template<typename OtherElementType, typename OtherExtents, typename OtherLayoutPolicy, typename OtherContainerPolicy, std::enable_if_t<is_copyable_from<mdbuffer<OtherElementType, OtherExtents, OtherLayoutPolicy, OtherContainerPolicy>>()>* = nullptr> inline mdbuffer(raft::resources const &res, mdbuffer<OtherElementType, OtherExtents, OtherLayoutPolicy, OtherContainerPolicy> const &other, std::optional<memory_type> specified_mem_type = std::nullopt)#: Construct an mdbuffer from an existing mdbuffer with arbitrary but compatible element type, extents, layout, and container policy. This constructor is used to coerce data to specific element types, layouts, or extents as well as specifying a memory type.

inline auto constexpr mem_type() const#: Return the memory type of the underlying data referenced by the mdbuffer.

inline auto constexpr is_owning() const#: Return a boolean indicating whether or not the mdbuffer owns its storage.

template<memory_type mem_type> inline auto view()#: Return an mdspan of the indicated memory type representing a view on the stored data. If the mdbuffer does not contain data of the indicated memory type, a std::bad_variant_access will be thrown.

template<memory_type mem_type> inline auto view() const#: Return an mdspan containing const elements of the indicated memory type representing a view on the stored data. If the mdbuffer does not contain data of the indicated memory type, a std::bad_variant_access will be thrown.

inline auto view()#

Return a std::variant representing the possible mdspan types that could be returned as views on the mdbuffer. The variant will contain the mdspan corresponding to its current memory type.

This method is useful for writing generic code to handle any memory type that might be contained in an mdbuffer at a particular point in a workflow. By performing a std::visit on the returned value, the caller can easily dispatch to the correct code path for the memory type.

inline auto view() const#

Return a std::variant representing the possible mdspan types that could be returned as const views on the mdbuffer. The variant will contain the mdspan corresponding to its current memory type.

template<typename T, typename = void> struct is_mdbuffer : public std::false_type#

template<typename T> struct is_mdbuffer<T, std::void_t<decltype(__takes_an_mdbuffer_ptr(std::declval<T*>()))>> : public std::true_type#

template<typename T, typename = void> struct is_input_mdbuffer : public std::false_type#

template<typename T> struct is_input_mdbuffer<T, std::void_t<decltype(__takes_an_mdbuffer_ptr(std::declval<T*>()))>> : public std::bool_constant<std::is_const_v<T::element_type>>#

template<typename T, typename = void> struct is_output_mdbuffer : public std::false_type#

template<typename T> struct is_output_mdbuffer<T, std::void_t<decltype(__takes_an_mdbuffer_ptr(std::declval<T*>()))>> : public std::bool_constant<not std::is_const_v<T::element_type>>#