[mlir][vulkan-runner] Avoid dependency on LLVM libraries

The Vulkan runtime wrapper will be compiled to a shared library
that are loaded by the JIT runner. Depending on LLVM libraries
means that LLVM symbols will be compiled into the shared library.
That can cause problems if we are using it with other shared
libraries depending on LLVM, notably Mesa, the open-source graphics
driver framework. The Vulkan API wrappers invoked by the JIT runner
links to the system libvulkan.so. If it's Mesa providing the
implementation, Mesa will normally try to load the system libLLVM.so
for its shader compilation. That causes issues because the JIT runner
already loaded the Vulkan runtime wrapper which has LLVM sybmols
compiled in. So system linker will instruct Mesa to use those symbols
instead.

Differential Revision: https://reviews.llvm.org/D79860

Author: antiagainst

mlir/tools/mlir-vulkan-runner/CMakeLists.txt

@@ -38,15 +38,12 @@ if (MLIR_VULKAN_RUNNER_ENABLED)
     VulkanRuntime.cpp
   )
 
-  target_include_directories(vulkan-runtime-wrappers
-    PRIVATE ${Vulkan_INCLUDE_DIR}
-    LLVMSupport
-  )
-
+  # *IMPORTANT*: This library cannot depend on LLVM libraries. Otherwise,
+  # it may cause LLVM version conflict when used together with other shared
+  # libraries depending on LLVM. Notably, Mesa, who implemnents Vulkan
+  # drivers on Linux, depends on the system libLLVM.so.
   target_link_libraries(vulkan-runtime-wrappers
     PUBLIC
-    LLVMCore
-    LLVMSupport
     ${Vulkan_LIBRARY}
   )
 

mlir/tools/mlir-vulkan-runner/VulkanRuntime.cpp

@@ -13,8 +13,22 @@
 
 #include "VulkanRuntime.h"
 
-#include "llvm/Support/Format.h"
 #include <chrono>
+#include <cstring>
+// TODO(antiagainst): It's generally bad to access stdout/stderr in a library.
+// Figure out a better way for error reporting.
+#include <iomanip>
+#include <iostream>
+
+inline void emitVulkanError(const char *api, VkResult error) {
+  std::cerr << " failed with error code " << error << " when executing " << api;
+}
+
+#define RETURN_ON_VULKAN_ERROR(result, api)                                    \
+  if ((result) != VK_SUCCESS) {                                                \
+    emitVulkanError(api, (result));                                            \
+    return failure();                                                          \
+  }
 
 using namespace mlir;
 
