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
|
#include <gtest/gtest.h>
#include <ATen/ops/tensor.h>
#include <torch/nativert/executor/ExecutionFrame.h>
namespace torch::nativert {
TEST(ExecutionFrameTest, CreateFrame) {
auto graph = stringToGraph(R"(
graph(%x, %y):
%a = foo(a=%x, b=%y)
%b = foo1(a=%x, b=%y)
%c = foo2(c=%a, d=%b)
return(%c)
)");
auto frame = ExecutionFrame(*graph);
for (auto* v : graph->values()) {
frame.setIValue(v->id(), c10::IValue(at::tensor({v->id()}, at::kInt)));
auto& frame_v = frame.getIValue(v->id());
EXPECT_EQ(frame_v.tagKind(), "Tensor");
}
auto outputs = frame.tryMoveUserOutputs();
EXPECT_EQ(outputs.size(), 1);
EXPECT_EQ(outputs[0].tagKind(), "Tensor");
EXPECT_EQ(outputs[0].toTensor().item().toInt(), graph->getValue("c")->id());
}
TEST(ExecutionFrameTest, TestSetBorrowedValue) {
auto graph = stringToGraph(R"(
graph(%x, %y):
%a = foo(a=%x, b=%y)
%b = foo1(a=%x, b=%y)
%c = foo2(c=%a, d=%b)
return(%c)
)");
auto x = c10::IValue(at::tensor({1}, at::kInt));
auto y = c10::IValue(at::tensor({2}, at::kInt));
{
auto frame = ExecutionFrame(*graph);
frame.setBorrowedIValue(
graph->getValue("x")->id(),
c10::MaybeOwnedTraits<c10::IValue>::createBorrow(x));
frame.setBorrowedIValue(
graph->getValue("y")->id(),
c10::MaybeOwnedTraits<c10::IValue>::createBorrow(y));
[[maybe_unused]] auto& w = frame.getIValue(graph->getValue("x")->id());
[[maybe_unused]] auto& z = frame.getIValue(graph->getValue("y")->id());
EXPECT_EQ(x.use_count(), 1);
EXPECT_EQ(y.use_count(), 1);
EXPECT_TRUE(c10::MaybeOwnedTraits<c10::IValue>{}.debugBorrowIsValid(
frame.getIValue(graph->getValue("x")->id())));
EXPECT_TRUE(c10::MaybeOwnedTraits<c10::IValue>{}.debugBorrowIsValid(
frame.getIValue(graph->getValue("y")->id())));
}
EXPECT_EQ(x.use_count(), 1);
EXPECT_EQ(y.use_count(), 1);
}
TEST(ExecutionFrameTest, TestPersistentValue) {
auto graph = stringToGraph(R"(
graph(%x, %y, %my_weight):
%a = foo(a=%x, b=%y)
%b = foo1(a=%x, b=%y)
%c = foo2(c=%a, d=%b)
return(%c)
)");
Weights weights(graph.get());
weights.setValue("my_weight", at::tensor({1}, at::kInt));
auto new_sig = graph->signature();
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
const_cast<std::vector<std::pair<std::string, std::string>>&>(
new_sig.inputsToWeights())
.emplace_back("my_weight", "my_weight");
graph->setSignature(new_sig);
auto frame = ExecutionFrame(*graph, weights);
EXPECT_EQ(frame.weightVersion(), 0);
auto wid = graph->getValue("my_weight")->id();
EXPECT_NO_THROW(frame.getTensor(wid));
// can't release persistent value
frame.releaseValueIfNeeded(wid);
EXPECT_FALSE(frame.getIValue(wid).isNone());
}
} // namespace torch::nativert
|
