1 | #include <algorithm> |
2 | #include <climits> |
3 | #include "triton/codegen/analysis/layout.h" |
4 | #include "triton/codegen/analysis/allocation.h" |
5 | #include "triton/codegen/analysis/liveness.h" |
6 | #include "triton/ir/utils.h" |
7 | |
8 | namespace triton{ |
9 | namespace codegen{ |
10 | namespace analysis{ |
11 | |
12 | |
13 | void allocation::run(ir::module &mod) { |
14 | using std::max; |
15 | using std::min; |
16 | typedef std::multimap<unsigned, segment> triples_map_type; |
17 | |
18 | std::vector<shared_layout*> I; |
19 | for(auto x: liveness_->get()) |
20 | I.push_back(x.first); |
21 | std::vector<shared_layout*> J = I; |
22 | |
23 | triples_map_type H; |
24 | H.insert({0, segment{0, INT_MAX}}); |
25 | |
26 | std::vector<shared_layout*> V; |
27 | std::map<shared_layout*, unsigned> starts; |
28 | while(!J.empty()){ |
29 | auto h_it = H.begin(); |
30 | unsigned w = h_it->first; |
31 | segment xh = h_it->second; |
32 | H.erase(h_it); |
33 | auto j_it = std::find_if(J.begin(), J.end(), [&](shared_layout* JJ){ |
34 | segment xj = liveness_->get(JJ); |
35 | bool res = xj.intersect(xh); |
36 | for(auto val: H) |
37 | res = res && !val.second.intersect(xj); |
38 | return res; |
39 | }); |
40 | if(j_it != J.end()){ |
41 | unsigned size = (*j_it)->get_size(); |
42 | segment xj = liveness_->get(*j_it); |
43 | starts[*j_it] = w; |
44 | H.insert({w + size, segment{max(xh.start, xj.start), min(xh.end, xj.end)}}); |
45 | if(xh.start < xj.start) |
46 | H.insert({w, segment{xh.start, xj.end}}); |
47 | if(xj.end < xh.end) |
48 | H.insert({w, segment{xj.start, xh.end}}); |
49 | V.push_back(*j_it); |
50 | J.erase(j_it); |
51 | } |
52 | } |
53 | // Build interference graph |
54 | std::map<shared_layout*, std::set<shared_layout*>> interferences; |
55 | for(shared_layout* x: V) |
56 | for(shared_layout* y: V){ |
57 | if(x == y) |
58 | continue; |
59 | unsigned X0 = starts[x], Y0 = starts[y]; |
60 | unsigned NX = x->get_size(); |
61 | unsigned NY = y->get_size(); |
62 | segment XS = {X0, X0 + NX}; |
63 | segment YS = {Y0, Y0 + NY}; |
64 | if(liveness_->get(x).intersect(liveness_->get(y)) |
65 | && XS.intersect(YS)) |
66 | interferences[x].insert(y); |
67 | } |
68 | // Initialize colors |
69 | std::map<shared_layout*, int> colors; |
70 | for(shared_layout* X: V) |
71 | colors[X] = (X==V[0])?0:-1; |
72 | // First-fit graph coloring |
73 | std::vector<bool> available(V.size()); |
74 | for(shared_layout* x: V){ |
75 | // Non-neighboring colors are available |
76 | std::fill(available.begin(), available.end(), true); |
77 | for(shared_layout* Y: interferences[x]){ |
78 | int color = colors[Y]; |
79 | if(color >= 0) |
80 | available[color] = false; |
81 | } |
82 | // Assigns first available color |
83 | auto It = std::find(available.begin(), available.end(), true); |
84 | colors[x] = std::distance(available.begin(), It); |
85 | } |
86 | // Finalize allocation |
87 | for(shared_layout* x: V){ |
88 | unsigned Adj = 0; |
89 | for(shared_layout* y: interferences[x]) |
90 | Adj = std::max<unsigned>(Adj, starts[y] + y->get_size()); |
91 | offsets_[x] = starts[x] + colors[x] * Adj; |
92 | } |
93 | // Save maximum size of induced memory space |
94 | allocated_size_ = 0; |
95 | for(shared_layout* x: V){ |
96 | allocated_size_ = std::max<size_t>(allocated_size_, starts[x] + x->get_size()); |
97 | // std::cout << "start: " << starts[x] << " | end: " << starts[x] + x->get_size() << std::endl; |
98 | } |
99 | } |
100 | |
101 | } |
102 | } |
103 | } |
104 | |