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 /// Fills the given 1D float memref with the given float value.
139 void _mlir_ciface_fillResource1DFloat(MemRefDescriptor<float, 1> *ptr, // NOLINT
140                                       float value) {
141   std::fill_n(ptr->allocated, ptr->sizes[0], value);
142 }
143 
144 /// Fills the given 2D float memref with the given float value.
145 void _mlir_ciface_fillResource2DFloat(MemRefDescriptor<float, 2> *ptr, // NOLINT
146                                       float value) {
147   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
148 }
149 
150 /// Fills the given 3D float memref with the given float value.
151 void _mlir_ciface_fillResource3DFloat(MemRefDescriptor<float, 3> *ptr, // NOLINT
152                                       float value) {
153   std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
154               value);
155 }
156 }
157