OptionalDeviceGuard Class — pytorch Architecture

Source Code

c10/core/DeviceGuard.h lines 130–186

class OptionalDeviceGuard {
 public:
  /// Create an uninitialized guard.  Set the guard later using reset_device.
  explicit OptionalDeviceGuard() = default;

  /// Initialize the guard, setting the current device to the passed Device.
  explicit OptionalDeviceGuard(Device device) : guard_(device) {}

  /// Initialize the guard if a Device is passed; otherwise leave the
  /// guard uninitialized.
  explicit OptionalDeviceGuard(std::optional<Device> device) : guard_(device) {}

  /// Constructor for testing only.
  explicit OptionalDeviceGuard(
      Device device,
      const impl::DeviceGuardImplInterface* impl)
      : guard_(device, impl) {}

  ~OptionalDeviceGuard() = default;
  /// Copy is disallowed
  OptionalDeviceGuard(const OptionalDeviceGuard&) = delete;
  OptionalDeviceGuard& operator=(const OptionalDeviceGuard&) = delete;

  /// Move is disallowed
  /// See Note [Explicit initialization of optional fields]
  /// and // Note [Move construction for RAII guards is tricky]
  /// for rationale.
  OptionalDeviceGuard(OptionalDeviceGuard&& other) = delete;
  OptionalDeviceGuard& operator=(OptionalDeviceGuard&& other) = delete;

  /// Sets the device to the given one.  The specified device must be consistent
  /// with the device type originally specified during guard construction.
  void reset_device(at::Device device) {
    guard_.reset_device(device);
  }

  /// For testing only
  void reset_device(
      at::Device device,
      const impl::DeviceGuardImplInterface* impl) {
    guard_.reset_device(device, impl);
  }

  /// Returns the device that was set at the time the guard was constructed.
  std::optional<Device> original_device() const {
    return guard_.original_device();
  }

  /// Returns the most recent device that was set using this device guard,
  /// either from construction, or via reset_device.
  std::optional<Device> current_device() const {
    return guard_.current_device();
  }

 private:
  impl::InlineOptionalDeviceGuard<impl::VirtualGuardImpl> guard_;
};