开发者指南

概览

请先阅读贡献指南。

性能

1. 在代码的性能关键部分，优先使用 cudaDeviceGetAttribute 而不是 cudaDeviceGetProperties。请参阅相应的 CUDA 开发者博客此处了解更多信息。

2. 如果某个算法需要您在多个 CUDA 流中启动 GPU 工作，请不要为每个此类工作流创建多个 raft::resources 对象。而是使用给定 raft::resources 实例上配置的流池 raft::resources::get_stream_from_stream_pool() 来选择正确的 CUDA 流。有关更多详细信息，请参阅关于CUDA 资源的章节和关于线程的章节。提示：使用 raft::resources::get_stream_pool_size() 可以了解有多少此类流可供您使用。

本地开发

为 RAFT 库本身开发功能和修复 bug 是很直接的，只需要构建和安装相关的 RAFT artifact 即可。

处理可能跨越 RAFT 和一个或多个消费库的 CUDA/C++ 功能的过程可能略有不同，具体取决于消费项目是否依赖于源代码构建（如BUILD文档中所述）。在这种情况下，可以将选项 CPM_raft_SOURCE=/path/to/raft/source 传递给消费项目的 cmake，以便从源代码构建本地 RAFT。包含对消费项目相关更改的 PR 也可以通过在调用 find_and_configure_raft 时显式更改 FORK 和 PINNED_TAG 参数为包含其更改的 RAFT 分支来临时固定 RAFT 版本。在更改合并到 RAFT 项目后、合并到下游依赖项目之前，应撤销该固定。

如果构建跨项目的功能且未在 cmake 中使用源代码构建，则需要先将 RAFT 更改（包括 C++ 和 Python）安装到消费项目的环境中，然后才能使用。将 RAFT 集成到消费项目中的理想方式是仅在这种情况下启用消费项目中的源代码构建，否则依赖更稳定的打包方式（例如 conda 打包）。

线程模型

除了 raft::resources 之外，RAFT 算法应保持线程安全，并且通常假定为单线程。这意味着只要使用不同的 raft::resources 实例，它们就可以从多个主机线程调用。

可以利用多个主机线程内的多个 CUDA 流以超额订阅或增加单个 GPU 占用率的算法例外。在这些情况下，RAFT 算法中多个主机线程的使用仅应用于维护底层 CUDA 流的并发性。应谨慎使用多个主机线程，限制其数量，并应避免执行 CPU 密集型计算。

一个可接受的在 RAFT 算法中使用主机线程的良好示例可能如下所示

#include <raft/core/resources.hpp>
#include <raft/core/resource/cuda_stream.hpp>
#include <raft/core/resource/cuda_stream_pool.hpp>
raft::resources res;

...

sync_stream(res);

...

int n_streams = get_stream_pool_size(res);

#pragma omp parallel for num_threads(n_threads)
for(int i = 0; i < n; i++) {
    int thread_num = omp_get_thread_num() % n_threads;
    cudaStream_t s = get_stream_from_stream_pool(res, thread_num);
    ... possible light cpu pre-processing ...
    my_kernel1<<<b, tpb, 0, s>>>(...);
    ...
    ... some possible async d2h / h2d copies ...
    my_kernel2<<<b, tpb, 0, s>>>(...);
    ...
    sync_stream(res, s);
    ... possible light cpu post-processing ...
}

在上面的示例中，如果在 for 循环开始时没有 CPU 预处理，可以在 for 循环内的每个流中注册一个事件，使它们等待来自句柄的流。如果在每个 for 循环迭代结束时没有 CPU 后处理，则可以将 sync_stream(res, s) 替换为 for 循环后的单个 sync_stream_pool(res)。

为了避免不同线程模型之间的兼容性问题，RAFT 中唯一允许的线程编程是 OpenMP。尽管 RAFT 默认启用 OpenMP 构建，但即使禁用 OpenMP，RAFT 算法也应能正常运行。如果上面的示例中不需要 CPU 预处理和后处理，则不需要 OpenMP。

