Create a gpu.module operation for the GPU Dialect.

[tpopp]

Summary

This is based on the use of code constantly checking for an attribute on
a model and instead represents the distinct operaion with a different
op. Instead, this op can be used to provide better filtering.

Reverts "Revert "[mlir] Create a gpu.module operation for the GPU Dialect.""

This reverts commit ac446302ca4145cdc89f377c0c364c29ee303be5 after
fixing internal Google issues.

This additionally updates ROCDL lowering to use the new gpu.module.

[merge_guards_bot]

Unit tests: pass. 61953 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[mehdi_amini]

Can you expand on the original reason for the revert? What has been fix in the patch?

[tpopp]

Run clang-format with LLVM style instead of google style.

[tpopp]

In D72921#1828330, @mehdi_amini wrote:

Can you expand on the original reason for the revert? What has been fix in the patch?

In the commit message or just in general? An XLA:MLIR test was failing when Google tried to update LLVM, so this was reverted. It took some time to figure out what patch is needed on the Tensorflow side to successfully integrate this patch because the style of failure was only caught by integration tests in a code base I didn't know.

[merge_guards_bot]

Unit tests: pass. 61953 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[herhut]

Thanks.

[mehdi_amini]

In D72921#1829009, @tpopp wrote:

In D72921#1828330, @mehdi_amini wrote:

Can you expand on the original reason for the revert? What has been fix in the patch?

In the commit message or just in general? An XLA:MLIR test was failing when Google tried to update LLVM, so this was reverted. It took some time to figure out what patch is needed on the Tensorflow side to successfully integrate this patch because the style of failure was only caught by integration tests in a code base I didn't know.

OK, I expect that there is a strong suspicion of an upstream bug when reverting a commit ; "causing problem downstream" by itself isn't enough.

