1 #include <torch/csrc/jit/passes/utils/check_alias_annotation.h> 7 IValue deepCopy(
const IValue&
self) {
9 if (!
self.isPtrType()) {
14 if (
self.isTensor()) {
15 return IValue(
self.toTensor().clone());
17 if (
self.isTensorList()) {
18 std::vector<at::Tensor> newList;
19 for (
const auto& oldTensor :
self.toTensorListRef()) {
20 newList.push_back(oldTensor.clone());
26 if (
self.isGenericList()) {
27 std::vector<IValue> newList;
28 for (
const auto& value :
self.toGenericListRef()) {
29 newList.push_back(deepCopy(value));
35 if (
self.isIntList()) {
36 return IValue(
self.toIntListRef());
37 }
else if (
self.isDoubleList()) {
38 return IValue(
self.toDoubleListRef());
39 }
else if (
self.isBoolList()) {
40 return IValue(
self.toBoolListRef());
41 }
else if (
self.isString()) {
42 return IValue(
self.toStringRef());
50 Stack deepCopy(
const Stack& stack) {
52 ret.reserve(stack.size());
53 for (
const auto& v : stack) {
54 ret.push_back(deepCopy(v));
59 bool deepEquals(
const IValue& lhs,
const IValue& rhs) {
60 if (lhs.isInt() && rhs.isInt()) {
61 return lhs.toInt() == rhs.toInt();
62 }
else if (lhs.isDouble() && rhs.isDouble()) {
63 return lhs.toDouble() == rhs.toDouble();
64 }
else if (lhs.isNone() && rhs.isNone()) {
66 }
else if (lhs.isIntList() && rhs.isIntList()) {
67 return lhs.toIntList()->elements() == rhs.toIntList()->elements();
68 }
else if (lhs.isTensor() && rhs.isTensor()) {
69 return lhs.toTensor().equal(rhs.toTensor());
72 throw std::runtime_error(
"Deep equals not implemented for type");
75 struct AliasAndIValue {
77 : aliasInfo(
std::move(aliasInfo)), iValue(
std::move(iValue)) {}
84 void checkInputPreconditions(
const Stack& inputs) {
85 for (
size_t i = 0; i < inputs.size(); i++) {
86 for (
size_t j = 0; j < inputs.size(); j++) {
90 const auto& lhs = inputs.at(i);
91 const auto& rhs = inputs.at(j);
92 AT_ASSERT(!lhs.isAliasOf(rhs));
99 const std::vector<AliasAndIValue>& inputs,
100 const std::vector<AliasAndIValue>& outputs) {
101 for (
const auto& output : outputs) {
103 for (
const auto& input : inputs) {
104 if (output.iValue.isAliasOf(input.iValue)) {
105 const auto inputSet = input.aliasInfo;
106 const auto outputSet = output.aliasInfo;
107 AT_ASSERT(inputSet && outputSet);
109 for (
const auto&
set : inputSet->beforeSets()) {
110 if (outputSet->beforeSets().count(
set)) {
124 const std::vector<AliasAndIValue>& inputs,
125 const std::vector<IValue>& deepCopiedInputs) {
126 AT_ASSERT(inputs.size() == deepCopiedInputs.size());
127 for (
size_t i = 0; i < inputs.size(); i++) {
128 const auto& input = inputs[i];
129 const auto& deepCopiedInput = deepCopiedInputs[i];
130 if (!input.aliasInfo || !input.aliasInfo->isWrite()) {
131 AT_ASSERT(deepEquals(input.iValue, deepCopiedInput));
136 const Node* findNodeForOp(
138 const std::string& unqualifiedOpName) {
139 const auto opName = Symbol::fromQualString(
"aten::" + unqualifiedOpName);
140 for (
const auto node : g.nodes()) {
141 if (node->kind() == opName) {
152 if (v->node()->kind() == prim::ListConstruct) {
153 std::vector<IValue> genericList;
154 for (
auto input : v->node()->inputs()) {
155 if (
auto elem = toIValue(input)) {
156 genericList.push_back(*elem);
164 auto listType = v->node()->output()->type();
165 auto containedType = listType->containedTypes().at(0);
166 if (containedType == IntType::get()) {
167 return fmap(genericList, [](
const IValue& v) {
return v.toInt(); });
168 }
else if (containedType == FloatType::get()) {
169 return fmap(genericList, [](
const IValue& v) {
return v.toDouble(); });
170 }
else if (containedType->isSubtypeOf(TensorType::get())) {
171 return fmap(genericList, [](
const IValue& v) {
return v.toTensor(); });
177 if (v->node()->kind() == prim::Float) {
178 auto op = getOperation(v->node());
179 if (
auto input = toIValue(v->node()->input())) {
180 auto op = getOperation(v->node());
193 void checkAliasAnnotation(
194 const std::shared_ptr<Graph>& graph,
195 std::vector<IValue> pythonInputs,
196 const std::string& unqualifiedOpName) {
198 const auto node = findNodeForOp(*graph, unqualifiedOpName);
202 for (
const auto input : node->inputs()) {
203 if (input->node() == graph->param_node()) {
205 push(stack, pythonInputs.at(input->offset()));
208 auto inputValue = toIValue(input);
210 inputValue = toIValueProp(input);
214 push(stack, *inputValue);
216 AT_ASSERT(input->type()->kind() == TypeKind::OptionalType);
217 push(stack, IValue());
224 checkInputPreconditions(stack);
226 const auto schema = node->schema();
228 std::vector<AliasAndIValue> inputsToCheck;
229 for (
size_t i = 0; i < schema.arguments().size(); i++) {
230 inputsToCheck.emplace_back(
231 schema.arguments().at(i).alias_info(), stack.at(i));
236 const auto inputsDeepCopy = deepCopy(stack);
239 getOperation(node)(stack);
241 const auto outputs = std::move(stack);
243 std::vector<AliasAndIValue> outputsToCheck;
244 for (
size_t i = 0; i < schema.returns().size(); i++) {
245 outputsToCheck.emplace_back(
246 schema.returns().at(i).alias_info(), outputs.at(i));
250 checkAliases(inputsToCheck, outputsToCheck);
253 checkWrites(inputsToCheck, inputsDeepCopy);