mlir/tools/mlir-vulkan-runner/vulkan-runtime-wrappers.cpp

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