[MLIR][SPIRV] Added optional name to SPIR-V module

[georgemitenkov]

Summary

This patch adds an optional name to SPIR-V module.
This will help with lowering from GPU dialect (so that we
can pass the kernel module name) and will be more naturally
aligned with ModuleOp.

[mravishankar]

I am not totally opposed to it though. Its one way of making the name optional, but still using the symbol table traits.

@rriddle any thoughts here?

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

422

It might be better to just add a SymbolTable trait to this (similar to this). THen the spv.module will also be uniqued which is probably more robust.

That would make it a bigger change though. Is this something that is needed for the SPIRVToLLVM conversion? For now there it is not handling the case of multiple spv.module translation anyway, so we can just skip this for now if this takes long. I would assume you have to update a LOT of tests....

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2288

Nit: Better to add { .. } since the body spans multiple lines.

[rriddle]

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

364

You should add the symbol interface, as it includes support for optional symbols, and provide a bool isOptionalSymbol() { return true; } method:
https://github.com/llvm/llvm-project/blob/05777ab941063192b9ccb1775358a83a2700ccc1/mlir/include/mlir/IR/SymbolInterfaces.td#L118

413

This needs to be sym_name.

439

This can just be Optional<StringRef> getName() { return sym_name(); }.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2289

nit: You can do state.attributes.append(mlir::SymbolTable::getSymbolAttrName(), ...)

2305

Same as above.

2316

nit: Drop the mlir::.

[georgemitenkov]

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

422

It makes sense to have spv.module be uniqued by name. I am fine with both ways and I think that updating tests should not be too long :)

This is needed for passing the GPU module name to the SPIR-V module (and later is used in mlir-spirv-cpu-runner to find the necessary SPIR-V module simply by looking it up in a symbol table). This is not 100% needed for the runner, but is a much nicer way to implement SPIR-V module lookup.

[mravishankar]

Making this request changes. I think after addressing comments this is ready to land.

[georgemitenkov]

Addressed comments.

[mravishankar]

Thanks George! Please fix the clang-tidy errors before submitting if possible.

[georgemitenkov]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2296

@mravishankar In this case is it fine to leave like that? That's how it was previously so I suppose the snake case is fine here?

[georgemitenkov]

Fixed clang-tidy error.

[georgemitenkov]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2296

I decided to change the names to camel case instead.

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

