注意

RAFT 中的向量搜索和聚类算法正在迁移到一个新的专门用于向量搜索的库，称为 cuVS。在此迁移期间，我们将继续支持 RAFT 中的向量搜索算法，但在 RAPIDS 24.06（六月）版本之后将不再更新。我们计划在 RAPIDS 24.10（十月）版本之前完成迁移，并在 24.12（十二月）版本中将它们从 RAFT 中完全移除。

mdbuffer: 多维可选择拥有容器#

#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>#: 根据给定的内存类型从 variant 中检索类型。

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)>>#: 用于为缓冲区构建默认容器策略的、每种内存类型的容器策略 variant。

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> 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> bool is_input_mdbuffer_v = std::conjunction_v<is_input_mdbuffer_t<Tn...>>#

template<typename ...Tn> 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)#

检索与给定内存类型关联的标准索引。

对于基于内存类型的 variant，此索引可用于帮助保持 variant 中内存类型的一致顺序。

inline auto constexpr memory_type_from_variant_index(std::underlying_type_t<raft::memory_type> index)#: 检索与标准索引关联的内存类型。

template<typename ElementType, typename Extents, typename LayoutPolicy, typename ContainerPolicy> void __takes_an_mdbuffer_ptr(mdbuffer<ElementType, Extents, LayoutPolicy, ContainerPolicy>*)#: \brief 用于确定类型 T 是否为 mdbuffer 或派生类型的模板检查和助手

template<typename ElementType, typename ContainerPolicyVariant = default_container_policy_variant<std::remove_cv_t<ElementType>>> struct default_buffer_container_policy#: 一个模板，用于将底层 mdarray 容器策略的 variant 转换为 mdbuffer 可以使用的容器策略。

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

表示多维数据的一种类型，该数据可能拥有也可能不拥有其底层存储。raft::mdbuffer 用于方便地执行数据复制，仅当需要确保数据在所需的内存空间和格式中可访问时才执行复制。

在开发与 GPU 交互的函数时，通常需要确保数据位于特定的内存空间（例如，设备、主机、托管、Pinned），但这些函数可能在调用时使用的数据可能已经或尚未位于所需的内存空间中。例如，在某个工作流程中调用时，数据可能之前已被传输到设备，从而无需复制。在另一个工作流程中，该函数可能直接对主机数据进行操作。

即使只处理主机内存，也经常需要确保数据采用特定布局以实现高效访问（例如，列优先 vs 行优先），或者即使我们希望使用另一种兼容类型（例如 float）的数据调用函数，也要确保数据是特定类型（例如 double）。

mdbuffer 是一个工具，用于确保数据以恰好所需的格式和位置表示，同时灵活支持可能尚未处于该格式或位置的数据。它通过提供对已处于所需形式的数据的非拥有视图来实现这一点，但仅在必要时才会分配（拥有的）内存并执行复制。

使用示例

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

请注意，在此示例中，foo_device 模板可以正确地为主机和设备 mdspans 实例化。同样，我们可以使用 mdbuffer 将数据强制转换为特定的内存布局和数据类型，如下例所示

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

请注意，在此示例中，foo_device 模板可以接受任何布局和任何内存类型的、可转换为 float 的任何类型的数据，并将其强制转换为所需的设备可访问表示形式。

由于 mdspan 类型可以隐式转换为 mdbuffer，因此甚至可以通过直接将 mdbuffer 作为参数接受来避免多次模板实例化，如下例所示

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

在此示例中，foo_device 现在可以接受任何行优先的 float mdspan，无论其内存类型如何，而无需为每种类型进行单独的模板实例化。

虽然 view 方法采用可选的编译时内存类型参数，但省略此参数将返回 mdspan 类型的 std::variant。这允许使用 std::visit 基于内存类型进行直接的运行时分派，如下例所示

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

为方便起见，运行时内存类型分派也可以在使用 raft::memory_type_dispatcher 时不显式使用 mdbuffer，如基于内存类型的分派函数对象中所述。请参阅该函数模板的完整文档，以更详细地讨论其多种用法。然而，为了说明它与 mdbuffer 的联系，请考虑以下示例，该示例执行与上述 std::visit 调用类似的任务

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

在此处，memory_type_dispatcher 从输入隐式构造一个 mdbuffer，并在将输入传递给 foo_device 之前执行任何必要的转换。虽然 mdbuffer 不需要使用 memory_type_dispatcher，但在许多常见用例中，可以使用 memory_type_dispatcher 省略对 mdbuffer 的显式调用。

