[mlir] 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.

[antiagainst]

Cool stuff! :)

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

22

Nit: move this after template <typename OpT> class OpPassBase to keep fwd declared symbols in mlir together? Same for the following file.

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

590

We can directly use IsolatedFromAbove here.

594

It would be nice to be more clear regarding the purpose of this module, how it involves the host/device modelling, what's the invariants before jumping to the structural representation.

597

It's a bit confusing for me here saying that gpu.launch_func is to launch a gpu.module... I'd expect a module to contain some entry functions that a gpu.launch_func can launch. Seems like we either need to have a gpu.launch_module to be consistent, which needs to define the semantics first for sure... Maybe consider revising the wording here at least?

605

The code should start a new line.

611

Should we skipDefaultBuilders given that the default generated ones cannot guarantee terminating with gpu.module_end?

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

67

KernelGPU seems saying the same thing with different word... What about either KernelModule or GPUModule?

317

Can we wrap this in namespace { ... }?

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.td

9

Nit: SPIR-V

[rriddle]

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

587

Why are there extra newlines here?

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.td

3

LLVM -> MLIR

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

813

nit: GpuModuleOp -> GPUModuleOp

822

Why is this method necessary?

850

nit: Merge this with the line above.

[herhut]

Thanks for adding more structure! This needs a little more work on the description but otherwise looks good to me.

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

597

I think the underlying problem is that we have no notion of qualified symbols implemented in MLIR yet. So the launch_func actually gets a module and function. I agree the description is misleading.

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

822

It is the equivalent of ModuleOp::create and FuncOp::create to allow creation of modules outside the context of a surrounding module. We could refactor the code here to use a builder, as the module is inserted into a surrounding module but in general that might not be the case. Or is there a way to have create methods generated?

[tpopp]

Rename GpuModule to GPUModule and improve description.

This was responding to various comments on the code review.

[tpopp]

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

822

We need GPUModuleOp::build because we have to ensure a terminator, we can't use a generated create method as I understand it. The create method is to be consistent with similar ops like those that Stephan referenced.

[rriddle]

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

822

ModuleOp has a create because it is the top level operation(also partially legacy), FuncOp only has one because it was a part of the transition from separate class to operation that was never removed. Realistically, I don't see why GPUModuleOp needs one given that it is not any more difficult for a user to just have:

OpBuilder builder(context);
builder.create<GPUModuleOp>(...);

[antiagainst]

SPIR-V side LGTM. I'd wait @herhut to give the final approval.

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.td

4

Only need to change the first one. See an example here: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/ADT/TypeSwitch.h

[tpopp]

Remove GpuModuleOp::create.

This is because GPUModuleOp is inside of another Module always, so it does not
need a static create method and can use OpBuilders.

[herhut]

LGTM from my side if rriddle@ is happy with the latest changes.

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

181

Can you use kernelFunc.getLoc() here?

183–186

Maybe SymbolTable should also provide a create<T> abstraction to hide these details of construction an operation. Not for this change, though.

[rriddle]

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

183–186

That would be fine to add in a followup.

[tpopp]

Change which location information is used at one point.

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

 #include <string>
 #include <vector>
 
 namespace mlir {
 
 class Location;
 class ModuleOp;
 
+namespace gpu {
+class GpuModuleOp;
+}
+
 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;
+namespace gpu {
+class GpuModuleOp;
+}
+
 template <typename OpT> class OpPassBase;
 
 /// 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/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", [
+  NativeOpTrait<"IsIsolatedFromAbove">, SymbolTable, Symbol,
+  SingleBlockImplicitTerminator<"ModuleEndOp">
+]> {
+  let summary = "A top level compilation unit to be run on a GPU.";
+  let description = [{
+    This op defines a GPU module. The module's top-level scope is modeled by a
+    single region with a single block. This block is expected to contain
+    operations that define a symbol. 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.module_end
+    }
+
+    ```
+  }];
+  let builders = [OpBuilder<"Builder *builder, OperationState &result, "
+                            "StringRef name">];
+  let extraClassDeclaration = [{
+    static GpuModuleOp create(Location loc, StringRef name);
+  }];
+  let parser = [{ return ::parseGpuModuleOp(parser, result); }];
+  let printer = [{ return ::print(p, *this); }];
+  let regions = (region SizedRegion<1>:$body);
+}
+
+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()))
+            *llvmModule, module.getLoc(), module.getName()))
       module.setAttr(kCubinAnnotation, cubinAttr);
     else
       signalPassFailure();
   }
 
 private:
   static OwnedCubin compilePtxToCubinForTesting(const std::string &ptx,
                                                 Location, StringRef);

 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::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(
patterns
               GPUIndexIntrinsicOpLowering<gpu::GridDimOp, NVVM::GridDimXOp,
                                           NVVM::GridDimYOp, NVVM::GridDimZOp>,
               GPUAllReduceOpLowering, GPUShuffleOpLowering, GPUFuncOpLowering,
               GPUReturnOpLowering>(converter);
   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/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 KernelGpuModuleConversion 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 KernelGpuModuleConversion::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.
 //===----------------------------------------------------------------------===//
 
+#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<
-      GPUReturnOpConversion, ForOpConversion, KernelModuleConversion,
-      KernelModuleTerminatorConversion,
+      GPUReturnOpConversion, ForOpConversion, KernelGpuModuleConversion,
       LaunchConfigConversion<gpu::BlockDimOp, spirv::BuiltIn::WorkgroupSize>,
       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
       LaunchConfigConversion<gpu::ThreadIdOp,
                              spirv::BuiltIn::LocalInvocationId>>(context,
                                                                  typeConverter);
 }

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp

 } // namespace
 
 void GPUToSPIRVPass::runOnModule() {
   auto context = &getContext();
   auto 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);
 
   ConversionTarget target(*context);

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.td

+//===-- GPUToSPIRV.td - GPU to SPIR-V Dialect Lowerings ----*- tablegen -*-===//
+//
+// Part of the LLVM 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 SPIRV 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)));
+}
+
+GpuModuleOp GpuModuleOp::create(Location loc, StringRef name) {
+  OperationState state(loc, getOperationName());
+  Builder builder(loc->getContext());
+  GpuModuleOp::build(&builder, state, name);
+  return cast<GpuModuleOp>(Operation::create(state));
+}
+
+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 << ' ';
+  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) {
-    auto context = getModule().getContext();
-    Builder builder(context);
     auto kernelModule =
-        ModuleOp::create(builder.getUnknownLoc(), kernelFunc.getName());
-    kernelModule.setAttr(gpu::GPUDialect::getKernelModuleAttrName(),
-                         builder.getUnitAttr());
+        gpu::GpuModuleOp::create(kernelFunc.getLoc(), kernelFunc.getName());
     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) {
           StringRef symbolName =
               symbolUse.getSymbolRef().cast<FlatSymbolRefAttr>().getValue();
           if (symbolTable.lookup(symbolName))
             continue;
 

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_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/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/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);
 }