[clang] [MLIR] Update convert-gpu-to-spirv pass to prepare using GPU compilat… (PR #69941)

[Lei Zhang via cfe-commits]

================
@@ -54,22 +55,52 @@ void GPUToSPIRVPass::runOnOperation() {
 
   SmallVector<Operation *, 1> gpuModules;
   OpBuilder builder(context);
+
+  auto targetEnvSupportsKernelCapability = [](gpu::GPUModuleOp moduleOp) {
+    Operation *gpuModule = moduleOp.getOperation();
+    auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
+    spirv::TargetEnv targetEnv(targetAttr);
+    return targetEnv.allows(spirv::Capability::Kernel);
+  };
+
   module.walk([&](gpu::GPUModuleOp moduleOp) {
     // Clone each GPU kernel module for conversion, given that the GPU
     // launch op still needs the original GPU kernel module.
-    builder.setInsertionPoint(moduleOp.getOperation());
+    // For Vulkan Shader capabilities, we insert the newly converted SPIR-V
+    // module right after the original GPU module, as that's the expectation of
+    // the in-tree Vulkan runner.
+    // For OpenCL Kernel capabilities, we insert the newly converted SPIR-V
+    // module inside the original GPU module, as that's the expectaion of the
+    // normal GPU compilation pipeline.
+    if (targetEnvSupportsKernelCapability(moduleOp)) {
+      builder.setInsertionPoint(moduleOp.getBody(),
+                                moduleOp.getBody()->begin());
+    } else {
+      builder.setInsertionPoint(moduleOp.getOperation());
+    }
     gpuModules.push_back(builder.clone(*moduleOp.getOperation()));
   });
 
   // Run conversion for each module independently as they can have different
   // TargetEnv attributes.
   for (Operation *gpuModule : gpuModules) {
+    spirv::TargetEnvAttr targetAttr =
+        spirv::lookupTargetEnvOrDefault(gpuModule);
+
     // Map MemRef memory space to SPIR-V storage class first if requested.
     if (mapMemorySpace) {
+      spirv::TargetEnv targetEnv(targetAttr);
+      FailureOr<spirv::MemoryModel> memoryModel =
+          spirv::getMemoryModel(targetEnv);
+      if (failed(memoryModel))
+        return signalPassFailure();
+
       std::unique_ptr<ConversionTarget> target =
           spirv::getMemorySpaceToStorageClassTarget(*context);
       spirv::MemorySpaceToStorageClassMap memorySpaceMap =
-          spirv::mapMemorySpaceToVulkanStorageClass;
+          (memoryModel == spirv::MemoryModel::OpenCL)
----------------
antiagainst wrote:

Here we can also use the above `targetEnvSupportsKernelCapability` no? Something like:

```c++
spirv::MemorySpaceToStorageClassMap memorySpaceMap =
  targetEnvSupportsKernelCapability(gpuModule) ? ...
```

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