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
|
#include <gtest/gtest.h>
#include <c10/util/Exception.h>
#include <torch/csrc/lazy/core/cache.h>
#include <torch/csrc/lazy/core/hash.h>
#include <torch/csrc/lazy/core/ir.h>
#include <torch/csrc/lazy/core/shape.h>
#include <torch/csrc/lazy/ts_backend/ts_node.h>
namespace torch {
namespace lazy {
class CacheNode : public Node {
public:
explicit CacheNode(const std::string& str)
: Node(OpKind(), /* num_outputs */ 1), hash_(Hash(str)), str_(str) {}
~CacheNode() override = default;
const std::vector<Output>& operands() const override {
TORCH_INTERNAL_ASSERT(false, "Can't access operands of test node");
}
const Output& operand(size_t i) const override {
TORCH_INTERNAL_ASSERT(false, "Can't access operand[i] of test node");
}
hash_t hash() const override {
return hash_;
}
hash_t shapeHash() const override {
return hash_;
}
private:
hash_t hash_;
std::string str_;
};
TEST(CacheTest, BasicTest) {
std::shared_ptr<CacheNode> a = std::make_shared<CacheNode>("a");
std::shared_ptr<CacheNode> b = std::make_shared<CacheNode>("b");
std::shared_ptr<CacheNode> c = std::make_shared<CacheNode>("c");
Cache<hash_t, CacheNode, HashReducer> cache(2);
cache.Add(a->hash(), a);
EXPECT_EQ(cache.Get(a->hash()), a);
EXPECT_EQ(cache.Get(b->hash()), nullptr);
EXPECT_EQ(cache.Get(c->hash()), nullptr);
cache.Add(b->hash(), b);
EXPECT_EQ(cache.Get(a->hash()), a);
EXPECT_EQ(cache.Get(b->hash()), b);
EXPECT_EQ(cache.Get(c->hash()), nullptr);
cache.Add(c->hash(), c);
EXPECT_EQ(cache.Get(a->hash()), nullptr); // a has been evicted
EXPECT_EQ(cache.Get(b->hash()), b);
EXPECT_EQ(cache.Get(c->hash()), c);
cache.Erase(c->hash());
EXPECT_EQ(cache.Get(a->hash()), nullptr);
EXPECT_EQ(cache.Get(b->hash()), b);
EXPECT_EQ(cache.Get(c->hash()), nullptr); // c has been removed
cache.Clear();
EXPECT_EQ(cache.Get(a->hash()), nullptr);
EXPECT_EQ(cache.Get(b->hash()), nullptr);
EXPECT_EQ(cache.Get(c->hash()), nullptr);
}
class CacheNodeWithShape : public TsNode {
public:
explicit CacheNodeWithShape(const Shape& shape)
: TsNode(OpKind(), shape, /* num_outputs */ 1, /* seed */ 0) {}
};
TEST(CacheTest, ShapeCacheTestForDynamicShape) {
// enable dynamic shape
FLAGS_ltc_enable_dynamic_shapes = true;
CacheNodeWithShape nodes[] = {
CacheNodeWithShape(Shape(c10::kFloat, {2, 4})),
CacheNodeWithShape(Shape(c10::kFloat, {4, 2}))};
/*
* Make sure the cached shape for node (2, 4) is not used for node (4, 2)
*/
for (auto& node : nodes) {
EXPECT_EQ(node.shape(), node.computeShape([&]() { return node.shape(); }));
}
// reset the flag
FLAGS_ltc_enable_dynamic_shapes = false;
}
} // namespace lazy
} // namespace torch
|
