1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file relay/backend/token_allocator.h
22 * \brief Token allocation classes for backend
23 */
24#ifndef TVM_RELAY_BACKEND_TOKEN_ALLOCATOR_H_
25#define TVM_RELAY_BACKEND_TOKEN_ALLOCATOR_H_
26
27#include <tvm/relay/type.h>
28#include <tvm/target/virtual_device.h>
29
30#include <map>
31#include <string>
32#include <unordered_map>
33#include <unordered_set>
34#include <vector>
35
36#include "../../runtime/texture.h"
37
38namespace tvm {
39namespace relay {
40
41/*! A representation of a block of memory required at runtime on some device. */
42struct StorageToken {
43 /*! \brief Reference counter */
44 int ref_counter{0};
45 /*! \brief number of bytes */
46 size_t max_bytes{0};
47 /*! \brief The corresponding tensor type. */
48 TensorType ttype{nullptr};
49 /*! \brief VirtualDevice on which the memory will reside. */
50 VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained();
51 /*! \brief The storage id */
52 int64_t storage_id{-1};
53
54 bool is_valid() const { return !virtual_device->IsFullyUnconstrained(); }
55
56 bool is_compatible(const StorageToken& that) const {
57 return virtual_device == that.virtual_device;
58 }
59
60 std::string ToString() const {
61 std::ostringstream os;
62 os << "{storage_id: " << storage_id << ", max_bytes: " << max_bytes
63 << ", ttype: " << PrettyPrint(ttype) << ", virtual_device: " << virtual_device << "}";
64 return os.str();
65 }
66};
67
68/**
69 * @brief Memory manager for flattened 1d memory (buffers)
70 */
71class TokenAllocator1D {
72 public:
73 /*!
74 * \brief ceil(size/word_size) to get number of words.
75 * \param size The original size.
76 * \param word_size The element size.
77 */
78 static size_t DivRoundUp(size_t size, size_t word_size) {
79 return (size + word_size - 1) / word_size;
80 }
81
82 /*!
83 * \brief Get the memory requirement.
84 * \param prototype The prototype token.
85 * \return The required memory size.
86 *
87 * TODO(mbs): Gf GetMemorySizeBytes in aot_executor_codegen.cc,
88 * CalculateRelayExprSizeBytes in utils.cc
89 */
90 size_t GetMemorySize(StorageToken* prototype);
91 /*!
92 * \brief Request a storage token for a given prototype.
93 * \param prototype. The prototype storage token.
94 * \return The result token.
95 */
96 StorageToken* Request(StorageToken* prototype);
97 /*!
98 * \brief Alloacte a storage token by consuming prototype
99 * \param prototype The prototype token.
100 * \param size The size of memory being requested.
101 */
102 StorageToken* Alloc(StorageToken* prototype, int64_t storage_id);
103 /*!
104 * \brief Check if we can release token.
105 * \param tok The token to be released.
106 */
107 void CheckForRelease(StorageToken* tok);
108
109 private:
110 // scale used for rough match
111 const size_t match_range_{16};
112 // free list of storage entry
113 std::multimap<size_t, StorageToken*> free_;
114 // all the storage resources available
115 std::vector<StorageToken*> data_;
116};
117
118/**
119 * @brief Memory manager for 2d memory (textures)
120 */
121class TokenAllocator2D {
122 public:
123 /*!
124 * \brief Request a storage token for a given prototype.
125 * \param prototype. The prototype storage token.
126 * \return The result token.
127 */
128 StorageToken* Request(StorageToken* prototype);
129 /*!
130 * \brief Alloacte a storage token by consuming prototype
131 * \param prototype The prototype token.
132 * \param size The size of memory being requested.
133 */
134 StorageToken* Alloc(StorageToken* prototype, int64_t storage_id);
135 /*!
136 * \brief Check if we can release token.
137 * \param tok The token to be released.
138 */
139 void CheckForRelease(StorageToken* tok);
140 /*!
141 * \brief Get the texture 2d size requirement
142 * \param prototype The prototype token.
143 * \return The required texture 2d memory size in (width, height, channel).
144 */
145 runtime::Texture2DShape<int64_t> GetSize2D(StorageToken* prototype);
146
147 protected:
148 struct MemBlock {
149 StorageToken* token_;
150 int64_t x_;
151 int64_t y_;
152 };
153
154 std::unordered_map<int64_t, MemBlock> blocks_;
155 std::unordered_set<int64_t> free_list_;
156};
157
158} // namespace relay
159} // namespace tvm
160
161#endif // TVM_RELAY_BACKEND_TOKEN_ALLOCATOR_H_
162