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 // Explicitly export entry points to the vulkan-runtime-wrapper.
20 #define VULKAN_WRAPPER_SYMBOL_EXPORT __attribute__((visibility("default")))
21 
22 namespace {
23 
24 class VulkanRuntimeManager {
25 public:
26   VulkanRuntimeManager() = default;
27   VulkanRuntimeManager(const VulkanRuntimeManager &) = delete;
28   VulkanRuntimeManager operator=(const VulkanRuntimeManager &) = delete;
29   ~VulkanRuntimeManager() = default;
30 
31   void setResourceData(DescriptorSetIndex setIndex, BindingIndex bindIndex,
32                        const VulkanHostMemoryBuffer &memBuffer) {
33     std::lock_guard<std::mutex> lock(mutex);
34     vulkanRuntime.setResourceData(setIndex, bindIndex, memBuffer);
35   }
36 
37   void setEntryPoint(const char *entryPoint) {
38     std::lock_guard<std::mutex> lock(mutex);
39     vulkanRuntime.setEntryPoint(entryPoint);
40   }
41 
42   void setNumWorkGroups(NumWorkGroups numWorkGroups) {
43     std::lock_guard<std::mutex> lock(mutex);
44     vulkanRuntime.setNumWorkGroups(numWorkGroups);
45   }
46 
47   void setShaderModule(uint8_t *shader, uint32_t size) {
48     std::lock_guard<std::mutex> lock(mutex);
49     vulkanRuntime.setShaderModule(shader, size);
50   }
51 
52   void runOnVulkan() {
53     std::lock_guard<std::mutex> lock(mutex);
54     if (failed(vulkanRuntime.initRuntime()) || failed(vulkanRuntime.run()) ||
55         failed(vulkanRuntime.updateHostMemoryBuffers()) ||
56         failed(vulkanRuntime.destroy())) {
57       std::cerr << "runOnVulkan failed";
58     }
59   }
60 
61 private:
62   VulkanRuntime vulkanRuntime;
63   std::mutex mutex;
64 };
65 
66 } // namespace
67 
68 template <typename T, int N> struct MemRefDescriptor {
69   T *allocated;
70   T *aligned;
71   int64_t offset;
72   int64_t sizes[N];
73   int64_t strides[N];
74 };
75 
76 template <typename T, uint32_t S>
77 void bindMemRef(void *vkRuntimeManager, DescriptorSetIndex setIndex,
78                 BindingIndex bindIndex, MemRefDescriptor<T, S> *ptr) {
79   uint32_t size = sizeof(T);
80   for (unsigned i = 0; i < S; i++)
81     size *= ptr->sizes[i];
82   VulkanHostMemoryBuffer memBuffer{ptr->allocated, size};
83   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
84       ->setResourceData(setIndex, bindIndex, memBuffer);
85 }
86 
87 extern "C" {
88 /// Initializes `VulkanRuntimeManager` and returns a pointer to it.
89 VULKAN_WRAPPER_SYMBOL_EXPORT void *initVulkan() {
90   return new VulkanRuntimeManager();
91 }
92 
93 /// Deinitializes `VulkanRuntimeManager` by the given pointer.
94 VULKAN_WRAPPER_SYMBOL_EXPORT void deinitVulkan(void *vkRuntimeManager) {
95   delete reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager);
96 }
97 
98 VULKAN_WRAPPER_SYMBOL_EXPORT void runOnVulkan(void *vkRuntimeManager) {
99   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)->runOnVulkan();
100 }
101 
102 VULKAN_WRAPPER_SYMBOL_EXPORT void setEntryPoint(void *vkRuntimeManager,
103                                                 const char *entryPoint) {
104   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
105       ->setEntryPoint(entryPoint);
106 }
107 
108 VULKAN_WRAPPER_SYMBOL_EXPORT void
109 setNumWorkGroups(void *vkRuntimeManager, uint32_t x, uint32_t y, uint32_t z) {
110   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
111       ->setNumWorkGroups({x, y, z});
112 }
113 
114 VULKAN_WRAPPER_SYMBOL_EXPORT void
115 setBinaryShader(void *vkRuntimeManager, uint8_t *shader, uint32_t size) {
116   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
117       ->setShaderModule(shader, size);
118 }
119 
120 /// Binds the given memref to the given descriptor set and descriptor
121 /// index.
122 #define DECLARE_BIND_MEMREF(size, type, typeName)                              \
123   VULKAN_WRAPPER_SYMBOL_EXPORT void bindMemRef##size##D##typeName(             \
124       void *vkRuntimeManager, DescriptorSetIndex setIndex,                     \
125       BindingIndex bindIndex, MemRefDescriptor<type, size> *ptr) {             \
126     bindMemRef<type, size>(vkRuntimeManager, setIndex, bindIndex, ptr);        \
127   }
128 
129 DECLARE_BIND_MEMREF(1, float, Float)
130 DECLARE_BIND_MEMREF(2, float, Float)
131 DECLARE_BIND_MEMREF(3, float, Float)
132 DECLARE_BIND_MEMREF(1, int32_t, Int32)
133 DECLARE_BIND_MEMREF(2, int32_t, Int32)
134 DECLARE_BIND_MEMREF(3, int32_t, Int32)
135 DECLARE_BIND_MEMREF(1, int16_t, Int16)
136 DECLARE_BIND_MEMREF(2, int16_t, Int16)
137 DECLARE_BIND_MEMREF(3, int16_t, Int16)
138 DECLARE_BIND_MEMREF(1, int8_t, Int8)
139 DECLARE_BIND_MEMREF(2, int8_t, Int8)
140 DECLARE_BIND_MEMREF(3, int8_t, Int8)
141 DECLARE_BIND_MEMREF(1, int16_t, Half)
142 DECLARE_BIND_MEMREF(2, int16_t, Half)
143 DECLARE_BIND_MEMREF(3, int16_t, Half)
144 
145 /// Fills the given 1D float memref with the given float value.
146 VULKAN_WRAPPER_SYMBOL_EXPORT void
147 _mlir_ciface_fillResource1DFloat(MemRefDescriptor<float, 1> *ptr, // NOLINT
148                                  float value) {
149   std::fill_n(ptr->allocated, ptr->sizes[0], value);
150 }
151 
152 /// Fills the given 2D float memref with the given float value.
153 VULKAN_WRAPPER_SYMBOL_EXPORT void
154 _mlir_ciface_fillResource2DFloat(MemRefDescriptor<float, 2> *ptr, // NOLINT
155                                  float value) {
156   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
157 }
158 
159 /// Fills the given 3D float memref with the given float value.
160 VULKAN_WRAPPER_SYMBOL_EXPORT void
161 _mlir_ciface_fillResource3DFloat(MemRefDescriptor<float, 3> *ptr, // NOLINT
162                                  float value) {
163   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
164               value);
165 }
166 
167 /// Fills the given 1D int memref with the given int value.
168 VULKAN_WRAPPER_SYMBOL_EXPORT void
169 _mlir_ciface_fillResource1DInt(MemRefDescriptor<int32_t, 1> *ptr, // NOLINT
170                                int32_t value) {
171   std::fill_n(ptr->allocated, ptr->sizes[0], value);
172 }
173 
174 /// Fills the given 2D int memref with the given int value.
175 VULKAN_WRAPPER_SYMBOL_EXPORT void
176 _mlir_ciface_fillResource2DInt(MemRefDescriptor<int32_t, 2> *ptr, // NOLINT
177                                int32_t value) {
178   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
179 }
180 
181 /// Fills the given 3D int memref with the given int value.
182 VULKAN_WRAPPER_SYMBOL_EXPORT void
183 _mlir_ciface_fillResource3DInt(MemRefDescriptor<int32_t, 3> *ptr, // NOLINT
184                                int32_t value) {
185   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
186               value);
187 }
188 
189 /// Fills the given 1D int memref with the given int8 value.
190 VULKAN_WRAPPER_SYMBOL_EXPORT void
191 _mlir_ciface_fillResource1DInt8(MemRefDescriptor<int8_t, 1> *ptr, // NOLINT
192                                 int8_t value) {
193   std::fill_n(ptr->allocated, ptr->sizes[0], value);
194 }
195 
196 /// Fills the given 2D int memref with the given int8 value.
197 VULKAN_WRAPPER_SYMBOL_EXPORT void
198 _mlir_ciface_fillResource2DInt8(MemRefDescriptor<int8_t, 2> *ptr, // NOLINT
199                                 int8_t value) {
200   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
201 }
202 
203 /// Fills the given 3D int memref with the given int8 value.
204 VULKAN_WRAPPER_SYMBOL_EXPORT void
205 _mlir_ciface_fillResource3DInt8(MemRefDescriptor<int8_t, 3> *ptr, // NOLINT
206                                 int8_t value) {
207   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
208               value);
209 }
210 }
211