1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
#pragma once
#include <memory>
#include <ostream>
#include <string>
#include <ATen/Tensor.h>
#include <c10/macros/Export.h>
#include <c10/util/Deprecated.h>
#include <optional>
namespace c10 {
struct Device;
}
namespace torch::lazy {
// Backend should extend it and define their own supported hardware types.
struct TORCH_API BackendDeviceType {
int8_t type{(int8_t)at::kCPU};
// Note: previous default value was '0', which actually maps to at::kCPU, at
// least now it is explicit, we may want to make default/undefined semantics
// more clear though
BackendDeviceType() : type((int8_t)at::kCPU) {}
BackendDeviceType(int8_t type) : type(type) {}
virtual ~BackendDeviceType() = default;
virtual std::string toString() const {
return "Unknown";
}
};
class TORCH_API BackendDevice {
public:
// The default constructor will set both the device type and ordinal
// to backend specific defaults.
BackendDevice();
BackendDevice(std::shared_ptr<BackendDeviceType>&& type, int64_t ordinal);
int8_t type() const;
int64_t ordinal() const {
return ordinal_;
}
bool operator==(const BackendDevice& other) const {
return compare(other) == 0;
}
bool operator!=(const BackendDevice& other) const {
return compare(other) != 0;
}
bool operator<(const BackendDevice& rhs) const {
return compare(rhs) < 0;
}
std::string toString() const;
private:
int compare(const BackendDevice& rhs) const;
// Use shared_ptr instead of unique_ptr so that BackendDevice can be copied.
std::shared_ptr<BackendDeviceType> type_;
int64_t ordinal_;
};
TORCH_API std::ostream& operator<<(
std::ostream& os,
const BackendDevice& device);
// Helpers for converting a c10::Device to BackendDevice and vice versa.
TORCH_API BackendDevice atenDeviceToBackendDevice(const c10::Device& device);
TORCH_API c10::Device backendDeviceToAtenDevice(const BackendDevice& device);
// Tries to extract the backend device out of the lazy tensor. Returns nullopt
// if the input is not a lazy tensor.
TORCH_API std::optional<BackendDevice> GetBackendDevice(
const at::ITensorListRef tensors);
TORCH_API std::optional<BackendDevice> GetBackendDevice(
const at::TensorList tensors);
TORCH_API std::optional<BackendDevice> GetBackendDevice(
const at::Tensor& tensor);
TORCH_API std::optional<BackendDevice> GetBackendDevice(
const std::optional<c10::Device>& device);
// For variadic template.
TORCH_API std::optional<BackendDevice> GetBackendDevice();
C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Winfinite-recursion")
template <typename T, typename... Args>
std::optional<BackendDevice> GetBackendDevice(
const T& tensor,
const Args&... forward_tensors) {
auto optional_device = GetBackendDevice(tensor);
if (optional_device) {
return optional_device;
}
return GetBackendDevice(forward_tensors...);
}
C10_DIAGNOSTIC_POP()
} // namespace torch::lazy
|
