PyTorch C++ 中的 CUDA Guards 全解析：设备与流的安全 RAII 封装

CUDA guards 是 PyTorch C++ 中一组基于 RAII 的上下文管理工具，它们把“设置当前 CUDA 设备 / 当前 CUDA 流”与“离开作用域时自动恢复原上下文”两件事封装成类，避免手工调用 cudaSetDevice / cudaSetDevice + 恢复逻辑带来的异常安全问题与代码噪音。本文以 guards.md 为骨架，完整梳理 c10::cuda::CUDAGuard、CUDAStreamGuard、OptionalCUDAGuard、OptionalCUDAStreamGuard 与 at::cuda::CUDAMultiStreamGuard 五类守卫的构造方式、成员方法、适用场景，并结合 CUDAGuard.h 与 InlineDeviceGuard.h 等源码剖析其底层实现，帮助读者在多卡、多流编程中写出正确且可复用的设备/流管理代码。

什么是 CUDA Guard：核心语义与设计动机

在 CUDA 编程中，“当前设备”和“当前流”是驱动层维护的线程本地状态，任何不显式指定设备/流的操作都会落到它们上。手写切换代码通常长这样：先记录旧设备，再切到新设备，执行完再切回去——一旦中间抛出异常或提前 return，恢复逻辑就可能被跳过，后续代码会继续跑在错误的设备上。

CUDA guards 用 RAII 解决这个问题：构造时切换上下文，析构时自动恢复进入作用域前的上下文。因此原文档对它的定义可以精确概括为：

CUDA guards are RAII wrappers that set a CUDA device or stream as the current context and automatically restore the previous context when the guard goes out of scope.

全部五个类型都声明在单一头文件 CUDAGuard.h 中，使用时只需：

#include <c10/cuda/CUDAGuard.h>

（少数基于 at::cuda 的类型如示例中 at::cuda::CUDAStream 另见 CUDAStream.h。）这些类在 c10::cuda 命名空间下，它们继承自 c10 通用设备/流守卫机制，仅在"能直接链接 CUDA 运行时的代码"中可用，因为实现会退化为直截了当的 cudaSetDevice / cudaGetDevice 调用。

CUDAGuard：面向 CUDA 设备的守卫

c10::cuda::CUDAGuard 是泛型 DeviceGuard 的 CUDA 特化版本。它接收整数设备索引（解释为 CUDA 设备号），相比通用 DeviceGuard 略高效，编译为直截了当的 cudaSetDevice/cudaGetDevice 调用序列。

构造与核心方法

从 CUDAGuard.h 可见其接口：

成员	说明
CUDAGuard(DeviceIndex) / CUDAGuard(Device)	构造时把当前 CUDA 设备切到指定设备；传入非 CUDA 设备会报错
set_device(Device)	把当前设备切到给定 Device，传入非 CUDA 设备报错
reset_device(Device)	先恢复到原设备，再设置到给定设备（与泛型 DeviceGuard 保持接口一致）
set_index(DeviceIndex)	按设备索引设置当前设备
original_device()	返回构造时记录的原设备
current_device()	返回最近一次经 set_device 设置的设备；若无则返回构造时传入的设备

注意：该类型没有默认构造函数，也不允许拷贝或移动——设计上它必须携带"进入作用域前的设备"这一状态，移动语义在同时存在多个守卫的场景下无法确定谁先恢复（源码注释 Note [Move construction for RAII guards is tricky] 有详细讨论），因此被显式删除。

示例：作用域内切换到 device 1

原文档给出的最小示例可以直接编译运行：

#include <c10/cuda/CUDAGuard.h>

{
    c10::cuda::CUDAGuard guard(1);  // 切换到 device 1
    // 此作用域内所有 CUDA 操作都运行在 device 1 上
    auto tensor = torch::zeros({2, 2}, torch::device(torch::kCUDA));
}
// 离开作用域，自动恢复进入前的设备

与 Python 侧写法对应，若想用完整 Device 构造：

c10::cuda::CUDAGuard guard(torch::kCUDA);     // 默认 device（索引 -1 表示当前设备）
c10::cuda::CUDAGuard guard(torch::Device(torch::kCUDA, 1));

CUDAStreamGuard：设备 + 流的双重守卫

CUDAStreamGuard 是泛型 StreamGuard 的 CUDA 特化，它的行为比 CUDAGuard 更进一步：构造时同时把当前设备切到 stream 所属设备、并把该设备上的当前流设为传入 stream；析构时恢复原 stream 与原设备。见 CUDAGuard.h。

