3 #include <c10/core/DeviceType.h> 4 #include <c10/macros/Macros.h> 5 #include <c10/util/Exception.h> 31 using Type = DeviceType;
36 : type_(type), index_(index) {
45 Device(
const std::string& device_string);
50 return this->type_ == other.type_ && this->index_ == other.index_;
56 return !(*
this == other);
65 DeviceType
type() const noexcept {
81 return type_ == DeviceType::CUDA;
86 return type_ == DeviceType::CPU;
95 C10_API std::ostream& operator<<(
103 struct hash<
c10::Device> {
107 static_assert(
sizeof(c10::DeviceType) == 2,
"DeviceType is not 16-bit");
119 static_cast<uint32_t
>(
static_cast<uint16_t
>(d.
type())) << 16
120 |
static_cast<uint32_t
>(
static_cast<uint16_t
>(d.
index()));
121 return std::hash<uint32_t>{}(bits);
bool has_index() const noexcept
Returns true if the device has a non-default index.
bool operator!=(const Device &other) const noexcept
Returns true if the type or index of this Device differs from that of other.
bool operator==(const Device &other) const noexcept
Returns true if the type and index of this Device matches that of other.
bool is_cuda() const noexcept
Return true if the device is of CUDA type.
TensorOptions device(Device device)
Convenience function that returns a TensorOptions object with the device set to the given one...
Represents a a compute device on which a tensor is located.
int16_t DeviceIndex
An index representing a specific device; e.g., the 1 in GPU 1.
bool is_cpu() const noexcept
Return true if the device is of CPU type.
To register your own kernel for an operator, do in one (!) cpp file: C10_REGISTER_KERNEL(OperatorHand...
Device(DeviceType type, DeviceIndex index=-1)
Constructs a new Device from a DeviceType and an optional device index.
DeviceIndex index() const noexcept
Returns the optional index.
void set_index(DeviceIndex index)
Sets the device index.
DeviceType type() const noexcept
Returns the type of device this is.