[mlir][spirv] GPUToSPIRVPass: support case when `TargetEnv` attribute attached to the `gpu.module`

[Hardcode84]

Summary

Previously, only case when TargetEnv was attached to the top level ModuleOp was supported.

[Hardcode84]

fix use-after-free bug

[antiagainst]

Thanks! Generally LGTM; I just have one question actually.

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

342

The spirv.target_env should only be needed when _converting to_ SPIR-V? Once that's done we are in SPIR-V land already; it should not be necessary anymore. So curious why would you need to copy it over?

[Hardcode84]

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

342

Some of the patterns (specifically GPUFuncOpConversion) calls lookupTargetEnvOrDefault themselves during conversion and they won't find it if they run after GPUModuleConversion. Also we will need to preserve it if we ever decide to switch to progressive lowering.

[antiagainst]

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

342

Good points. Would mind to put the above as comments so later we know why this is needed? Thanks!

[Hardcode84]

update comment

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

auto spvModule = rewriter.create<spirv::ModuleOp>(
       StringRef(spvModuleName));
 
   // Move the region from the module op into the SPIR-V module.
   Region &spvModuleRegion = spvModule.getRegion();
   rewriter.inlineRegionBefore(moduleOp.getBodyRegion(), spvModuleRegion,
                               spvModuleRegion.begin());
   // The spirv.module build method adds a block. Remove that.
   rewriter.eraseBlock(&spvModuleRegion.back());
+
+  // Some of the patterns call `lookupTargetEnv` during conversion and they
+  // will fail if called after GPUModuleConversion and we don't preserve
+  // `TargetEnv` attribute.
+  // Copy TargetEnvAttr only if it is attached directly to the GPUModuleOp.
+  if (auto attr = moduleOp->getAttrOfType<spirv::TargetEnvAttr>(
+          spirv::getTargetEnvAttrName()))
+    spvModule->setAttr(spirv::getTargetEnvAttrName(), attr);
+
   rewriter.eraseOp(moduleOp);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // GPU return inside kernel functions to SPIR-V return.
 //===----------------------------------------------------------------------===//
 

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp

void GPUToSPIRVPass::runOnOperation() {
   OpBuilder builder(context);
   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());
     gpuModules.push_back(builder.clone(*moduleOp.getOperation()));
   });
 
+  // Run conversion for each module independently as they can have different
+  // TargetEnv attributes.
+  for (Operation *gpuModule : gpuModules) {
     // Map MemRef memory space to SPIR-V storage class first if requested.
     if (mapMemorySpace) {
       std::unique_ptr<ConversionTarget> target =
           spirv::getMemorySpaceToStorageClassTarget(*context);
       spirv::MemorySpaceToStorageClassMap memorySpaceMap =
           spirv::mapMemorySpaceToVulkanStorageClass;
       spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);
 
       RewritePatternSet patterns(context);
       spirv::populateMemorySpaceToStorageClassPatterns(converter, patterns);
 
-    if (failed(applyFullConversion(gpuModules, *target, std::move(patterns))))
+      if (failed(applyFullConversion(gpuModule, *target, std::move(patterns))))
         return signalPassFailure();
     }
 
-  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+    auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
     std::unique_ptr<ConversionTarget> target =
         SPIRVConversionTarget::get(targetAttr);
 
     SPIRVTypeConverter typeConverter(targetAttr);
     RewritePatternSet patterns(context);
     populateGPUToSPIRVPatterns(typeConverter, patterns);
 
     // TODO: Change SPIR-V conversion to be progressive and remove the following
     // patterns.
     mlir::arith::populateArithToSPIRVPatterns(typeConverter, patterns);
     populateMemRefToSPIRVPatterns(typeConverter, patterns);
     populateFuncToSPIRVPatterns(typeConverter, patterns);
 
-  if (failed(applyFullConversion(gpuModules, *target, std::move(patterns))))
+    if (failed(applyFullConversion(gpuModule, *target, std::move(patterns))))
       return signalPassFailure();
   }
+}
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertGPUToSPIRVPass(bool mapMemorySpace) {
   return std::make_unique<GPUToSPIRVPass>(mapMemorySpace);
 }

mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -convert-gpu-to-spirv -verify-diagnostics %s -o - | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -convert-gpu-to-spirv -verify-diagnostics -split-input-file %s -o - | FileCheck %s
 
 module attributes {
   gpu.container_module,
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kernel, Addresses], []>, #spirv.resource_limits<>>
 } {
   gpu.module @kernels {
     // CHECK-LABEL: spirv.module @{{.*}} Physical64 OpenCL
     //       CHECK:   spirv.func
func.func @main() {
     %1 = "op"() : () -> (memref<12xf32, #spirv.storage_class<CrossWorkgroup>>)
     %cst = arith.constant 1 : index
     gpu.launch_func @kernels::@basic_module_structure
         blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
         args(%0 : f32, %1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>)
     return
   }
 }
+
+// -----
+
+module attributes {
+  gpu.container_module
+} {
+  gpu.module @kernels attributes {
+    spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kernel, Addresses], []>, #spirv.resource_limits<>>
+  } {
+    // CHECK-LABEL: spirv.module @{{.*}} Physical64 OpenCL
+    //  CHECK-SAME: spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kernel, Addresses], []>, #spirv.resource_limits<>>
+    //       CHECK:   spirv.func
+    //  CHECK-SAME:     {{%.*}}: f32
+    //   CHECK-NOT:     spirv.interface_var_abi
+    //  CHECK-SAME:     {{%.*}}: !spirv.ptr<!spirv.array<12 x f32>, CrossWorkgroup>
+    //   CHECK-NOT:     spirv.interface_var_abi
+    //  CHECK-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
+    gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>) kernel
+        attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]>: vector<3xi32>>} {
+      gpu.return
+    }
+  }
+
+  func.func @main() {
+    %0 = "op"() : () -> (f32)
+    %1 = "op"() : () -> (memref<12xf32, #spirv.storage_class<CrossWorkgroup>>)
+    %cst = arith.constant 1 : index
+    gpu.launch_func @kernels::@basic_module_structure
+        blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
+        args(%0 : f32, %1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>)
+    return
+  }
+}