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 <iostream>
14 #include <mutex>
15 #include <numeric>
16 
17 #include "VulkanRuntime.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       std::cerr << "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 template <typename T, uint32_t S>
75 void bindMemRef(void *vkRuntimeManager, DescriptorSetIndex setIndex,
76                 BindingIndex bindIndex, MemRefDescriptor<T, S> *ptr) {
77   uint32_t size = sizeof(T);
78   for (unsigned i = 0; i < S; i++)
79     size *= ptr->sizes[i];
80   VulkanHostMemoryBuffer memBuffer{ptr->allocated, size};
81   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
82       ->setResourceData(setIndex, bindIndex, memBuffer);
83 }
84 
85 extern "C" {
86 /// Initializes `VulkanRuntimeManager` and returns a pointer to it.
87 void *initVulkan() { return new VulkanRuntimeManager(); }
88 
89 /// Deinitializes `VulkanRuntimeManager` by the given pointer.
90 void deinitVulkan(void *vkRuntimeManager) {
91   delete reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager);
92 }
93 
94 void runOnVulkan(void *vkRuntimeManager) {
95   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)->runOnVulkan();
96 }
97 
98 void setEntryPoint(void *vkRuntimeManager, const char *entryPoint) {
99   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
100       ->setEntryPoint(entryPoint);
101 }
102 
103 void setNumWorkGroups(void *vkRuntimeManager, uint32_t x, uint32_t y,
104                       uint32_t z) {
105   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
106       ->setNumWorkGroups({x, y, z});
107 }
108 
109 void setBinaryShader(void *vkRuntimeManager, uint8_t *shader, uint32_t size) {
110   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
111       ->setShaderModule(shader, size);
112 }
113 
114 /// Binds the given memref to the given descriptor set and descriptor
115 /// index.
116 #define DECLARE_BIND_MEMREF(size, type, typeName)                              \
117   void bindMemRef##size##D##typeName(                                          \
118       void *vkRuntimeManager, DescriptorSetIndex setIndex,                     \
119       BindingIndex bindIndex, MemRefDescriptor<type, size> *ptr) {             \
120     bindMemRef<type, size>(vkRuntimeManager, setIndex, bindIndex, ptr);        \
121   }
122 
123 DECLARE_BIND_MEMREF(1, float, Float)
124 DECLARE_BIND_MEMREF(2, float, Float)
125 DECLARE_BIND_MEMREF(3, float, Float)
126 DECLARE_BIND_MEMREF(1, int32_t, Int32)
127 DECLARE_BIND_MEMREF(2, int32_t, Int32)
128 DECLARE_BIND_MEMREF(3, int32_t, Int32)
129 DECLARE_BIND_MEMREF(1, int16_t, Int16)
130 DECLARE_BIND_MEMREF(2, int16_t, Int16)
131 DECLARE_BIND_MEMREF(3, int16_t, Int16)
132 DECLARE_BIND_MEMREF(1, int8_t, Int8)
133 DECLARE_BIND_MEMREF(2, int8_t, Int8)
134 DECLARE_BIND_MEMREF(3, int8_t, Int8)
135 DECLARE_BIND_MEMREF(1, int16_t, Half)
136 DECLARE_BIND_MEMREF(2, int16_t, Half)
137 DECLARE_BIND_MEMREF(3, int16_t, Half)
138 
139 /// Fills the given 1D float memref with the given float value.
140 void _mlir_ciface_fillResource1DFloat(MemRefDescriptor<float, 1> *ptr, // NOLINT
141                                       float value) {
142   std::fill_n(ptr->allocated, ptr->sizes[0], value);
143 }
144 
145 /// Fills the given 2D float memref with the given float value.
146 void _mlir_ciface_fillResource2DFloat(MemRefDescriptor<float, 2> *ptr, // NOLINT
147                                       float value) {
148   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
149 }
150 
151 /// Fills the given 3D float memref with the given float value.
152 void _mlir_ciface_fillResource3DFloat(MemRefDescriptor<float, 3> *ptr, // NOLINT
153                                       float value) {
154   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
155               value);
156 }
157 
158 /// Fills the given 1D int memref with the given int value.
159 void _mlir_ciface_fillResource1DInt(MemRefDescriptor<int32_t, 1> *ptr, // NOLINT
160                                     int32_t value) {
161   std::fill_n(ptr->allocated, ptr->sizes[0], value);
162 }
163 
164 /// Fills the given 2D int memref with the given int value.
165 void _mlir_ciface_fillResource2DInt(MemRefDescriptor<int32_t, 2> *ptr, // NOLINT
166                                     int32_t value) {
167   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
168 }
169 
170 /// Fills the given 3D int memref with the given int value.
171 void _mlir_ciface_fillResource3DInt(MemRefDescriptor<int32_t, 3> *ptr, // NOLINT
172                                     int32_t value) {
173   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
174               value);
175 }
176 
177 /// Fills the given 1D int memref with the given int8 value.
178 void _mlir_ciface_fillResource1DInt8(MemRefDescriptor<int8_t, 1> *ptr, // NOLINT
179                                      int8_t value) {
180   std::fill_n(ptr->allocated, ptr->sizes[0], value);
181 }
182 
183 /// Fills the given 2D int memref with the given int8 value.
184 void _mlir_ciface_fillResource2DInt8(MemRefDescriptor<int8_t, 2> *ptr, // NOLINT
185                                      int8_t value) {
186   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
187 }
188 
189 /// Fills the given 3D int memref with the given int8 value.
190 void _mlir_ciface_fillResource3DInt8(MemRefDescriptor<int8_t, 3> *ptr, // NOLINT
191                                      int8_t value) {
192   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
193               value);
194 }
195 }
196