最后，我们应该注意 mdbuffer 几乎不应作为 const 引用传递。为了表示底层数据的 const-ness，应使用 const 内存类型构造 mdbuffer，但 mdbuffer 本身通常应在函数参数中作为右值引用传递。使用本身就是 const 的 mdbuffer 并非严格意义上的错误，但这表明可能误用了该类型。

模板参数:

ElementType – 存储在缓冲区中的元素类型
Extents – 指定维度数量及其大小
LayoutPolicy – 指定数据在内存中的布局方式
ContainerPolicy – 指定在必要时如何分配数据以及如何访问数据。这通常很少需要自定义。对于必须自定义的情况，建议使用 std::variant 容器策略（每种内存类型对应一个）实例化 default_buffer_container_policy。请注意，每种容器策略 variant 的 accessor policy 用作相应内存类型的缓冲区的 mdspan 视图的 accessor policy。

公共函数

constexpr mdbuffer() = default#: 构造一个空的、未初始化的缓冲区。

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)#: 构造一个包装现有 mdspan 的 mdbuffer。生成的 mdbuffer 将是非拥有的，并且匹配该 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)#: 构造一个 const 元素的 mdbuffer，包装现有具有非 const 元素的 mdspan。生成的 mdbuffer 将是非拥有的，并且匹配该 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)#: 构造一个 mdbuffer 以容纳现有的 mdarray 右值。该 mdarray 将被移动到 mdbuffer 中，并且该 mdbuffer 将是拥有的。

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)#: 从现有的 mdarray 左值构造一个 mdbuffer。将从该 mdarray 中获取一个 mdspan 视图来构造 mdbuffer，并且该 mdbuffer 将是非拥有的。

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

从另一个具有匹配元素类型、extents、布局和容器策略的 mdbuffer 右值构造一个 mdbuffer。

如果现有 mdbuffer 是拥有的且具有正确的内存类型，则新的 mdbuffer 将取得底层内存的所有权（防止查看已被移动的对象的内存）。可以显式指定新 mdbuffer 的内存类型，在这种情况下，仅在必要时才会执行复制。

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

从另一个具有匹配元素类型、extents、布局和容器策略的 mdbuffer 左值构造一个 mdbuffer。

与从右值构造不同，新的 mdbuffer 在可能的情况下将采用非拥有视图，因为假定调用者将管理左值输入的生命周期。请注意，此处传递的 mdbuffer 本身必须是非 const 的，以便允许此构造函数提供底层数据的等效视图。为了指示底层数据的 const-ness，mdbuffers 应该使用 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)#: 从具有任意但兼容的元素类型、extents、布局和容器策略的现有 mdbuffer 构造 mdbuffer。此构造函数用于将数据强制转换为特定元素类型、布局或 extents，并指定内存类型。

inline auto constexpr mem_type() const#: 返回 mdbuffer 引用的底层数据的内存类型。

inline auto constexpr is_owning() const#: 返回一个布尔值，指示 mdbuffer 是否拥有其存储。

template<memory_type mem_type> inline auto view()#: 返回一个指定内存类型的 mdspan，表示对存储数据的视图。如果 mdbuffer 不包含指定内存类型的数据，则会抛出 std::bad_variant_access。

template<memory_type mem_type> inline auto view() const#: 返回一个包含指定内存类型的 const 元素的 mdspan，表示对存储数据的视图。如果 mdbuffer 不包含指定内存类型的数据，则会抛出 std::bad_variant_access。

inline auto view()#

返回一个 std::variant，表示可能作为 mdbuffer 视图返回的 mdspan 类型。该 variant 将包含与其当前内存类型对应的 mdspan。

此方法对于编写通用代码以处理工作流程中特定点可能包含在 mdbuffer 中的任何内存类型非常有用。通过对返回值执行 std::visit，调用者可以轻松地分派到针对该内存类型的正确代码路径。

inline auto view() const#

返回一个 std::variant，表示可能作为 mdbuffer const 视图返回的 mdspan 类型。该 variant 将包含与其当前内存类型对应的 mdspan。

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::false_type, 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::false_type, 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::false_type, public std::bool_constant<not std::is_const_v<T::element_type>>#

mdbuffer: 多维可选择拥有容器#

本页