@@ -32,7 +46,7 @@ void VulkanRuntime::setResourceData(
     const VulkanHostMemoryBuffer &hostMemBuffer) {
   resourceData[desIndex][bindIndex] = hostMemBuffer;
   resourceStorageClassData[desIndex][bindIndex] =
-      spirv::StorageClass::StorageBuffer;
+      SPIRVStorageClass::StorageBuffer;
 }
 
 void VulkanRuntime::setEntryPoint(const char *entryPointName) {
@@ -49,33 +63,27 @@ void VulkanRuntime::setShaderModule(uint8_t *shader, uint32_t size) {
 }
 
 LogicalResult VulkanRuntime::mapStorageClassToDescriptorType(
-    spirv::StorageClass storageClass, VkDescriptorType &descriptorType) {
+    SPIRVStorageClass storageClass, VkDescriptorType &descriptorType) {
   switch (storageClass) {
-  case spirv::StorageClass::StorageBuffer:
+  case SPIRVStorageClass::StorageBuffer:
     descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
     break;
-  case spirv::StorageClass::Uniform:
+  case SPIRVStorageClass::Uniform:
     descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
     break;
-  default:
-    llvm::errs() << "unsupported storage class";
-    return failure();
   }
   return success();
 }
 
 LogicalResult VulkanRuntime::mapStorageClassToBufferUsageFlag(
-    spirv::StorageClass storageClass, VkBufferUsageFlagBits &bufferUsage) {
+    SPIRVStorageClass storageClass, VkBufferUsageFlagBits &bufferUsage) {
   switch (storageClass) {
-  case spirv::StorageClass::StorageBuffer:
+  case SPIRVStorageClass::StorageBuffer:
     bufferUsage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
     break;
-  case spirv::StorageClass::Uniform:
+  case SPIRVStorageClass::Uniform:
     bufferUsage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
     break;
-  default:
-    llvm::errs() << "unsupported storage class";
-    return failure();
   }
   return success();
 }
@@ -87,8 +95,7 @@ LogicalResult VulkanRuntime::countDeviceMemorySize() {
       if (resourceDataBindingPair.second.size) {
         memorySize += resourceDataBindingPair.second.size;
       } else {
-        llvm::errs()
-            << "expected buffer size greater than zero for resource data";
+        std::cerr << "expected buffer size greater than zero for resource data";
         return failure();
       }
     }
@@ -98,11 +105,11 @@ LogicalResult VulkanRuntime::countDeviceMemorySize() {
 
 LogicalResult VulkanRuntime::initRuntime() {
   if (!resourceData.size()) {
-    llvm::errs() << "Vulkan runtime needs at least one resource";
+    std::cerr << "Vulkan runtime needs at least one resource";
     return failure();
   }
   if (!binarySize || !binary) {
-    llvm::errs() << "binary shader size must be greater than zero";
+    std::cerr << "binary shader size must be greater than zero";
     return failure();
   }
   if (failed(countDeviceMemorySize())) {
@@ -130,7 +137,7 @@ LogicalResult VulkanRuntime::destroy() {
   vkDestroyDescriptorPool(device, descriptorPool, nullptr);
   vkDestroyPipeline(device, pipeline, nullptr);
   vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
-  for (auto &descriptorSetLayout: descriptorSetLayouts) {
+  for (auto &descriptorSetLayout : descriptorSetLayouts) {
     vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);
   }
   vkDestroyShaderModule(device, shaderModule, nullptr);
@@ -199,12 +206,12 @@ LogicalResult VulkanRuntime::run() {
             VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT),
         "vkGetQueryPoolResults");
     float microsec = (timestamps[1] - timestamps[0]) * timestampPeriod / 1000;
-    llvm::outs() << "Compute shader execution time: "
-                 << llvm::format("%0.3fus\n", microsec);
+    std::cout << "Compute shader execution time: " << std::setprecision(3)
+              << microsec << "us\n";
   }
 
-  llvm::outs() << "Command buffer submit time: " << submitDuration.count()
-               << "us\nWait idle time: " << execDuration.count() << "us\n";
+  std::cout << "Command buffer submit time: " << submitDuration.count()
+            << "us\nWait idle time: " << execDuration.count() << "us\n";
 
   return success();
 }
@@ -240,7 +247,7 @@ LogicalResult VulkanRuntime::createDevice() {
       vkEnumeratePhysicalDevices(instance, &physicalDeviceCount, 0),
       "vkEnumeratePhysicalDevices");
 
-  llvm::SmallVector<VkPhysicalDevice, 1> physicalDevices(physicalDeviceCount);
+  std::vector<VkPhysicalDevice> physicalDevices(physicalDeviceCount);
   RETURN_ON_VULKAN_ERROR(vkEnumeratePhysicalDevices(instance,
                                                     &physicalDeviceCount,
                                                     physicalDevices.data()),
@@ -313,7 +320,7 @@ LogicalResult VulkanRuntime::getBestComputeQueue() {
   vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice,
                                            &queueFamilyPropertiesCount, 0);
 
-  SmallVector<VkQueueFamilyProperties, 1> familyProperties(
+  std::vector<VkQueueFamilyProperties> familyProperties(
       queueFamilyPropertiesCount);
   vkGetPhysicalDeviceQueueFamilyProperties(
       physicalDevice, &queueFamilyPropertiesCount, familyProperties.data());
@@ -339,14 +346,14 @@ LogicalResult VulkanRuntime::getBestComputeQueue() {
     }
   }
 
-  llvm::errs() << "cannot find valid queue";
+  std::cerr << "cannot find valid queue";
   return failure();
 }
 
 LogicalResult VulkanRuntime::createMemoryBuffers() {
   // For each descriptor set.
   for (const auto &resourceDataMapPair : resourceData) {
-    llvm::SmallVector<VulkanDeviceMemoryBuffer, 1> deviceMemoryBuffers;
+    std::vector<VulkanDeviceMemoryBuffer> deviceMemoryBuffers;
     const auto descriptorSetIndex = resourceDataMapPair.first;
     const auto &resourceDataMap = resourceDataMapPair.second;
 
@@ -362,7 +369,7 @@ LogicalResult VulkanRuntime::createMemoryBuffers() {
       const auto resourceStorageClassMapIt =
           resourceStorageClassData.find(descriptorSetIndex);
       if (resourceStorageClassMapIt == resourceStorageClassData.end()) {
-        llvm::errs()
+        std::cerr
             << "cannot find storage class for resource in descriptor set: "
             << descriptorSetIndex;
         return failure();
@@ -373,7 +380,7 @@ LogicalResult VulkanRuntime::createMemoryBuffers() {
       const auto resourceStorageClassIt =
           resourceStorageClassMap.find(resourceDataBindingPair.first);
       if (resourceStorageClassIt == resourceStorageClassMap.end()) {
-        llvm::errs()
+        std::cerr
             << "cannot find storage class for resource with descriptor index: "
             << resourceDataBindingPair.first;
         return failure();
@@ -384,10 +391,10 @@ LogicalResult VulkanRuntime::createMemoryBuffers() {
                                                  descriptorType)) ||
           failed(mapStorageClassToBufferUsageFlag(resourceStorageClassBinding,
                                                   bufferUsage))) {
-        llvm::errs() << "storage class for resource with descriptor binding: "
-                     << resourceDataBindingPair.first
-                     << " in the descriptor set: " << descriptorSetIndex
-                     << " is not supported ";
+        std::cerr << "storage class for resource with descriptor binding: "
+                  << resourceDataBindingPair.first
+                  << " in the descriptor set: " << descriptorSetIndex
+                  << " is not supported ";
         return failure();
       }
 
@@ -464,7 +471,7 @@ LogicalResult VulkanRuntime::createShaderModule() {
 
 void VulkanRuntime::initDescriptorSetLayoutBindingMap() {
   for (const auto &deviceMemoryBufferMapPair : deviceMemoryBufferMap) {
-    SmallVector<VkDescriptorSetLayoutBinding, 1> descriptorSetLayoutBindings;
+    std::vector<VkDescriptorSetLayoutBinding> descriptorSetLayoutBindings;
     const auto &deviceMemoryBuffers = deviceMemoryBufferMapPair.second;
     const auto descriptorSetIndex = deviceMemoryBufferMapPair.first;
 
@@ -495,8 +502,8 @@ LogicalResult VulkanRuntime::createDescriptorSetLayout() {
         descriptorSetLayoutBindingMap.find(descriptorSetIndex);
 
     if (descriptorSetLayoutBindingIt == descriptorSetLayoutBindingMap.end()) {
-      llvm::errs() << "cannot find layout bindings for the set with number: "
-                   << descriptorSetIndex;
+      std::cerr << "cannot find layout bindings for the set with number: "
+                << descriptorSetIndex;
       return failure();
     }
 
@@ -573,7 +580,7 @@ LogicalResult VulkanRuntime::createComputePipeline() {
 }
 
 LogicalResult VulkanRuntime::createDescriptorPool() {
-  llvm::SmallVector<VkDescriptorPoolSize, 1> descriptorPoolSizes;
+  std::vector<VkDescriptorPoolSize> descriptorPoolSizes;
   for (const auto &descriptorSetInfo : descriptorSetInfoPool) {
     // For each descriptor set populate descriptor pool size.
     VkDescriptorPoolSize descriptorPoolSize = {};
@@ -616,7 +623,7 @@ LogicalResult VulkanRuntime::allocateDescriptorSets() {
 
 LogicalResult VulkanRuntime::setWriteDescriptors() {
   if (descriptorSets.size() != descriptorSetInfoPool.size()) {
-    llvm::errs() << "Each descriptor set must have descriptor set information";
+    std::cerr << "Each descriptor set must have descriptor set information";
     return failure();
   }
   // For each descriptor set.

mlir/tools/mlir-vulkan-runner/VulkanRuntime.h

@@ -13,14 +13,10 @@
 #ifndef VULKAN_RUNTIME_H
 #define VULKAN_RUNTIME_H
 
-#include "mlir/Dialect/SPIRV/SPIRVOps.h"
-#include "mlir/Dialect/SPIRV/Serialization.h"
-#include "mlir/IR/Module.h"
 #include "mlir/Support/LogicalResult.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/StringExtras.h"
-#include "llvm/Support/ToolOutputFile.h"
 
+#include <unordered_map>
+#include <vector>
 #include <vulkan/vulkan.h>
 
 using namespace mlir;
@@ -64,25 +60,23 @@ struct DescriptorSetInfo {
 };
 
 /// VulkanHostMemoryBuffer mapped into a descriptor set and a binding.
-using ResourceData =
-    llvm::DenseMap<DescriptorSetIndex,
-                   llvm::DenseMap<BindingIndex, VulkanHostMemoryBuffer>>;
+using ResourceData = std::unordered_map<
+    DescriptorSetIndex,
+    std::unordered_map<BindingIndex, VulkanHostMemoryBuffer>>;
+
+/// SPIR-V storage classes.
+/// Note that this duplicates spirv::StorageClass but it keeps the Vulkan
+/// runtime library detached from SPIR-V dialect, so we can avoid pick up lots
+/// of dependencies.
+enum class SPIRVStorageClass {
+  Uniform = 2,
+  StorageBuffer = 12,
+};
 
 /// StorageClass mapped into a descriptor set and a binding.
 using ResourceStorageClassBindingMap =
-    llvm::DenseMap<DescriptorSetIndex,
-                   llvm::DenseMap<BindingIndex, mlir::spirv::StorageClass>>;
-
-inline void emitVulkanError(const llvm::Twine &message, VkResult error) {
-  llvm::errs()
-      << message.concat(" failed with error code ").concat(llvm::Twine{error});
-}
-
-#define RETURN_ON_VULKAN_ERROR(result, msg)                                    \
-  if ((result) != VK_SUCCESS) {                                                \
-    emitVulkanError(msg, (result));                                            \
-    return failure();                                                          \
-  }
+    std::unordered_map<DescriptorSetIndex,
+                       std::unordered_map<BindingIndex, SPIRVStorageClass>>;
 
 /// Vulkan runtime.
 /// The purpose of this class is to run SPIR-V compute shader on Vulkan
@@ -150,12 +144,12 @@ class VulkanRuntime {
 
   /// Maps storage class to a descriptor type.
   LogicalResult
-  mapStorageClassToDescriptorType(spirv::StorageClass storageClass,
+  mapStorageClassToDescriptorType(SPIRVStorageClass storageClass,
                                   VkDescriptorType &descriptorType);
 
   /// Maps storage class to buffer usage flags.
   LogicalResult
-  mapStorageClassToBufferUsageFlag(spirv::StorageClass storageClass,
+  mapStorageClassToBufferUsageFlag(SPIRVStorageClass storageClass,
                                    VkBufferUsageFlagBits &bufferUsage);
 
   LogicalResult countDeviceMemorySize();
@@ -170,28 +164,27 @@ class VulkanRuntime {
   VkQueue queue{VK_NULL_HANDLE};
 
   /// Specifies VulkanDeviceMemoryBuffers divided into sets.
-  llvm::DenseMap<DescriptorSetIndex,
-                 llvm::SmallVector<VulkanDeviceMemoryBuffer, 1>>
+  std::unordered_map<DescriptorSetIndex, std::vector<VulkanDeviceMemoryBuffer>>
       deviceMemoryBufferMap;
 
   /// Specifies shader module.
   VkShaderModule shaderModule{VK_NULL_HANDLE};
 
   /// Specifies layout bindings.
-  llvm::DenseMap<DescriptorSetIndex,
-                 llvm::SmallVector<VkDescriptorSetLayoutBinding, 1>>
+  std::unordered_map<DescriptorSetIndex,
+                     std::vector<VkDescriptorSetLayoutBinding>>
       descriptorSetLayoutBindingMap;
 
   /// Specifies layouts of descriptor sets.
-  llvm::SmallVector<VkDescriptorSetLayout, 1> descriptorSetLayouts;
+  std::vector<VkDescriptorSetLayout> descriptorSetLayouts;
   VkPipelineLayout pipelineLayout{VK_NULL_HANDLE};
 
   /// Specifies descriptor sets.
-  llvm::SmallVector<VkDescriptorSet, 1> descriptorSets;
+  std::vector<VkDescriptorSet> descriptorSets;
 
   /// Specifies a pool of descriptor set info, each descriptor set must have
   /// information such as type, index and amount of bindings.
-  llvm::SmallVector<DescriptorSetInfo, 1> descriptorSetInfoPool;
+  std::vector<DescriptorSetInfo> descriptorSetInfoPool;
   VkDescriptorPool descriptorPool{VK_NULL_HANDLE};
 
   /// Timestamp query.
@@ -202,7 +195,7 @@ class VulkanRuntime {
   /// Computation pipeline.
   VkPipeline pipeline{VK_NULL_HANDLE};
   VkCommandPool commandPool{VK_NULL_HANDLE};
-  llvm::SmallVector<VkCommandBuffer, 1> commandBuffers;
+  std::vector<VkCommandBuffer> commandBuffers;
 
   //===--------------------------------------------------------------------===//
   // Vulkan memory context.

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

@@ -10,11 +10,11 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <iostream>
 #include <mutex>
 #include <numeric>
 
 #include "VulkanRuntime.h"
-#include "llvm/Support/raw_ostream.h"
 
 namespace {
 
@@ -51,7 +51,7 @@ class VulkanRuntimeManager {
     if (failed(vulkanRuntime.initRuntime()) || failed(vulkanRuntime.run()) ||
         failed(vulkanRuntime.updateHostMemoryBuffers()) ||
         failed(vulkanRuntime.destroy())) {
-      llvm::errs() << "runOnVulkan failed";
+      std::cerr << "runOnVulkan failed";
     }
   }