允许在第三方库中使用线程，但它们仍应避免依赖特定的 OpenMP 运行时。

公共接口

通用准则

通过 C++ API 公开的函数必须是无状态的。可以在接口上公开的事项

1. 任何POD（Plain Old Data）- 请参阅std::is_pod作为 C++11 POD 类型的参考。

2. raft::resources - 因为它存储与资源相关的状态，与模型/算法状态无关。

3. 避免使用指向 POD 类型的指针（尽管可以将其视为 POD，但在此明确说明），而是通过引用传递结构体。在 C++ API 内部，这些无状态函数可以自由使用自己的临时类，只要它们不在接口上公开即可。

4. 接受单维 (raft::span) 和多维视图 (raft::mdspan)，并在可能的情况下验证其元数据。

5. 对于任何可选参数，优先使用 std::optional（例如，不接受 nullptr）

6. 所有公共 API 都应该是围绕 detail 命名空间内部私有 API 调用的轻量级封装。

API 稳定性

由于 RAFT 是一个拥有多个消费者的核心库，因此公共 API 保持跨版本的稳定性非常重要，对其进行的任何更改都应谨慎进行，必要时在几个版本中添加新函数并弃用旧函数。

无状态 C++ API

以 IVF-PQ 算法为例，根据本节的准则，以下公开其 API 的方式是错误的，因为它在 C++ API 中公开了非 POD 的 C++ 类对象

template <typename value_t, typename idx_t>
class ivf_pq {
  ivf_pq_params params_;
  raft::resources const& res_;

public:
  ivf_pq(raft::resources const& res);
  void train(raft::device_matrix<value_t, idx_t, raft::row_major> dataset);
  void search(raft::device_matrix<value_t, idx_t, raft::row_major> queries,
              raft::device_matrix<value_t, idx_t, raft::row_major> out_inds,
              raft::device_matrix<value_t, idx_t, raft::row_major> out_dists);
};

另一种正确的公开方式可以是

namespace raft::ivf_pq {

template<typename value_t, typename value_idx>
void ivf_pq_train(raft::resources const& res, const raft::ivf_pq_params &params, raft::ivf_pq_index &index,
raft::device_matrix<value_t, idx_t, raft::row_major> dataset);

template<typename value_t, typename value_idx>
void ivf_pq_search(raft::resources const& res, raft::ivf_pq_params const&params, raft::ivf_pq_index const & index,
raft::device_matrix<value_t, idx_t, raft::row_major> queries,
raft::device_matrix<value_t, idx_t, raft::row_major> out_inds,
raft::device_matrix<value_t, idx_t, raft::row_major> out_dists);
}

其他状态相关函数

这些准则还意味着 C++ API 的责任是公开加载和存储（又称封送）此类数据结构的方法。进一步延续 IVF-PQ 示例，以下方法可以实现这一点

namespace raft::ivf_pq {
   void save(raft::ivf_pq_index const& model, std::ostream &os);
   void load(raft::ivf_pq_index& model, std::istream &is);
}

编码风格

代码格式化

使用 pre-commit 钩子

RAFT 使用pre-commit来执行所有代码 linter 和格式化工具。这些工具确保整个项目的代码格式一致。使用 pre-commit 可确保所有开发人员的 linter 版本和选项保持一致。此外，还有一个 CI 检查来强制执行已提交的代码符合我们的标准。

要使用 pre-commit，请通过 conda 或 pip 安装

conda install -c conda-forge pre-commit

pip install pre-commit

然后在提交代码之前运行 pre-commit 钩子

pre-commit run

默认情况下，pre-commit 在暂存文件上运行（仅包括将要提交的更改和新增内容）。要在所有文件上运行 pre-commit 检查，请执行

pre-commit run --all-files

或者，您可以设置 pre-commit 钩子以便在您进行 git commit 时自动运行。这可以通过运行以下命令实现

pre-commit install

现在，每次提交更改时都会运行代码 linter 和格式化工具。

