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 texture.h
22 * \brief Texture utilities
23 */
24#ifndef TVM_RUNTIME_TEXTURE_H_
25#define TVM_RUNTIME_TEXTURE_H_
26
27#include <tvm/runtime/device_api.h>
28
29#include <memory>
30#include <string>
31#include <vector>
32
33namespace tvm {
34namespace runtime {
35
36/*! \brief Structure to represent flattened texture shape */
37template <typename T>
38struct Texture2DShape {
39 T width;
40 T height;
41 T channel;
42};
43
44/*!
45 * \param shape_rank Rank N of the Nd-shape
46 * \param convention Storage scope convention to use for flattening
47 * \return The axis separator that defines the Nd shape partitioning in 2d
48 */
49inline size_t DefaultTextureLayoutSeparator(size_t shape_rank,
50 std::string convention = "global.texture") {
51 // Texture activation:
52 // e.g. [N,C,H,W,c] -> Texture2d[N*C*H, W, c]
53 // Texture weight:
54 // e.g. [O,I,H,W,c] -> Texture2d[O, I*H*W, c]
55 size_t separator = 0;
56 if (convention == "global.texture") {
57 separator = shape_rank - 2;
58 } else if (convention == "global.texture-weight") {
59 separator = 1;
60 } else if (convention == "global.texture-nhwc") {
61 if (shape_rank == 3) {
62 separator = 1;
63 } else {
64 separator = 2;
65 }
66 } else {
67 LOG(FATAL) << "Encountered unknown texture lowering convention: " << convention;
68 }
69 return separator;
70}
71
72/*!
73 * \param shape Nd shape
74 * \param rank Number of dimensions N of the Nd shape
75 * \param axis The axis separator that splits the Nd axes into two sets
76 * \return Width and height of the 2d shape
77 */
78template <typename T, typename S>
79Texture2DShape<T> ApplyTexture2DFlattening(const S& shape, size_t rank, size_t axis) {
80 ICHECK(axis < rank)
81 << "Number of axes to flatten into rows must be less than shape rank for 2d flattening";
82 Texture2DShape<T> texture{1, 1, shape[rank - 1]};
83 for (size_t i = 0; i < rank - 1; i++) {
84 if (i < axis) {
85 texture.height *= shape[i];
86 } else {
87 texture.width *= shape[i];
88 }
89 }
90 return texture;
91}
92
93inline bool IsTextureStorage(std::string scope) {
94 return scope.find("texture") != std::string::npos;
95}
96
97class TVM_DLL Pool2D {
98 public:
99 Pool2D() = default;
100 void* Alloc(Device dev, DeviceAPI* device, size_t width, size_t height, DLDataType type_hint);
101 void Free(void* data);
102 // Release all resources immediately
103 void Release(Device dev, DeviceAPI* device);
104
105 protected:
106 struct Entry {
107 void* data;
108 size_t x;
109 size_t y;
110 DLDataType type;
111 };
112 std::vector<Entry> free_list_;
113 std::vector<Entry> allocated_;
114};
115
116/*!
117 * \brief A two dimensional storage pool that recycles temporal workspace
118 * allocations for dynamically allocated texture. See AllocTexture docstring
119 * for approach to allocation and reuse.
120 */
121class TVM_DLL TexturePool {
122 public:
123 /*!
124 * \brief Create pool with specific device type and device.
125 * \param device_type The device type.
126 * \param device_api The device API.
127 */
128 TexturePool(DLDeviceType device_type, DeviceAPI* device_api);
129 /*! \brief destructor */
130 ~TexturePool();
131
132 /*!
133 * \brief Allocate a two dimensional temporal texture workspace on device
134 *
135 * \note Two dimensional texture workspaces will be grown and reused
136 * according to the following strategy:
137 * - Choose the workspace which minimizes the amount of memory required to
138 * grow the workspace to fit the request.
139 * - If a set of workspaces exist that fit the current request without
140 * expansion, choose the workspace of that set which most closely
141 * matches the request size, minimizing wasted space.
142 *
143 * \param dev The context of allocation.
144 * \param width The width of the 2d texture to be allocated.
145 * \param height The height of the 2d texture to be allocated.
146 * \param type_hint The type of elements.
147 */
148 void* AllocTexture(Device dev, size_t width, size_t height, DLDataType type_hint);
149 /*!
150 * \brief Free temporal texture in backend execution.
151 *
152 * \param dev The context of allocation.
153 * \param ptr The pointer to be freed.
154 */
155 void FreeTexture(Device dev, void* ptr);
156
157 private:
158 /*! \brief pool of device local array */
159 std::vector<Pool2D*> array_;
160 /*! \brief device type this pool support */
161 DLDeviceType device_type_;
162 /*! \brief The device API */
163 DeviceAPI* device_;
164};
165
166} // namespace runtime
167} // namespace tvm
168#endif // TVM_RUNTIME_TEXTURE_H_
169