Device.h

#pragma once

#include <c10/core/DeviceType.h>
#include <c10/macros/Macros.h>
#include <c10/util/Exception.h>

#include <cstddef>
#include <functional>
#include <iosfwd>
#include <string>

namespace c10 {

using DeviceIndex = int16_t;

struct C10_API Device final {
  using Type = DeviceType;

  /* implicit */ Device(DeviceType type, DeviceIndex index = -1)
      : type_(type), index_(index) {
    validate();
  }

  /* implicit */ Device(const std::string& device_string);

  bool operator==(const Device& other) const noexcept {
    return this->type_ == other.type_ && this->index_ == other.index_;
  }

  bool operator!=(const Device& other) const noexcept {
    return !(*this == other);
  }

  void set_index(DeviceIndex index) {
    index_ = index;
  }

  DeviceType type() const noexcept {
    return type_;
  }

  DeviceIndex index() const noexcept {
    return index_;
  }

  bool has_index() const noexcept {
    return index_ != -1;
  }

  bool is_cuda() const noexcept {
    return type_ == DeviceType::CUDA;
  }

  bool is_cpu() const noexcept {
    return type_ == DeviceType::CPU;
  }

 private:
  DeviceType type_;
  DeviceIndex index_ = -1;
  void validate();
};

C10_API std::ostream& operator<<(
    std::ostream& stream,
    const Device& device);

} // namespace c10

namespace std {
template <>
struct hash<c10::Device> {
  size_t operator()(c10::Device d) const noexcept {
    // Are you here because this static assert failed?  Make sure you ensure
    // that the bitmasking code below is updated accordingly!
    static_assert(sizeof(c10::DeviceType) == 2, "DeviceType is not 16-bit");
    static_assert(sizeof(c10::DeviceIndex) == 2, "DeviceIndex is not 16-bit");
    // Note [Hazard when concatenating signed integers]
    // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    // We must first convert to a same-sized unsigned type, before promoting to
    // the result type, to prevent sign extension when any of the values is -1.
    // If sign extension occurs, you'll clobber all of the values in the MSB
    // half of the resulting integer.
    //
    // Technically, by C/C++ integer promotion rules, we only need one of the
    // uint32_t casts to the result type, but we put in both for explicitness's sake.
    uint32_t bits =
        static_cast<uint32_t>(static_cast<uint16_t>(d.type())) << 16
      | static_cast<uint32_t>(static_cast<uint16_t>(d.index()));
    return std::hash<uint32_t>{}(bits);
  }
};
} // namespace std