您可以使用 git commit --no-verify 或简写版本 git commit -n 跳过这些检查。

pre-commit 钩子总结

以下部分描述了仓库使用的一些核心 pre-commit 钩子。请参阅 .pre-commit-config.yaml 获取完整列表。

C++/CUDA 使用clang-format进行格式化。

RAFT 依赖 clang-format 来强制执行所有 C++ 和 CUDA 源代码的代码风格。编码风格基于Google 风格指南。与此风格唯一的偏差如下。

1. 不要拆分空函数/记录/命名空间。

2. 所有地方都使用两个空格缩进，包括行续行。

3. 禁用注释的重排。这些偏离 Google 风格指南的原因在此处的注释中给出。

使用doxygen作为文档生成器，也作为文档 linter。为了将 doxygen 作为 linter 运行在 C++/CUDA 代码上，请运行

./ci/checks/doxygen.sh

Python 代码运行多个 linter，包括Black、isort和flake8。

RAFT 还使用codespell查找拼写错误，并且此检查作为 pre-commit 钩子运行。要应用建议的拼写更正，您可以在仓库根目录运行 codespell -i 3 -w .。这将弹出一个交互式提示，让您选择应用哪些拼写更正。

#include 风格

使用include_checker.py来强制执行 #include 风格，如下所示

1. #include "..." 仅应用于引用本地文件。可以用于引用同一算法的子文件夹/父文件夹中的文件，但绝不应包含其他算法中的文件，或算法与原语或其他依赖项之间的文件。

2. #include <...> 应用于引用其他所有内容

手动运行以下命令批量修复 #include 风格问题

python ./cpp/scripts/include_checker.py --inplace [cpp/include cpp/tests ... list of folders which you want to fix]

版权头

RAPIDSpre-commit-hooks检查所有 git 修改文件的版权头。

手动地，您可以运行以下命令批量修复仓库中所有文件的版权头

pre-commit run -a verify-copyright

请记住，这仅适用于 git 跟踪且已修改的文件。

错误处理

通过提供的辅助宏 RAFT_CUDA_TRY、RAFT_CUBLAS_TRY 和 RAFT_CUSOLVER_TRY 调用 CUDA API。这些宏负责检查所用 API 调用的返回值，并在命令不成功时生成异常。如果您需要避免异常，例如在析构函数内部，请使用 RAFT_CUDA_TRY_NO_THROW、RAFT_CUBLAS_TRY_NO_THROW 和 RAFT_CUSOLVER_TRY_NO_THROW。这些宏会记录错误但不会抛出异常。

日志记录

介绍

关于日志记录的一切都定义在logger.hpp中。它底层使用了spdlog，但此信息对所有人都是透明的。

用法

#include <raft/core/logger.hpp>

// Inside your method or function, use any of these macros
RAFT_LOG_TRACE("Hello %s!", "world");
RAFT_LOG_DEBUG("Hello %s!", "world");
RAFT_LOG_INFO("Hello %s!", "world");
RAFT_LOG_WARN("Hello %s!", "world");
RAFT_LOG_ERROR("Hello %s!", "world");
RAFT_LOG_CRITICAL("Hello %s!", "world");

更改日志级别

共有 7 个日志级别，每个连续级别都变得更安静

1. RAFT_LEVEL_TRACE

2. RAFT_LEVEL_DEBUG

3. RAFT_LEVEL_INFO

4. RAFT_LEVEL_WARN

5. RAFT_LEVEL_ERROR

6. RAFT_LEVEL_CRITICAL

7. RAFT_LEVEL_OFF 根据您的需要，将其中一个传递给 set_level() 方法，如下所示

raft::default_logger().set_level(RAFT_LEVEL_WARN);
// From now onwards, this will print only WARN and above kind of messages

更改日志模式

传递格式字符串如下，以便使用与默认不同的日志模式。

raft::default_logger().set_pattern(YourFavoriteFormat);

也可以使用相应的 get_pattern() 方法来了解当前格式。