mlir/include/mlir/Conversion/GPUToCUDA/GPUToCUDAPass.h

 #include <string>
 #include <vector>
 
 namespace mlir {
 
 class Location;
 class ModuleOp;
 
+template <typename T> class OpPassBase;
+
+namespace gpu {
+class GPUModuleOp;
+} // namespace gpu
+
 namespace LLVM {
 class LLVMDialect;
 } // namespace LLVM
 
-template <typename T> class OpPassBase;
-
 using OwnedCubin = std::unique_ptr<std::vector<char>>;
 using CubinGenerator =
     std::function<OwnedCubin(const std::string &, Location, StringRef)>;
 
 /// Creates a pass to convert kernel functions into CUBIN blobs.
 ///
 /// This transformation takes the body of each function that is annotated with
 /// the 'nvvm.kernel' attribute, copies it to a new LLVM module, compiles the
 /// module with help of the nvptx backend to PTX and then invokes the provided
 /// cubinGenerator to produce a binary blob (the cubin). Such blob is then
 /// attached as a string attribute named 'nvvm.cubin' to the kernel function.
 /// After the transformation, the body of the kernel function is removed (i.e.,
 /// it is turned into a declaration).
-std::unique_ptr<OpPassBase<ModuleOp>>
+std::unique_ptr<OpPassBase<gpu::GPUModuleOp>>
 createConvertGPUKernelToCubinPass(CubinGenerator cubinGenerator);
 
 /// Creates a pass to convert a gpu.launch_func operation into a sequence of
 /// CUDA calls.
 ///
 /// This pass does not generate code to call CUDA directly but instead uses a
 /// small wrapper library that exports a stable and conveniently typed ABI
 /// on top of CUDA.
 std::unique_ptr<OpPassBase<ModuleOp>>
 createConvertGpuLaunchFuncToCudaCallsPass();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUTOCUDA_GPUTOCUDAPASS_H_

mlir/include/mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h

 //===- GPUToNVVMPass.h - Convert GPU kernel to NVVM dialect -----*- C++ -*-===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_CONVERSION_GPUTONVVM_GPUTONVVMPASS_H_
 #define MLIR_CONVERSION_GPUTONVVM_GPUTONVVMPASS_H_
 
 #include <memory>
 
 namespace mlir {
 class LLVMTypeConverter;
 class OwningRewritePatternList;
 
-class ModuleOp;
 template <typename OpT> class OpPassBase;
 
+namespace gpu {
+class GPUModuleOp;
+}
+
 /// Collect a set of patterns to convert from the GPU dialect to NVVM.
 void populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter,
                                          OwningRewritePatternList &patterns);
 
 /// Creates a pass that lowers GPU dialect operations to NVVM counterparts.
-std::unique_ptr<OpPassBase<ModuleOp>> createLowerGpuOpsToNVVMOpsPass();
+std::unique_ptr<OpPassBase<gpu::GPUModuleOp>> createLowerGpuOpsToNVVMOpsPass();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUTONVVM_GPUTONVVMPASS_H_

mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h

 //===- GPUToROCDLPass.h - Convert GPU kernel to ROCDL dialect ---*- C++ -*-===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_CONVERSION_GPUTOROCDL_GPUTOROCDLPASS_H_
 #define MLIR_CONVERSION_GPUTOROCDL_GPUTOROCDLPASS_H_
 
 #include <memory>
 
 namespace mlir {
 
-class ModuleOp;
+namespace gpu {
+class GPUModuleOp;
+} // namespace gpu
 template <typename OpT> class OpPassBase;
 
 /// Creates a pass that lowers GPU dialect operations to ROCDL counterparts.
-std::unique_ptr<OpPassBase<ModuleOp>> createLowerGpuOpsToROCDLOpsPass();
+std::unique_ptr<OpPassBase<gpu::GPUModuleOp>> createLowerGpuOpsToROCDLOpsPass();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUTOROCDL_GPUTOROCDLPASS_H_

mlir/include/mlir/Dialect/GPU/GPUOps.td

let description = [{
 
     Either none or all work items of a workgroup need to execute this op
     in convergence.
   }];
   let parser = [{ return success(); }];
   let printer = [{ p << getOperationName(); }];
 }
 
+def GPU_GPUModuleOp : GPU_Op<"module", [
+  IsolatedFromAbove, SymbolTable, Symbol,
+  SingleBlockImplicitTerminator<"ModuleEndOp">
+]> {
+  let summary = "A top level compilation unit containing code to be run on a GPU.";
+  let description = [{
+    GPU module contains code that is intended to be run on a GPU. A host device
+    can launch this code through a gpu.launc_func that creates a fully
+    qualified symbol through the gpu.module's symbol and a gpu.func symbol
+    contained in the gpu.module.
+
+    The module's top-level scope is modeled by a single region with a single
+    block. GPU modules are required to have a name that is used for symbol
+    resolution by the gpu.launch_func operation.
+
+    Using an op with a region to define a GPU module enables "embedding" GPU
+    modules with SIMT execution models in other dialects in a clean manner and
+    allows filtering of code regions to execute passes on only code intended to
+    or not intended to be run on the separate device.
+
+    ```
+      gpu.module @symbol_name {
+      gpu.func {}
+        ...
+      gpu.module_end
+    }
+
+    ```
+  }];
+  let builders = [OpBuilder<"Builder *builder, OperationState &result, "
+                            "StringRef name">];
+  let parser = [{ return ::parseGPUModuleOp(parser, result); }];
+  let printer = [{ return ::print(p, *this); }];
+  let regions = (region SizedRegion<1>:$body);
+
+  // We need to ensure the block inside the region is properly terminated;
+  // the auto-generated builders do not guarantee that.
+  let skipDefaultBuilders = 1;
+}
+
+def GPU_ModuleEndOp : GPU_Op<"module_end", [
+  Terminator, HasParent<"GPUModuleOp">
+]> {
+  let summary = "A pseudo op that marks the end of a gpu.module.";
+  let description = [{
+    This op terminates the only block inside the only region of a `gpu.module`.
+  }];
+
+  let parser = [{ return success(); }];
+  let printer = [{ p << getOperationName(); }];
+}
+
 #endif // GPU_OPS

mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp

 static constexpr const char *kCubinAnnotation = "nvvm.cubin";
 
 /// A pass converting tagged kernel modules to cubin blobs.
 ///
 /// If tagged as a kernel module, each contained function is translated to NVVM
 /// IR and further to PTX. A user provided CubinGenerator compiles the PTX to
 /// GPU binary code, which is then attached as an attribute to the function. The
 /// function body is erased.
-class GpuKernelToCubinPass : public ModulePass<GpuKernelToCubinPass> {
+class GpuKernelToCubinPass
+    : public OperationPass<GpuKernelToCubinPass, gpu::GPUModuleOp> {
 public:
   GpuKernelToCubinPass(
       CubinGenerator cubinGenerator = compilePtxToCubinForTesting)
       : cubinGenerator(cubinGenerator) {}
 
-  void runOnModule() override {
-    ModuleOp module = getModule();
-    if (!module.getAttrOfType<UnitAttr>(
-            gpu::GPUDialect::getKernelModuleAttrName()) ||
-        !module.getName())
-      return;
+  void runOnOperation() override {
+    gpu::GPUModuleOp module = getOperation();
 
     // Make sure the NVPTX target is initialized.
     LLVMInitializeNVPTXTarget();
     LLVMInitializeNVPTXTargetInfo();
     LLVMInitializeNVPTXTargetMC();
     LLVMInitializeNVPTXAsmPrinter();
 
     auto llvmModule = translateModuleToNVVMIR(module);
     if (!llvmModule)
       return signalPassFailure();
 
     // Translate the module to CUBIN and attach the result as attribute to the
     // module.
-    if (auto cubinAttr = translateGpuModuleToCubinAnnotation(
-            *llvmModule, module.getLoc(), *module.getName()))
+    if (auto cubinAttr = translateGPUModuleToCubinAnnotation(
+            *llvmModule, module.getLoc(), module.getName()))
       module.setAttr(kCubinAnnotation, cubinAttr);
     else
       signalPassFailure();
   }
 
 private:
   static OwnedCubin compilePtxToCubinForTesting(const std::string &ptx,
                                                 Location, StringRef);
 
   std::string translateModuleToPtx(llvm::Module &module,
                                    llvm::TargetMachine &target_machine);
 
   /// Converts llvmModule to cubin using the user-provided generator. Location
   /// is used for error reporting and name is forwarded to the CUBIN generator
   /// to use in its logging mechanisms.
   OwnedCubin convertModuleToCubin(llvm::Module &llvmModule, Location loc,
                                   StringRef name);
 
   /// Translates llvmModule to cubin and returns the result as attribute.
-  StringAttr translateGpuModuleToCubinAnnotation(llvm::Module &llvmModule,
+  StringAttr translateGPUModuleToCubinAnnotation(llvm::Module &llvmModule,
                                                  Location loc, StringRef name);
 
   CubinGenerator cubinGenerator;
 };
 
 } // anonymous namespace
 
 std::string GpuKernelToCubinPass::translateModuleToPtx(
OwnedCubin GpuKernelToCubinPass::convertModuleToCubin(llvm::Module &llvmModule,
   // Set the data layout of the llvm module to match what the ptx target needs.
   llvmModule.setDataLayout(targetMachine->createDataLayout());
 
   auto ptx = translateModuleToPtx(llvmModule, *targetMachine);
 
   return cubinGenerator(ptx, loc, name);
 }
 
-StringAttr GpuKernelToCubinPass::translateGpuModuleToCubinAnnotation(
+StringAttr GpuKernelToCubinPass::translateGPUModuleToCubinAnnotation(
     llvm::Module &llvmModule, Location loc, StringRef name) {
   auto cubin = convertModuleToCubin(llvmModule, loc, name);
   if (!cubin)
     return {};
   return StringAttr::get({cubin->data(), cubin->size()}, loc->getContext());
 }
 
-std::unique_ptr<OpPassBase<ModuleOp>>
+std::unique_ptr<OpPassBase<gpu::GPUModuleOp>>
 mlir::createConvertGPUKernelToCubinPass(CubinGenerator cubinGenerator) {
   return std::make_unique<GpuKernelToCubinPass>(cubinGenerator);
 }
 
 static PassRegistration<GpuKernelToCubinPass>
     pass("test-kernel-to-cubin",
          "Convert all kernel functions to CUDA cubin blobs");

mlir/lib/Conversion/GPUToCUDA/ConvertLaunchFuncToCudaCalls.cpp

void runOnModule() override {
     initializeCachedTypes();
 
     getModule().walk([this](mlir::gpu::LaunchFuncOp op) {
       translateGpuLaunchCalls(op);
     });
 
     // GPU kernel modules are no longer necessary since we have a global
     // constant with the CUBIN data.
-    for (auto m : llvm::make_early_inc_range(getModule().getOps<ModuleOp>()))
-      if (m.getAttrOfType<UnitAttr>(gpu::GPUDialect::getKernelModuleAttrName()))
-        m.erase();
+    for (auto m :
+         llvm::make_early_inc_range(getModule().getOps<gpu::GPUModuleOp>()))
   }
 
 private:
   LLVM::LLVMDialect *llvmDialect;
   LLVM::LLVMType llvmVoidType;
   LLVM::LLVMType llvmPointerType;
   LLVM::LLVMType llvmPointerPointerType;
   LLVM::LLVMType llvmInt8Type;
void GpuLaunchFuncToCudaCallsPass::translateGpuLaunchCalls(
   OpBuilder builder(launchOp);
   Location loc = launchOp.getLoc();
   declareCudaFunctions(loc);
 
   auto zero = builder.create<LLVM::ConstantOp>(loc, getInt32Type(),
                                                builder.getI32IntegerAttr(0));
   // Create an LLVM global with CUBIN extracted from the kernel annotation and
   // obtain a pointer to the first byte in it.
-  auto kernelModule =
-      getModule().lookupSymbol<ModuleOp>(launchOp.getKernelModuleName());
+  auto kernelModule = getModule().lookupSymbol<gpu::GPUModuleOp>(
+      launchOp.getKernelModuleName());
   assert(kernelModule && "expected a kernel module");
 
   auto cubinAttr = kernelModule.getAttrOfType<StringAttr>(kCubinAnnotation);
   if (!cubinAttr) {
     kernelModule.emitOpError()
         << "missing " << kCubinAnnotation << " attribute";
     return signalPassFailure();
   }
 
-  assert(kernelModule.getName() && "expected a named module");
-  SmallString<128> nameBuffer(*kernelModule.getName());
+  SmallString<128> nameBuffer(kernelModule.getName());
   nameBuffer.append(kCubinStorageSuffix);
   Value data = LLVM::createGlobalString(
       loc, builder, nameBuffer.str(), cubinAttr.getValue(),
       LLVM::Linkage::Internal, getLLVMDialect());
 
   // Emit the load module call to load the module data. Error checking is done
   // in the called helper function.
   auto cuModule = allocatePointer(builder, loc);

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

private:
   ///     return %result
   ///
   Value createBlockReduce(Location loc, Value operand,
                           AccumulatorFactory &accumFactory,
                           ConversionPatternRewriter &rewriter) const {
     auto type = operand.getType().cast<LLVM::LLVMType>();
 
     // Create shared memory array to store the warp reduction.
-    auto module = operand.getDefiningOp()->getParentOfType<ModuleOp>();
+    auto module = operand.getDefiningOp()->getParentOfType<gpu::GPUModuleOp>();
     assert(module && "op must belong to a module");
     Value sharedMemPtr =
         createSharedMemoryArray(loc, module, type, kWarpSize, rewriter);
 
     Value zero = rewriter.create<LLVM::ConstantOp>(
         loc, int32Type, rewriter.getI32IntegerAttr(0u));
     Value laneId = rewriter.create<NVVM::LaneIdOp>(loc, int32Type);
     Value isFirstLane = rewriter.create<LLVM::ICmpOp>(
createIf(
             value = accumFactory(loc, value, shflValue, rewriter);
           }
           return SmallVector<Value, 1>{value};
         });
     return rewriter.getInsertionBlock()->getArgument(0);
   }
 
   /// Creates a global array stored in shared memory.
-  Value createSharedMemoryArray(Location loc, ModuleOp module,
+  Value createSharedMemoryArray(Location loc, gpu::GPUModuleOp module,
                                 LLVM::LLVMType elementType, int numElements,
                                 ConversionPatternRewriter &rewriter) const {
-    OpBuilder builder(module.getBodyRegion());
+    OpBuilder builder(module.body());
 
     auto arrayType = LLVM::LLVMType::getArrayTy(elementType, numElements);
     StringRef name = "reduce_buffer";
     auto globalOp = builder.create<LLVM::GlobalOp>(
         loc, arrayType.cast<LLVM::LLVMType>(),
         /*isConstant=*/false, LLVM::Linkage::Internal, name,
         /*value=*/Attribute(), gpu::GPUDialect::getWorkgroupAddressSpace());
 

 /// Import the GPU Ops to NVVM Patterns.
 #include "GPUToNVVM.cpp.inc"
 
 /// A pass that replaces all occurrences of GPU device operations with their
 /// corresponding NVVM equivalent.
 ///
 /// This pass only handles device code and is not meant to be run on GPU host
 /// code.
-class LowerGpuOpsToNVVMOpsPass : public ModulePass<LowerGpuOpsToNVVMOpsPass> {
+class LowerGpuOpsToNVVMOpsPass
+    : public OperationPass<LowerGpuOpsToNVVMOpsPass, gpu::GPUModuleOp> {
 public:
-  void runOnModule() override {
-    ModuleOp m = getModule();
-    if (!m.getAttrOfType<UnitAttr>(gpu::GPUDialect::getKernelModuleAttrName()))
-      return;
-
+  void runOnOperation() override {
+    gpu::GPUModuleOp m = getOperation();
     OwningRewritePatternList patterns;
     NVVMTypeConverter converter(m.getContext());
     populateStdToLLVMConversionPatterns(converter, patterns);
     populateGpuToNVVMConversionPatterns(converter, patterns);
     ConversionTarget target(getContext());
     target.addIllegalDialect<gpu::GPUDialect>();
     target.addIllegalOp<LLVM::FAbsOp, LLVM::FCeilOp, LLVM::CosOp,
                         LLVM::ExpOp>();
     target.addIllegalOp<FuncOp>();
     target.addLegalDialect<LLVM::LLVMDialect>();
     target.addLegalDialect<NVVM::NVVMDialect>();
     // TODO(csigg): Remove once we support replacing non-root ops.
-    target.addLegalOp<gpu::YieldOp>();
+    target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>();
     if (failed(applyPartialConversion(m, target, patterns, &converter)))
       signalPassFailure();
   }
 };
 
 } // anonymous namespace
 
 void mlir::populateGpuToNVVMConversionPatterns(
void mlir::populateGpuToNVVMConversionPatterns(
   patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__nv_ceilf",
                                                "__nv_ceil");
   patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__nv_cosf",
                                                "__nv_cos");
   patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__nv_expf",
                                                "__nv_exp");
 }
 
-std::unique_ptr<OpPassBase<ModuleOp>> mlir::createLowerGpuOpsToNVVMOpsPass() {
+std::unique_ptr<OpPassBase<gpu::GPUModuleOp>>
+mlir::createLowerGpuOpsToNVVMOpsPass() {
   return std::make_unique<LowerGpuOpsToNVVMOpsPass>();
 }
 
 static PassRegistration<LowerGpuOpsToNVVMOpsPass>
     pass("convert-gpu-to-nvvm", "Generate NVVM operations for gpu operations");

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

 
 namespace {
 
 // A pass that replaces all occurrences of GPU device operations with their
 // corresponding ROCDL equivalent.
 //
 // This pass only handles device code and is not meant to be run on GPU host
 // code.
-class LowerGpuOpsToROCDLOpsPass : public ModulePass<LowerGpuOpsToROCDLOpsPass> {
+class LowerGpuOpsToROCDLOpsPass
+    : public OperationPass<LowerGpuOpsToROCDLOpsPass, gpu::GPUModuleOp> {
 public:
-  void runOnModule() override {
-    ModuleOp m = getModule();
-    if (!m.getAttrOfType<UnitAttr>(gpu::GPUDialect::getKernelModuleAttrName()))
-      return;
+  void runOnOperation() override {
+    gpu::GPUModuleOp m = getOperation();
 
     OwningRewritePatternList patterns;
     LLVMTypeConverter converter(m.getContext());
     populateStdToLLVMConversionPatterns(converter, patterns);
     patterns.insert<
         GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp,
                                     ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>,
         GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp,
target.addDynamicallyLegalOp<FuncOp>(
         [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
     if (failed(applyPartialConversion(m, target, patterns, &converter)))
       signalPassFailure();
   }
 };
 
 } // anonymous namespace
 
-std::unique_ptr<OpPassBase<ModuleOp>> mlir::createLowerGpuOpsToROCDLOpsPass() {
+std::unique_ptr<OpPassBase<gpu::GPUModuleOp>>
+mlir::createLowerGpuOpsToROCDLOpsPass() {
   return std::make_unique<LowerGpuOpsToROCDLOpsPass>();
 }
 
 static PassRegistration<LowerGpuOpsToROCDLOpsPass>
     pass("convert-gpu-to-rocdl",
          "Generate ROCDL operations for gpu operations");

mlir/lib/Conversion/GPUToSPIRV/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS GPUToSPIRV.td)
+mlir_tablegen(GPUToSPIRV.cpp.inc -gen-rewriters)
+add_public_tablegen_target(MLIRGPUToSPIRVIncGen)
+
 add_llvm_library(MLIRGPUtoSPIRVTransforms
   ConvertGPUToSPIRV.cpp
   ConvertGPUToSPIRVPass.cpp
   )
 
+add_dependencies(MLIRGPUtoSPIRVTransforms
+  MLIRGPUToSPIRVIncGen)
+
 target_link_libraries(MLIRGPUtoSPIRVTransforms
   MLIRGPU
   MLIRIR
   MLIRPass
   MLIRSPIRV
   MLIRStandardOps
   MLIRStandardToSPIRVTransforms
   MLIRSupport
   MLIRTransforms
   )

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

public:
   PatternMatchResult
   matchAndRewrite(gpu::GPUFuncOp funcOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 
 private:
   SmallVector<int32_t, 3> workGroupSizeAsInt32;
 };
 
-/// Pattern to convert a module with gpu.kernel_module attribute to a
-/// spv.module.
-class KernelModuleConversion final : public SPIRVOpLowering<ModuleOp> {
+/// Pattern to convert a gpu.module to a spv.module.
+class GPUModuleConversion final : public SPIRVOpLowering<gpu::GPUModuleOp> {
 public:
-  using SPIRVOpLowering<ModuleOp>::SPIRVOpLowering;
+  using SPIRVOpLowering<gpu::GPUModuleOp>::SPIRVOpLowering;
 
   PatternMatchResult
-  matchAndRewrite(ModuleOp moduleOp, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override;
-};
-
-/// Pattern to convert a module terminator op to a terminator of spv.module op.
-// TODO: Move this into DRR, but that requires ModuleTerminatorOp to be defined
-// in ODS.
-class KernelModuleTerminatorConversion final
-    : public SPIRVOpLowering<ModuleTerminatorOp> {
-public:
-  using SPIRVOpLowering<ModuleTerminatorOp>::SPIRVOpLowering;
-
-  PatternMatchResult
-  matchAndRewrite(ModuleTerminatorOp terminatorOp, ArrayRef<Value> operands,
+  matchAndRewrite(gpu::GPUModuleOp moduleOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 /// Pattern to convert a gpu.return into a SPIR-V return.
 // TODO: This can go to DRR when GPU return has operands.
 class GPUReturnOpConversion final : public SPIRVOpLowering<gpu::ReturnOp> {
 public:
   using SPIRVOpLowering<gpu::ReturnOp>::SPIRVOpLowering;
if (!newFuncOp) {
     return matchFailure();
   }
   newFuncOp.removeAttr(Identifier::get(gpu::GPUDialect::getKernelFuncAttrName(),
                                        rewriter.getContext()));
   return matchSuccess();
 }
 
 //===----------------------------------------------------------------------===//
-// ModuleOp with gpu.kernel_module.
+// ModuleOp with gpu.module.
 //===----------------------------------------------------------------------===//
 
-PatternMatchResult KernelModuleConversion::matchAndRewrite(
-    ModuleOp moduleOp, ArrayRef<Value> operands,
+PatternMatchResult GPUModuleConversion::matchAndRewrite(
+    gpu::GPUModuleOp moduleOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
-  if (!moduleOp.getAttrOfType<UnitAttr>(
-          gpu::GPUDialect::getKernelModuleAttrName())) {
-    return matchFailure();
-  }
   // TODO : Generalize this to account for different extensions,
   // capabilities, extended_instruction_sets, other addressing models
   // and memory models.
   auto spvModule = rewriter.create<spirv::ModuleOp>(
       moduleOp.getLoc(), spirv::AddressingModel::Logical,
       spirv::MemoryModel::GLSL450, spirv::Capability::Shader,
       spirv::Extension::SPV_KHR_storage_buffer_storage_class);
   // Move the region from the module op into the SPIR-V module.
-  Region &spvModuleRegion = spvModule.getOperation()->getRegion(0);
-  rewriter.inlineRegionBefore(moduleOp.getBodyRegion(), spvModuleRegion,
+  Region &spvModuleRegion = spvModule.body();
+  rewriter.inlineRegionBefore(moduleOp.body(), spvModuleRegion,
                               spvModuleRegion.begin());
   // The spv.module build method adds a block with a terminator. Remove that
   // block. The terminator of the module op in the remaining block will be
   // legalized later.
   spvModuleRegion.back().erase();
   rewriter.eraseOp(moduleOp);
   return matchSuccess();
 }
 
 //===----------------------------------------------------------------------===//
-// ModuleTerminatorOp for gpu.kernel_module.
-//===----------------------------------------------------------------------===//
-
-PatternMatchResult KernelModuleTerminatorConversion::matchAndRewrite(
-    ModuleTerminatorOp terminatorOp, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
-  rewriter.replaceOpWithNewOp<spirv::ModuleEndOp>(terminatorOp);
-  return matchSuccess();
-}
-
-//===----------------------------------------------------------------------===//
 // GPU return inside kernel functions to SPIR-V return.
 //===----------------------------------------------------------------------===//
 
 PatternMatchResult GPUReturnOpConversion::matchAndRewrite(
     gpu::ReturnOp returnOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   if (!operands.empty())
     return matchFailure();
 
   rewriter.replaceOpWithNewOp<spirv::ReturnOp>(returnOp);
   return matchSuccess();
 }
 
 //===----------------------------------------------------------------------===//
 // GPU To SPIRV Patterns.
 //===----------------------------------------------------------------------===//
 
+namespace {
+#include "GPUToSPIRV.cpp.inc"
+}
+
 void mlir::populateGPUToSPIRVPatterns(MLIRContext *context,
                                       SPIRVTypeConverter &typeConverter,
                                       OwningRewritePatternList &patterns,
                                       ArrayRef<int64_t> workGroupSize) {
+  populateWithGenerated(context, &patterns);
   patterns.insert<KernelFnConversion>(context, typeConverter, workGroupSize);
   patterns.insert<
       ForOpConversion, GPUReturnOpConversion, IfOpConversion,
-      KernelModuleConversion, KernelModuleTerminatorConversion,
+      GPUModuleConversion,
+      GPUReturnOpConversion, ForOpConversion, GPUModuleConversion,
       LaunchConfigConversion<gpu::BlockDimOp, spirv::BuiltIn::WorkgroupSize>,
       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
       LaunchConfigConversion<gpu::ThreadIdOp,
                              spirv::BuiltIn::LocalInvocationId>,
       TerminatorOpConversion>(context, typeConverter);
 }

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp

 } // namespace
 
 void GPUToSPIRVPass::runOnModule() {
   MLIRContext *context = &getContext();
   ModuleOp module = getModule();
 
   SmallVector<Operation *, 1> kernelModules;
   OpBuilder builder(context);
-  module.walk([&builder, &kernelModules](ModuleOp moduleOp) {
-    if (moduleOp.getAttrOfType<UnitAttr>(
-            gpu::GPUDialect::getKernelModuleAttrName())) {
-      // For each kernel module (should be only 1 for now, but that is not a
-      // requirement here), clone the module for conversion because the
-      // gpu.launch function still needs the kernel module.
-      builder.setInsertionPoint(moduleOp.getOperation());
-      kernelModules.push_back(builder.clone(*moduleOp.getOperation()));
-    }
+  module.walk([&builder, &kernelModules](gpu::GPUModuleOp moduleOp) {
   });
 
   SPIRVTypeConverter typeConverter;
   OwningRewritePatternList patterns;
   populateGPUToSPIRVPatterns(context, typeConverter, patterns, workGroupSize);
   populateStandardToSPIRVPatterns(context, typeConverter, patterns);
 
   std::unique_ptr<ConversionTarget> target = spirv::SPIRVConversionTarget::get(

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.td

+//===-- GPUToSPIRV.td - GPU to SPIR-V Dialect Lowerings ----*- tablegen -*-===//
+//
+// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains patterns to lower GPU dialect ops to to SPIR-V ops.
+//
+//===----------------------------------------------------------------------===//
+
+
+#ifndef CONVERT_GPU_TO_SPIRV
+#define CONVERT_GPU_TO_SPIRV
+
+include "mlir/Dialect/GPU/GPUOps.td"
+include "mlir/Dialect/SPIRV/SPIRVStructureOps.td"
+
+def : Pat<(GPU_ModuleEndOp), (SPV_ModuleEndOp)>;
+
+#endif // CONVERT_GPU_TO_SPIRV

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

auto walkResult = module.walk([&module](LaunchFuncOp launchOp) -> WalkResult {
     if (!launchOp.getAttrOfType<StringAttr>(
             LaunchFuncOp::getKernelAttrName()) ||
         !launchOp.getAttrOfType<SymbolRefAttr>(
             LaunchFuncOp::getKernelModuleAttrName()))
       return success();
 
     // Check that `launch_func` refers to a well-formed GPU kernel module.
     StringRef kernelModuleName = launchOp.getKernelModuleName();
-    auto kernelModule = module.lookupSymbol<ModuleOp>(kernelModuleName);
+    auto kernelModule = module.lookupSymbol<GPUModuleOp>(kernelModuleName);
     if (!kernelModule)
       return launchOp.emitOpError()
              << "kernel module '" << kernelModuleName << "' is undefined";
-    if (!kernelModule.getAttrOfType<UnitAttr>(
-            GPUDialect::getKernelModuleAttrName()))
-      return launchOp.emitOpError("module '")
-             << kernelModuleName << "' is missing the '"
-             << GPUDialect::getKernelModuleAttrName() << "' attribute";
 
     // Check that `launch_func` refers to a well-formed kernel function.
     StringRef kernelName = launchOp.kernel();
     Operation *kernelFunc = kernelModule.lookupSymbol(kernelName);
     auto kernelGPUFunction = dyn_cast_or_null<gpu::GPUFuncOp>(kernelFunc);
     auto kernelLLVMFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(kernelFunc);
     if (!kernelGPUFunction && !kernelLLVMFunction)
       return launchOp.emitOpError("kernel function '")
void LaunchFuncOp::build(Builder *builder, OperationState &result,
                          Value gridSizeZ, Value blockSizeX, Value blockSizeY,
                          Value blockSizeZ, ValueRange kernelOperands) {
   // Add grid and block sizes as op operands, followed by the data operands.
   result.addOperands(
       {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
   result.addOperands(kernelOperands);
   result.addAttribute(getKernelAttrName(),
                       builder->getStringAttr(kernelFunc.getName()));
-  auto kernelModule = kernelFunc.getParentOfType<ModuleOp>();
-  if (Optional<StringRef> kernelModuleName = kernelModule.getName())
-    result.addAttribute(getKernelModuleAttrName(),
-                        builder->getSymbolRefAttr(*kernelModuleName));
+  auto kernelModule = kernelFunc.getParentOfType<GPUModuleOp>();
+                      builder->getSymbolRefAttr(kernelModule.getName()));
 }
 
 void LaunchFuncOp::build(Builder *builder, OperationState &result,
                          GPUFuncOp kernelFunc, KernelDim3 gridSize,
                          KernelDim3 blockSize, ValueRange kernelOperands) {
   build(builder, result, kernelFunc, gridSize.x, gridSize.y, gridSize.z,
         blockSize.x, blockSize.y, blockSize.z, kernelOperands);
 }
if (failed(verifyAttributions(getOperation(), getWorkgroupAttributions(),
                                 GPUDialect::getWorkgroupAddressSpace())) ||
       failed(verifyAttributions(getOperation(), getPrivateAttributions(),
                                 GPUDialect::getPrivateAddressSpace())))
     return failure();
 
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// GPUModuleOp
+//===----------------------------------------------------------------------===//
+
+void GPUModuleOp::build(Builder *builder, OperationState &result,
+                        StringRef name) {
+  ensureTerminator(*result.addRegion(), *builder, result.location);
+  result.attributes.push_back(builder->getNamedAttr(
+      ::mlir::SymbolTable::getSymbolAttrName(), builder->getStringAttr(name)));
+}
+
+static ParseResult parseGPUModuleOp(OpAsmParser &parser,
+                                    OperationState &result) {
+  StringAttr nameAttr;
+  if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
+                             result.attributes))
+    return failure();
+
+  // If module attributes are present, parse them.
+  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
+    return failure();
+
+  // Parse the module body.
+  auto *body = result.addRegion();
+  if (parser.parseRegion(*body, None, None))
+    return failure();
+
+  // Ensure that this module has a valid terminator.
+  GPUModuleOp::ensureTerminator(*body, parser.getBuilder(), result.location);
+  return success();
+}
+
+static void print(OpAsmPrinter &p, GPUModuleOp op) {
+  p << op.getOperationName() << ' ';
+  p.printSymbolName(op.getName());
+  p.printOptionalAttrDictWithKeyword(op.getAttrs(),
+                                     {SymbolTable::getSymbolAttrName()});
+  p.printRegion(op.getOperation()->getRegion(0), /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/false);
+}
+
 // Namespace avoids ambiguous ReturnOpOperandAdaptor.
 namespace mlir {
 namespace gpu {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
 } // namespace gpu
 } // namespace mlir

mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp

 namespace {
 
 /// Pass that moves the kernel of each LaunchOp into its separate nested module.
 ///
 /// This pass moves the kernel code of each LaunchOp into a function created
 /// inside a nested module. It also creates an external function of the same
 /// name in the parent module.
 ///
-/// The kernel modules are intended to be compiled to a cubin blob independently
-/// in a separate pass. The external functions can then be annotated with the
+/// The gpu.modules are intended to be compiled to a cubin blob independently in
+/// a separate pass. The external functions can then be annotated with the
 /// symbol of the cubin accessor function.
 class GpuKernelOutliningPass : public ModulePass<GpuKernelOutliningPass> {
 public:
   void runOnModule() override {
     SymbolTable symbolTable(getModule());
     bool modified = false;
     for (auto func : getModule().getOps<FuncOp>()) {
       // Insert just after the function.
void runOnModule() override {
     // If any new module was inserted in this module, annotate this module as
     // a container module.
     if (modified)
       getModule().setAttr(gpu::GPUDialect::getContainerModuleAttrName(),
                           UnitAttr::get(&getContext()));
   }
 
 private:
-  // Returns a module containing kernelFunc and all callees (recursive).
-  ModuleOp createKernelModule(gpu::GPUFuncOp kernelFunc,
+  // Returns a gpu.module containing kernelFunc and all callees (recursive).
+  gpu::GPUModuleOp createKernelModule(gpu::GPUFuncOp kernelFunc,
                                       const SymbolTable &parentSymbolTable) {
+    // TODO: This code cannot use an OpBuilder because it must be inserted into
+    // a SymbolTable by the caller. SymbolTable needs to be refactored to
+    // prevent manual building of Ops with symbols in code using SymbolTables
+    // and then this needs to use the OpBuilder.
     auto context = getModule().getContext();
     Builder builder(context);
-    auto kernelModule =
-        ModuleOp::create(builder.getUnknownLoc(), kernelFunc.getName());
-    kernelModule.setAttr(gpu::GPUDialect::getKernelModuleAttrName(),
-                         builder.getUnitAttr());
+    OperationState state(kernelFunc.getLoc(),
+                         gpu::GPUModuleOp::getOperationName());
+    gpu::GPUModuleOp::build(&builder, state, kernelFunc.getName());
+    auto kernelModule = cast<gpu::GPUModuleOp>(Operation::create(state));
     SymbolTable symbolTable(kernelModule);
     symbolTable.insert(kernelFunc);
 
     SmallVector<Operation *, 8> symbolDefWorklist = {kernelFunc};
     while (!symbolDefWorklist.empty()) {
       if (Optional<SymbolTable::UseRange> symbolUses =
               SymbolTable::getSymbolUses(symbolDefWorklist.pop_back_val())) {
         for (SymbolTable::SymbolUse symbolUse : *symbolUses) {

mlir/test/Conversion/GPUToCUDA/lower-launch-func-to-cuda.mlir

 // RUN: mlir-opt %s --launch-func-to-cuda | FileCheck %s
 
 module attributes {gpu.container_module} {
 
   // CHECK: llvm.mlir.global internal constant @[[kernel_name:.*]]("kernel\00")
   // CHECK: llvm.mlir.global internal constant @[[global:.*]]("CUBIN")
 
-  module @kernel_module attributes {gpu.kernel_module, nvvm.cubin = "CUBIN"} {
+  gpu.module @kernel_module attributes {nvvm.cubin = "CUBIN"} {
     gpu.func @kernel(%arg0: !llvm.float, %arg1: !llvm<"float*">) attributes {gpu.kernel} {
       gpu.return
     }
   }
 
   llvm.func @foo() {
     %0 = "op"() : () -> (!llvm.float)
     %1 = "op"() : () -> (!llvm<"float*">)

mlir/test/Conversion/GPUToCUDA/lower-nvvm-kernel-to-cubin.mlir

 // RUN: mlir-opt %s --test-kernel-to-cubin -split-input-file | FileCheck %s
 
-// CHECK: attributes {gpu.kernel_module, nvvm.cubin = "CUBIN"}
-module @foo attributes {gpu.kernel_module} {
+// CHECK: attributes {nvvm.cubin = "CUBIN"}
+gpu.module @foo {
   llvm.func @kernel(%arg0 : !llvm.float, %arg1 : !llvm<"float*">)
     // CHECK: attributes  {gpu.kernel}
     attributes  { gpu.kernel } {
     llvm.return
   }
 }
 
 // -----
 
-module @bar attributes {gpu.kernel_module} {
+gpu.module @bar {
   // CHECK: func @kernel_a
   llvm.func @kernel_a()
     attributes  { gpu.kernel } {
     llvm.return
   }
 
   // CHECK: func @kernel_b
   llvm.func @kernel_b()
     attributes  { gpu.kernel } {
     llvm.return
   }
 }

mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir

 // RUN: mlir-opt %s -convert-gpu-to-nvvm -split-input-file | FileCheck %s
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK-LABEL: func @gpu_index_ops()
   func @gpu_index_ops()
       attributes { gpu.kernel } {
     // CHECK: = nvvm.read.ptx.sreg.tid.x : !llvm.i32
     %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
     // CHECK: = nvvm.read.ptx.sreg.tid.y : !llvm.i32
     %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index)
     // CHECK: = nvvm.read.ptx.sreg.tid.z : !llvm.i32
func @gpu_index_ops()
     %gDimZ = "gpu.grid_dim"() {dimension = "z"} : () -> (index)
 
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK-LABEL: func @gpu_all_reduce_op()
   func @gpu_all_reduce_op()
       attributes { gpu.kernel } {
     %arg0 = constant 1.0 : f32
     // TODO(csigg): Check full IR expansion once lowering has settled.
     // CHECK: nvvm.shfl.sync.bfly
     // CHECK: nvvm.barrier0
     // CHECK: llvm.fadd
     %result = "gpu.all_reduce"(%arg0) ({}) {op = "add"} : (f32) -> (f32)
 
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK-LABEL: func @gpu_all_reduce_region()
   func @gpu_all_reduce_region()
       attributes { gpu.kernel } {
     %arg0 = constant 1 : i32
     // TODO(csigg): Check full IR expansion once lowering has settled.
     // CHECK: nvvm.shfl.sync.bfly
     // CHECK: nvvm.barrier0
     %result = "gpu.all_reduce"(%arg0) ({
     ^bb(%lhs : i32, %rhs : i32):
       %xor = xor %lhs, %rhs : i32
       "gpu.yield"(%xor) : (i32) -> ()
     }) : (i32) -> (i32)
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK-LABEL: func @gpu_shuffle()
   func @gpu_shuffle()
       attributes { gpu.kernel } {
     // CHECK: %[[#VALUE:]] = llvm.mlir.constant(1.000000e+00 : f32) : !llvm.float
     %arg0 = constant 1.0 : f32
     // CHECK: %[[#OFFSET:]] = llvm.mlir.constant(4 : i32) : !llvm.i32
     %arg1 = constant 4 : i32
     // CHECK: %[[#WIDTH:]] = llvm.mlir.constant(23 : i32) : !llvm.i32
     %arg2 = constant 23 : i32
     // CHECK: %[[#ONE:]] = llvm.mlir.constant(1 : i32) : !llvm.i32
     // CHECK: %[[#SHL:]] = llvm.shl %[[#ONE]], %[[#WIDTH]] : !llvm.i32
     // CHECK: %[[#MASK:]] = llvm.sub %[[#SHL]], %[[#ONE]] : !llvm.i32
     // CHECK: %[[#CLAMP:]] = llvm.sub %[[#WIDTH]], %[[#ONE]] : !llvm.i32
     // CHECK: %[[#SHFL:]] = nvvm.shfl.sync.bfly %[[#MASK]], %[[#VALUE]], %[[#OFFSET]], %[[#CLAMP]] : !llvm<"{ float, i1 }">
     // CHECK: llvm.extractvalue %[[#SHFL]][0 : index] : !llvm<"{ float, i1 }">
     // CHECK: llvm.extractvalue %[[#SHFL]][1 : index] : !llvm<"{ float, i1 }">
     %shfl, %pred = "gpu.shuffle"(%arg0, %arg1, %arg2) { mode = "xor" } : (f32, i32, i32) -> (f32, i1)
 
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK-LABEL: func @gpu_sync()
   func @gpu_sync()
       attributes { gpu.kernel } {
     // CHECK: nvvm.barrier0
     gpu.barrier
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK: llvm.func @__nv_fabsf(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__nv_fabs(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_fabs
   func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) {
     %result32 = std.absf %arg_f32 : f32
     // CHECK: llvm.call @__nv_fabsf(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.absf %arg_f64 : f64
     // CHECK: llvm.call @__nv_fabs(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK: llvm.func @__nv_ceilf(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__nv_ceil(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_ceil
   func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) {
     %result32 = std.ceilf %arg_f32 : f32
     // CHECK: llvm.call @__nv_ceilf(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.ceilf %arg_f64 : f64
     // CHECK: llvm.call @__nv_ceil(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK: llvm.func @__nv_cosf(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__nv_cos(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_cos
   func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) {
     %result32 = std.cos %arg_f32 : f32
     // CHECK: llvm.call @__nv_cosf(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.cos %arg_f64 : f64
     // CHECK: llvm.call @__nv_cos(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   // CHECK: llvm.func @__nv_expf(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__nv_exp(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_exp
   func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
     %exp_f32 = std.exp %arg_f32 : f32
     // CHECK: llvm.call @__nv_expf(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result_f32 = std.exp %exp_f32 : f32
     // CHECK: llvm.call @__nv_expf(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.exp %arg_f64 : f64
     // CHECK: llvm.call @__nv_exp(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 // -----
 
 // Test that we handled properly operation with SymbolTable other than module op
-module attributes {gpu.kernel_module} {
+gpu.module @test_module {
   "test.symbol_scope"() ({
   // CHECK: test.symbol_scope
   // CHECK: llvm.func @__nv_expf(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__nv_exp(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_exp
     func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
       %exp_f32 = std.exp %arg_f32 : f32
       // CHECK: llvm.call @__nv_expf(%{{.*}}) : (!llvm.float) -> !llvm.float

mlir/test/Conversion/GPUToNVVM/memory-attrbution.mlir

 // RUN: mlir-opt --convert-gpu-to-nvvm --split-input-file %s | FileCheck %s
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel {
   // CHECK-LABEL:  llvm.func @private
   gpu.func @private(%arg0: f32) private(%arg1: memref<4xf32, 5>) {
     // Allocate private memory inside the function.
     // CHECK: %[[size:.*]] = llvm.mlir.constant(4 : i64) : !llvm.i64
     // CHECK: %[[raw:.*]] = llvm.alloca %[[size]] x !llvm.float : (!llvm.i64) -> !llvm<"float*">
 
     // Populate the memref descriptor.
     // CHECK: %[[descr1:.*]] = llvm.mlir.undef : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
gpu.func @private(%arg0: f32) private(%arg1: memref<4xf32, 5>) {
     store %arg0, %arg1[%c0] : memref<4xf32, 5>
 
     "terminator"() : () -> ()
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel {
   // Workgroup buffers are allocated as globals.
   // CHECK: llvm.mlir.global internal @[[buffer:.*]]()
   // CHECK-SAME:  addr_space = 3
   // CHECK-SAME:  !llvm<"[4 x float]">
 
   // CHECK-LABEL: llvm.func @workgroup
   // CHECK-SAME: {
   gpu.func @workgroup(%arg0: f32) workgroup(%arg1: memref<4xf32, 3>) {
gpu.func @workgroup(%arg0: f32) workgroup(%arg1: memref<4xf32, 3>) {
     store %arg0, %arg1[%c0] : memref<4xf32, 3>
 
     "terminator"() : () -> ()
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel {
   // Check that the total size was computed correctly.
   // CHECK: llvm.mlir.global internal @[[buffer:.*]]()
   // CHECK-SAME:  addr_space = 3
   // CHECK-SAME:  !llvm<"[48 x float]">
 
   // CHECK-LABEL: llvm.func @workgroup3d
   gpu.func @workgroup3d(%arg0: f32) workgroup(%arg1: memref<4x2x6xf32, 3>) {
     // Get the address of the first element in the global array.
gpu.func @workgroup3d(%arg0: f32) workgroup(%arg1: memref<4x2x6xf32, 3>) {
     %c0 = constant 0 : index
     store %arg0, %arg1[%c0,%c0,%c0] : memref<4x2x6xf32, 3>
     "terminator"() : () -> ()
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel {
   // Check that several buffers are defined.
   // CHECK: llvm.mlir.global internal @[[buffer1:.*]]()
   // CHECK-SAME:  !llvm<"[1 x float]">
   // CHECK: llvm.mlir.global internal @[[buffer2:.*]]()
   // CHECK-SAME:  !llvm<"[2 x float]">
 
   // CHECK-LABEL: llvm.func @multiple
   gpu.func @multiple(%arg0: f32)

mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

 // RUN: mlir-opt %s -convert-gpu-to-rocdl -split-input-file | FileCheck %s
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel_module {
   // CHECK-LABEL: func @gpu_index_ops()
   func @gpu_index_ops()
       attributes { gpu.kernel } {
     // CHECK: rocdl.workitem.id.x : !llvm.i32
     %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
     // CHECK: rocdl.workitem.id.y : !llvm.i32
     %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index)
     // CHECK: rocdl.workitem.id.z : !llvm.i32
func @gpu_index_ops()
     %gDimZ = "gpu.grid_dim"() {dimension = "z"} : () -> (index)
 
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel_module {
   // CHECK: llvm.func @_ocml_fabs_f32(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @_ocml_fabs_f64(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_fabs
   func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) {
     %result32 = std.absf %arg_f32 : f32
     // CHECK: llvm.call @_ocml_fabs_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.absf %arg_f64 : f64
     // CHECK: llvm.call @_ocml_fabs_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel_module {
   // CHECK: llvm.func @_ocml_ceil_f32(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @_ocml_ceil_f64(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_ceil
   func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) {
     %result32 = std.ceilf %arg_f32 : f32
     // CHECK: llvm.call @_ocml_ceil_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.ceilf %arg_f64 : f64
     // CHECK: llvm.call @_ocml_ceil_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel_module {
   // CHECK: llvm.func @_ocml_cos_f32(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @_ocml_cos_f64(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_cos
   func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) {
     %result32 = std.cos %arg_f32 : f32
     // CHECK: llvm.call @_ocml_cos_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.cos %arg_f64 : f64
     // CHECK: llvm.call @_ocml_cos_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 // -----
 
-module attributes {gpu.kernel_module} {
+gpu.module @kernel_module {
   // CHECK: llvm.func @_ocml_exp_f32(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @_ocml_exp_f64(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_exp
   func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
     %exp_f32 = std.exp %arg_f32 : f32
     // CHECK: llvm.call @_ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result_f32 = std.exp %exp_f32 : f32
     // CHECK: llvm.call @_ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.exp %arg_f64 : f64
     // CHECK: llvm.call @_ocml_exp_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 
 // -----
 
 // Test that we handled properly operation with SymbolTable other than module op
-module attributes {gpu.kernel_module} {
+gpu.module @kernel_module {
   "test.symbol_scope"() ({
     // CHECK: test.symbol_scope
     // CHECK: llvm.func @_ocml_exp_f32(!llvm.float) -> !llvm.float
     // CHECK: llvm.func @_ocml_exp_f64(!llvm.double) -> !llvm.double
     // CHECK-LABEL: func @gpu_exp
     func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
       %exp_f32 = std.exp %arg_f32 : f32
       // CHECK: llvm.call @_ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

 // RUN: mlir-opt -split-input-file -convert-gpu-to-spirv %s -o - | FileCheck %s
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_workgroup_id_x", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL:  spv.module "Logical" "GLSL450"
   // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @builtin_workgroup_id_x()
       attributes {gpu.kernel} {
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[WORKGROUPID]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
       %0 = "gpu.block_id"() {dimension = "x"} : () -> index
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_workgroup_id_y", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL:  spv.module "Logical" "GLSL450"
   // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @builtin_workgroup_id_y()
       attributes {gpu.kernel} {
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[WORKGROUPID]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}1 : i32{{\]}}
       %0 = "gpu.block_id"() {dimension = "y"} : () -> index
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_workgroup_id_z", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL:  spv.module "Logical" "GLSL450"
   // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @builtin_workgroup_id_z()
       attributes {gpu.kernel} {
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[WORKGROUPID]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}2 : i32{{\]}}
       %0 = "gpu.block_id"() {dimension = "z"} : () -> index
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_workgroup_size_x", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL:  spv.module "Logical" "GLSL450"
   // CHECK: spv.globalVariable [[WORKGROUPSIZE:@.*]] built_in("WorkgroupSize")
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @builtin_workgroup_size_x()
       attributes {gpu.kernel} {
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[WORKGROUPSIZE]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
       %0 = "gpu.block_dim"() {dimension = "x"} : () -> index
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_local_id_x", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL:  spv.module "Logical" "GLSL450"
   // CHECK: spv.globalVariable [[LOCALINVOCATIONID:@.*]] built_in("LocalInvocationId")
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @builtin_local_id_x()
       attributes {gpu.kernel} {
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[LOCALINVOCATIONID]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
       %0 = "gpu.thread_id"() {dimension = "x"} : () -> index
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_num_workgroups_x", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL:  spv.module "Logical" "GLSL450"
   // CHECK: spv.globalVariable [[NUMWORKGROUPS:@.*]] built_in("NumWorkgroups")
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @builtin_num_workgroups_x()
       attributes {gpu.kernel} {
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[NUMWORKGROUPS]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
       %0 = "gpu.grid_dim"() {dimension = "x"} : () -> index
       gpu.return
     }
   }
 }

mlir/test/Conversion/GPUToSPIRV/if.mlir

 // RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
 
 module attributes {gpu.container_module} {
   func @main(%arg0 : memref<10xf32>, %arg1 : i1) {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0, %arg0, %arg1) { kernel = "kernel_simple_selection", kernel_module = @kernels} : (index, index, index, index, index, index, memref<10xf32>, i1) -> ()
     return
   }
 
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     // CHECK-LABEL: @kernel_simple_selection
     gpu.func @kernel_simple_selection(%arg2 : memref<10xf32>, %arg3 : i1)
     attributes {gpu.kernel} {
       %value = constant 0.0 : f32
       %i = constant 0 : index
 
       // CHECK:       spv.selection {
       // CHECK-NEXT:    spv.BranchConditional {{%.*}}, [[TRUE:\^.*]], [[MERGE:\^.*]]

mlir/test/Conversion/GPUToSPIRV/load-store.mlir

func @load_store(%arg0: memref<12x4xf32>, %arg1: memref<12x4xf32>, %arg2: memref<12x4xf32>) {
     %1 = subi %c4, %c0_0 : index
     %c1_1 = constant 1 : index
     %c1_2 = constant 1 : index
     "gpu.launch_func"(%0, %c1_2, %c1_2, %1, %c1_2, %c1_2, %arg0, %arg1, %arg2, %c0, %c0_0, %c1, %c1_1) {kernel = "load_store_kernel", kernel_module = @kernels} : (index, index, index, index, index, index, memref<12x4xf32>, memref<12x4xf32>, memref<12x4xf32>, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL: spv.module "Logical" "GLSL450"
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     // CHECK-DAG: spv.globalVariable [[WORKGROUPSIZEVAR:@.*]] built_in("WorkgroupSize") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-DAG: spv.globalVariable [[NUMWORKGROUPSVAR:@.*]] built_in("NumWorkgroups") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-DAG: spv.globalVariable [[LOCALINVOCATIONIDVAR:@.*]] built_in("LocalInvocationId") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-DAG: spv.globalVariable [[WORKGROUPIDVAR:@.*]] built_in("WorkgroupId") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-LABEL:    func @load_store_kernel
     // CHECK-SAME: [[ARG0:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
     // CHECK-SAME: [[ARG1:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 1 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
     // CHECK-SAME: [[ARG2:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 2 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}

mlir/test/Conversion/GPUToSPIRV/loop.mlir

 // RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
 
 module attributes {gpu.container_module} {
   func @loop(%arg0 : memref<10xf32>, %arg1 : memref<10xf32>) {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0, %arg0, %arg1) { kernel = "loop_kernel", kernel_module = @kernels} : (index, index, index, index, index, index, memref<10xf32>, memref<10xf32>) -> ()
     return
   }
 
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @loop_kernel(%arg2 : memref<10xf32>, %arg3 : memref<10xf32>)
     attributes {gpu.kernel} {
       // CHECK: [[LB:%.*]] = spv.constant 4 : i32
       %lb = constant 4 : index
       // CHECK: [[UB:%.*]] = spv.constant 42 : i32
       %ub = constant 42 : index
       // CHECK: [[STEP:%.*]] = spv.constant 2 : i32
       %step = constant 2 : index

mlir/test/Conversion/GPUToSPIRV/simple.mlir

 // RUN: mlir-opt -pass-pipeline='convert-gpu-to-spirv{workgroup-size=32,4}' %s -o - | FileCheck %s
 
 module attributes {gpu.container_module} {
 
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     // CHECK:       spv.module "Logical" "GLSL450" {
     // CHECK-LABEL: func @kernel_1
     // CHECK-SAME: {{%.*}}: f32 {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
     // CHECK-SAME: {{%.*}}: !spv.ptr<!spv.struct<!spv.array<12 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 1 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
     // CHECK-SAME: spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>}
     gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<12xf32>) attributes {gpu.kernel} {
       // CHECK: spv.Return
       gpu.return

mlir/test/Dialect/GPU/invalid.mlir

 
 // -----
 
 module attributes {gpu.container_module} {
   module @kernels {
   }
 
   func @launch_func_missing_module_attribute(%sz : index) {
-    // expected-error@+1 {{module 'kernels' is missing the 'gpu.kernel_module' attribute}}
+    // expected-error@+1 {{kernel module 'kernels' is undefined}}
     "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz)
     { kernel = "kernel_1", kernel_module = @kernels }
         : (index, index, index, index, index, index) -> ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
-  module @kernels attributes {gpu.kernel_module} {
-  }
+  gpu.module @kernels { }
 
   func @launch_func_undefined_function(%sz : index) {
     // expected-error@+1 {{kernel function 'kernel_1' is undefined}}
     "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz)
     { kernel = "kernel_1", kernel_module = @kernels }
         : (index, index, index, index, index, index) -> ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @kernel_1(%arg1 : !llvm<"float*">) kernel {
       gpu.return
     }
   }
 
   func @launch_func_missing_kernel_attr(%sz : index, %arg : !llvm<"float*">) {
     // xpected-error@+1 {{kernel function is missing the 'gpu.kernel' attribute}}
     "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg)
     {kernel = "kernel_1", kernel_module = @kernels}
         : (index, index, index, index, index, index, !llvm<"float*">) -> ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @kernel_1(%arg1 : !llvm<"float*">) attributes { gpu.kernel } {
       gpu.return
     }
   }
 
   func @launch_func_kernel_operand_size(%sz : index, %arg : !llvm<"float*">) {
     // expected-error@+1 {{got 2 kernel operands but expected 1}}
     "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg, %arg)
         {kernel = "kernel_1", kernel_module = @kernels}
         : (index, index, index, index, index, index, !llvm<"float*">,
            !llvm<"float*">) -> ()
     return
   }
 }
 
 // -----
 
-module @kernels attributes {gpu.kernel_module} {
+gpu.module @kernels {
   gpu.func @kernel_1(%arg1 : !llvm<"float*">) attributes { gpu.kernel } {
     gpu.return
   }
 }
 
 // Due to the ordering of the current impl of lowering and LLVMLowering, type
 // checks need to be temporarily disabled.
 // TODO(ntv,zinenko,herhut): reactivate checks once "changing gpu.launchFunc

mlir/test/Dialect/GPU/ops.mlir

gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %sz, %grid_y = %sz, %grid_z = %sz)
         }) : () -> ()
       }) : () -> ()
       // CHECK: gpu.return
       gpu.return
     }
     return
   }
 
-  module @kernels attributes {gpu.kernel_module} {
+  gpu.module @kernels {
     gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) attributes {gpu.kernel} {
       %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
       %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index)
       %tIdZ = "gpu.thread_id"() {dimension = "z"} : () -> (index)
 
       %bDimX = "gpu.block_dim"() {dimension = "x"} : () -> (index)
       %bDimY = "gpu.block_dim"() {dimension = "y"} : () -> (index)
       %bDimZ = "gpu.block_dim"() {dimension = "z"} : () -> (index)

mlir/test/Dialect/GPU/outlining.mlir

func @device_function() {
   gpu.return
 }
 
 func @recursive_device_function() {
   call @recursive_device_function() : () -> ()
   gpu.return
 }
 
-// CHECK: module @function_call_kernel attributes {gpu.kernel_module} {
+// CHECK: gpu.module @function_call_kernel {
 // CHECK:   gpu.func @function_call_kernel()
 // CHECK:     call @device_function() : () -> ()
 // CHECK:     call @device_function() : () -> ()
 // CHECK:     llvm.mlir.addressof @global : !llvm<"i64*">
 //
 // CHECK:   llvm.mlir.global internal @global(42 : i64) : !llvm.i64
 //
 // CHECK:   func @device_function()
 // CHECK:   func @recursive_device_function()
 // CHECK-NOT:   func @device_function

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

OwnedCubin compilePtxToCubin(const std::string ptx, Location loc,
   return result;
 }
 
 static LogicalResult runMLIRPasses(ModuleOp m) {
   PassManager pm(m.getContext());
   applyPassManagerCLOptions(pm);
 
   pm.addPass(createGpuKernelOutliningPass());
-  auto &kernelPm = pm.nest<ModuleOp>();
+  auto &kernelPm = pm.nest<gpu::GPUModuleOp>();
   kernelPm.addPass(createLowerGpuOpsToNVVMOpsPass());
   kernelPm.addPass(createConvertGPUKernelToCubinPass(&compilePtxToCubin));
   pm.addPass(createLowerToLLVMPass());
   pm.addPass(createConvertGpuLaunchFuncToCudaCallsPass());
 
   return pm.run(m);
 }
 
 int main(int argc, char **argv) {
   registerPassManagerCLOptions();
   return mlir::JitRunnerMain(argc, argv, &runMLIRPasses);
 }