[Mlir-commits] [llvm] [mlir] [mlir][spirv] Integrate `convert-to-spirv` into `mlir-vulkan-runner` (PR #106082)

[Jakub Kuderski]

================
@@ -64,31 +94,32 @@ struct ConvertToSPIRVPass final
     RewritePatternSet patterns(context);
     ScfToSPIRVContext scfToSPIRVContext;
 
-    // Map MemRef memory space to SPIR-V storage class.
-    spirv::TargetEnv targetEnv(targetAttr);
-    bool targetEnvSupportsKernelCapability =
-        targetEnv.allows(spirv::Capability::Kernel);
-    spirv::MemorySpaceToStorageClassMap memorySpaceMap =
-        targetEnvSupportsKernelCapability
-            ? spirv::mapMemorySpaceToOpenCLStorageClass
-            : spirv::mapMemorySpaceToVulkanStorageClass;
-    spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);
-    spirv::convertMemRefTypesAndAttrs(op, converter);
-
-    // Populate patterns for each dialect.
-    arith::populateCeilFloorDivExpandOpsPatterns(patterns);
-    arith::populateArithToSPIRVPatterns(typeConverter, patterns);
-    populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
-    populateFuncToSPIRVPatterns(typeConverter, patterns);
-    populateGPUToSPIRVPatterns(typeConverter, patterns);
-    index::populateIndexToSPIRVPatterns(typeConverter, patterns);
-    populateMemRefToSPIRVPatterns(typeConverter, patterns);
-    populateVectorToSPIRVPatterns(typeConverter, patterns);
-    populateSCFToSPIRVPatterns(typeConverter, scfToSPIRVContext, patterns);
-    ub::populateUBToSPIRVConversionPatterns(typeConverter, patterns);
-
-    if (failed(applyPartialConversion(op, *target, std::move(patterns))))
-      return signalPassFailure();
+    if (runOnGPUModules) {
+      SmallVector<Operation *, 1> gpuModules;
+      OpBuilder builder(context);
+      op->walk([&](gpu::GPUModuleOp gpuModule) {
+        builder.setInsertionPoint(gpuModule);
+        gpuModules.push_back(builder.clone(*gpuModule));
+      });
+      // Run conversion for each module independently as they can have
+      // different TargetEnv attributes.
+      for (Operation *gpuModule : gpuModules) {
+        spirv::TargetEnvAttr targetAttr =
+            spirv::lookupTargetEnvOrDefault(gpuModule);
+        mapToMemRef(gpuModule, targetAttr);
+        populateConvertToSPIRVPatterns(typeConverter, scfToSPIRVContext,
+                                       patterns);
+        if (failed(
+                applyFullConversion(gpuModule, *target, std::move(patterns))))
+          return signalPassFailure();
+      }
+    } else {
----------------
kuhar wrote:

Prefer early exits: https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code

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