[mlir][spirv] Fix a few issues in ModuleCombiner

- Fixed symbol insertion into `symNameToModuleMap`. Insertion
  needs to happen whether symbols are renamed or not.
- Added check for the VCE triple and avoid dropping it.
- Disabled function deduplication. It requires more careful
  rules. Right now it can remove different functions.
- Added tests for symbol rename listener.
- And some other code/comment cleanups.

Reviewed By: ergawy

Differential Revision: https://reviews.llvm.org/D106886

Author: antiagainst

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.h

@@ -13,6 +13,7 @@
 #ifndef MLIR_DIALECT_SPIRV_IR_SPIRVOPS_H_
 #define MLIR_DIALECT_SPIRV_IR_SPIRVOPS_H_
 
+#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOpTraits.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/IR/BuiltinOps.h"

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

@@ -467,8 +467,9 @@ def SPV_ModuleOp : SPV_Op<"module",
   let builders = [
     OpBuilder<(ins CArg<"Optional<StringRef>", "llvm::None">:$name)>,
     OpBuilder<(ins "spirv::AddressingModel":$addressing_model,
-      "spirv::MemoryModel":$memory_model,
-      CArg<"Optional<StringRef>", "llvm::None">:$name)>
+                   "spirv::MemoryModel":$memory_model,
+                   CArg<"Optional<spirv::VerCapExtAttr>", "llvm::None">:$vce_triple,
+                   CArg<"Optional<StringRef>", "llvm::None">:$name)>
   ];
 
   // We need to ensure the block inside the region is properly terminated;

mlir/include/mlir/Dialect/SPIRV/Linking/ModuleCombiner.h

@@ -22,53 +22,54 @@ class OpBuilder;
 namespace spirv {
 class ModuleOp;
 
-/// To combine a number of MLIR SPIR-V modules, we move all the module-level ops
+/// The listener function to receive symbol renaming events.
+///
+/// `originalModule` is the input spirv::ModuleOp that contains the renamed
+/// symbol. `oldSymbol` and `newSymbol` are the original and renamed symbol.
+/// Note that it's the responsibility of the caller to properly retain the
+/// storage underlying the passed StringRefs if the listener callback outlives
+/// this function call.
+using SymbolRenameListener = function_ref<void(
+    spirv::ModuleOp originalModule, StringRef oldSymbol, StringRef newSymbol)>;
+
+/// Combines a list of SPIR-V `inputModules` into one. Returns the combined
+/// module on success; returns a null module otherwise.
+//
+/// \param inputModules the list of modules to combine. They won't be modified.
+/// \param combinedMdouleBuilder an OpBuilder for building the combined module.
+/// \param symbRenameListener a listener that gets called everytime a symbol in
+///                           one of the input modules is renamed.
+///
+/// To combine multiple SPIR-V modules, we move all the module-level ops
 /// from all the input modules into one big combined module. To that end, the
 /// combination process proceeds in 2 phases:
 ///
-///   (1) resolve conflicts between pairs of ops from different modules
-///   (2) deduplicate equivalent ops/sub-ops in the merged module.
+/// 1. resolve conflicts between pairs of ops from different modules,
+/// 2. deduplicate equivalent ops/sub-ops in the merged module.
 ///
 /// For the conflict resolution phase, the following rules are employed to
 /// resolve such conflicts:
 ///
-///   - If 2 spv.func's have the same symbol name, then rename one of the
+/// - If 2 spv.func's have the same symbol name, then rename one of the
 ///   functions.
-///   - If an spv.func and another op have the same symbol name, then rename the
+/// - If an spv.func and another op have the same symbol name, then rename the
 ///   other symbol.
-///   - If none of the 2 conflicting ops are spv.func, then rename either.
+/// - If none of the 2 conflicting ops are spv.func, then rename either.
 ///
 /// For deduplication, the following 3 cases are taken into consideration:
 ///
-///   - If 2 spv.GlobalVariable's have either the same descriptor set + binding
+/// - If 2 spv.GlobalVariable's have either the same descriptor set + binding
 ///   or the same build_in attribute value, then replace one of them using the
 ///   other.
-///   - If 2 spv.SpecConstant's have the same spec_id attribute value, then
+/// - If 2 spv.SpecConstant's have the same spec_id attribute value, then
 ///   replace one of them using the other.
-///   - If 2 spv.func's are identical replace one of them using the other.
+/// - Deduplicating functions are not supported right now.
 ///
 /// In all cases, the references to the updated symbol (whether renamed or
 /// deduplicated) are also updated to reflect the change.
-///
-/// \param modules the list of modules to combine. Input modules are not
-/// modified.
-/// \param combinedMdouleBuilder an OpBuilder to be used for
-//                               building up the combined module.
-/// \param symbRenameListener a listener that gets called everytime a symbol in
-///                           one of the input modules is renamed. The arguments
-///                           passed to the listener are: the input
-///                           spirv::ModuleOp that contains the renamed symbol,
-///                           a StringRef to the old symbol name, and a
-///                           StringRef to the new symbol name. Note that it is
-///                           the responsibility of the caller to properly
-///                           retain the storage underlying the passed
-///                           StringRefs if the listener callback outlives this
-///                           function call.
-///
-/// \return the combined module.
-OwningOpRef<spirv::ModuleOp>
-combine(MutableArrayRef<ModuleOp> modules, OpBuilder &combinedModuleBuilder,
-        function_ref<void(ModuleOp, StringRef, StringRef)> symbRenameListener);
+OwningOpRef<spirv::ModuleOp> combine(ArrayRef<spirv::ModuleOp> inputModules,
+                                     OpBuilder &combinedModuleBuilder,
+                                     SymbolRenameListener symRenameListener);
 } // namespace spirv
 } // namespace mlir
 

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

@@ -310,7 +310,7 @@ LogicalResult GPUModuleConversion::matchAndRewrite(
   // Add a keyword to the module name to avoid symbolic conflict.
   std::string spvModuleName = (kSPIRVModule + moduleOp.getName()).str();
   auto spvModule = rewriter.create<spirv::ModuleOp>(
-      moduleOp.getLoc(), addressingModel, memoryModel.getValue(),
+      moduleOp.getLoc(), addressingModel, memoryModel.getValue(), llvm::None,
       StringRef(spvModuleName));
 
   // Move the region from the module op into the SPIR-V module.

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

@@ -2540,6 +2540,7 @@ void spirv::ModuleOp::build(OpBuilder &builder, OperationState &state,
 void spirv::ModuleOp::build(OpBuilder &builder, OperationState &state,
                             spirv::AddressingModel addressingModel,
                             spirv::MemoryModel memoryModel,
+                            Optional<VerCapExtAttr> vceTriple,
                             Optional<StringRef> name) {
   state.addAttribute(
       "addressing_model",
@@ -2548,10 +2549,11 @@ void spirv::ModuleOp::build(OpBuilder &builder, OperationState &state,
                                          static_cast<int32_t>(memoryModel)));
   OpBuilder::InsertionGuard guard(builder);
   builder.createBlock(state.addRegion());
-  if (name) {
-    state.attributes.append(mlir::SymbolTable::getSymbolAttrName(),
-                            builder.getStringAttr(*name));
-  }
+  if (vceTriple)
+    state.addAttribute(getVCETripleAttrName(), *vceTriple);
+  if (name)
+    state.addAttribute(mlir::SymbolTable::getSymbolAttrName(),
+                       builder.getStringAttr(*name));
 }
 
 static ParseResult parseModuleOp(OpAsmParser &parser, OperationState &state) {