临时更改日志模式

有时，我们需要临时更改日志模式（例如：报告决策树结构）。这可以通过类似 RAII 的方法实现，如下所示

{
  PatternSetter _(MyNewTempFormat);
  // new log format is in effect from here onwards
  doStuff();
  // once the above temporary object goes out-of-scope, the old format will be restored
}

技巧

• 不要以换行符结束您的日志消息！它由 spdlog 自动添加。

• 出于性能原因，RAFT_LOG_TRACE() 默认情况下由于 RAFT_ACTIVE_LEVEL 宏设置而未编译。如果您需要启用它，请在编译时相应地更改此宏

常见设计考虑事项

1. 对于可以使用 gcc 针对 CUDA 运行时编译的文件，使用 hpp 扩展名。对于需要使用 nvcc 编译的文件，使用 cuh 扩展名。仅当有适当的检查来移除未使用 nvcc 编译时的 __device__ 指定时，hpp 也可以用于标记为 __host__ __device__ 的函数。

2. 当需要在公共 API 中使用额外的类、结构体或一般的 POCO（Plain Old CLR Object）类型来表示数据时，将它们放在一个名为 <primitive_name>_types.hpp 的新文件中。这告诉用户可以安全地在他们自己的公共 API 上公开这些类型，而无需引入设备代码。至少，这些类型的定义不应需要 nvcc。一般来说，这些类应仅存储非常简单的状态，而不应执行自己的计算。相反，公共 API 上应公开接受这些对象的新函数，根据需要读取或更新其状态。

3. 公共 API 的文档应清晰，易于使用，并且强烈建议包含使用说明。

4. 在创建新的原语之前，请检查是否已存在。如果存在，但 API 不够灵活无法满足您的用例，请首先考虑重构现有原语。如果这不可能在没有大量更改的情况下实现，请考虑如何使公共 API 更灵活。如果新原语与所有现有原语足够不同，请考虑现有公共 API 是否可以作为选项或参数调用新原语。如果新原语与已有的足够不同，请将新公共 API 函数的头文件添加到相应的子目录和命名空间。

昂贵函数模板的头文件组织

RAFT 是一个 heavily templated 的库。一些核心函数编译起来很昂贵，我们希望防止重复编译此功能。为了限制构建时间，RAFT 提供了一个预编译库（libraft.so），其中为最常用的模板参数实例化了昂贵的函数模板。为了防止 (1) 这些模板的意外实例化和 (2) 对这些模板内部实现的不必要依赖，我们使用了一种拆分头文件结构并定义了宏来控制模板实例化。本节描述了这些宏和头文件结构。

宏。 我们定义了宏 RAFT_COMPILED 和 RAFT_EXPLICIT_INSTANTIATE_ONLY。RAFT_COMPILED 宏由 CMake 在编译代码时定义，该代码 (1) 是 libraft.so 的一部分或 (2) 与 libraft.so 链接。它表示在运行时存在一个预编译的 libraft.so。

RAFT_EXPLICIT_INSTANTIATE_ONLY 宏由 CMake 在编译 libraft.so 本身时定义。当定义时，它表示禁止昂贵函数模板的隐式实例化（它们会导致编译器错误）。在 RAFT 项目中，我们还在编译测试和基准测试时定义此宏。

下面，我们总结了 RAFT_COMPILED 和 RAFT_EXPLICIT_INSTANTIATE_ONLY 在实践中使用的组合以及这些组合的效果。

RAFT_COMPILED	RAFT_EXPLICIT_INSTANTIATE_ONLY	哪些目标
已定义	已定义	raft::compiled, RAFT 测试, RAFT 基准测试
已定义		依赖于 libraft 的下游库，如 cuML, cuGraph。
		依赖于 libraft-headers 的下游库，如 cugraph-ops。

