Line data Source code
1 : // SPDX-License-Identifier: Apache-2.0
2 : #pragma once
3 :
4 : #include "kompute/Core.hpp"
5 : #include "logger/Logger.hpp"
6 : #include <memory>
7 : #include <string>
8 :
9 : namespace kp {
10 :
11 : /**
12 : * Structured data used in GPU operations.
13 : *
14 : * Tensors are the base building block in Kompute to perform operations across
15 : * GPUs. Each tensor would have a respective Vulkan memory and buffer, which
16 : * would be used to store their respective data. The tensors can be used for GPU
17 : * data storage or transfer.
18 : */
19 : class Tensor
20 : {
21 : public:
22 : /**
23 : * Type for tensors created: Device allows memory to be transferred from
24 : * staging buffers. Staging are host memory visible. Storage are device
25 : * visible but are not set up to transfer or receive data (only for shader
26 : * storage).
27 : */
28 : enum class TensorTypes
29 : {
30 : eDevice = 0, ///< Type is device memory, source and destination
31 : eHost = 1, ///< Type is host memory, source and destination
32 : eStorage = 2, ///< Type is Device memory (only)
33 : };
34 : enum class TensorDataTypes
35 : {
36 : eBool = 0,
37 : eInt = 1,
38 : eUnsignedInt = 2,
39 : eFloat = 3,
40 : eDouble = 4,
41 : };
42 :
43 : static std::string toString(TensorDataTypes dt);
44 : static std::string toString(TensorTypes dt);
45 :
46 : /**
47 : * Constructor with data provided which would be used to create the
48 : * respective vulkan buffer and memory.
49 : *
50 : * @param physicalDevice The physical device to use to fetch properties
51 : * @param device The device to use to create the buffer and memory from
52 : * @param data Non-zero-sized vector of data that will be used by the
53 : * tensor
54 : * @param tensorTypes Type for the tensor which is of type TensorTypes
55 : */
56 : Tensor(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
57 : std::shared_ptr<vk::Device> device,
58 : void* data,
59 : uint32_t elementTotalCount,
60 : uint32_t elementMemorySize,
61 : const TensorDataTypes& dataType,
62 : const TensorTypes& tensorType = TensorTypes::eDevice);
63 :
64 : /**
65 : * Destructor which is in charge of freeing vulkan resources unless they
66 : * have been provided externally.
67 : */
68 : virtual ~Tensor();
69 :
70 : /**
71 : * Function to trigger reinitialisation of the tensor buffer and memory with
72 : * new data as well as new potential device type.
73 : *
74 : * @param data Vector of data to use to initialise vector from
75 : * @param tensorType The type to use for the tensor
76 : */
77 : void rebuild(void* data,
78 : uint32_t elementTotalCount,
79 : uint32_t elementMemorySize);
80 :
81 : /**
82 : * Destroys and frees the GPU resources which include the buffer and memory.
83 : */
84 : void destroy();
85 :
86 : /**
87 : * Check whether tensor is initialized based on the created gpu resources.
88 : *
89 : * @returns Boolean stating whether tensor is initialized
90 : */
91 : bool isInit();
92 :
93 : /**
94 : * Retrieve the tensor type of the Tensor
95 : *
96 : * @return Tensor type of tensor
97 : */
98 : TensorTypes tensorType();
99 :
100 : /**
101 : * Records a copy from the memory of the tensor provided to the current
102 : * thensor. This is intended to pass memory into a processing, to perform
103 : * a staging buffer transfer, or to gather output (between others).
104 : *
105 : * @param commandBuffer Vulkan Command Buffer to record the commands into
106 : * @param copyFromTensor Tensor to copy the data from
107 : */
108 : void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
109 : std::shared_ptr<Tensor> copyFromTensor);
110 :
111 : /**
112 : * Records a copy from the internal staging memory to the device memory
113 : * using an optional barrier to wait for the operation. This function would
114 : * only be relevant for kp::Tensors of type eDevice.
115 : *
116 : * @param commandBuffer Vulkan Command Buffer to record the commands into
117 : */
118 : void recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer);
119 :
120 : /**
121 : * Records a copy from the internal device memory to the staging memory
122 : * using an optional barrier to wait for the operation. This function would
123 : * only be relevant for kp::Tensors of type eDevice.
124 : *
125 : * @param commandBuffer Vulkan Command Buffer to record the commands into
126 : */
127 : void recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer);
128 :
129 : /**
130 : * Records the buffer memory barrier into the primary buffer and command
131 : * buffer which ensures that relevant data transfers are carried out
132 : * correctly.
133 : *
134 : * @param commandBuffer Vulkan Command Buffer to record the commands into
135 : * @param srcAccessMask Access flags for source access mask
136 : * @param dstAccessMask Access flags for destination access mask
137 : * @param scrStageMask Pipeline stage flags for source stage mask
138 : * @param dstStageMask Pipeline stage flags for destination stage mask
139 : */
140 : void recordPrimaryBufferMemoryBarrier(
141 : const vk::CommandBuffer& commandBuffer,
142 : vk::AccessFlagBits srcAccessMask,
143 : vk::AccessFlagBits dstAccessMask,
144 : vk::PipelineStageFlagBits srcStageMask,
145 : vk::PipelineStageFlagBits dstStageMask);
146 : /**
147 : * Records the buffer memory barrier into the staging buffer and command
148 : * buffer which ensures that relevant data transfers are carried out
149 : * correctly.
150 : *
151 : * @param commandBuffer Vulkan Command Buffer to record the commands into
152 : * @param srcAccessMask Access flags for source access mask
153 : * @param dstAccessMask Access flags for destination access mask
154 : * @param scrStageMask Pipeline stage flags for source stage mask
155 : * @param dstStageMask Pipeline stage flags for destination stage mask
156 : */
157 : void recordStagingBufferMemoryBarrier(
158 : const vk::CommandBuffer& commandBuffer,
159 : vk::AccessFlagBits srcAccessMask,
160 : vk::AccessFlagBits dstAccessMask,
161 : vk::PipelineStageFlagBits srcStageMask,
162 : vk::PipelineStageFlagBits dstStageMask);
163 :
164 : /**
165 : * Constructs a vulkan descriptor buffer info which can be used to specify
166 : * and reference the underlying buffer component of the tensor without
167 : * exposing it.
168 : *
169 : * @return Descriptor buffer info with own buffer
170 : */
171 : vk::DescriptorBufferInfo constructDescriptorBufferInfo();
172 :
173 : /**
174 : * Returns the size/magnitude of the Tensor, which will be the total number
175 : * of elements across all dimensions
176 : *
177 : * @return Unsigned integer representing the total number of elements
178 : */
179 : uint32_t size();
180 :
181 : /**
182 : * Returns the total size of a single element of the respective data type
183 : * that this tensor holds.
184 : *
185 : * @return Unsigned integer representing the memory of a single element of
186 : * the respective data type.
187 : */
188 : uint32_t dataTypeMemorySize();
189 :
190 : /**
191 : * Returns the total memory size of the data contained by the Tensor object
192 : * which would equate to (this->size() * this->dataTypeMemorySize())
193 : *
194 : * @return Unsigned integer representing the memory of a single element of
195 : * the respective data type.
196 : */
197 : uint32_t memorySize();
198 :
199 : /**
200 : * Retrieve the data type of the tensor (host, device, storage)
201 : *
202 : * @return Data type of tensor of type kp::Tensor::TensorDataTypes
203 : */
204 : TensorDataTypes dataType();
205 :
206 : /**
207 : * Retrieve the raw data via the pointer to the memory that contains the raw
208 : * memory of this current tensor. This tensor gets changed to a nullptr when
209 : * the Tensor is removed.
210 : *
211 : * @return Pointer to raw memory containing raw bytes data of Tensor.
212 : */
213 : void* rawData();
214 :
215 : /**
216 : * Sets / resets the data of the tensor which is directly done on the GPU
217 : * host visible memory available by the tensor.
218 : */
219 : void setRawData(const void* data);
220 :
221 : /**
222 : * Template to return the pointer data converted by specific type, which
223 : * would be any of the supported types including float, double, int32,
224 : * uint32 and bool.
225 : *
226 : * @return Pointer to raw memory containing raw bytes data of Tensor.
227 : */
228 : template<typename T>
229 : T* data()
230 : {
231 : return (T*)this->mRawData;
232 : }
233 :
234 : /**
235 : * Template to get the data of the current tensor as a vector of specific
236 : * type, which would be any of the supported types including float, double,
237 : * int32, uint32 and bool.
238 : *
239 : * @return Vector of type provided by template.
240 : */
241 : template<typename T>
242 0 : std::vector<T> vector()
243 : {
244 0 : return { (T*)this->mRawData, ((T*)this->mRawData) + this->size() };
245 : }
246 :
247 : protected:
248 : // -------------- ALWAYS OWNED RESOURCES
249 : TensorTypes mTensorType;
250 : TensorDataTypes mDataType;
251 : uint32_t mSize;
252 : uint32_t mDataTypeMemorySize;
253 : void* mRawData;
254 :
255 : private:
256 : // -------------- NEVER OWNED RESOURCES
257 : std::shared_ptr<vk::PhysicalDevice> mPhysicalDevice;
258 : std::shared_ptr<vk::Device> mDevice;
259 :
260 : // -------------- OPTIONALLY OWNED RESOURCES
261 : std::shared_ptr<vk::Buffer> mPrimaryBuffer;
262 : bool mFreePrimaryBuffer = false;
263 : std::shared_ptr<vk::Buffer> mStagingBuffer;
264 : bool mFreeStagingBuffer = false;
265 : std::shared_ptr<vk::DeviceMemory> mPrimaryMemory;
266 : bool mFreePrimaryMemory = false;
267 : std::shared_ptr<vk::DeviceMemory> mStagingMemory;
268 : bool mFreeStagingMemory = false;
269 :
270 : void allocateMemoryCreateGPUResources(); // Creates the vulkan buffer
271 : void createBuffer(std::shared_ptr<vk::Buffer> buffer,
272 : vk::BufferUsageFlags bufferUsageFlags);
273 : void allocateBindMemory(std::shared_ptr<vk::Buffer> buffer,
274 : std::shared_ptr<vk::DeviceMemory> memory,
275 : vk::MemoryPropertyFlags memoryPropertyFlags);
276 : void recordCopyBuffer(const vk::CommandBuffer& commandBuffer,
277 : std::shared_ptr<vk::Buffer> bufferFrom,
278 : std::shared_ptr<vk::Buffer> bufferTo,
279 : vk::DeviceSize bufferSize,
280 : vk::BufferCopy copyRegion);
281 : void recordBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
282 : const vk::Buffer& buffer,
283 : vk::AccessFlagBits srcAccessMask,
284 : vk::AccessFlagBits dstAccessMask,
285 : vk::PipelineStageFlagBits srcStageMask,
286 : vk::PipelineStageFlagBits dstStageMask);
287 :
288 : // Private util functions
289 : vk::BufferUsageFlags getPrimaryBufferUsageFlags();
290 : vk::MemoryPropertyFlags getPrimaryMemoryPropertyFlags();
291 : vk::BufferUsageFlags getStagingBufferUsageFlags();
292 : vk::MemoryPropertyFlags getStagingMemoryPropertyFlags();
293 :
294 : void mapRawData();
295 : void unmapRawData();
296 : };
297 :
298 : template<typename T>
299 : class TensorT : public Tensor
300 : {
301 :
302 : public:
303 97 : TensorT(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
304 : std::shared_ptr<vk::Device> device,
305 : const std::vector<T>& data,
306 : const TensorTypes& tensorType = TensorTypes::eDevice)
307 : : Tensor(physicalDevice,
308 : device,
309 97 : (void*)data.data(),
310 97 : data.size(),
311 : sizeof(T),
312 98 : this->dataType(),
313 196 : tensorType)
314 : {
315 288 : KP_LOG_DEBUG("Kompute TensorT constructor with data size {}",
316 : data.size());
317 96 : }
318 :
319 384 : ~TensorT() { KP_LOG_DEBUG("Kompute TensorT destructor"); }
320 :
321 6012 : T* data() { return (T*)this->mRawData; }
322 :
323 67 : std::vector<T> vector()
324 : {
325 67 : return { (T*)this->mRawData, ((T*)this->mRawData) + this->size() };
326 : }
327 :
328 : T& operator[](int index) { return *(((T*)this->mRawData) + index); }
329 :
330 : void setData(const std::vector<T>& data)
331 : {
332 :
333 : KP_LOG_DEBUG("Kompute TensorT setting data with data size {}",
334 : data.size());
335 :
336 : if (data.size() != this->mSize) {
337 : throw std::runtime_error(
338 : "Kompute TensorT Cannot set data of different sizes");
339 : }
340 :
341 : Tensor::setRawData(data.data());
342 : }
343 :
344 : TensorDataTypes dataType();
345 : };
346 :
347 : } // End namespace kp
|
