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
|
#include <c10/core/Device.h>
#include <c10/macros/Macros.h>
#include <c10/util/Exception.h>
#include <algorithm>
#include <array>
#include <exception>
#include <ostream>
#include <string>
#include <tuple>
#include <vector>
#include <regex>
// Check if compiler has working std::regex implementation
//
// Test below is adapted from https://stackoverflow.com/a/41186162
#if defined(_MSVC_LANG) && _MSVC_LANG >= 201103L
// Compiler has working regex. MSVC has erroneous __cplusplus.
#elif __cplusplus >= 201103L && \
(!defined(__GLIBCXX__) || (__cplusplus >= 201402L) || \
(defined(_GLIBCXX_REGEX_DFS_QUANTIFIERS_LIMIT) || \
defined(_GLIBCXX_REGEX_STATE_LIMIT) || \
(defined(_GLIBCXX_RELEASE) && \
_GLIBCXX_RELEASE > 4)))
// Compiler has working regex.
#else
static_assert(false, "Compiler does not have proper regex support.");
#endif
namespace c10 {
namespace {
DeviceType parse_type(const std::string& device_string) {
static const std::array<std::pair<std::string, DeviceType>, 10> types = {{
{"cpu", DeviceType::CPU},
{"cuda", DeviceType::CUDA},
{"mkldnn", DeviceType::MKLDNN},
{"opengl", DeviceType::OPENGL},
{"opencl", DeviceType::OPENCL},
{"ideep", DeviceType::IDEEP},
{"hip", DeviceType::HIP},
{"fpga", DeviceType::FPGA},
{"msnpu", DeviceType::MSNPU},
{"xla", DeviceType::XLA},
}};
auto device = std::find_if(
types.begin(),
types.end(),
[device_string](const std::pair<std::string, DeviceType>& p) {
return p.first == device_string;
});
if (device != types.end()) {
return device->second;
}
AT_ERROR(
"Expected one of cpu, cuda, mkldnn, opengl, opencl, ideep, hip, msnpu, xla device type at start of device string: ", device_string);
}
} // namespace
Device::Device(const std::string& device_string) : Device(Type::CPU) {
TORCH_CHECK(!device_string.empty(), "Device string must not be empty");
// We assume gcc 5+, so we can use proper regex.
static const std::regex regex("([a-zA-Z_]+)(?::([1-9]\\d*|0))?");
std::smatch match;
TORCH_CHECK(
std::regex_match(device_string, match, regex),
"Invalid device string: '", device_string, "'");
type_ = parse_type(match[1].str());
if (match[2].matched) {
try {
index_ = c10::stoi(match[2].str());
} catch (const std::exception &) {
AT_ERROR(
"Could not parse device index '", match[2].str(),
"' in device string '", device_string, "'");
}
}
validate();
}
std::string Device::str() const {
std::string str = DeviceTypeName(type(), /* lower case */ true);
if (has_index()) {
str.push_back(':');
str.append(to_string(index()));
}
return str;
}
std::ostream& operator<<(std::ostream& stream, const Device& device) {
stream << device.str();
return stream;
}
} // namespace c10
|
