注意

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

资源#

所有特定于计算环境（如主机或设备）的资源都包含在 raft::resources 中并由其管理。这种设计通过默认使 API 不透明，但允许基于用户偏好进行自定义，从而简化了 API 并减轻了用户负担。

术语#

#include <raft/core/resource/resource_types.hpp>

namespace raft::resource

enum resource_type#

资源类型可以应用于任何资源，不必是主机或设备特有的。

取值

enumerator CUBLAS_HANDLE#
enumerator CUSOLVER_DN_HANDLE#
enumerator CUSOLVER_SP_HANDLE#
enumerator CUSPARSE_HANDLE#
enumerator CUDA_STREAM_VIEW#
enumerator CUDA_STREAM_POOL#
enumerator CUDA_STREAM_SYNC_EVENT#
enumerator COMMUNICATOR#
enumerator SUB_COMMUNICATOR#
enumerator DEVICE_PROPERTIES#
enumerator DEVICE_ID#
enumerator STREAM_VIEW#
enumerator THRUST_POLICY#
enumerator WORKSPACE_RESOURCE#
enumerator CUBLASLT_HANDLE#
enumerator CUSTOM#
enumerator LARGE_WORKSPACE_RESOURCE#
enumerator NCCL_CLIQUE#
enumerator LAST_KEY#
class resource#
#include <resource_types.hpp>

资源构建并包含某种预定对象类型的实例，并通过通用 API 为该对象提供外观。

派生类 raft::resource::comms_resource, raft::resource::cublas_resource, raft::resource::cublaslt_resource, raft::resource::cuda_event_resource, raft::resource::cuda_stream_pool_resource, raft::resource::cuda_stream_resource, raft::resource::cusolver_dn_resource, raft::resource::cusolver_sp_resource, raft::resource::cusparse_resource, raft::resource::custom_resource, raft::resource::device_id_resource, raft::resource::device_memory_resource, raft::resource::device_properties_resource, raft::resource::empty_resource, raft::resource::limiting_memory_resource, raft::resource::nccl_clique_resource, raft::resource::stream_view_resource, raft::resource::sub_comms_resource, raft::resource::thrust_policy_resource

class empty_resource : public raft::resource::resource#
#include <resource_types.hpp>
class resource_factory#
#include <resource_types.hpp>

资源工厂知道如何构建特定 raft::resource::resource 的实例。

派生类 raft::resource::comms_resource_factory, raft::resource::cublas_resource_factory, raft::resource::cublaslt_resource_factory, raft::resource::cuda_stream_pool_resource_factory, raft::resource::cuda_stream_resource_factory, raft::resource::cusolver_dn_resource_factory, raft::resource::cusolver_sp_resource_factory, raft::resource::cusparse_resource_factory, raft::resource::custom_resource_factory, raft::resource::device_id_resource_factory, raft::resource::device_properties_resource_factory, raft::resource::empty_resource_factory, raft::resource::large_workspace_resource_factory, raft::resource::nccl_clique_resource_factory, raft::resource::stream_view_resource_factory, raft::resource::sub_comms_resource_factory, raft::resource::thrust_policy_resource_factory, raft::resource::workspace_resource_factory

公有函数

virtual resource_type get_resource_type() = 0#

返回与当前工厂关联的 resource_type。

返回值:

resource_type 对应于当前工厂

virtual resource *make_resource() = 0#

构建工厂底层资源的实例。

返回值:

资源实例

class empty_resource_factory : public raft::resource::resource_factory#
#include <resource_types.hpp>

资源工厂知道如何构建特定 raft::resource::resource 的实例。

公有函数

inline virtual resource_type get_resource_type() override#

返回与当前工厂关联的 resource_type。

返回值:

resource_type 对应于当前工厂

inline virtual resource *make_resource() override#

构建工厂底层资源的实例。

返回值:

资源实例

设备资源#

raft::device_resources 是对直接使用 raft::resources 的一种便利封装。它提供了访问器方法来检索资源，例如 CUDA 流、流池以及各种 CUDA 数学库（如 cuBLAS 和 cuSOLVER）的句柄。

#include <raft/core/device_resources.hpp>

namespace raft::core

class device_resources : public raft::resources#

主要的资源容器对象，存储用于调用必要的设备函数、CUDA 内核和/或库所需的所有资源。

公有函数

inline device_resources(rmm::cuda_stream_view stream_view = rmm::cuda_stream_per_thread, std::shared_ptr<rmm::cuda_stream_pool> stream_pool = {nullptr}, std::shared_ptr<rmm::mr::device_memory_resource> workspace_resource = {nullptr}, std::optional<std::size_t> allocation_limit = std::nullopt)#

构造一个带有流视图和流池的资源实例。