def SPV_GlobalVariableOp : SPV_Op<"globalVariable", [InModuleScope, Symbol]> {
   }];
 }
 
 // -----
 
 def SPV_ModuleOp : SPV_Op<"module",
                           [IsolatedFromAbove,
                            SingleBlockImplicitTerminator<"ModuleEndOp">,
-                           SymbolTable]> {
+                           SymbolTable, Symbol]> {
   let summary = "The top-level op that defines a SPIR-V module";
 
   let description = [{
     This op defines a SPIR-V module using a MLIR region. The region contains
     one block. Module-level operations, including functions definitions,
     are all placed in this block.
 
     Using an op with a region to define a SPIR-V module enables "embedding"
spv.module Logical Vulkan
       }
     }
     ```
   }];
 
   let arguments = (ins
     SPV_AddressingModelAttr:$addressing_model,
     SPV_MemoryModelAttr:$memory_model,
-    OptionalAttr<SPV_VerCapExtAttr>:$vce_triple
+    OptionalAttr<SPV_VerCapExtAttr>:$vce_triple,
+    OptionalAttr<StrAttr>:$sym_name
   );
 
   let results = (outs);
 
   let regions = (region SizedRegion<1>:$body);
 
   let builders = [
-    OpBuilder<[{OpBuilder &, OperationState &state}]>,
+    OpBuilder<[{OpBuilder &, OperationState &state,
+                Optional<StringRef> name = llvm::None}]>,
     OpBuilder<[{OpBuilder &, OperationState &state,
                 spirv::AddressingModel addressing_model,
-                spirv::MemoryModel memory_model}]>
+                spirv::MemoryModel memory_model,
+                Optional<StringRef> name = llvm::None}]>
   ];
 
   // We need to ensure the block inside the region is properly terminated;
   // the auto-generated builders do not guarantee that.
   let skipDefaultBuilders = 1;
 
   let hasOpcode = 0;
 
   let autogenSerialization = 0;
 
   let extraClassDeclaration = [{
+
+    bool isOptionalSymbol() { return true; }
+
+    Optional<StringRef> getName() { return sym_name(); }
+
     static StringRef getVCETripleAttrName() { return "vce_triple"; }
 
     Block& getBlock() {
       return this->getOperation()->getRegion(0).front();
     }
   }];
 }
 

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

return mergeOp.emitOpError(
         "can only be used in the last block of 'spv.selection' or 'spv.loop'");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // spv.module
 //===----------------------------------------------------------------------===//
 
-void spirv::ModuleOp::build(OpBuilder &builder, OperationState &state) {
+void spirv::ModuleOp::build(OpBuilder &builder, OperationState &state,
+                            Optional<StringRef> name) {
   ensureTerminator(*state.addRegion(), builder, state.location);
+  if (name) {
+    state.attributes.append(mlir::SymbolTable::getSymbolAttrName(),
+                            builder.getStringAttr(*name));
+  }
 }
 
 void spirv::ModuleOp::build(OpBuilder &builder, OperationState &state,
                             spirv::AddressingModel addressing_model,
-                            spirv::MemoryModel memory_model) {
+                            spirv::MemoryModel memory_model,
+                            Optional<StringRef> name) {
   state.addAttribute(
       "addressing_model",
       builder.getI32IntegerAttr(static_cast<int32_t>(addressing_model)));
   state.addAttribute("memory_model", builder.getI32IntegerAttr(
                                          static_cast<int32_t>(memory_model)));
   ensureTerminator(*state.addRegion(), builder, state.location);
+  if (name) {
+    state.attributes.append(mlir::SymbolTable::getSymbolAttrName(),
+                            builder.getStringAttr(*name));
+  }
 }
 
 static ParseResult parseModuleOp(OpAsmParser &parser, OperationState &state) {
   Region *body = state.addRegion();
 
+  // If the name is present, parse it.
+  StringAttr nameAttr;
+  parser.parseOptionalSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
+                                 state.attributes);
+
   // Parse attributes
   spirv::AddressingModel addrModel;
   spirv::MemoryModel memoryModel;
   if (parseEnumKeywordAttr(addrModel, parser, state) ||
       parseEnumKeywordAttr(memoryModel, parser, state))
     return failure();
 
   if (succeeded(parser.parseOptionalKeyword("requires"))) {
static ParseResult parseModuleOp(OpAsmParser &parser, OperationState &state) {
 
   spirv::ModuleOp::ensureTerminator(*body, parser.getBuilder(), state.location);
   return success();
 }
 
 static void print(spirv::ModuleOp moduleOp, OpAsmPrinter &printer) {
   printer << spirv::ModuleOp::getOperationName();
 
+  if (Optional<StringRef> name = moduleOp.getName()) {
+    printer << ' ';
+    printer.printSymbolName(*name);
+  }
+
   SmallVector<StringRef, 2> elidedAttrs;
 
   printer << " " << spirv::stringifyAddressingModel(moduleOp.addressing_model())
           << " " << spirv::stringifyMemoryModel(moduleOp.memory_model());
   auto addressingModelAttrName = spirv::attributeName<spirv::AddressingModel>();
   auto memoryModelAttrName = spirv::attributeName<spirv::MemoryModel>();
-  elidedAttrs.assign({addressingModelAttrName, memoryModelAttrName});
+  elidedAttrs.assign({addressingModelAttrName, memoryModelAttrName,
+                      SymbolTable::getSymbolAttrName()});
 
   if (Optional<spirv::VerCapExtAttr> triple = moduleOp.vce_triple()) {
     printer << " requires " << *triple;
     elidedAttrs.push_back(spirv::ModuleOp::getVCETripleAttrName());
   }
 
   printer.printOptionalAttrDictWithKeyword(moduleOp.getAttrs(), elidedAttrs);
   printer.printRegion(moduleOp.body(), /*printEntryBlockArgs=*/false,

mlir/test/Dialect/SPIRV/structure-ops.mlir

 //===----------------------------------------------------------------------===//
 // spv.module
 //===----------------------------------------------------------------------===//
 
 // Module without capability and extension
 // CHECK: spv.module Logical GLSL450
 spv.module Logical GLSL450 { }
 
+// Module with a name
+// CHECK: spv.module @{{.*}} Logical GLSL450
+spv.module @name Logical GLSL450 { }
 
 // Module with (version, capabilities, extensions) triple
 // CHECK: spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], [SPV_KHR_16bit_storage]>
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], [SPV_KHR_16bit_storage]> { }
 
 // Module with additional attributes
 // CHECK: spv.module Logical GLSL450 attributes {foo = "bar"}
 spv.module Logical GLSL450 attributes {foo = "bar"} { }