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/IR/IRGen.h" |
18 | |
19 | #include "llvm/ADT/StringExtras.h" |
20 | #include "llvm/Support/Casting.h" |
21 | |
22 | #include <unordered_map> |
23 | #include <unordered_set> |
24 | |
25 | //===----------------------------------------------------------------------===// |
26 | // IRGen visitor - the code that generates the IR. |
27 | //===----------------------------------------------------------------------===// |
28 | |
29 | using namespace glow; |
30 | |
31 | using llvm::cast; |
32 | using llvm::dyn_cast; |
33 | using llvm::isa; |
34 | |
35 | #define DECORATE_NODE_NAME(Node, ...) \ |
36 | llvm::join_items("_", Node->getName(), __VA_ARGS__) |
37 | |
38 | /// Helper function that \returns the number of times the same consecutive |
39 | /// NodeValue in \p inputs is found, starting from index \p i. |
40 | static size_t getConsecutiveSameNodeCount(NodeValueArrayRef inputs, |
41 | const size_t i) { |
42 | assert(i < inputs.size() && "Index must fit inside the size of the inputs." ); |
43 | for (size_t j = i, e = inputs.size(); j < e; j++) { |
44 | if (inputs[i] != inputs[j]) { |
45 | return j - i; |
46 | } |
47 | } |
48 | return inputs.size() - i; |
49 | } |
50 | |
51 | bool IRGenVisitor::shouldVisit(Node *parent, Node *N) { |
52 | // Don't revisit nodes that we've already processed. |
53 | return !visited_.count(N); |
54 | } |
55 | |
56 | /// \returns the generated instruction for the node \p N. |
57 | Value *IRGenVisitor::valueForNode(NodeValue N) { |
58 | if (auto *V = dyn_cast<Storage>(N)) { |
59 | auto &map = F_->getVariableMap(); |
60 | return map[V]; |
61 | } |
62 | auto it = generatedNodeDest_.find(N); |
63 | assert(it != generatedNodeDest_.end() && "IR was not generated for the node" ); |
64 | return it->second; |
65 | } |
66 | /// Saves the generated IR in \p v for the node \p N. |
67 | void IRGenVisitor::registerIR(NodeValue N, Value *v) { |
68 | if (auto *V = dyn_cast<Storage>(N)) { |
69 | auto &map = F_->getVariableMap(); |
70 | map[V] = v; |
71 | return; |
72 | } |
73 | assert(!generatedNodeDest_.count(N) && |
74 | "Already generated code for this node" ); |
75 | assert(isa<AllocActivationInst>(v) && "The value must be an activation" ); |
76 | generatedNodeDest_[N] = v; |
77 | } |
78 | |
79 | /// Adds to Node \p N --> Instruction \p inst map. |
80 | void IRGenVisitor::setNodeToIR(Node *N, Instruction *inst) { |
81 | nodeToInstr_[N] = inst; |
82 | } |
83 | |
84 | /// Return Instruction that is mapped to Node \p N. |
85 | /// If mapping doesn't exists returns nullptr. |
86 | Instruction *IRGenVisitor::getNodeToIR(Node *N) { |
87 | Instruction *retNode = nullptr; |
88 | auto iterInst = nodeToInstr_.find(N); |
89 | if (iterInst != nodeToInstr_.end()) |
90 | retNode = iterInst->second; |
91 | |
92 | return retNode; |
93 | } |
94 | |
95 | void IRGenVisitor::post(Node *parent, Node *N) { |
96 | visited_.insert(N); |
97 | |
98 | // Allows backend to generate their custom instrution IR. |
99 | if (B_.generateInst(N, *this)) { |
100 | return; |
101 | } |
102 | |
103 | switch (N->getKind()) { |
104 | default: |
105 | llvm_unreachable("Unhandled node; perhaps the node should have been " |
106 | "lowered, or the backend should have specified an IRGen " |
107 | "case for this node to a backend-specific Instr." ); |
108 | break; |
109 | |
110 | // Include all automatically generated cases: |
111 | #include "glow/AutoGenIRGen.h" |
112 | |
113 | case glow::Kinded::Kind::ReshapeNodeKind: { |
114 | auto *RN = cast<ReshapeNode>(N); |
115 | |
116 | auto *inVal = valueForNode(RN->getInput()); |
117 | std::vector<dim_t> offsets(inVal->getType()->dims().size(), 0); |
118 | auto *TVI = builder_.createTensorViewInst( |
119 | DECORATE_NODE_NAME(N, "tensorview" ), inVal, RN->getResult().getType(), |
120 | offsets); |
121 | auto *dest = builder_.createAllocActivationInst( |
122 | DECORATE_NODE_NAME(N, "res" ), RN->getResult().getType()); |
123 | builder_.createCopyInst(DECORATE_NODE_NAME(N, "copy" ), dest, TVI); |
124 | registerIR(N, dest); |
125 | break; |
126 | } |
127 | case glow::Kinded::Kind::ConvolutionGradNodeKind: { |
128 | auto *CG = cast<ConvolutionGradNode>(N); |
129 | |
130 | auto *input = valueForNode(CG->getInput()); |
131 | auto *filter = valueForNode(CG->getFilter()); |
132 | auto *bias = valueForNode(CG->getBias()); |
133 | |
134 | auto *outGrad = valueForNode(CG->getGradOfOriginalOutputNamedResult()); |
135 | |
136 | auto *inG = builder_.createAllocActivationInst( |
137 | DECORATE_NODE_NAME(N, "input" , "grad" ), input->getType()); |
138 | auto *biasG = builder_.createAllocActivationInst( |
139 | DECORATE_NODE_NAME(N, "bias" , "grad" ), bias->getType()); |
140 | auto *filterG = builder_.createAllocActivationInst( |
141 | DECORATE_NODE_NAME(N, "filter" , "grad" ), filter->getType()); |
142 | |
143 | builder_.createConvolutionGradInst( |
144 | N->getName(), input, filter, outGrad, inG, filterG, biasG, |
145 | CG->getKernels(), CG->getStrides(), CG->getPads(), CG->getGroup(), |
146 | CG->getDilation(), CG->getLayout(), CG->getFusedActivation(), |
147 | CG->getFusedActivationArgs()); |
148 | |
149 | registerIR(CG->getGradOfInputNamedInput(), inG); |
150 | registerIR(CG->getGradOfInputNamedFilter(), filterG); |
151 | registerIR(CG->getGradOfInputNamedBias(), biasG); |
152 | break; |
153 | } |
154 | case glow::Kinded::Kind::MaxPoolNodeKind: { |
155 | auto *P = cast<MaxPoolNode>(N); |
156 | auto *in = valueForNode(P->getInput()); |
157 | auto argMax = P->getArgmax(); |
158 | auto *V = builder_.createMaxPoolWithArgmaxOp( |
159 | N->getName(), in, P->getKernels(), P->getStrides(), P->getPads(), |
160 | P->getLayout(), argMax.getElementType()); |
161 | Value *dest = V->getDest(); |
162 | Value *argmax = V->getArgmax(); |
163 | nodeToInstr_[N] = V; |
164 | registerIR(P->getResult(), dest); |
165 | registerIR(P->getArgmax(), argmax); |
166 | break; |
167 | } |
168 | case glow::Kinded::Kind::MaxPoolGradNodeKind: { |
169 | auto *PG = cast<MaxPoolGradNode>(N); |
170 | |
171 | auto poolIn = PG->getInput(); |
172 | auto poolOut = PG->getOriginalOutputForResult(); |
173 | auto *inW = valueForNode(poolIn); |
174 | auto *outW = valueForNode(poolOut); |
175 | auto *outG = valueForNode(PG->getGradOfOriginalOutputNamedResult()); |
176 | |
177 | auto *inG = builder_.createAllocActivationInst( |
178 | DECORATE_NODE_NAME(N, "outG" ), PG->getInput().getType()); |
179 | |
180 | // Find the original pool instruction. |
181 | assert(nodeToInstr_.count(poolOut) && "Pool IRgen did not register itself" ); |
182 | auto *PI = cast<MaxPoolWithArgmaxInst>(nodeToInstr_[poolOut.getNode()]); |
183 | |
184 | builder_.createMaxPoolWithArgmaxGradInst( |
185 | N->getName(), outW, inW, PI->getArgmax(), outG, inG, PG->getKernels(), |
186 | PG->getStrides(), PG->getPads(), PG->getLayout()); |
187 | registerIR(PG->getGradOfInputNamedInput(), inG); |
188 | break; |
189 | } |
190 | case glow::Kinded::Kind::AvgPoolGradNodeKind: { |
191 | auto *PG = cast<AvgPoolGradNode>(N); |
192 | |
193 | auto poolIn = PG->getInput(); |
194 | auto poolOut = PG->getOriginalOutputForResult(); |
195 | auto *inW = valueForNode(poolIn); |
196 | auto *outW = valueForNode(poolOut); |
197 | auto *outG = valueForNode(PG->getGradOfOriginalOutputNamedResult()); |
198 | |
199 | auto *inG = builder_.createAllocActivationInst( |
200 | DECORATE_NODE_NAME(N, "outG" ), PG->getInput().getType()); |
201 | |
202 | builder_.createAvgPoolGradInst( |
203 | N->getName(), outW, inW, outG, inG, PG->getKernels(), PG->getStrides(), |
204 | PG->getPads(), PG->getLayout(), PG->getCountIncludePads()); |
205 | registerIR(PG->getGradOfInputNamedInput(), inG); |
206 | break; |
207 | } |
208 | case glow::Kinded::Kind::AdaptiveAvgPoolGradNodeKind: { |
209 | auto *PG = cast<AdaptiveAvgPoolGradNode>(N); |
210 | |
211 | auto poolOut = PG->getOriginalOutputForResult(); |
212 | auto *outW = valueForNode(poolOut); |
213 | auto *outG = valueForNode(PG->getGradOfOriginalOutputNamedResult()); |
214 | |
215 | auto *inG = builder_.createAllocActivationInst( |
216 | DECORATE_NODE_NAME(N, "outG" ), PG->getInput().getType()); |
217 | |
218 | builder_.createAdaptiveAvgPoolGradInst(N->getName(), outW, outG, inG); |
219 | registerIR(PG->getGradOfInputNamedInput(), inG); |
220 | break; |
221 | } |
222 | case glow::Kinded::Kind::SoftMaxGradNodeKind: { |
223 | auto *SMG = cast<SoftMaxGradNode>(N); |
224 | // Original inputs: |
225 | auto *origIn = valueForNode(SMG->getInput()); |
226 | auto *origSelect = valueForNode(SMG->getSelected()); |
227 | // Values related to the output of the node. |
228 | auto *outGrad = valueForNode(SMG->getGradOfOriginalOutputNamedResult()); |
229 | auto originalNodeResult = SMG->getOriginalOutputForResult(); |
230 | assert(nodeToInstr_.count(originalNodeResult.getNode()) && |
231 | "Unknown original node" ); |
232 | auto *origOut = valueForNode(originalNodeResult); |
233 | auto *srcGrad = builder_.createAllocActivationInst( |
234 | DECORATE_NODE_NAME(N, "res" ), outGrad->getType()); |
235 | auto *SMGI = builder_.createSoftMaxGradInst(N->getName(), origOut, origIn, |
236 | origSelect, srcGrad); |
237 | registerIR(SMG->getGradOfInputNamedInput(), SMGI->getSrcGrad()); |
238 | break; |
239 | } |
240 | case glow::Kinded::Kind::CrossEntropyLossNodeKind: { |
241 | auto *CELoss = cast<CrossEntropyLossNode>(N); |
242 | auto *P = valueForNode(CELoss->getP()); |
243 | auto *Labels = valueForNode(CELoss->getLabels()); |
244 | auto *V = builder_.createCrossEntropyLossOp(N->getName(), P, Labels); |
245 | registerIR(N, V->getCE()); |
246 | nodeToInstr_[N] = V; |
247 | break; |
248 | } |
249 | case glow::Kinded::Kind::CrossEntropyLossGradNodeKind: { |
250 | auto *CELossG = cast<CrossEntropyLossGradNode>(N); |
251 | // Forward pass inputs. |
252 | auto *P = valueForNode(CELossG->getP()); |
253 | auto *Y = valueForNode(CELossG->getLabels()); |
254 | // Backward pass gradient dL/dY. |
255 | auto *dY = valueForNode(CELossG->getGradOfOriginalOutputNamedCE()); |
256 | auto *pGrad = builder_.createAllocActivationInst( |
257 | DECORATE_NODE_NAME(N, "p" , "grad" ), P->getType()); |
258 | auto *yGrad = builder_.createAllocActivationInst( |
259 | DECORATE_NODE_NAME(N, "labels" , "grad" ), Y->getType()); |
260 | auto *CELossGI = builder_.createCrossEntropyLossGradInst( |
261 | N->getName(), dY, P, Y, pGrad, yGrad); |
262 | registerIR(CELossG->getGradOfInputNamedP(), CELossGI->getPgrad()); |
263 | registerIR(CELossG->getGradOfInputNamedLabels(), CELossGI->getLabelsgrad()); |
264 | break; |
265 | } |
266 | case glow::Kinded::Kind::ConcatNodeKind: { |
267 | auto *CC = cast<ConcatNode>(N); |
268 | |
269 | auto *dest = builder_.createAllocActivationInst(CC->getName(), |
270 | CC->getResult().getType()); |
271 | // Mark the buffer as initialized, this is safe since the InsertTensors |
272 | // below will fully overwrite the buffer. |
273 | builder_.createTouchInst(CC->getName(), dest); |
274 | auto inputs = CC->getInputs(); |
275 | |
276 | // We start inserting to the shape at (0,0, ... ). |
277 | std::vector<dim_t> offsets(CC->getResult().dims().size(), 0); |
278 | unsigned dim = CC->getDim(); |
279 | |
280 | for (size_t i = 0, e = inputs.size(); i < e;) { |
281 | // Look for a series of the same Node being concated consecutively many |
282 | // times. We can wrap n such consecutive repeats into a single insert |
283 | // with count n along the dim axis. |
284 | const size_t consecutiveCount = getConsecutiveSameNodeCount(inputs, i); |
285 | |
286 | // Create the new InsertTensor instruction given the input node, along |
287 | // with the number of times to insert the node and the axis (dim) we are |
288 | // inserting in. |
289 | builder_.createInsertTensorInst( |
290 | DECORATE_NODE_NAME(CC, inputs[i].getNode()->getName()), dest, |
291 | valueForNode(inputs[i]), offsets, consecutiveCount, dim); |
292 | |
293 | // We are stacking the tensors along a specific dimension. This means |
294 | // that we increase the size of the tensor along this dimension, count |
295 | // times. |
296 | offsets[dim] += inputs[i].dims()[dim] * consecutiveCount; |
297 | |
298 | // Increment i by the number of the same nodes that were found in a row, |
299 | // which were all wrapped into a single InsertTensorInst. |
300 | i += consecutiveCount; |
301 | } |
302 | registerIR(N, dest); |
303 | break; |
304 | } |
305 | case glow::Kinded::Kind::CollectRpnProposalsNodeKind: { |
306 | auto *CRPN = llvm::cast<CollectRpnProposalsNode>(N); |
307 | |
308 | std::string allocName = std::string(CRPN->getName()) + ".res" ; |
309 | auto *dest = builder_.createAllocActivationInst( |
310 | allocName, CRPN->getResult().getType()); |
311 | |
312 | auto *inst = builder_.createCollectRpnProposalsInst( |
313 | CRPN->getName(), dest, CRPN->getRpnMaxLevel(), CRPN->getRpnMinLevel(), |
314 | CRPN->getRpnPostNmsTopN()); |
315 | |
316 | // Adding inputs to instruction |
317 | for (auto &in : CRPN->getRoisIn()) { |
318 | inst->pushOperand({valueForNode(in), OperandKind::In}); |
319 | } |
320 | |
321 | for (auto &in : CRPN->getRoisProbsIn()) { |
322 | inst->pushOperand({valueForNode(in), OperandKind::In}); |
323 | } |
324 | |
325 | registerIR(CRPN->getResult(), dest); |
326 | break; |
327 | } |
328 | case glow::Kinded::Kind::SliceNodeKind: { |
329 | auto *SL = cast<SliceNode>(N); |
330 | auto start = SL->getStart(); |
331 | auto *in = valueForNode(SL->getInput()); |
332 | auto *dest = builder_.createAllocActivationInst(SL->getName(), |
333 | SL->getResult().getType()); |
334 | builder_.createExtractTensorInst(SL->getName(), dest, in, start); |
335 | registerIR(N, dest); |
336 | break; |
337 | } |
338 | case glow::Kinded::Kind::InsertTensorNodeKind: { |
339 | auto *IT = cast<InsertTensorNode>(N); |
340 | auto start = IT->getStart(); |
341 | auto count = IT->getCount(); |
342 | auto axis = IT->getAxis(); |
343 | auto *big = valueForNode(IT->getBig()); |
344 | auto *small = valueForNode(IT->getSmall()); |
345 | auto *dest = builder_.createAllocActivationInst(IT->getName(), |
346 | IT->getResult().getType()); |
347 | if (small->getSizeInBytes() * count < big->getSizeInBytes()) { |
348 | builder_.createCopyInst(DECORATE_NODE_NAME(N, "copy" ), dest, big); |
349 | } else { |
350 | // Small tensor completely fills the big tensor, thus no need to |
351 | // initialize the destination. |
352 | builder_.createTouchInst(DECORATE_NODE_NAME(N, "init" ), dest); |
353 | } |
354 | builder_.createInsertTensorInst(IT->getName(), dest, small, start, count, |
355 | axis); |
356 | |
357 | registerIR(N, dest); |
358 | break; |
359 | } |
360 | case glow::Kinded::Kind::ScatterDataNodeKind: { |
361 | auto *SDI = cast<ScatterDataNode>(N); |
362 | auto *dataTensor = valueForNode(SDI->getData()); |
363 | auto *indicesTensor = valueForNode(SDI->getIndices()); |
364 | auto *slicesTensor = valueForNode(SDI->getSlices()); |
365 | auto *dest = builder_.createAllocActivationInst(SDI->getName(), |
366 | SDI->getResult().getType()); |
367 | builder_.createCopyInst(DECORATE_NODE_NAME(N, "copy" ), dest, dataTensor); |
368 | builder_.createScatterDataInst(SDI->getName(), dest, indicesTensor, |
369 | slicesTensor, SDI->getCumulative()); |
370 | registerIR(N, dest); |
371 | break; |
372 | } |
373 | case glow::Kinded::Kind::LocalResponseNormalizationNodeKind: { |
374 | auto *LR = cast<LocalResponseNormalizationNode>(N); |
375 | auto *in = valueForNode(LR->getInput()); |
376 | auto *V = builder_.createLocalResponseNormalizationOp( |
377 | N->getName(), in, LR->getHalfWindowSize(), LR->getAlpha(), |
378 | LR->getBeta(), LR->getK()); |
379 | nodeToInstr_[N] = V; |
380 | registerIR(N, V->getDest()); |
381 | break; |
382 | } |
383 | |
384 | case glow::Kinded::Kind::LocalResponseNormalizationGradNodeKind: { |
385 | auto *LRG = cast<LocalResponseNormalizationGradNode>(N); |
386 | auto *origIn = valueForNode(LRG->getInput()); |
387 | |
388 | auto originalNodeResult = LRG->getOriginalOutputForResult(); |
389 | assert(nodeToInstr_.count(originalNodeResult.getNode()) && |
390 | "Unknown original node" ); |
391 | auto *LRI = |
392 | cast<LocalResponseNormalizationInst>(nodeToInstr_[originalNodeResult]); |
393 | |
394 | auto *srcGrad = builder_.createAllocActivationInst( |
395 | DECORATE_NODE_NAME(N, "res" , "grad" ), origIn->getType()); |
396 | |
397 | builder_.createLocalResponseNormalizationGradInst( |
398 | N->getName(), valueForNode(LRG->getOriginalOutputForResult()), |
399 | valueForNode(LRG->getInput()), LRI->getScale(), |
400 | valueForNode(LRG->getGradOfOriginalOutputNamedResult()), srcGrad, |
401 | LRG->getHalfWindowSize(), LRG->getAlpha(), LRG->getBeta(), LRG->getK()); |
402 | |
403 | registerIR(LRG->getGradOfInputNamedInput(), srcGrad); |
404 | break; |
405 | } |
406 | case glow::Kinded::Kind::SaveNodeKind: { |
407 | auto *R = cast<SaveNode>(N); |
408 | auto *src = valueForNode(R->getInput()); |
409 | auto *dest = valueForNode(R->getOutput()); |
410 | builder_.createCopyInst(N->getName(), dest, src); |
411 | break; |
412 | } |
413 | case glow::Kinded::Kind::ConstantKind: { |
414 | auto *V = cast<Constant>(N); |
415 | auto *W = builder_.createWeightVar(V->getType(), V->getName(), |
416 | WeightVar::MutabilityKind::Constant); |
417 | registerIR(N, W); |
418 | break; |
419 | } |
420 | case glow::Kinded::Kind::PlaceholderKind: { |
421 | auto *P = cast<Placeholder>(N); |
422 | auto *W = builder_.createWeightVar(P->getType(), P->getName(), |
423 | WeightVar::MutabilityKind::Mutable); |
424 | registerIR(N, W); |
425 | break; |
426 | } |
427 | case glow::Kinded::Kind::QuantizationProfileNodeKind: { |
428 | auto *QPN = cast<QuantizationProfileNode>(N); |
429 | auto *inputTensor = valueForNode(QPN->getInput()); |
430 | auto *histogram = valueForNode(QPN->getHistogramPlaceholder()); |
431 | auto *computationInfo = valueForNode(QPN->getComputationInfoPlaceholder()); |
432 | builder_.createQuantizationProfileInst(QPN->getName(), inputTensor, |
433 | histogram, computationInfo); |
434 | break; |
435 | } |
436 | case glow::Kinded::Kind::TopKNodeKind: { |
437 | auto *TKN = cast<TopKNode>(N); |
438 | auto *inputTensor = valueForNode(TKN->getInput()); |
439 | auto k = TKN->getK(); |
440 | auto *V = builder_.createTopKOp(N->getName(), inputTensor, k, |
441 | TKN->getIndices().getElementType()); |
442 | registerIR(TKN->getValues(), V->getValues()); |
443 | registerIR(TKN->getIndices(), V->getIndices()); |
444 | break; |
445 | } |
446 | case glow::Kinded::Kind::TraceEventNodeKind: { |
447 | auto *TEN = cast<TraceEventNode>(N); |
448 | auto *dataTensor = valueForNode(TEN->getData()); |
449 | builder_.createTraceEventInst(TEN->getName(), dataTensor, TEN->getIndex()); |
450 | break; |
451 | } |
452 | case glow::Kinded::Kind::SparseLengthsSumGradNodeKind: { |
453 | auto *SLSG = cast<SparseLengthsSumGradNode>(N); |
454 | |
455 | auto *data = valueForNode(SLSG->getData()); |
456 | auto *indices = valueForNode(SLSG->getIndices()); |
457 | auto *lengths = valueForNode(SLSG->getLengths()); |
458 | |
459 | auto *destGrad = valueForNode(SLSG->getGradOfOriginalOutputNamedResult()); |
460 | auto *dataGrad = builder_.createAllocActivationInst( |
461 | DECORATE_NODE_NAME(N, "dataG" ), |
462 | SLSG->getGradOfInputNamedData().getType()); |
463 | |
464 | builder_.createSparseLengthsSumGradInst( |
465 | N->getName(), data, indices, lengths, destGrad, dataGrad, |
466 | SLSG->getLengthsMode(), SLSG->getAvgLength()); |
467 | |
468 | registerIR(SLSG->getGradOfInputNamedData(), dataGrad); |
469 | break; |
470 | } |
471 | case glow::Kinded::Kind::SparseLengthsWeightedSumGradNodeKind: { |
472 | auto *SLWSG = cast<SparseLengthsWeightedSumGradNode>(N); |
473 | |
474 | auto *data = valueForNode(SLWSG->getData()); |
475 | auto *weights = valueForNode(SLWSG->getWeights()); |
476 | auto *indices = valueForNode(SLWSG->getIndices()); |
477 | auto *lengths = valueForNode(SLWSG->getLengths()); |
478 | |
479 | auto *destGrad = valueForNode(SLWSG->getGradOfOriginalOutputNamedResult()); |
480 | auto *dataGrad = builder_.createAllocActivationInst( |
481 | DECORATE_NODE_NAME(N, "dataG" ), |
482 | SLWSG->getGradOfInputNamedData().getType()); |
483 | auto *weightsGrad = builder_.createAllocActivationInst( |
484 | DECORATE_NODE_NAME(N, "weightsG" ), |
485 | SLWSG->getGradOfInputNamedWeights().getType()); |
486 | |
487 | builder_.createSparseLengthsWeightedSumGradInst( |
488 | N->getName(), data, weights, indices, lengths, destGrad, dataGrad, |
489 | weightsGrad, SLWSG->getLengthsMode(), SLWSG->getAvgLength()); |
490 | |
491 | registerIR(SLWSG->getGradOfInputNamedData(), dataGrad); |
492 | registerIR(SLWSG->getGradOfInputNamedWeights(), weightsGrad); |
493 | break; |
494 | } |
495 | case glow::Kinded::Kind::BatchedPairwiseDotProductNodeKind: { |
496 | auto *BPDPN = llvm::cast<BatchedPairwiseDotProductNode>(N); |
497 | auto firstInput = BPDPN->getInputs()[0]; |
498 | |
499 | std::string allocName = std::string(BPDPN->getName()) + ".res" ; |
500 | auto *dest = builder_.createAllocActivationInst( |
501 | allocName, BPDPN->getResult().getType()); |
502 | |
503 | auto *inst = builder_.createBatchedPairwiseDotProductInst( |
504 | BPDPN->getName(), dest, BPDPN->getInputs().size(), |
505 | firstInput.getType()->dims()[1]); |
506 | |
507 | // First instruction operand is the buffer to write the dot products, the |
508 | // rest are all inputs. |
509 | for (auto &in : BPDPN->getInputs()) { |
510 | inst->pushOperand({valueForNode(in), OperandKind::In}); |
511 | } |
512 | |
513 | registerIR(BPDPN->getResult(), dest); |
514 | break; |
515 | } |
516 | |
517 | case glow::Kinded::Kind::BatchedPairwiseDotProductGradNodeKind: { |
518 | auto *BPDPGN = llvm::cast<BatchedPairwiseDotProductGradNode>(N); |
519 | |
520 | auto *in0 = valueForNode(BPDPGN->getOriginalInputs()[0]); |
521 | auto *outputGrad = valueForNode(BPDPGN->getOutputGrad()); |
522 | |
523 | // First, create alloc instructions for all of the gradients. This needs to |
524 | // be done first so that these instructions precede the first use of the |
525 | // buffers they create. |
526 | std::vector<Value *> dests; |
527 | for (unsigned i = 0, e = BPDPGN->getNumResults(); i < e; ++i) { |
528 | NodeValue res = BPDPGN->getNthResult(i); |
529 | std::string allocName = |
530 | std::string(BPDPGN->getName()) + ".res." + std::to_string(i); |
531 | auto *dest = builder_.createAllocActivationInst(allocName, res.getType()); |
532 | dests.emplace_back(dest); |
533 | } |
534 | |
535 | auto *inst = builder_.createBatchedPairwiseDotProductGradInst( |
536 | BPDPGN->getName(), outputGrad, BPDPGN->getOriginalInputs().size(), |
537 | in0->dims()[1]); |
538 | |
539 | // Operands 1 -> numInputs are gradients. |
540 | for (unsigned i = 0, e = BPDPGN->getNumResults(); i < e; ++i) { |
541 | NodeValue res = BPDPGN->getNthResult(i); |
542 | inst->pushOperand({dests[i], OperandKind::Out}); |
543 | registerIR(res, dests[i]); |
544 | } |
545 | |
546 | // Operands numInputs + 1 -> 2 * numInputs are original inputs. |
547 | for (auto &in : BPDPGN->getOriginalInputs()) { |
548 | inst->pushOperand({valueForNode(in), OperandKind::In}); |
549 | } |
550 | break; |
551 | } |
552 | case glow::Kinded::Kind::ExternalFunctionCallNodeKind: { |
553 | auto *EFCN = llvm::cast<ExternalFunctionCallNode>(N); |
554 | std::string externalCallType = std::string(EFCN->getName()); |
555 | std::string allocName = std::string(EFCN->getName()) + ".res" ; |
556 | auto *dest = builder_.createAllocActivationInst( |
557 | allocName, EFCN->getResult().getType()); |
558 | |
559 | auto *inst = builder_.createExternalFunctionCallInst( |
560 | EFCN->getName(), dest, EFCN->getFunctionName(), EFCN->getFunctionImpl(), |
561 | EFCN->getFunctionKind()); |
562 | |
563 | // First instruction operand is the buffer for the result, the |
564 | // rest are all inputs. |
565 | for (auto &in : EFCN->getInputs()) { |
566 | inst->pushOperand({valueForNode(in), OperandKind::In}); |
567 | } |
568 | registerIR(EFCN->getResult(), dest); |
569 | break; |
570 | } |
571 | } |
572 | } |
573 | |
574 | void IRFunction::generateIR(const Backend &B) { |
575 | assert(getGraph()->verify(&B) && "Invalid function" ); |
576 | // Schedule the nodes. |
577 | NodesPtrList ScheduledNodes; |
578 | scheduleGraph(ScheduledNodes); |
579 | IRGenVisitor irgen(this, B); |
580 | |
581 | for (auto &N : ScheduledNodes) { |
582 | N->visit(nullptr, &irgen); |
583 | } |
584 | |
585 | if (!B.verify(*this)) { |
586 | EXIT_ON_ERR( |
587 | MAKE_ERR(ErrorValue::ErrorCode::COMPILE_UNSUPPORTED_IR_AFTER_GENERATE, |
588 | "Unsupported instruction(s) found after generating IR " + |
589 | getName().str() + " for backend " + B.getBackendName())); |
590 | } |
591 | } |
592 | |