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/Graph/Graph.h" |
18 | |
19 | #include "glow/IR/IRBuilder.h" |
20 | #include "glow/IR/IRUtils.h" |
21 | |
22 | using namespace glow; |
23 | using llvm::dyn_cast; |
24 | using llvm::isa; |
25 | |
26 | IRBuilder::~IRBuilder() { deallocateActiveInstrs(); } |
27 | |
28 | static bool hasDeallocas(AllocActivationInst *AA) { |
29 | for (auto &U : AA->getUsers()) { |
30 | if (isa<DeallocActivationInst>(U.get())) { |
31 | return true; |
32 | } |
33 | } |
34 | return false; |
35 | } |
36 | |
37 | void IRBuilder::deallocateActiveInstrs() { |
38 | auto &instrs = F_->getInstrs(); |
39 | // Inserts dealloc instructions for all instructions that don't have |
40 | // 'dealloc' as one of their users. |
41 | for (auto &I : instrs) { |
42 | auto AA = dyn_cast<AllocActivationInst>(&I); |
43 | if (!AA) { |
44 | continue; |
45 | } |
46 | |
47 | if (hasDeallocas(AA)) { |
48 | continue; |
49 | } |
50 | |
51 | createDeallocActivationInst("dealloc." + AA->getName().str(), AA); |
52 | } |
53 | } |
54 | |
55 | //===----------------------------------------------------------------------===// |
56 | // High level operators. |
57 | //===----------------------------------------------------------------------===// |
58 | MaxPoolWithArgmaxInst *IRBuilder::createMaxPoolWithArgmaxOp( |
59 | llvm::StringRef name, Value *input, llvm::ArrayRef<unsigned_t> kernels, |
60 | llvm::ArrayRef<unsigned_t> strides, llvm::ArrayRef<unsigned_t> pads, |
61 | unsigned_t layout, ElemKind argMaxIndicesTy) { |
62 | TypeRef outTy{nullptr}; |
63 | Value *argmax{nullptr}; |
64 | |
65 | if (layout == NHWC) { |
66 | ShapeNHWC idim = ShapeNHWC(input->dims()); |
67 | |
68 | auto outSz = |
69 | calculateConvPoolOutputDims(idim.h, idim.w, kernels, strides, pads); |
70 | |
71 | // Allocate storage for flattened NCHW index of max element. |
72 | argmax = |
73 | createAllocActivationInst(name.str() + ".argmax" , argMaxIndicesTy, |
74 | {idim.n, outSz.first, outSz.second, idim.c}); |
75 | |
76 | outTy = F_->getParent()->uniqueTypeWithNewShape( |
77 | input->getType(), {idim.n, outSz.first, outSz.second, idim.c}); |
78 | } else { |
79 | ShapeNCHW idim(input->dims()); |
80 | |
81 | auto outSz = |
82 | calculateConvPoolOutputDims(idim.h, idim.w, kernels, strides, pads); |
83 | |
84 | // Allocate storage for flattened NCHW index of max element. |
85 | argmax = |
86 | createAllocActivationInst(name.str() + ".argmax" , argMaxIndicesTy, |
87 | {idim.n, idim.c, outSz.first, outSz.second}); |
88 | |
89 | outTy = F_->getParent()->uniqueTypeWithNewShape( |
90 | input->getType(), {idim.n, idim.c, outSz.first, outSz.second}); |
91 | } |
92 | |
93 | Value *dest = createAllocActivationInst(name.str() + ".res" , outTy); |
94 | |
95 | return createMaxPoolWithArgmaxInst(name, dest, input, argmax, kernels, |
96 | strides, pads, layout); |
97 | } |
98 | |
99 | ArgMaxInst *IRBuilder::createArgMaxOp(llvm::StringRef name, Value *input, |
100 | unsigned_t axis, bool keepDims, |
101 | ElemKind outIndicesTy) { |
102 | ShapeVector odim = reduceDims(input->dims(), {axis}, keepDims); |
103 | Value *argmax = |
104 | createAllocActivationInst(name.str() + ".argmax" , outIndicesTy, odim); |
105 | return createArgMaxInst(name, argmax, input, axis, keepDims); |
106 | } |
107 | |
108 | AvgPoolInst *IRBuilder::createAvgPoolOp(Value *input, |
109 | llvm::ArrayRef<unsigned_t> kernels, |
110 | llvm::ArrayRef<unsigned_t> strides, |
111 | llvm::ArrayRef<unsigned_t> pads, |
112 | unsigned_t layout, |
113 | bool countIncludePads) { |
114 | |
115 | TypeRef outTy; |
116 | |
117 | if (layout == NHWC) { |
118 | ShapeNHWC idim(input->dims()); |
119 | auto outSz = |
120 | calculateConvPoolOutputDims(idim.h, idim.w, kernels, strides, pads); |
121 | outTy = F_->getParent()->uniqueTypeWithNewShape( |
122 | input->getType(), {idim.n, outSz.first, outSz.second, idim.c}); |
123 | } else { |
124 | ShapeNCHW idim(input->dims()); |
125 | auto outSz = |
126 | calculateConvPoolOutputDims(idim.h, idim.w, kernels, strides, pads); |
127 | outTy = F_->getParent()->uniqueTypeWithNewShape( |
128 | input->getType(), {idim.n, idim.c, outSz.first, outSz.second}); |
129 | } |
130 | |
131 | Value *dest = createAllocActivationInst("pool.res" , outTy); |
132 | return createAvgPoolInst("pool" , dest, input, kernels, strides, pads, layout, |
133 | countIncludePads); |
134 | } |
135 | |
136 | CrossEntropyLossInst *IRBuilder::createCrossEntropyLossOp(llvm::StringRef name, |
137 | Value *p, |
138 | Value *labels) { |
139 | auto *res = |
140 | createAllocActivationInst(name.str() + ".res" , ElemKind::FloatTy, {1}); |
141 | return createCrossEntropyLossInst(name, p, labels, res); |
142 | } |
143 | |
144 | /// Creates a tensorview instruction with the following parameters: |
145 | /// \param elemKind the type of elements in a tensor |
146 | /// \param dims dimensions of the view, such that the number of elements |
147 | /// in the view is the same as the number of elements in the source tensor |
148 | /// \p src |
149 | /// \param src the source tensor used to create the unowned tensor. |
150 | /// \param offsets is a vector of offsets into the Tensor for this view of the |
151 | /// Tensor. |
152 | TensorViewInst *IRBuilder::createTensorView(ElemKind elemKind, |
153 | llvm::ArrayRef<dim_t> dims, |
154 | Value *src, llvm::StringRef name, |
155 | llvm::ArrayRef<dim_t> offsets) { |
156 | auto ty = getIRFunction().getParent()->uniqueType(Type(elemKind, dims)); |
157 | return createTensorViewInst( |
158 | name, src, ty, |
159 | (offsets.size() |
160 | ? offsets |
161 | : llvm::ArrayRef<dim_t>(std::vector<dim_t>(dims.size(), 0)))); |
162 | } |
163 | |
164 | LocalResponseNormalizationInst *IRBuilder::createLocalResponseNormalizationOp( |
165 | llvm::StringRef name, Value *input, size_t halfWindowSize, float alpha, |
166 | float beta, float k) { |
167 | auto ty = input->getType(); |
168 | auto *scale = createAllocActivationInst(name.str() + ".scale" , ty); |
169 | |
170 | // The output tensor is of the same shape as the input tensor. |
171 | auto *res = createAllocActivationInst(name.str() + ".res" , ty); |
172 | return createLocalResponseNormalizationInst(name, res, input, scale, |
173 | halfWindowSize, alpha, beta, k); |
174 | } |
175 | |
176 | TopKInst *IRBuilder::createTopKOp(llvm::StringRef name, Value *input, size_t k, |
177 | ElemKind outIndicesTy) { |
178 | auto inDims = input->dims(); |
179 | assert(inDims.size() > 0); |
180 | assert(k <= inDims.back()); |
181 | assert(outIndicesTy == ElemKind::Int32ITy || |
182 | outIndicesTy == ElemKind::Int64ITy); |
183 | ShapeVector outDims(inDims.begin(), inDims.end()); |
184 | outDims.back() = k; |
185 | auto outTy = |
186 | F_->getParent()->uniqueTypeWithNewShape(input->getType(), outDims); |
187 | auto *values = createAllocActivationInst(name.str() + ".values" , outTy); |
188 | auto *indices = |
189 | createAllocActivationInst(name.str() + ".indices" , outIndicesTy, outDims); |
190 | return createTopKInst(name.str(), values, indices, input, k); |
191 | } |
192 | |
193 | Value *IRBuilder::createReturnOp(Value *input) { |
194 | auto *W = createWeightVar(input->getType(), "result" , |
195 | WeightVar::MutabilityKind::Mutable); |
196 | createCopyInst("return" , W, input); |
197 | return W; |
198 | } |
199 | |
200 | //===----------------------------------------------------------------------===// |
201 | // Low level instructions. |
202 | //===----------------------------------------------------------------------===// |
203 | |
204 | WeightVar *IRBuilder::createWeightVar(ElemKind elemTy, |
205 | llvm::ArrayRef<dim_t> dims, |
206 | llvm::StringRef name, |
207 | WeightVar::MutabilityKind m) { |
208 | auto T = F_->getParent()->uniqueType(elemTy, dims); |
209 | return createWeightVar(T, name, m); |
210 | } |
211 | |
212 | WeightVar *IRBuilder::createWeightVar(TypeRef T, llvm::StringRef name, |
213 | WeightVar::MutabilityKind m) { |
214 | auto *A = new WeightVar(uniqueName(name), T, m); |
215 | F_->getWeights().push_back(A); |
216 | return A; |
217 | } |
218 | |
219 | WeightVar *IRBuilder::createWeightVar(ElemKind elemTy, |
220 | llvm::ArrayRef<dim_t> dims, float scale, |
221 | int32_t offset, llvm::StringRef name, |
222 | WeightVar::MutabilityKind m) { |
223 | auto T = F_->getParent()->uniqueType(elemTy, dims, scale, offset); |
224 | return createWeightVar(T, name, m); |
225 | } |
226 | |
227 | AllocActivationInst * |
228 | IRBuilder::createAllocActivationInst(llvm::StringRef name, ElemKind elemTy, |
229 | llvm::ArrayRef<dim_t> dims) { |
230 | auto T = F_->getParent()->uniqueType(elemTy, dims); |
231 | return createAllocActivationInst(name, T); |
232 | } |
233 | |