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/IRUtils.h" |
18 | #include "glow/IR/Instrs.h" |
19 | |
20 | #include "llvm/Support/Casting.h" |
21 | |
22 | using namespace glow; |
23 | using llvm::dyn_cast; |
24 | using llvm::isa; |
25 | |
26 | bool glow::isTensorView(glow::Value *v) { return isa<TensorViewInst>(v); } |
27 | |
28 | size_t glow::calculateTensorViewOffset(const TensorViewInst *TVI) { |
29 | // Pop tensor views off repeatedly until we reach the origin, in case there |
30 | // are multiple stacked together, to calculate the total offset. |
31 | const TensorViewInst *currTVI = TVI; |
32 | size_t totalOffsetLength = 0; |
33 | do { |
34 | // Get the offset into the current base Tensor in bytes. Aggregate all |
35 | // offsets from stacked TVIs into totalOffsetLength. |
36 | totalOffsetLength += |
37 | getFlattenedOffset(currTVI->getSrc()->getType()->strides(), |
38 | currTVI->getOffsets()) * |
39 | currTVI->getType()->getElementSize(); |
40 | } while ((currTVI = dyn_cast<TensorViewInst>(currTVI->getSrc()))); |
41 | |
42 | return totalOffsetLength; |
43 | } |
44 | |
45 | Value *glow::getAllocationOrigin(Value *V) { |
46 | while (true) { |
47 | if (auto *AI = dyn_cast<AllocActivationInst>(V)) |
48 | return AI; |
49 | if (auto *TVI = dyn_cast<TensorViewInst>(V)) { |
50 | V = TVI->getSrc(); |
51 | continue; |
52 | } |
53 | return nullptr; |
54 | } |
55 | return nullptr; |
56 | } |
57 | |
58 | Value *glow::getOrigin(Value *V) { |
59 | return const_cast<Value *>(getOrigin(const_cast<const Value *>(V))); |
60 | } |
61 | |
62 | const Value *glow::getOrigin(const Value *V) { |
63 | while (true) { |
64 | auto *TVI = dyn_cast<TensorViewInst>(V); |
65 | if (!TVI) |
66 | return V; |
67 | V = TVI->getSrc(); |
68 | } |
69 | return V; |
70 | } |
71 | |
72 | size_t glow::getOriginOffset(Value *V) { |
73 | size_t off = 0; |
74 | |
75 | // Since each TensorView can either maintain or decrease the size of a buffer, |
76 | // the linearized offset of a TensorView into its underlying buffer can be |
77 | // computed by adding together the linearized offsets of intermediate |
78 | // TensorViews relative to their sources. |
79 | while (true) { |
80 | auto *TVI = dyn_cast<TensorViewInst>(V); |
81 | if (!TVI) { |
82 | return off; |
83 | } |
84 | |
85 | llvm::ArrayRef<dim_t> offsets = TVI->getOffsets(); |
86 | llvm::ArrayRef<dim_t> srcDims = TVI->getSrc()->getType()->dims(); |
87 | |
88 | size_t numSrcDims = srcDims.size(); |
89 | |
90 | // Iterate backwards in order to figure out how big the slices corresponding |
91 | // to each offset element are. |
92 | for (size_t i = 0, j = numSrcDims - 1; i < numSrcDims; ++i, --j) { |
93 | // For each offset, add into the linearized offset the offset value |
94 | // multiplied by the slice size corresponding to that offset. |
95 | off += offsets[j] * TVI->getSrc()->getType()->getSliceSize(j); |
96 | } |
97 | |
98 | // Move on to the src of the TensorView that was just processed. |
99 | V = TVI->getSrc(); |
100 | } |
101 | |
102 | return off; |
103 | } |
104 | |