参数:

  • stream_view – 默认流（如果未指定，则具有默认的每线程流）

  • stream_pool – [in] 使用的流池（如果未指定，默认为 nullptr）

  • workspace_resource – [in] 一些函数用于分配临时工作空间的可选资源。

  • allocation_limit – [in] 可用于临时工作空间资源的内存总字节数。

inline virtual ~device_resources()#

销毁所有持有的资源。

inline void sync_stream(rmm::cuda_stream_view stream) const#

同步当前容器上的流

inline void sync_stream() const#

同步当前容器上的主流程

inline rmm::cuda_stream_view get_stream() const#

返回当前容器上的主流

inline bool is_stream_pool_initialized() const#

返回流池是否在当前容器上已初始化

inline const rmm::cuda_stream_pool &get_stream_pool() const#

返回当前容器上的流池

inline rmm::cuda_stream_view get_stream_from_stream_pool() const#

从流池返回流

inline rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const#

从流池中返回指定索引处的流

inline rmm::cuda_stream_view get_next_usable_stream() const#

如果流池大小 > 0，则从流池返回流；否则返回当前容器上的主流

inline rmm::cuda_stream_view get_next_usable_stream(std::size_t stream_idx) const#

如果流池大小 > 0，则从流池中返回指定索引处的流；否则返回当前容器上的主流

参数:

stream_idx – [in] 如果可用，流在流池中的所需索引

inline void sync_stream_pool() const#

同步当前容器上的流池

inline void sync_stream_pool(const std::vector<std::size_t> stream_indices) const#

同步流池的子集

参数:

stream_indices – [in] 要同步的流在流池中的索引

inline void wait_stream_pool_on_stream() const#

要求流池等待主流中的最后一个事件

设备资源管理器#

虽然 raft::device_resources 为一系列 RAFT 调用提供了访问设备相关资源的便利方式，但在整个应用程序中限制这些资源有时也很有用。例如，在高度多线程的应用程序中，限制流的总数而不是依赖于默认的每线程流会很有帮助。raft::device_resources_manager 提供了一种方式来访问从有限的底层设备资源池中获取的 raft::device_resources 实例。

#include <raft/core/device_resources_manager.hpp>

namespace raft::core

警告

doxygenclass: Cannot find class “raft::device_resources_manager” in doxygen xml output for project “RAFT” from directory: ../../cpp/doxygen/_xml/

资源函数#

Comms#

#include <raft/core/resource/comms.hpp>

namespace raft::resource

inline bool comms_initialized(resources const &res)#
inline comms::comms_t const &get_comms(resources const &res)#
inline void set_comms(resources const &res, std::shared_ptr<comms::comms_t> communicator)#

cuBLAS 句柄#

#include <raft/core/resource/cublas_handle.hpp>

namespace raft::resource

inline cublasHandle_t get_cublas_handle(resources const &res)#

从 raft res 加载一个 cublasHandle_t，如果存在则添加并返回。

参数:

res – [in] raft 资源对象

返回值:

cublas 句柄

cuBLASLt 句柄#

#include <raft/core/resource/cublaslt_handle.hpp>

namespace raft::resource

inline cublasLtHandle_t get_cublaslt_handle(resources const &res)#

从 raft res 加载一个 cublasLtHandle_t，如果存在则添加并返回。

参数:

res – [in] raft 资源对象

返回值:

cublasLt 句柄

CUDA 流#

#include <raft/core/resource/cuda_stream.hpp>

namespace raft::resource

inline rmm::cuda_stream_view get_cuda_stream(resources const &res)#

从资源实例加载一个 rmm::cuda_stream_view（如果需要，将其填充到 res 上）。

参数:

res – 用于管理资源的 raft res 对象

返回值:

inline void set_cuda_stream(resources const &res, rmm::cuda_stream_view stream_view)#

从资源实例加载一个 rmm::cuda_stream_view（如果需要，将其填充到 res 上）。

参数:

  • res – [in] 用于管理资源的 raft resources 对象

  • stream_view – cuda 流视图

inline void sync_stream(const resources &res, rmm::cuda_stream_view stream)#

同步特定流

参数:

  • res – [in] raft 资源对象

  • stream – [in] 要同步的流

inline void sync_stream(const resources &res)#

同步资源实例上的主流程

CUDA 流池#

#include <raft/core/resource/cuda_stream_pool.hpp>

namespace raft::resource

inline const rmm::cuda_stream_pool &get_cuda_stream_pool(const resources &res)#

加载一个 cuda_stream_pool，如果不存在则创建一个新的

参数:

res – 用于管理资源的 raft res 对象

返回值:

inline void set_cuda_stream_pool(const resources &res, std::shared_ptr<rmm::cuda_stream_pool> stream_pool)#

在当前 res 上显式设置一个流池。注意，这将覆盖 res 上现有的流池。

