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
|
#include <c10/util/Enumerate.h>
#include <gtest/gtest.h>
#include <torch/nativert/executor/memory/Bump.h>
#include <torch/nativert/executor/memory/DisjointStorageGroups.h>
#include <torch/nativert/executor/memory/GreedyBySize.h>
using namespace ::testing;
using namespace torch::nativert;
std::vector<AllocationSpec> create_test_allocation_specs() {
std::vector<AllocationSpec> specs;
const std::vector<std::tuple<size_t, size_t, size_t>> test_cases = {
{0, 1, 32},
{1, 4, 28},
{2, 5, 36},
{3, 5, 16},
{4, 5, 8},
{5, 7, 64},
{6, 8, 10},
{7, 8, 40},
};
specs.reserve(test_cases.size());
for (const auto& [l_start, l_end, size] : test_cases) {
specs.push_back(AllocationSpec{AllocationLifetime(l_start, l_end), size});
};
return specs;
}
// figure 6 -- https://arxiv.org/pdf/2001.03288
TEST(LayoutPlannerAlgorithmTests, TestGreedyBySize) {
auto result = GreedyBySizeAllocationPlanner(create_test_allocation_specs());
EXPECT_EQ(result.total_size, 124);
auto& allocations = result.allocations;
EXPECT_EQ(allocations[0].offset, 0);
EXPECT_EQ(allocations[1].offset, 32);
EXPECT_EQ(allocations[2].offset, 64);
EXPECT_EQ(allocations[3].offset, 100);
EXPECT_EQ(allocations[4].offset, 116);
EXPECT_EQ(allocations[5].offset, 0);
EXPECT_EQ(allocations[6].offset, 104);
EXPECT_EQ(allocations[7].offset, 64);
}
TEST(LayoutPlannerAlgorithmTests, TestBump) {
auto specs = create_test_allocation_specs();
auto result = BumpAllocationPlanner(create_test_allocation_specs());
auto& allocations = result.allocations;
size_t offset = 0;
for (auto&& [i, spec] : c10::enumerate(specs)) {
EXPECT_EQ(allocations[i].offset, offset);
offset += spec.size;
}
EXPECT_EQ(result.total_size, offset);
}
TEST(LayoutPlannerAlgorithmTests, TestStorageGroup) {
auto specs = create_test_allocation_specs();
auto result = DisjointStorageGroupsPlanner(create_test_allocation_specs());
auto& allocations = result.allocations;
EXPECT_EQ(allocations[0].offset, 0);
EXPECT_EQ(allocations[1].offset, 36);
EXPECT_EQ(allocations[2].offset, 0);
EXPECT_EQ(allocations[3].offset, 100);
EXPECT_EQ(allocations[4].offset, 140);
EXPECT_EQ(allocations[5].offset, 36);
EXPECT_EQ(allocations[6].offset, 140);
EXPECT_EQ(allocations[7].offset, 100);
for (auto&& [i, spec] : c10::enumerate(specs)) {
EXPECT_EQ(allocations[i].size, spec.size);
}
EXPECT_EQ(result.total_size, 150);
}
|