#include <c10/cuda/CUDAGuard.h>

auto stream = c10::cuda::getStreamFromPool();
{
    c10::cuda::CUDAStreamGuard guard(stream);
    // 此作用域内的操作使用指定的 stream
}
// 离开作用域，恢复之前的 stream

CUDAStreamGuard 还提供一个值得注意的方法 reset_stream(Stream)：它会先把当前设置的流/设备恢复为原始值，再切到新传入的 stream。源码注释特别警示：reset_stream 不会保留此前在其他设备上设置的流，若需要在多个设备上同时设置多个流，应改用 CUDAMultiStreamGuard。

查询接口返回的是 CUDAStream 而非裸 Stream：

• original_stream()：构造时刻被记录的原 stream；
• current_stream()：最近一次经守卫设置的 stream；
• current_device() / original_device()：与设备相关的查询。

getStreamFromPool() 从 c10 的全局流池中按 round-robin 方式取流，创建成本近似为零；它支持两个重载，见 CUDAStream.h：

// 从默认优先级流池取流（每个设备 32 条，轮转复用）
c10::cuda::CUDAStream getStreamFromPool(bool isHighPriority = false, DeviceIndex device = -1);
// 按显式优先级取流
c10::cuda::CUDAStream getStreamFromPool(int priority, DeviceIndex device = -1);

其中 device = -1 表示取当前设备的流池。

OptionalCUDAGuard：可延迟/可条件初始化的设备守卫

OptionalCUDAGuard 解决一个真实痛点：许多情况下"是否需要切换设备"取决于运行时条件，若用普通 CUDAGuard 就必须提前知道目标设备。Optional 变体允许先构造一个未初始化的守卫，稍后再通过 set_device/set_index 按需初始化；未初始化就析构则什么都不做。

其接口与 CUDAGuard 一一对应，但返回值改为 std::optional<Device>：

方法	说明
OptionalCUDAGuard()	默认构造，处于未初始化状态
OptionalCUDAGuard(std::optional<Device> / std::optional<DeviceIndex>)	传入非空 optional 即初始化并切换设备
set_device / reset_device / set_index	若未初始化则先初始化再切换；已初始化则直接切换
original_device() / current_device()	未初始化时返回 std::nullopt
reset()	立即恢复原设备并把守卫重置为未初始化状态

原文档的示例完整呈现了"条件守卫"的用法：

c10::cuda::OptionalCUDAGuard guard;
if (use_cuda) {
    guard.set_device(0);
}
// 只有调用过 set_device 才会切换设备

从实现上看，OptionalCUDAGuard 内部持有 std::optional<InlineDeviceGuard<CUDAGuardImpl>>（见 InlineDeviceGuard.h），未初始化即 optional 为空，自然不触发任何设备切换与恢复。

OptionalCUDAStreamGuard：可延迟的流守卫

OptionalCUDAStreamGuard 是 CUDAStreamGuard 的 optional 版本，语义与 OptionalCUDAGuard 平行：构造时可传入一个 Stream，或传入 std::optional<Stream>（为空则保持未初始化），并支持 reset_stream(Stream) 按需初始化。查询接口 original_stream() / current_stream() 在未初始化时返回 std::nullopt，而 reset() 会恢复原设备与流并将守卫重置为未初始化。见 CUDAGuard.h。

c10::cuda::OptionalCUDAStreamGuard guard;
if (use_stream) {
    guard.reset_stream(c10::cuda::getStreamFromPool());
}

CUDAMultiStreamGuard：多设备多流的一次性守卫

CUDAMultiStreamGuard 面向“不同设备上需要同时运行不同流”的场景。它与单流守卫的关键差异是：它在每个传入 stream 所属的设备上分别把该 stream 设为当前流。原文档示例基于 at::cuda（即 c10::cuda 的别名）API：

at::cuda::CUDAStream stream0 = at::cuda::getStreamFromPool(false, 0);
at::cuda::CUDAStream stream1 = at::cuda::getStreamFromPool(false, 1);

{
    at::cuda::CUDAMultiStreamGuard multi_guard({stream0, stream1});
    // device 0 上当前流是 stream0，device 1 上当前流是 stream1
}
// 离开作用域，两个设备的流都恢复原状

构造函数接收 ArrayRef<CUDAStream>，内部先把 CUDAStream 解包成通用 Stream 列表，再交由 InlineMultiStreamGuard 逐条执行 exchangeStream（见 CUDAGuard.h）。其底层实现 InlineStreamGuard.h 会校验所有 stream 属于同一 device type（此处即均为 CUDA），否则抛出值错误；析构时按逆序把所有设备恢复为各自的原始流。

