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#ifndef GLOW_LLVMIRCODEGEN_ALLOCATIONSINFO_H
17#define GLOW_LLVMIRCODEGEN_ALLOCATIONSINFO_H
18
19#include "glow/Backend/BackendUtils.h"
20#include "glow/CodeGen/MemoryAllocator.h"
21#include "glow/Graph/Nodes.h"
22#include "llvm/IR/Module.h"
23
24#include <functional>
25
26namespace glow {
27class Value;
28class IRFunction;
29class WeightVar;
30class Constant;
31class PlaceholderBindings;
32
33/// Information about allocations for activations, constant weight variables
34/// and mutable weight variables.
35class AllocationsInfo {
36public:
37 /// Different kinds of values that need to be allocated.
38 enum class ValueKind { ConstantWeight, MutableWeight, Activation };
39 using KindAndNumber = std::pair<ValueKind, size_t>;
40 /// Map Values in the module to their numbers.
41 llvm::DenseMap<const Kinded *, KindAndNumber> valueNumbers_;
42 /// To get the offset of a given value simply use
43 /// numberOffsets_[valueNumbers_[v]]
44
45 /// Maps Values in the module to their offsets.
46 llvm::DenseMap<const Kinded *, uint64_t> allocatedAddress_;
47 /// Amount of memory to be allocated for constant WeightVars.
48 size_t constantWeightVarsMemSize_{0};
49 /// Amount of memory to be allocated for mutable WeightVars.
50 size_t mutableWeightVarsMemSize_{0};
51 /// Amount of memory to be allocated for activations.
52 size_t activationsMemSize_{0};
53 /// Base address of stored constant weights.
54 uint8_t *baseConstantWeightVarsStore_{nullptr};
55
56 /// Ctor.
57 AllocationsInfo()
58 : constantWeightVarsAllocator_("ConstantWeights", 0),
59 mutableWeightVarsAllocator_("MutableWeights", 0),
60 activationsAllocator_("Activations", 0) {}
61 virtual ~AllocationsInfo() = default;
62 /// Assign offsets to all of the variables in the module \p M and to the
63 /// placeholders.
64 virtual void allocateWeightVars(const IRFunction *F);
65 /// Assign offsets to all activations.
66 /// No actual memory allocation is performed. All the allocations should be
67 /// performed by the client based on the information provided by the
68 /// AllocationsInfo or RuntimeBundle.
69 virtual void allocateActivations(const IRFunction *F);
70 /// Assign offsets to all tensorviews.
71 /// No memory allocation is performed. Sets up all offsets into already
72 /// defined offsets for WeightVars and AllocActivations. Assumes the weight
73 /// vars and alloc activations have already been added to allocatedAddress_.
74 virtual void allocateTensorViews(const IRFunction *F);
75 /// Number all allocations and weight variables by assigning them unique
76 /// numbers.
77 virtual void numberValues(const IRFunction *F);
78 /// Getters for allocators.
79 MemoryAllocator &getConstantWeightVarsAllocator() {
80 return constantWeightVarsAllocator_;
81 }
82 MemoryAllocator &getMutableWeightVarsAllocator() {
83 return mutableWeightVarsAllocator_;
84 }
85 MemoryAllocator &getActivationsAllocator() { return activationsAllocator_; }
86
87 const glow::runtime::SymbolTableTy &getSymbolTable() const {
88 return symbolTable_;
89 }
90
91protected:
92 /// Index to be used for a new value.
93 size_t valueIdx_{0};
94 /// Use two different allocators, because constant weights and mutable weights
95 /// may use different memory blocks.
96 MemoryAllocator constantWeightVarsAllocator_;
97 MemoryAllocator mutableWeightVarsAllocator_;
98 /// Use a memory allocator with no upper bound on how much memory we can
99 /// allocate.
100 MemoryAllocator activationsAllocator_;
101 /// Symbol table for the allocated symbols.
102 glow::runtime::SymbolTableTy symbolTable_;
103};
104
105} // namespace glow
106#endif // GLOW_LLVMIRCODEGEN_ALLOCATIONSINFO_H
107