DeviceGuard.h

#pragma once

#include <c10/core/impl/InlineDeviceGuard.h>

namespace c10 {

class DeviceGuard {
public:
  explicit DeviceGuard() = delete;

  explicit DeviceGuard(Device device) : guard_(device) {}

  explicit DeviceGuard(Device device, const impl::DeviceGuardImplInterface* impl) : guard_(device, impl) {}

  DeviceGuard(const DeviceGuard&) = delete;
  DeviceGuard& operator=(const DeviceGuard&) = delete;

  DeviceGuard(DeviceGuard&& other) = delete;
  DeviceGuard& operator=(DeviceGuard&& other) = delete;

  void reset_device(at::Device device) {
    guard_.reset_device(device);
  }

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

  void set_index(DeviceIndex index) {
    guard_.set_index(index);
  }

  Device original_device() const {
    return guard_.original_device();
  }

  Device current_device() const {
    return guard_.current_device();
  }

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

class OptionalDeviceGuard {
public:
  explicit OptionalDeviceGuard() : guard_() {}

  explicit OptionalDeviceGuard(Device device) : guard_(device) {}

  explicit OptionalDeviceGuard(optional<Device> device) : guard_(device) {}

  explicit OptionalDeviceGuard(Device device, const impl::DeviceGuardImplInterface* impl) : guard_(device, impl) {}

  OptionalDeviceGuard(const OptionalDeviceGuard&) = delete;
  OptionalDeviceGuard& operator=(const OptionalDeviceGuard&) = delete;

  OptionalDeviceGuard(OptionalDeviceGuard&& other) = delete;
  OptionalDeviceGuard& operator=(OptionalDeviceGuard&& other) = delete;

  void reset_device(at::Device device) {
    guard_.reset_device(device);
  }

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

  optional<Device> original_device() const {
    return guard_.original_device();
  }

  optional<Device> current_device() const {
    return guard_.current_device();
  }

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

// Note [Whither the DeviceGuard boilerplate]
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// Design note: in principle, we could avoid these wrappers using:
//
// using DeviceGuard = impl::InlineDeviceGuard<impl::VirtualGuardImpl>;
// using OptionalDeviceGuard = impl::InlineOptionalDeviceGuard<impl::VirtualGuardImpl>;
//
// But the error messages are worse, and our users can't just look at the
// header file to find out what's going on.  Furthermore, for specializations
// like CUDAStreamGuard, it can be profitable to replace some interfaces with
// refined types (e.g., return CUDAStream instead of Stream).  So, we eat
// the boilerplate and write out the API explicitly.

} // namespace c10