RAFT_COMPILED	RAFT_EXPLICIT_INSTANTIATE_ONLY	效果
已定义	已定义	模板已预编译。意外实例化昂贵函数模板时会产生编译器错误。
已定义		模板已预编译。允许隐式实例化。
		未预编译任何内容。允许隐式实例化。
	已定义	避免此情况：未预编译任何内容。任何昂贵函数模板实例化时都会产生编译器错误。

头文件组织。 任何定义昂贵函数模板的头文件（例如 expensive.cuh）应拆分为三个部分：expensive.cuh、expensive-inl.cuh 和 expensive-ext.cuh。expensive-inl.cuh 文件（“inl”表示“inline”）包含模板定义，即实际代码。expensive.cuh 文件根据 RAFT_COMPILED 和 RAFT_EXPLICIT_INSTANTIATE_ONLY 宏的值，包含其他两个文件中的一个或两个。expensive-ext.cuh 文件包含 extern template 实例化。此外，如果设置了 RAFT_EXPLICIT_INSTANTIATE_ONLY，它还包含模板定义，以确保在意外实例化时引发编译器错误。

expensive.cuh 的分派如下进行

#ifndef RAFT_EXPLICIT_INSTANTIATE_ONLY
// If implicit instantiation is allowed, include template definitions.
#include "expensive-inl.cuh"
#endif

#ifdef RAFT_COMPILED
// Include extern template instantiations when RAFT is compiled.
#include "expensive-ext.cuh"
#endif

不变的 expensive-inl.cuh 文件

namespace raft {
template <typename T>
void expensive(T arg) {
  // .. function body
}
} // namespace raft

expensive-ext.cuh 文件包含以下内容

#include <raft/util/raft_explicit.cuh> // RAFT_EXPLICIT

#ifdef RAFT_EXPLICIT_INSTANTIATE_ONLY
namespace raft {
// (1) define templates to raise an error in case of accidental instantiation
template <typename T> void expensive(T arg) RAFT_EXPLICIT;
} // namespace raft
#endif //RAFT_EXPLICIT_INSTANTIATE_ONLY

// (2) Provide extern template instantiations.
extern template void raft::expensive<int>(int);
extern template void raft::expensive<float>(float);

此头文件有两个职责：(1) 定义模板以在意外实例化时引发错误，以及 (2) 提供 extern template 实例化。首先，如果设置了 RAFT_EXPLICIT_INSTANTIATE_ONLY，则定义 expensive。这样做有两个原因：(1) 提供定义，因为 expensive-inl.cuh 中的定义被跳过，以及 (2) 通过使用 RAFT_EXPLICIT 宏标记来表明模板应显式实例化。此宏定义了函数体，并确保在隐式实例化错误发生时生成信息丰富的错误消息。最后，列出 extern template 实例化。

为了实际生成模板实例的代码，文件 src/expensive.cu 包含以下内容。请注意，expensive-ext.cuh 中的 extern 模板实例化与这些行之间的唯一区别是移除了单词 extern

#include <raft/expensive-inl.cuh>

template void raft::expensive<int>(int);
template void raft::expensive<float>(float);

设计考虑事项:

1. 在 -ext.cuh 头文件中，不要包含实现头文件。只包含函数参数类型和用于实例化模板的类型。如果原语接受自定义参数类型，请在名为 <primitive_name>_types.hpp 的单独头文件中定义它们。（请参阅常见设计考虑事项）。

2. 将文档字符串保留在 -inl.cuh 头文件中，因为它更接近代码。从 -ext.cuh 头文件中的模板定义中移除文档字符串。确保在 RAFT API 文档中显式包含公共 API。也就是说，将 #include <raft/expensive.cuh> 添加到 docs/source/cpp_api/expensive.rst 中的文档（而不是 #include <raft/expensive-inl.cuh>）。

3. 在 expensive.cuh 中包含的顺序至关重要。如果未定义 RAFT_EXPLICIT_INSTANTIATE_ONLY，但定义了 RAFT_COMPILED，则我们必须在 extern template 实例化之前包含模板定义。

