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

[Lei Zhang]

================
@@ -54,22 +55,67 @@ void GPUToSPIRVPass::runOnOperation() {
 
   SmallVector<Operation *, 1> gpuModules;
   OpBuilder builder(context);
+
+  auto getTargetEnvFromGPUModuleOp = [*this](gpu::GPUModuleOp moduleOp) {
+    Operation *gpuModule = moduleOp.getOperation();
+    auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
+    std::unique_ptr<ConversionTarget> target =
+        SPIRVConversionTarget::get(targetAttr);
+
+    SPIRVConversionOptions options;
+    options.use64bitIndex = this->use64bitIndex;
+    SPIRVTypeConverter typeConverter(targetAttr, options);
+    const spirv::TargetEnv &targetEnv = typeConverter.getTargetEnv();
+    return targetEnv;
+  };
+
   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());
+    // SPIRV module insertion point by is after original GPU module.
+    // This works fine for Vulkan shader that has a dedicated runner.
+    // But OpenCL kernel needs SPIRV module placed inside original GPU module as
+    // OpenCL uses GPU compilation pipeline.
+    const mlir::spirv::TargetEnv &targetEnv =
+        getTargetEnvFromGPUModuleOp(moduleOp);
+    FailureOr<spirv::MemoryModel> memoryModel =
+        spirv::getMemoryModel(targetEnv);
+    if (failed(memoryModel))
+      return signalPassFailure();
+    if (memoryModel == spirv::MemoryModel::OpenCL) {
+      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) {
+    mlir::spirv::TargetEnvAttr targetAttr =
----------------
antiagainst wrote:

Can we use the lambda function in the above? This is duplicating the logic there.

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