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
34using namespace glow;
35using llvm::cast;
36using llvm::dyn_cast;
37using llvm::isa;
38
39void 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
110void 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
153void 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
181void 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