4. 如果一个头文件定义了多个昂贵的模板，可能其中一个没有被实例化。在这种情况下，请务必在 -ext 头文件中使用 RAFT_EXPLICIT 定义该模板。这样，当模板被实例化时，开发者会收到有用的错误消息，而不是令人困惑的“函数未找到”。

此头文件结构是在issue #1416中提出的，其中包含有关此结构的动机以及 C++ 模板实例化机制的更多背景信息。

测试

RAFT 保持公共 API 的高测试覆盖率非常重要，以便最大程度地减少下游项目因更改而遇到意外构建或运行时行为的可能性。

良好定义的公共 API 有助于保持编译时稳定性，但这意味着应更多地关注测试功能需求以及验证 RAFT 本身各种边缘情况下的执行。理想情况下，应能够在独立于消费项目的情况下对 RAFT 进行 bug 修复和新功能开发。

文档

公共 API 总是需要文档，因为它们将直接暴露给用户。对于 C++，我们使用doxygen；对于 Python/cython，我们使用pydoc。除了总结公共 API 中每个类/函数的作用之外，还应记录参数（和相关模板）以及简要的使用示例。

异步操作和流排序

所有 RAFT 算法都应尽可能异步，避免使用默认流（即 NULL 或 0 流）。只需要一个 CUDA Stream 的实现应使用来自 raft::resources 的流

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

void foo(const raft::resources& res, ...)
{
    cudaStream_t stream = get_cuda_stream(res);
}

当需要多个流时，例如管理一个流水线，使用 raft::resources 中可用的内部流（参见CUDA 资源）。如果使用了多个流，所有操作仍然必须按照 raft::resource::get_cuda_stream()（来自 raft/core/resource/cuda_stream.hpp）进行排序。在任何内部 CUDA 流中的任何操作开始之前，raft::resource::get_cuda_stream() 中的所有先前工作必须已完成。RAFT 函数返回后在 raft::resource::get_cuda_stream() 中排队的任何工作都不应在内部流中排队的所有工作完成之前开始。例如，如果 RAFT 算法像这样调用

#include <raft/core/resources.hpp>
#include <raft/core/resource/cuda_stream.hpp>
void foo(const double* srcdata, double* result)
{
    cudaStream_t stream;
    CUDA_RT_CALL( cudaStreamCreate( &stream ) );
    raft::resources res;
    set_cuda_stream(res, stream);

    ...

    RAFT_CUDA_TRY( cudaMemcpyAsync( srcdata, h_srcdata.data(), n*sizeof(double), cudaMemcpyHostToDevice, stream ) );

    raft::algo(raft::resources, dopredict, srcdata, result, ... );

    RAFT_CUDA_TRY( cudaMemcpyAsync( h_result.data(), result, m*sizeof(int), cudaMemcpyDeviceToHost, stream ) );

    ...
}

在调用 raft::algo 之前启动的 stream 中的 cudaMemcpyAsync 完成之前，raft::algo 中任何流中的工作都不应开始。并且 raft::algo 中使用的所有流中的所有工作都应在调用 raft::algo 后启动的 stream 中的 cudaMemcpyAsync 开始之前完成。

这可以通过使用 CUDA 事件和 cudaStreamWaitEvent 引入流间依赖关系来确保。为了方便起见，头文件 raft/core/device_resources.hpp 提供了 raft::stream_syncer 类，该类在其构造函数和析构函数中让所有 raft::resources 内部 CUDA 流等待 raft::resource::get_cuda_stream()，并让 raft::resource::get_cuda_stream() 等待在 raft::resources 内部 CUDA 流中排队的所有工作。预期的用法是在公共 RAFT API 的入口函数中首先创建一个 raft::stream_syncer 对象

namespace raft {
   void algo(const raft::resources& res, ...)
   {
       raft::streamSyncer _(res);
   }
}

这确保了上述的流排序行为。

使用 Thrust