参数:

  • res –

  • stream_pool –

inline std::size_t get_stream_pool_size(const resources &res)#
inline rmm::cuda_stream_view get_stream_from_stream_pool(const resources &res)#

从流池返回流

inline rmm::cuda_stream_view get_stream_from_stream_pool(const resources &res, std::size_t stream_idx)#

从流池中返回指定索引处的流

inline rmm::cuda_stream_view get_next_usable_stream(const resources &res)#

如果流池大小 > 0，则从流池返回流；否则返回 res 上的主流

inline rmm::cuda_stream_view get_next_usable_stream(const resources &res, std::size_t stream_idx)#

如果流池大小 > 0，则从流池中返回指定索引处的流；否则返回 res 上的主流

参数:

  • res – [in] raft 资源对象

  • stream_idx – [in] 如果可用，流在流池中的所需索引

inline void sync_stream_pool(const resources &res)#

同步 res 上的流池

参数:

res – [in] raft 资源对象

inline void sync_stream_pool(const resources &res, const std::vector<std::size_t> stream_indices)#

同步流池的子集

参数:

  • res – [in] raft 资源对象

  • stream_indices – [in] 要同步的流在流池中的索引

inline void wait_stream_pool_on_stream(const resources &res)#

要求流池等待主流中的最后一个事件

参数:

res – [in] raft 资源对象

cuSolverDn 句柄#

#include <raft/core/resource/cusolver_dn_handle.hpp> namespace raft::resource

inline cusolverDnHandle_t get_cusolver_dn_handle(resources const &res)#

从 raft res 加载一个 cusolverSpres_t，如果存在则添加并返回。

参数:

res – [in] raft 资源对象

返回值:

cusolver dn 句柄

class cusolver_dn_resource_factory : public raft::resource::resource_factory#
#include <cusolver_dn_handle.hpp>

知道如何构建特定 raft::resource 以填充 res_t 的工厂。

公有函数

inline virtual resource_type get_resource_type() override#

返回与当前工厂关联的 resource_type。

返回值:

resource_type 对应于当前工厂

inline virtual resource *make_resource() override#

构建工厂底层资源的实例。

返回值:

资源实例

cuSolverSp 句柄#

#include <raft/core/resource/cusolver_sp_handle.hpp>

namespace raft::resource

inline cusolverSpHandle_t get_cusolver_sp_handle(resources const &res)#

从 raft res 加载一个 cusolverSpres_t，如果存在则添加并返回。

参数:

res – [in] raft 资源对象

返回值:

cusolver sp 句柄

cuSparse 句柄#

#include <raft/core/resource/cusparse_handle.hpp>

namespace raft::resource

inline cusparseHandle_t get_cusparse_handle(resources const &res)#

从 raft res 加载一个 cusparseres_t，如果存在则添加并返回。

参数:

res – [in] raft 资源对象

返回值:

cusparse 句柄

设备 ID#

#include <raft/core/resource/device_id.hpp>

namespace raft::resource

inline int get_device_id(resources const &res)#

从 res 加载设备 ID（如果需要，将其填充到 res 上）。

参数:

res – 用于管理资源的 raft res 对象

返回值:

设备 ID

设备内存资源#

#include <raft/core/resource/device_memory_resource.hpp>

namespace raft::resource

inline rmm::mr::limiting_resource_adaptor<rmm::mr::device_memory_resource> *get_workspace_resource(resources const &res)#

从资源实例加载临时工作空间资源（如果需要，将其填充到 res 上）。

参数:

res – 用于管理资源的 raft resources 对象

返回值:

设备内存资源对象

inline size_t get_workspace_total_bytes(resources const &res)#

获取工作区资源的总大小。

inline size_t get_workspace_used_bytes(resources const &res)#

获取工作区资源已分配的大小。

inline size_t get_workspace_free_bytes(resources const &res)#

获取工作区资源的可用大小。

inline void set_workspace_resource(resources const &res, std::shared_ptr<rmm::mr::device_memory_resource> mr = {nullptr}, std::optional<std::size_t> allocation_limit = std::nullopt, std::optional<std::size_t> alignment = std::nullopt)#

在 resources 实例上设置一个临时工作区资源。

参数:

  • res – 用于管理资源的 raft resources 对象

  • mr – 可选的 RMM device_memory_resource

  • allocation_limit – 可用于临时工作区资源的总内存量（字节）。

  • alignment – 传递给 RMM 分配的可选对齐要求

inline void set_workspace_to_pool_resource(resources const &res, std::optional<std::size_t> allocation_limit = std::nullopt)#

将临时工作区资源设置为全局内存资源 (rmm::mr::get_current_device_resource()) 之上的一个池。

