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
8namespace triton{
9namespace codegen{
10namespace analysis{
11
12
13void 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