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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
|
#pragma once
#include <ATen/core/Tensor.h>
#include <string>
#include <vector>
#include <ATen/core/jit_type_base.h>
#include <ATen/core/symbol.h>
#include <torch/csrc/Export.h>
namespace torch {
namespace jit {
using ::c10::Symbol;
constexpr int max_tensor_display_size = 10;
enum class AttributeKind {
f,
fs,
c,
cs,
i,
is,
s,
ss,
t,
ts,
g,
gs,
ty,
tys,
ival
};
static inline const char* toString(AttributeKind kind) {
// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
static const char* names[] = {
"f",
"c",
"cs",
"fs",
"i",
"is",
"s",
"ss",
"t",
"ts",
"g",
"gs",
"ty",
"tys",
"ival"};
AT_ASSERT(size_t(kind) < sizeof(names) / sizeof(*names));
return names[int(kind)];
}
struct AttributeValue {
AttributeValue(Symbol name) : name(name) {}
using Ptr = std::unique_ptr<AttributeValue>;
Symbol name;
virtual AttributeKind kind() const = 0;
virtual Ptr clone() const = 0;
virtual ~AttributeValue() = default;
};
template <typename T, AttributeKind Kind>
struct ScalarAttributeValue : public AttributeValue {
using ConstructorType = T;
using ValueType = T;
ScalarAttributeValue(Symbol name, ConstructorType value_)
: AttributeValue(name), value_(std::move(value_)) {}
ValueType& value() {
return value_;
}
Ptr clone() const override {
return Ptr(new ScalarAttributeValue(name, value_));
}
AttributeKind kind() const override {
return Kind;
}
private:
ValueType value_;
};
template <typename T, AttributeKind Kind>
struct VectorAttributeValue : public AttributeValue {
using ConstructorType = std::vector<T>;
using ValueType = std::vector<T>;
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
VectorAttributeValue(Symbol name, ConstructorType value_)
: AttributeValue(name), value_(std::move(value_)) {}
ValueType& value() {
return value_;
}
AttributeKind kind() const override {
return Kind;
}
std::unique_ptr<AttributeValue> clone() const override {
auto copy = value_;
return Ptr(new VectorAttributeValue(name, std::move(copy)));
}
private:
ValueType value_;
};
using ComplexAttr =
ScalarAttributeValue<c10::complex<double>, AttributeKind::c>;
using ComplexValsAttr =
VectorAttributeValue<c10::complex<double>, AttributeKind::cs>;
using FloatAttr = ScalarAttributeValue<double, AttributeKind::f>;
using FloatsAttr = VectorAttributeValue<double, AttributeKind::fs>;
using IntAttr = ScalarAttributeValue<int64_t, AttributeKind::i>;
using IntsAttr = VectorAttributeValue<int64_t, AttributeKind::is>;
using StringAttr = ScalarAttributeValue<std::string, AttributeKind::s>;
using StringsAttr = VectorAttributeValue<std::string, AttributeKind::ss>;
using TensorAttr = ScalarAttributeValue<at::Tensor, AttributeKind::t>;
using TensorsAttr = VectorAttributeValue<at::Tensor, AttributeKind::ts>;
using TypeAttr = ScalarAttributeValue<c10::TypePtr, AttributeKind::ty>;
using TypesAttr = VectorAttributeValue<c10::TypePtr, AttributeKind::tys>;
using IValueAttr = ScalarAttributeValue<at::IValue, AttributeKind::ival>;
struct Graph;
// We special case Graph attributes like this because we want to ensure that
// Graph::copy() is called when we clone() these attributes.
struct TORCH_API GraphAttr : public AttributeValue {
using ConstructorType = std::shared_ptr<Graph>;
using ValueType = std::shared_ptr<Graph>;
GraphAttr(Symbol name, ConstructorType value_)
: AttributeValue(name), value_(std::move(value_)) {}
ValueType& value() {
return value_;
}
Ptr clone() const override;
AttributeKind kind() const override {
return AttributeKind::g;
}
private:
std::shared_ptr<Graph> value_;
};
struct TORCH_API GraphsAttr : public AttributeValue {
using ConstructorType = std::vector<std::shared_ptr<Graph>>;
using ValueType = std::vector<std::shared_ptr<Graph>>;
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
GraphsAttr(Symbol name, ConstructorType value_)
: AttributeValue(name), value_(std::move(value_)) {}
ValueType& value() {
return value_;
}
AttributeKind kind() const override {
return AttributeKind::gs;
}
std::unique_ptr<AttributeValue> clone() const override;
private:
ValueType value_;
};
struct IRAttributeError : public std::exception {
IRAttributeError(Symbol name, bool defined) {
std::stringstream ss;
// NOLINTNEXTLINE(bugprone-branch-clone)
if (!defined) {
ss << "required keyword attribute '" << name.toUnqualString()
<< "' is undefined";
} else {
ss << "required keyword attribute '" << name.toUnqualString()
<< "' has the wrong type";
}
msg = ss.str();
}
const char* what() const noexcept override {
return msg.c_str();
}
private:
std::string msg;
};
} // namespace jit
} // namespace torch
|