为了确保 thrust 算法在预期的流中执行，应使用 thrust::cuda::par 执行策略。为了确保 thrust 算法通过提供的设备内存分配器分配临时内存，请使用 raft/core/resource/thrust_policy.hpp 中可用的 rmm::exec_policy，该策略可以通过 raft::resources 使用

#include <raft/core/resources.hpp>
#include <raft/core/resource/thrust_policy.hpp>
void foo(const raft::resources& res, ...)
{
    auto execution_policy = get_thrust_policy(res);
    thrust::for_each(execution_policy, ... );
}

资源管理

不要在 RAFT 算法的实现中直接创建可重用的 CUDA 资源。而是使用 raft::resources 中的现有资源，以避免不断创建和删除可重用资源，如 CUDA 流、CUDA 事件或库句柄。如果在 raft::resources 中缺少资源句柄，请提交功能请求。可以像这样获取资源

#include <raft/core/resources.hpp>
#include <raft/core/resource/cublas_handle.hpp>
#include <raft/core/resource/cuda_stream_pool.hpp>
void foo(const raft::resources& h, ...)
{
    cublasHandle_t cublasHandle = get_cublas_handle(h);
    const int num_streams       = get_stream_pool_size(h);
    const int stream_idx        = ...
    cudaStream_t stream         = get_stream_from_stream_pool(stream_idx);
    ...
}

下面的示例展示了如何使用 rmm::stream_pool 创建 n_stream 个内部 cuda 流，这些流稍后可供 RAFT 内部的算法使用。

#include <raft/core/resources.hpp>
#include <raft/core/resource/cuda_stream_pool.hpp>
#include <rmm/cuda_stream_pool.hpp>
int main(int argc, char** argv)
{
    int n_streams = argc > 1 ? atoi(argv[1]) : 0;
    raft::resources res;
    set_cuda_stream_pool(res, std::make_shared<rmm::cuda_stream_pool>(n_streams));

    foo(res, ...);
}

多 GPU

RAFT 的多 GPU 范式是 One Process per GPU (OPG)，即每个 GPU 一个进程。每个算法的实现方式应使其可以在单个 GPU 上运行，而不依赖于特定的通信库。多 GPU 实现应使用 [raft/core/comms.hpp] 中 raft::comms::comms_t 类提供的方法进行跨 rank/GPU 通信。cuML 用户有责任创建一个初始化的 raft::comms::comms_t 实例。

例如，对于支持 CUDA 的 MPI，RAFT 用户可以使用类似这样的代码将初始化的 raft::comms::mpi_comms 实例注入到 raft::resources 中

#include <mpi.h>
#include <raft/core/resources.hpp>
#include <raft/comms/mpi_comms.hpp>
#include <raft/algo.hpp>
...
int main(int argc, char * argv[])
{
    MPI_Init(&argc, &argv);
    int rank = -1;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    int local_rank = -1;
    {
        MPI_Comm local_comm;
        MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, rank, MPI_INFO_NULL, &local_comm);

        MPI_Comm_rank(local_comm, &local_rank);

        MPI_Comm_free(&local_comm);
    }

    cudaSetDevice(local_rank);

    mpi_comms raft_mpi_comms;
    MPI_Comm_dup(MPI_COMM_WORLD, &raft_mpi_comms);

    {
        raft::resources res;
        initialize_mpi_comms(res, raft_mpi_comms);

        ...

        raft::algo(res, ... );
    }

    MPI_Comm_free(&raft_mpi_comms);

    MPI_Finalize();
    return 0;
}

RAFT 开发者可以假定以下几点

• 一个 raft::comms::comms_t 实例已正确初始化。

• 作为 raft::comms::comms_t 一部分的所有进程都协作调用 RAFT 算法。

可以从 raft::resources 实例访问已初始化的 raft::comms::comms_t 实例

#include <raft/core/resources.hpp>
#include <raft/core/resource/comms.hpp>
void foo(const raft::resources& res, ...)
{
    const raft::comms_t& communicator = get_comms(res);
    const int rank = communicator.get_rank();
    const int size = communicator.get_size();
    ...
}