底层原理：从 CUDAGuardImpl 到 Inline Guards

五个守卫类型本身都很薄——它们真正的实现都委托给以 c10::cuda::impl::CUDAGuardImpl 为模板参数的通用内联守卫。这个设计在 InlineDeviceGuard.h 的注释中有明确说明：直接以具体 DeviceGuardImpl 实例化可以得到去虚化的直线代码（等价于手写 cudaGetDevice/cudaSetDevice），而以 VirtualGuardImpl 实例化则退化为按 DeviceType 注册表做虚函数分发（即通用 DeviceGuard 的实现路径）。

CUDAGuardImpl（见 CUDAGuardImpl.h）扮演"设备/流原语适配器"，实现 DeviceGuardImplInterface：

• type() 返回 DeviceType::CUDA，构造时用非 CUDA 类型实例化会触发 TORCH_CHECK；
• exchangeDevice(d) 调用 c10::cuda::ExchangeDevice(d.index()) 并返回旧设备；
• setDevice(d) / getDevice() 分别封装 cudaSetDevice / cudaGetDevice；
• exchangeStream(s) 封装 setCurrentCUDAStream，返回切换前的旧流；
• getStream(d) 返回该设备的 getCurrentCUDAStream，getDefaultStream 返回默认流，getNewStream/getStreamFromGlobalPool 从流池取流。

关键的 exchange 语义在 CUDAFunctions.cpp：它先取当前设备索引，仅在目标设备与当前不同时才真正调用 cudaSetDevice，从而把"切换"压缩成一次调用并返回旧值。InlineDeviceGuard 的析构函数则调用 impl_.uncheckedSetDevice(original_device_)（见 InlineDeviceGuard.h），并且该恢复路径做了容错——即使恢复失败也只告警而不 std::terminate（对应 CUDAGuardImpl.h 中 uncheckedSetDevice 的 try/catch 注释：设备携带 sticky error 时恢复设备可能抛出，需要降级为警告）。

理解这条调用链后，一个关键结论是：CUDAGuard 析构永远恢复 original_device 而非"最近一次手动设置"，因此嵌套守卫各自负责各自的恢复，作用域层级清晰。这也是为什么源码强调守卫不可拷贝/移动——移动会破坏“谁先创建谁负责恢复”的次序。

实际使用建议与常见误区

1. 统一头文件入口：五个守卫均通过 #include <c10/cuda/CUDAGuard.h> 引入；需要显式操纵流池、默认流时可同时引入 CUDAStream.h。
2. 把守卫放在尽可能小的作用域：RAII 恢复发生在析构时，作用域越大，被"钉住"的设备/流就越久，会阻塞其它线程/设备上的调度。
3. 不要拷贝或 move 守卫：源码显式 = delete 了拷贝与移动构造/赋值，这是有意为之的正确性约束而非遗漏。
4. 区分"设置当前流"与"切换当前设备"：setCurrentCUDAStream 只切换"该流所在设备上的当前流"，与当前设备无关（见 CUDAStream.h 的注释）；跨设备的完整切换请交给 CUDAStreamGuard 或 CUDAMultiStreamGuard。
5. 多设备并行流请用 CUDAMultiStreamGuard：CUDAStreamGuard::reset_stream 不会保留其它设备上的流设置，多设备场景必须使用多流守卫。
6. ROCm/HIP 构建同样可用：这些 API 位于 c10::cuda 命名空间，HIP 下通过别名与包装（c10::hip::getDefaultHIPStream 等，见 CUDAStream.h）复用同一套守卫基础设施。

总结

CUDA guards 是 PyTorch C++ 库向使用者暴露的最实用的设备/流管理设施之一：CUDAGuard 解决"单设备切换"，CUDAStreamGuard 解决"单设备 + 单流"，CUDAMultiStreamGuard 解决"多设备 + 多流"，两个 Optional 变体则覆盖"运行期条件决定是否切换"的常见模式。无论选择哪一个，其底层都统一收敛到 InlineDeviceGuard/InlineStreamGuard + CUDAGuardImpl 的组合上，换来的是直截了当的运行时开销与无需手写 try/catch 的异常安全。在编写自定义 CUDA 算子、多卡训练框架或跨库推理服务时，优先使用这些守卫代替裸的 cudaSetDevice 调用，是规避上下文泄漏类 bug 的最简单手段。