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 | |
33 | namespace tvm { |
34 | namespace runtime { |
35 | |
36 | /*! \brief Structure to represent flattened texture shape */ |
37 | template <typename T> |
38 | struct 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 | */ |
49 | inline 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 | */ |
78 | template <typename T, typename S> |
79 | Texture2DShape<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 | |
93 | inline bool IsTextureStorage(std::string scope) { |
94 | return scope.find("texture" ) != std::string::npos; |
95 | } |
96 | |
97 | class 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 | */ |
121 | class 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 | |