1 //===- vulkan-runtime-wrappers.cpp - MLIR Vulkan runner wrapper library ---===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Implements C runtime wrappers around the VulkanRuntime.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include <mutex>
14 #include <numeric>
15 
16 #include "VulkanRuntime.h"
17 #include "llvm/Support/raw_ostream.h"
18 
19 namespace {
20 
21 class VulkanRuntimeManager {
22 public:
23   VulkanRuntimeManager() = default;
24   VulkanRuntimeManager(const VulkanRuntimeManager &) = delete;
25   VulkanRuntimeManager operator=(const VulkanRuntimeManager &) = delete;
26   ~VulkanRuntimeManager() = default;
27 
28   void setResourceData(DescriptorSetIndex setIndex, BindingIndex bindIndex,
29                        const VulkanHostMemoryBuffer &memBuffer) {
30     std::lock_guard<std::mutex> lock(mutex);
31     vulkanRuntime.setResourceData(setIndex, bindIndex, memBuffer);
32   }
33 
34   void setEntryPoint(const char *entryPoint) {
35     std::lock_guard<std::mutex> lock(mutex);
36     vulkanRuntime.setEntryPoint(entryPoint);
37   }
38 
39   void setNumWorkGroups(NumWorkGroups numWorkGroups) {
40     std::lock_guard<std::mutex> lock(mutex);
41     vulkanRuntime.setNumWorkGroups(numWorkGroups);
42   }
43 
44   void setShaderModule(uint8_t *shader, uint32_t size) {
45     std::lock_guard<std::mutex> lock(mutex);
46     vulkanRuntime.setShaderModule(shader, size);
47   }
48 
49   void runOnVulkan() {
50     std::lock_guard<std::mutex> lock(mutex);
51     if (failed(vulkanRuntime.initRuntime()) || failed(vulkanRuntime.run()) ||
52         failed(vulkanRuntime.updateHostMemoryBuffers()) ||
53         failed(vulkanRuntime.destroy())) {
54       llvm::errs() << "runOnVulkan failed";
55     }
56   }
57 
58 private:
59   VulkanRuntime vulkanRuntime;
60   std::mutex mutex;
61 };
62 
63 } // namespace
64 
65 template <typename T, int N>
66 struct MemRefDescriptor {
67   T *allocated;
68   T *aligned;
69   int64_t offset;
70   int64_t sizes[N];
71   int64_t strides[N];
72 };
73 
74 extern "C" {
75 /// Initializes `VulkanRuntimeManager` and returns a pointer to it.
76 void *initVulkan() { return new VulkanRuntimeManager(); }
77 
78 /// Deinitializes `VulkanRuntimeManager` by the given pointer.
79 void deinitVulkan(void *vkRuntimeManager) {
80   delete reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager);
81 }
82 
83 void runOnVulkan(void *vkRuntimeManager) {
84   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)->runOnVulkan();
85 }
86 
87 void setEntryPoint(void *vkRuntimeManager, const char *entryPoint) {
88   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
89       ->setEntryPoint(entryPoint);
90 }
91 
92 void setNumWorkGroups(void *vkRuntimeManager, uint32_t x, uint32_t y,
93                       uint32_t z) {
94   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
95       ->setNumWorkGroups({x, y, z});
96 }
97 
98 void setBinaryShader(void *vkRuntimeManager, uint8_t *shader, uint32_t size) {
99   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
100       ->setShaderModule(shader, size);
101 }
102 
103 /// Binds the given 1D float memref to the given descriptor set and descriptor
104 /// index.
105 void bindMemRef1DFloat(void *vkRuntimeManager, DescriptorSetIndex setIndex,
106                        BindingIndex bindIndex,
107                        MemRefDescriptor<float, 1> *ptr) {
108   VulkanHostMemoryBuffer memBuffer{
109       ptr->allocated, static_cast<uint32_t>(ptr->sizes[0] * sizeof(float))};
110   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
111       ->setResourceData(setIndex, bindIndex, memBuffer);
112 }
113 
114 /// Binds the given 2D float memref to the given descriptor set and descriptor
115 /// index.
116 void bindMemRef2DFloat(void *vkRuntimeManager, DescriptorSetIndex setIndex,
117                        BindingIndex bindIndex,
118                        MemRefDescriptor<float, 2> *ptr) {
119   VulkanHostMemoryBuffer memBuffer{
120       ptr->allocated,
121       static_cast<uint32_t>(ptr->sizes[0] * ptr->sizes[1] * sizeof(float))};
122   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
123       ->setResourceData(setIndex, bindIndex, memBuffer);
124 }
125 
126 /// Binds the given 3D float memref to the given descriptor set and descriptor
127 /// index.
128 void bindMemRef3DFloat(void *vkRuntimeManager, DescriptorSetIndex setIndex,
129                        BindingIndex bindIndex,
130                        MemRefDescriptor<float, 3> *ptr) {
131   VulkanHostMemoryBuffer memBuffer{
132       ptr->allocated, static_cast<uint32_t>(ptr->sizes[0] * ptr->sizes[1] *
133                                             ptr->sizes[2] * sizeof(float))};
134   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
135       ->setResourceData(setIndex, bindIndex, memBuffer);
136 }
137 
138 /// Binds the given 1D int memref to the given descriptor set and descriptor
139 /// index.
140 void bindMemRef1DInt(void *vkRuntimeManager, DescriptorSetIndex setIndex,
141                      BindingIndex bindIndex,
142                      MemRefDescriptor<int32_t, 1> *ptr) {
143   VulkanHostMemoryBuffer memBuffer{
144       ptr->allocated, static_cast<uint32_t>(ptr->sizes[0] * sizeof(int32_t))};
145   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
146       ->setResourceData(setIndex, bindIndex, memBuffer);
147 }
148 
149 /// Binds the given 2D int memref to the given descriptor set and descriptor
150 /// index.
151 void bindMemRef2DInt(void *vkRuntimeManager, DescriptorSetIndex setIndex,
152                      BindingIndex bindIndex,
153                      MemRefDescriptor<int32_t, 2> *ptr) {
154   VulkanHostMemoryBuffer memBuffer{
155       ptr->allocated,
156       static_cast<uint32_t>(ptr->sizes[0] * ptr->sizes[1] * sizeof(int32_t))};
157   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
158       ->setResourceData(setIndex, bindIndex, memBuffer);
159 }
160 
161 /// Binds the given 3D int memref to the given descriptor set and descriptor
162 /// index.
163 void bindMemRef3DInt(void *vkRuntimeManager, DescriptorSetIndex setIndex,
164                      BindingIndex bindIndex,
165                      MemRefDescriptor<int32_t, 3> *ptr) {
166   VulkanHostMemoryBuffer memBuffer{
167       ptr->allocated, static_cast<uint32_t>(ptr->sizes[0] * ptr->sizes[1] *
168                                             ptr->sizes[2] * sizeof(int32_t))};
169   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
170       ->setResourceData(setIndex, bindIndex, memBuffer);
171 }
172 
173 /// Fills the given 1D float memref with the given float value.
174 void _mlir_ciface_fillResource1DFloat(MemRefDescriptor<float, 1> *ptr, // NOLINT
175                                       float value) {
176   std::fill_n(ptr->allocated, ptr->sizes[0], value);
177 }
178 
179 /// Fills the given 2D float memref with the given float value.
180 void _mlir_ciface_fillResource2DFloat(MemRefDescriptor<float, 2> *ptr, // NOLINT
181                                       float value) {
182   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
183 }
184 
185 /// Fills the given 3D float memref with the given float value.
186 void _mlir_ciface_fillResource3DFloat(MemRefDescriptor<float, 3> *ptr, // NOLINT
187                                       float value) {
188   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
189               value);
190 }
191 
192 /// Fills the given 1D int memref with the given int value.
193 void _mlir_ciface_fillResource1DInt(MemRefDescriptor<int32_t, 1> *ptr, // NOLINT
194                                     int32_t value) {
195   std::fill_n(ptr->allocated, ptr->sizes[0], value);
196 }
197 
198 /// Fills the given 2D int memref with the given int value.
199 void _mlir_ciface_fillResource2DInt(MemRefDescriptor<int32_t, 2> *ptr, // NOLINT
200                                     int32_t value) {
201   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
202 }
203 
204 /// Fills the given 3D int memref with the given int value.
205 void _mlir_ciface_fillResource3DInt(MemRefDescriptor<int32_t, 3> *ptr, // NOLINT
206                                     int32_t value) {
207   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
208               value);
209 }
210 }
211