1 | /** |
2 | * Copyright (c) Glow Contributors. See CONTRIBUTORS file. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | */ |
16 | |
17 | #include "glow/LLVMIRCodeGen/AllocationsInfo.h" |
18 | #include "glow/Backend/BackendUtils.h" |
19 | #include "glow/Backend/CompiledFunction.h" |
20 | #include "glow/CodeGen/MemoryAllocator.h" |
21 | #include "glow/Graph/Graph.h" |
22 | #include "glow/Graph/Nodes.h" |
23 | #include "glow/Graph/PlaceholderBindings.h" |
24 | #include "glow/IR/IRUtils.h" |
25 | #include "glow/IR/Instrs.h" |
26 | #include "glow/Support/Debug.h" |
27 | #include "glow/Support/Memory.h" |
28 | |
29 | #include "llvm/Support/Debug.h" |
30 | #include "llvm/Support/raw_ostream.h" |
31 | |
32 | #define DEBUG_TYPE "jit-allocations" |
33 | |
34 | using namespace glow; |
35 | using llvm::cast; |
36 | using llvm::dyn_cast; |
37 | using llvm::isa; |
38 | |
39 | void AllocationsInfo::allocateWeightVars(const IRFunction *F) { |
40 | // Compute the new offsets for all the weights, do not reuse their current |
41 | // addresses. Process all constant WeightVars first. |
42 | allocateConstants(F->findConstants(), constantWeightVarsAllocator_, |
43 | symbolTable_); |
44 | for (auto &c : F->findConstants()) { |
45 | auto name = std::string(c->getName()); |
46 | CHECK(symbolTable_.find(name) != symbolTable_.end()) |
47 | << "Expected to find " << name << " in symbol table" ; |
48 | auto *w = cast<WeightVar>(F->getWeightForNode(c)); |
49 | if (allocatedAddress_.count(c)) { |
50 | allocatedAddress_[w] = allocatedAddress_[c]; |
51 | continue; |
52 | } |
53 | auto &symb = symbolTable_[name]; |
54 | CHECK(valueNumbers_.count(c)) << "Unexpected uncounted constant" ; |
55 | symb.index = valueNumbers_[c].second; |
56 | auto addr = symb.offset; |
57 | allocatedAddress_[c] = addr; |
58 | allocatedAddress_[w] = addr; |
59 | } |
60 | |
61 | // Placeholders should be allocated in a order of |
62 | // intput|inputOutput|output|neither. |
63 | auto contiguousPlaceholders = |
64 | getContiguousPlaceHolder(F->findPlaceholders(), *F); |
65 | |
66 | allocatePlaceholders(contiguousPlaceholders, mutableWeightVarsAllocator_, |
67 | symbolTable_); |
68 | // Compute the offsets and total memory requirements for Placeholders. |
69 | for (auto it = contiguousPlaceholders.begin(); |
70 | it != contiguousPlaceholders.end(); it++) { |
71 | auto &v = it->addr; |
72 | auto name = std::string(v->getName()); |
73 | CHECK(symbolTable_.find(name) != symbolTable_.end()) |
74 | << "Expected to find " << name << " in symbol table" ; |
75 | auto *w = cast<WeightVar>(F->getWeightForNode(v)); |
76 | if (allocatedAddress_.count(w)) { |
77 | continue; |
78 | } |
79 | auto &symb = symbolTable_[name]; |
80 | CHECK(valueNumbers_.count(w)) << "Unexpected uncounted placeholder" ; |
81 | symb.index = valueNumbers_[w].second; |
82 | auto addr = symb.offset; |
83 | allocatedAddress_[w] = addr; |
84 | } |
85 | |
86 | // Remember that max required memory size for each kind of weights. |
87 | constantWeightVarsMemSize_ = constantWeightVarsAllocator_.getMaxMemoryUsage(); |
88 | mutableWeightVarsMemSize_ = mutableWeightVarsAllocator_.getMaxMemoryUsage(); |
89 | |
90 | DEBUG_GLOW(for (auto &A |
91 | : allocatedAddress_) { |
92 | if (isa<AllocActivationInst>(A.first) || isa<TensorViewInst>(A.first)) |
93 | continue; |
94 | (CHECK(valueNumbers_.count(A.first)) << "Unknown weight" ); |
95 | if (isa<Constant>(A.first)) |
96 | continue; |
97 | auto *weight = dyn_cast<WeightVar>(A.first); |
98 | llvm::StringRef kind = |
99 | valueNumbers_[weight].first == ValueKind::ConstantWeight |
100 | ? "constant weight" |
101 | : "mutable weight" ; |
102 | llvm::dbgs() << "Allocated " << kind << " " << weight->getName() |
103 | << " size: " << weight->getSizeInBytes() |
104 | << " address range: [" << allocatedAddress_[weight] << ", " |
105 | << allocatedAddress_[weight] + weight->getSizeInBytes() |
106 | << "]\n" ; |
107 | }); |
108 | } |
109 | |
110 | void AllocationsInfo::allocateActivations(const IRFunction *F) { |
111 | glow::allocateActivations(F->getInstrs(), activationsAllocator_, |
112 | symbolTable_); |
113 | |
114 | // Maps activations and views to some offset within the heap. |
115 | llvm::DenseMap<const Value *, uint64_t> activationAddr; |
116 | |
117 | // Assign device-space addresses to the activations. |
118 | for (const auto &I : F->getInstrs()) { |
119 | if (auto *A = dyn_cast<AllocActivationInst>(&I)) { |
120 | auto name = std::string(A->getName()); |
121 | CHECK(symbolTable_.find(name) != symbolTable_.end()) |
122 | << "Expected to find " << name << " in symbol table" ; |
123 | auto &symb = symbolTable_[name]; |
124 | CHECK(valueNumbers_.count(A)) << "Unexpected uncounted activation" ; |
125 | symb.index = valueNumbers_[A].second; |
126 | CHECK(!activationAddr.count(A)) << "Allocation already made!" ; |
127 | auto addr = symb.offset; |
128 | activationAddr[A] = addr; |
129 | } |
130 | } |
131 | |
132 | activationsMemSize_ = activationsAllocator_.getMaxMemoryUsage(); |
133 | |
134 | // Register specific addresses within the heap to activations. |
135 | for (auto &A : activationAddr) { |
136 | allocatedAddress_[A.first] = A.second; |
137 | } |
138 | DEBUG_GLOW(for (auto &A |
139 | : allocatedAddress_) { |
140 | if (!isa<AllocActivationInst>(A.first)) { |
141 | continue; |
142 | } |
143 | if (isa<Constant>(A.first)) |
144 | continue; |
145 | auto *act = dyn_cast<AllocActivationInst>(A.first); |
146 | llvm::dbgs() << "Allocated activation " << act->getName() |
147 | << " size: " << act->getSizeInBytes() << " address range: [" |
148 | << allocatedAddress_[act] << ", " |
149 | << allocatedAddress_[act] + act->getSizeInBytes() << "]\n" ; |
150 | }); |
151 | } |
152 | |
153 | void AllocationsInfo::allocateTensorViews(const IRFunction *F) { |
154 | for (const auto &I : F->getInstrs()) { |
155 | if (const auto *TVI = dyn_cast<TensorViewInst>(&I)) { |
156 | auto *viewOrigin = getOrigin(TVI); |
157 | CHECK(allocatedAddress_.count(viewOrigin)) |
158 | << "Did not find original WeightVar or AllocActivation for a " |
159 | << "TensorView." ; |
160 | size_t originAddr = allocatedAddress_[viewOrigin]; |
161 | |
162 | // Calculate the offset into the underlying alloc activation. |
163 | size_t offset = calculateTensorViewOffset(TVI); |
164 | |
165 | // Calculate the correct address using this offset into the alloc |
166 | // activation and map from the original TVI to it. |
167 | CHECK(!allocatedAddress_.count(TVI)) << "Allocation already made!" ; |
168 | allocatedAddress_[TVI] = originAddr + offset; |
169 | |
170 | auto name = std::string(TVI->getName()); |
171 | CHECK(symbolTable_.count(name)) << "Unexpected tensorview symbol" ; |
172 | auto &symb = symbolTable_[name]; |
173 | CHECK(valueNumbers_.count(TVI)) << "Unexpected uncounted tensorview" ; |
174 | symb.index = valueNumbers_[TVI].second; |
175 | |
176 | continue; |
177 | } |
178 | } |
179 | } |
180 | |
181 | void AllocationsInfo::numberValues(const IRFunction *F) { |
182 | // Assign numbers to all weights. |
183 | for (auto &v : F->findConstants()) { |
184 | CHECK(isa<WeightVar>(F->getWeightForNode(v))) |
185 | << "Expected a weight variable" ; |
186 | auto *w = cast<WeightVar>(F->getWeightForNode(v)); |
187 | if (valueNumbers_.count(v)) { |
188 | valueNumbers_[w] = valueNumbers_[v]; |
189 | continue; |
190 | } |
191 | valueNumbers_[v] = std::make_pair(ValueKind::ConstantWeight, valueIdx_); |
192 | valueNumbers_[w] = std::make_pair(ValueKind::ConstantWeight, valueIdx_++); |
193 | } |
194 | |
195 | // Assign numbers to all placeholders. |
196 | for (auto &v : F->findPlaceholders()) { |
197 | CHECK(isa<WeightVar>(F->getWeightForNode(v))) |
198 | << "Expected a weight variable" ; |
199 | auto *w = cast<WeightVar>(F->getWeightForNode(v)); |
200 | if (valueNumbers_.count(w)) { |
201 | continue; |
202 | } |
203 | valueNumbers_[w] = std::make_pair(ValueKind::MutableWeight, valueIdx_++); |
204 | } |
205 | |
206 | // Assign numbers to all activations and tensorviews. |
207 | for (const auto &I : F->getInstrs()) { |
208 | if (auto *A = dyn_cast<AllocActivationInst>(&I)) { |
209 | CHECK(!valueNumbers_.count(A)) |
210 | << "Activation should be defined only once" ; |
211 | valueNumbers_[A] = std::make_pair(ValueKind::Activation, valueIdx_++); |
212 | continue; |
213 | } |
214 | if (auto *A = dyn_cast<TensorViewInst>(&I)) { |
215 | auto *viewOrigin = getOrigin(A); |
216 | auto kind = ValueKind::Activation; |
217 | if (auto *w = dyn_cast<WeightVar>(viewOrigin)) { |
218 | kind = w->isConstant() ? ValueKind::ConstantWeight |
219 | : ValueKind::MutableWeight; |
220 | } |
221 | CHECK(!valueNumbers_.count(A)) |
222 | << "TensorView should be defined only once" ; |
223 | valueNumbers_[A] = std::make_pair(kind, valueIdx_++); |
224 | continue; |
225 | } |
226 | } |
227 | DEBUG_GLOW(for (auto &A |
228 | : valueNumbers_) { |
229 | if (isa<Constant>(A.first)) |
230 | continue; |
231 | auto *v = static_cast<const Value *>(A.first); |
232 | llvm::dbgs() << "Value number for " << v->getName() << ": " |
233 | << A.second.second << "\n" ; |
234 | }); |
235 | } |
236 | |