参数:

  • res – 用于管理资源的 raft resources 对象

  • allocation_limit – 可用于临时工作区资源的总内存量（字节）；如果未提供，将使用上次使用或默认的限制。

inline void set_workspace_to_global_resource(resources const &res, std::optional<std::size_t> allocation_limit = std::nullopt)#

将临时工作区资源设置为与全局内存资源 (rmm::mr::get_current_device_resource()) 相同。

注意，工作区资源始终是受限的；此处的限制定义了工作区分配可以消耗多少全局内存资源。

参数:

  • res – 用于管理资源的 raft resources 对象

  • allocation_limit – 可用于临时工作区资源的总内存量（字节）。

inline rmm::mr::device_memory_resource *get_large_workspace_resource(resources const &res)#
inline void set_large_workspace_resource(resources const &res, std::shared_ptr<rmm::mr::device_memory_resource> mr = {nullptr})#
class device_memory_resource : public raft::resource::resource#
#include <device_memory_resource.hpp>
class limiting_memory_resource : public raft::resource::resource#
#include <device_memory_resource.hpp>
class large_workspace_resource_factory : public raft::resource::resource_factory#
#include <device_memory_resource.hpp>

知道如何构造特定 raft::resource 来填充 resources 实例的工厂。

公有函数

inline virtual resource_type override get_resource_type ()

返回与当前工厂关联的 resource_type。

返回值:

resource_type 对应于当前工厂

inline virtual resource *override make_resource ()

构建工厂底层资源的实例。

返回值:

资源实例

class workspace_resource_factory : public raft::resource::resource_factory#
#include <device_memory_resource.hpp>

知道如何构造特定 raft::resource 来填充 resources 实例的工厂。

公有函数

inline virtual resource_type override get_resource_type ()

返回与当前工厂关联的 resource_type。

返回值:

resource_type 对应于当前工厂

inline virtual resource *override make_resource ()

构建工厂底层资源的实例。

返回值:

资源实例

公共静态函数

static inline std::shared_ptr<rmm::mr::device_memory_resource> default_pool_resource(std::size_t limit)#

构造一个合理的默认池内存资源。

static inline std::shared_ptr<rmm::mr::device_memory_resource> default_plain_resource()#

获取包装在非管理 shared_ptr（无删除器）中的全局内存资源。

注意：底层 rmm::mr::get_current_device_resource() 的生命周期在别处管理，因为它通过原始指针传递。因此，此 shared_ptr 包装器不允许在析构时删除指针。

设备属性#

#include <raft/core/resource/device_properties.hpp>

namespace raft::resource

inline cudaDeviceProp &get_device_properties(resources const &res)#

从 res 加载 cudaDeviceProp（如果需要，并在 res 上填充）。

参数:

res – 用于管理资源的 raft res 对象

返回值:

填充的 cuda device properties 实例

class device_properties_resource_factory : public raft::resource::resource_factory#
#include <device_properties.hpp>

知道如何构建特定 raft::resource 以填充 res_t 的工厂。

公有函数

inline virtual resource_type get_resource_type() override#

返回与当前工厂关联的 resource_type。

返回值:

resource_type 对应于当前工厂

inline virtual resource *make_resource() override#

构建工厂底层资源的实例。

返回值:

资源实例

子通信器#

#include <raft/core/resource/sub_comms.hpp>

namespace raft::resource

inline const comms::comms_t &get_subcomm(const resources &res, std::string key)#
inline void set_subcomm(resources const &res, std::string key, std::shared_ptr<comms::comms_t> subcomm)#

Thrust 执行策略#

#include <raft/core/resource/thrust_policy.hpp>

namespace raft::resource

inline rmm::exec_policy_nosync &get_thrust_policy(resources const &res)#

从 res 加载 thrust 策略（如果需要，并在 res 上填充）。

参数:

res – 用于管理资源的 raft res 对象

返回值:

thrust 执行策略

自定义运行时共享资源#

自定义资源是任意的默认可构造 C++ 类。API 的消费者可以将此类资源保存在 raft::resources 句柄中。例如，考虑一个预期重复调用并涉及昂贵内核配置的函数。可以将内核配置缓存在自定义资源中。访问它的成本是一次哈希映射查找。

#include <raft/core/resource/custom_resource.hpp>

namespace raft::resource

template<typename ResourceT>
ResourceT *get_custom_resource(resources const &res)#

如果自定义的默认可构造资源存在则获取，否则创建它。

注意：与其他硬编码资源不同，编译时没有关于自定义资源的信息。因此，自定义资源存储在哈希映射中并在运行时查找。这会导致访问时间略慢。

模板参数:

ResourceT – 资源的类型；它必须是完整的且默认可构造的。

参数:

res – [in] raft 资源对象

返回值:

指向自定义资源的指针。

在本页