[Mlir-commits] [mlir] [mlir][spirv] Add mgpu* wrappers for Vulkan runtime, migrate some tests (PR #123114)

[Andrea Faulds]

================
@@ -91,6 +124,88 @@ void bindMemRef(void *vkRuntimeManager, DescriptorSetIndex setIndex,
 }
 
 extern "C" {
+
+//===----------------------------------------------------------------------===//
+//
+// New wrappers, intended for mlir-cpu-runner. Calls to these are generated by
+// GPUToLLVMConversionPass.
+//
+//===----------------------------------------------------------------------===//
+
+VULKAN_WRAPPER_SYMBOL_EXPORT void *mgpuStreamCreate() {
+  return new VulkanRuntimeManager();
+}
+
+VULKAN_WRAPPER_SYMBOL_EXPORT void mgpuStreamDestroy(void *vkRuntimeManager) {
+  delete static_cast<VulkanRuntimeManager *>(vkRuntimeManager);
+}
+
+VULKAN_WRAPPER_SYMBOL_EXPORT void mgpuStreamSynchronize(void *) {
+  // Currently a no-op as the other operations are synchronous.
+}
+
+VULKAN_WRAPPER_SYMBOL_EXPORT void *mgpuModuleLoad(const void *data,
+                                                  size_t gpuBlobSize) {
+  return new VulkanModule(static_cast<const char *>(data), gpuBlobSize);
+}
+
+VULKAN_WRAPPER_SYMBOL_EXPORT void mgpuModuleUnload(void *vkModule) {
+  delete static_cast<VulkanModule *>(vkModule);
+}
+
+VULKAN_WRAPPER_SYMBOL_EXPORT void *mgpuModuleGetFunction(void *vkModule,
+                                                         const char *name) {
+  return static_cast<VulkanModule *>(vkModule)->getFunction(name);
+}
+
+VULKAN_WRAPPER_SYMBOL_EXPORT void
+mgpuLaunchKernel(void *vkKernel, size_t gridX, size_t gridY, size_t gridZ,
+                 size_t /*blockX*/, size_t /*blockY*/, size_t /*blockZ*/,
+                 size_t /*smem*/, void *vkRuntimeManager, void **params,
+                 void ** /*extra*/, size_t paramsCount) {
+  auto manager = static_cast<VulkanRuntimeManager *>(vkRuntimeManager);
+
+  // The non-bare-pointer memref ABI interacts badly with mgpuLaunchKernel's
+  // signature:
+  // - The memref descriptor struct gets split into several elements, each
+  //   passed as their own "param".
+  // - No metadata is provided as to the rank or element type/size of a memref.
+  //   Here we assume that all MemRefs have rank 1 and an element size of
+  //   4 bytes. This means each descriptor struct will have five members.
+  // TODO: Refactor the ABI/API of mgpuLaunchKernel to use a different ABI for
+  //       memrefs, so that other memref types can also be used. This will allow
+  //       migrating the remaining tests and removal of mlir-vulkan-runner.
+  const size_t paramsPerMemRef = 5;
+  if (paramsCount % paramsPerMemRef) {
+    abort();
+  }
+  const DescriptorSetIndex setIndex = 0;
+  BindingIndex bindIndex = 0;
+  for (size_t i = 0; i < paramsCount; i += paramsPerMemRef) {
+    auto memref = static_cast<MemRefDescriptor<uint32_t, 1> *>(params[i]);
+    bindMemRef<uint32_t, 1>(manager, setIndex, bindIndex, memref);
+    bindIndex++;
+  }
+
+  manager->setNumWorkGroups(NumWorkGroups{static_cast<uint32_t>(gridX),
+                                          static_cast<uint32_t>(gridY),
+                                          static_cast<uint32_t>(gridZ)});
+
+  auto function = static_cast<VulkanFunction *>(vkKernel);
+  manager->setShaderModule(function->module->blobData(),
+                           function->module->blobSize());
----------------
andfau-amd wrote:

Apparently bytes. 

https://github.com/llvm/llvm-project/pull/123114