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

[Jakub Kuderski]

================
@@ -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
----------------
kuhar wrote:

Also link to the issue here

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