[MLIR][SPIRV] Start module combiner.

[ergawy]

Summary

This commit adds a new library that merges/combines a number of spv
modules into a combined one. The library has a single entry point:
combine(...).

To combine a number of MLIR spv modules, we move all the module-level ops
from all the input modules into one big combined module. To that end, the
combination process can proceed in 2 phases:

(1) resolving conflicts between pairs of ops from different modules
(2) deduplicate equivalent ops/sub-ops in the merged module. (TODO)

This patch implements only the first phase.

[ergawy]

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

1

Oh, forgot to add copyright headers and such. Will do.

[ergawy]

• Add copyright headers.

[ergawy]

• Remove trivial braces.

[antiagainst]

This looks awesome! I love the extensive tests! :)

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

25

s/spv/SPIR-V/

We should prefer to use the formal "SPIR-V' in comments as we are not restricted by valid symbol characters in comments. :)

29

I'd use

resolve ... / deduplicate ...
or
resolving ... / deduplicating ...

for consistency.

35

We might be able to condense this comment to just saying something like

• If both of the two symbols are for spv.func, then ..
• One of the two symbols is for spv.func. then ..
• None of the two symbols is for spv.func. then ..

?

85

Could you also document what happened to the input modules (left as untouched)?

86

We can use SmallVectorImpl here to avoid forcing the caller to use 4.

86

This should return a bool to indicate whether it is successful ? Actually what about just return the newly combined module op (so the caller can perform additional work on it without explicitly rediscovering it again from combinedModuleBuilder). Return nullptr to indicate error?

87

We should also have a way to let the callers know what symbols is mapped to what? How about registering a llvm::function_ref listener here? Signature can be (spirv::ModuleOp sourceOp, StringRef oldSymbol) -> StringRef newSymbol.

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

22

Please avoid using namespace spirv as we can pull in symbols of the same name (like ModuleOp, etc.). (We are using it in the codebase but it should be fixed actually.)

25

s/nextConflictID/nextFreeID/ or something like that? nextConflictID reads to me that it's an ID that triggers confliction. :)

26

No need to use the const: in mlir ir constructs normally don't require const modifier. See https://mlir.llvm.org/docs/Rationale/UsageOfConst/

30

This appends the number to the symbol again (like _123456)? I'd assume we want to mutate the number suffix?

47

s/globalVarOp/op/ (given this is not for global variable specifically)

48

if (!...) return WalkResult::advance();

MLIR typically prefers early exit if possible.

55

return op.emitError(...) here. Also should we let the function to report an error properly? Silently ignoring the issue seems not good.

83

module.emitError. combinedModule is still something under construction.

90

What about going through all the symbols in the cloned module and renaming them in one pass? I think you can do something like

for (auto op: module.getBlock().without_terminator()) {
  if (auto symbolOp = dyn_cast<SymbolOpInterface>(op)) ...
}

This way we avoids multiple iterations.

122

You can use moduleClone.getBlock().without_terminator() to avoid checking the module end op.

mlir/test/lib/Dialect/SPIRV/TestModuleCombiner.cpp

37

No need to explicitly erase here?

[rriddle]

Drive by comment: Why name this ModuleCombiner instead of ModuleLinker?

[antiagainst]

In D90022#2354674, @rriddle wrote:

Drive by comment: Why name this ModuleCombiner instead of ModuleLinker?

Because it does not do the stuff related to what one would expect from a linker. It just merges the contents from multiple SPIR-V modules; no symbol resolution across modules, no weak/strong symbols, etc. These issues do not exist for Vulkan-flavored SPIR-V because each SPIR-V module is required to be standalone. But for OpenCL-flavored SPIR-V they do matter. I guess in the future this may evolve into that but until then keeping it clear seems preferable to me. :)

[rriddle]

In D90022#2354689, @antiagainst wrote:

In D90022#2354674, @rriddle wrote:

Drive by comment: Why name this ModuleCombiner instead of ModuleLinker?

Because it does not do the stuff related to what one would expect from a linker. It just merges the contents from multiple SPIR-V modules; no symbol resolution across modules, no weak/strong symbols, etc. These issues do not exist for Vulkan-flavored SPIR-V because each SPIR-V module is required to be standalone. But for OpenCL-flavored SPIR-V they do matter. I guess in the future this may evolve into that but until then keeping it clear seems preferable to me. :)

Thanks for the response Lei.

[ergawy]

@antiagainst thanks a lot for the review. Your comments simplify things by quite a bit.

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

87

I think the signature should rather be: (spirv::ModuleOp sourceOp, StringRef oldSymbol, StringRef newSymbol) -> void, right? This way the listener will have access to the new and old symbol as part of the callback invocation. Or did I miss your point?

Also, we need to properly test this at some point using GTest maybe! However, I prefer to add a TODO and continue with the rest of the functionality first if you don't mind. (I expect we would need some tiny infrastructure, e.g. mocking and building MLIR entities ourselves, in place in order to do so).

[ergawy]

• Handle review comments:

[ergawy]

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

87

One more thing, is it wise to pass 'StringRef' to those callbacks? I mean you never know whether the user-supplied code will outlive the scope that created the string underlying the ref.

[antiagainst]

Generally looks good to me! Just a few more nits. Also please make sure adding the listener.

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

50

Sorry actually ArrayRef<ModuleOp> here? :)

50

Should return OwningSPIRVModuleRef to handle ownership and avoid leaks.

87

You are right about the signature. It's fine to use StringRef; typically the caller should copy it if they want to retain. We can be explicit in the comment for the signature.

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

27

Have a limit over nextFreeID to avoid infinite loop? It's a bit arbitrary; but I'd assume 2^20 might be a good number for now.

28

Nit: by ++nextFreeID, it's actually lastUsedID here. ;-P So I'd expect nextFreeID++ here. (You can set the intial value of nextFreeID as 1 if wanting to be 1 based.)

31

can directly return possible?

93

s/dyn_cast/isa/

[ergawy]

• Handle review comments.

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

50

Sorry actually ArrayRef<ModuleOp> here? :)

I had to use MutableArrayRef<ModuleOp> since the ModuleOp API is non-const. Let me know if I misunderstood something.

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

28

Indeed. I renamed to lastUsedID (laziest solution that works 😁 )

31

I feel better having a single exit point here 😁 . Please let me know if you have a strong opinion against this.

[ergawy]

Squash all commits into one.

[antiagainst]

Nice! Thanks @ergawy !

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

120

This can also be guarded by && symRenameListener?

[ergawy]

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

120

Didn't do this to make sure the management of symNameToModuleMap is properly handled. I can also guard all code that manages the map with the same check but this is cleaner.

[ergawy]

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

120

but _thought_ this _would be_ cleaner.

[antiagainst]

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

120

No worries! I changed it before landing, plus a few other tweaks. :)

But I don't have full visibility into what your planned next batch of code so I may very well misunderstand things. If that's the case, please just feel free to change it back. :)

[mehdi_amini]

Sorry I had to revert this back because it broke the build: https://buildkite.com/mlir/mlir-core/builds/8983#511574bd-3519-4f34-afd7-ef5702305815
Seems like a missing dependency somewhere.

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

+#ifndef MLIR_DIALECT_SPIRV_MODULE_COMBINER_H_
+#define MLIR_DIALECT_SPIRV_MODULE_COMBINER_H_
+
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+class OpBuilder;
+
+namespace spirv {
+class ModuleOp;
+
+/// To combine a number of MLIR spv modules, we move all the module-level ops
+/// from all the input modules into one big combined module. To that end, the
+/// combination process can proceed in 2 phases:
+///
+///   (1) resolving conflicts between pairs of ops from different modules
+///   (2) deduplicate equivalent ops/sub-ops in the merged module. (TODO)
+///
+/// For the conflict resolution phase, the following rules are employed to
+/// resolve such conflicts:
+///
+///   =========================================================================
+///   FuncOp vs. FuncOp
+///   Conflict: Same symbol name
+///   -------------------------------------------------------------------------
+///   Rename one of the functions and update its refernces.
+///   =========================================================================
+///
+///   =========================================================================
+///   FuncOp vs. GlobalVariableOp
+///   Conflict: Same symbol name
+///   -------------------------------------------------------------------------
+///   Rename the global variable and update references to the renamed symbol.
+///   =========================================================================
+///
+///   =========================================================================
+///   FuncOp vs. SpecConstantOp
+///   FuncOp vs. SpecConstantCompositeOp
+///   Conflict: Same symbol name
+///   -------------------------------------------------------------------------
+///   Rename the spec constant and update references to the renamed symbol.
+///   =========================================================================
+///
+///   =========================================================================
+///   GlobalVariableOp vs. GlobalVariableOp
+///   Conflict: Same symbol name
+///   -------------------------------------------------------------------------
+///   Rename either of the global variables and update references to it.
+///   =========================================================================
+///
+///   =========================================================================
+///   GlobalVariableOp vs. SpecConstantOp
+///   FuncOp vs. SpecConstantCompositeOp
+///   Conflict: Same symbol name
+///   -------------------------------------------------------------------------
+///   Rename the global variable and update its references.
+///   =========================================================================
+///
+///   =========================================================================
+///   2 spec constants (scalar or composite)
+///   Conflict: Same symbol name
+///   -------------------------------------------------------------------------
+///   Rename either of the constants and update its references.
+///   =========================================================================
+///
+///   =========================================================================
+///   EntryPointOp vs. EntryPointOp
+///   Conflict: Same symbol name and execution model
+///   -------------------------------------------------------------------------
+///   No need to resolve this explicitly as it will be resolved as part of
+///   resolving the conflict between the 2 associated functions.
+///   =========================================================================
+void combine(llvm::SmallVector<ModuleOp, 4> modules,
+             OpBuilder &combinedModuleBuilder);
+} // namespace spirv
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPIRV_MODULE_COMBINER_H_

mlir/lib/Dialect/SPIRV/CMakeLists.txt

add_mlir_dialect_library(MLIRSPIRV
   MLIRControlFlowInterfaces
   MLIRIR
   MLIRParser
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRTransforms
   )
 
+add_subdirectory(Linking)
 add_subdirectory(Serialization)
 add_subdirectory(Transforms)

mlir/lib/Dialect/SPIRV/Linking/CMakeLists.txt

+add_subdirectory(ModuleCombiner)

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/CMakeLists.txt

+add_mlir_dialect_library(MLIRSPIRVModuleCombiner
+  ModuleCombiner.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SPIRV
+  )

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

+#include "mlir/Dialect/SPIRV/ModuleCombiner.h"
+
+#include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/StringExtras.h"
+
+using namespace mlir;
+using namespace spirv;
+
+static SmallString<64> renameSymbol(StringRef oldSymName,
+                                    unsigned &nextConflictID,
+                                    const spirv::ModuleOp combinedModule) {
+  SmallString<64> newSymName(oldSymName);
+  newSymName.push_back('_');
+  while (true) {
+    newSymName += llvm::utostr(++nextConflictID);
+
+    if (!SymbolTable::lookupSymbolIn(combinedModule, newSymName))
+      break;
+  }
+
+  return newSymName;
+}
+
+/// Walks the target module for operations of type OpTy. And for each such
+/// operation, checks if another operation in the source module has the same
+/// symbol. If this is the case, renames the visited/walked operation and
+/// updates its references.
+template <typename OpTy>
+static void updateSymbolAndAllUses(spirv::ModuleOp target,
+                                   spirv::ModuleOp source,
+                                   unsigned &nextConflictID) {
+  target.walk([&](OpTy globalVarOp) {
+    if (SymbolTable::lookupSymbolIn(source, globalVarOp.sym_name())) {
+      StringRef oldSymName = globalVarOp.sym_name();
+      SmallString<64> newSymName =
+          renameSymbol(oldSymName, nextConflictID, target);
+
+      if (failed(SymbolTable::replaceAllSymbolUses(globalVarOp, newSymName,
+                                                   target)))
+        globalVarOp.emitError("unable to update all symbol uses for ")
+            << oldSymName << " to " << newSymName;
+
+      SymbolTable::setSymbolName(globalVarOp, newSymName);
+    }
+  });
+}
+
+namespace mlir {
+namespace spirv {
+
+void combine(llvm::SmallVector<ModuleOp, 4> modules,
+             OpBuilder &combinedModuleBuilder) {
+  unsigned nextConflictID = 0;
+
+  if (modules.empty())
+    return;
+
+  auto addressingModel = modules[0].addressing_model();
+  auto memoryModel = modules[0].memory_model();
+
+  auto combinedModule = combinedModuleBuilder.create<spirv::ModuleOp>(
+      modules[0].getLoc(), addressingModel, memoryModel);
+  combinedModuleBuilder.setInsertionPointToStart(&*combinedModule.getBody());
+
+  for (auto module : modules) {
+    if (module.addressing_model() != addressingModel ||
+        module.memory_model() != memoryModel) {
+      combinedModule.emitError(
+          "input modules differ in addressing model and/or memory model");
+      return;
+    }
+
+    spirv::ModuleOp moduleClone = module.clone();
+
+    // A global variable is renamed if it conflicts with a function or a spec
+    // constant:
+    //   (1) Rename global variables from the current input module that are
+    //   conflicting with any other module-level op currently in the combined
+    //   module.
+    updateSymbolAndAllUses<GlobalVariableOp>(moduleClone, combinedModule,
+                                             nextConflictID);
+    //   (2) Rename global variables currently in the combined module that are
+    //   conflicting with any other module-level op in the current input module.
+    updateSymbolAndAllUses<GlobalVariableOp>(combinedModule, moduleClone,
+                                             nextConflictID);
+
+    // A spec constant is renamed if it conflicts with a function:
+    //   (1) Rename spec constants from the current input module that are
+    //   conflicting with functions currently in the combined module.
+    updateSymbolAndAllUses<SpecConstantOp>(moduleClone, combinedModule,
+                                           nextConflictID);
+    updateSymbolAndAllUses<SpecConstantCompositeOp>(moduleClone, combinedModule,
+                                                    nextConflictID);
+    //   (2) Rename spec constants currently in the combined module that are
+    //   conflicting with functions in the current input module.
+    updateSymbolAndAllUses<SpecConstantOp>(combinedModule, moduleClone,
+                                           nextConflictID);
+    updateSymbolAndAllUses<SpecConstantCompositeOp>(combinedModule, moduleClone,
+                                                    nextConflictID);
+
+    // Rename function in the current input modules that are conflicting with
+    // functions already in the combined module.
+    updateSymbolAndAllUses<FuncOp>(moduleClone, combinedModule, nextConflictID);
+
+    // Clone all the module's ops to the combined module.
+    for (auto &op : moduleClone.getOps()) {
+      if (dyn_cast<spirv::ModuleEndOp>(op))
+        continue;
+
+      combinedModuleBuilder.insert(op.clone());
+    }
+  }
+}
+
+} // namespace spirv
+} // namespace mlir

mlir/test/Dialect/SPIRV/Linking/ModuleCombiner/basic.mlir

+// RUN: mlir-opt -test-spirv-module-combiner -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.specConstant @m1_sc
+// CHECK-NEXT:     spv.specConstant @m2_sc
+// CHECK-NEXT:     spv.func @variable_init_spec_constant
+// CHECK-NEXT:       spv._reference_of @m2_sc
+// CHECK-NEXT:       spv.Variable init
+// CHECK-NEXT:       spv.Return
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.specConstant @m1_sc = 42.42 : f32
+}
+
+spv.module Logical GLSL450 {
+  spv.specConstant @m2_sc = 42 : i32
+  spv.func @variable_init_spec_constant() -> () "None" {
+    %0 = spv._reference_of @m2_sc : i32
+    %1 = spv.Variable init(%0) : !spv.ptr<i32, Function>
+    spv.Return
+  }
+}
+}
+
+// -----
+
+module {
+// expected-error @+1 {{input modules differ in addressing model and/or memory model}}
+spv.module Physical64 GLSL450 {
+}
+
+spv.module Logical GLSL450 {
+}
+}
+
+// -----
+
+module {
+// expected-error @+1 {{input modules differ in addressing model and/or memory model}}
+spv.module Logical Simple {
+}
+
+spv.module Logical GLSL450 {
+}
+}

mlir/test/Dialect/SPIRV/Linking/ModuleCombiner/conflict_resolution.mlir

+// RUN: mlir-opt -test-spirv-module-combiner -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// Test basic renaming of conflicting funcOps.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : f32) -> f32 "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+}
+}
+
+// -----
+
+// Test basic renaming of conflicting funcOps across 3 modules.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo_2
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : f32) -> f32 "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+}
+
+// -----
+
+// Test properly updating references to a renamed funcOp.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @bar
+// CHECK-NEXT:       spv.FunctionCall @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : f32) -> f32 "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+
+  spv.func @bar(%arg0 : f32) -> f32 "None" {
+    %0 = spv.FunctionCall @foo(%arg0) : (f32) ->  (f32)
+    spv.ReturnValue %0 : f32
+  }
+}
+}
+
+// -----
+
+// Test properly updating references to a renamed funcOp if the functionCallOp
+// preceeds the callee funcOp definition.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @bar
+// CHECK-NEXT:       spv.FunctionCall @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @bar(%arg0 : f32) -> f32 "None" {
+    %0 = spv.FunctionCall @foo(%arg0) : (f32) ->  (f32)
+    spv.ReturnValue %0 : f32
+  }
+
+  spv.func @foo(%arg0 : f32) -> f32 "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+}
+}
+
+// -----
+
+// Test properly updating entryPointOp and executionModeOp attached to renamed
+// funcOp.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.EntryPoint "GLCompute" @foo_1
+// CHECK-NEXT:     spv.ExecutionMode @foo_1 "ContractionOff"
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : f32) -> f32 "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+
+  spv.EntryPoint "GLCompute" @foo
+  spv.ExecutionMode @foo "ContractionOff"
+}
+}
+
+// -----
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.EntryPoint "GLCompute" @fo
+// CHECK-NEXT:     spv.ExecutionMode @foo "ContractionOff"
+
+// CHECK-NEXT:     spv.func @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.EntryPoint "GLCompute" @foo_1
+// CHECK-NEXT:     spv.ExecutionMode @foo_1 "ContractionOff"
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+  
+  spv.EntryPoint "GLCompute" @foo
+  spv.ExecutionMode @foo "ContractionOff"
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : f32) -> f32 "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+
+  spv.EntryPoint "GLCompute" @foo
+  spv.ExecutionMode @foo "ContractionOff"
+}
+}
+
+// -----
+
+// Resolve conflicting funcOp and globalVariableOp.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.globalVariable @foo_1
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+}
+}
+
+// -----
+
+// Resolve conflicting funcOp and globalVariableOp and update the global variable's
+// references.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.globalVariable @foo_1
+// CHECK-NEXT:     spv.func @bar
+// CHECK-NEXT:       spv._address_of @foo_1
+// CHECK-NEXT:       spv.Load
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+
+  spv.func @bar() -> f32 "None" {
+    %0 = spv._address_of @foo : !spv.ptr<f32, Input>
+    %1 = spv.Load "Input" %0 : f32
+    spv.ReturnValue %1 : f32
+  }
+}
+}
+
+// -----
+
+// Resolve conflicting globalVariableOp and funcOp and update the global variable's
+// references.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.globalVariable @foo_1
+// CHECK-NEXT:     spv.func @bar
+// CHECK-NEXT:       spv._address_of @foo_1
+// CHECK-NEXT:       spv.Load
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+
+  spv.func @bar() -> f32 "None" {
+    %0 = spv._address_of @foo : !spv.ptr<f32, Input>
+    %1 = spv.Load "Input" %0 : f32
+    spv.ReturnValue %1 : f32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+}
+
+// -----
+
+// Resolve conflicting funcOp and specConstantOp.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.specConstant @foo_1
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.specConstant @foo = -5 : i32
+}
+}
+
+// -----
+
+// Resolve conflicting funcOp and specConstantOp and update the spec constant's
+// references.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.specConstant @foo_1
+// CHECK-NEXT:     spv.func @bar
+// CHECK-NEXT:       spv._reference_of @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.specConstant @foo = -5 : i32
+
+  spv.func @bar() -> i32 "None" {
+    %0 = spv._reference_of @foo : i32 
+    spv.ReturnValue %0 : i32
+  }
+}
+}
+
+// -----
+
+// Resolve conflicting specConstantOp and funcOp and update the spec constant's
+// references.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.specConstant @foo_1
+// CHECK-NEXT:     spv.func @bar
+// CHECK-NEXT:       spv._reference_of @foo_1
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.specConstant @foo = -5 : i32
+
+  spv.func @bar() -> i32 "None" {
+    %0 = spv._reference_of @foo : i32
+    spv.ReturnValue %0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+}
+
+// -----
+
+// Resolve conflicting funcOp and specConstantCompositeOp.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.specConstant @bar
+// CHECK-NEXT:     spv.specConstantComposite @foo_1 (@bar, @bar)
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.specConstant @bar = -5 : i32
+  spv.specConstantComposite @foo (@bar, @bar) : !spv.array<2 x i32>
+}
+}
+
+// -----
+
+// Resolve conflicting funcOp and specConstantCompositeOp and update the spec
+// constant's references.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.specConstant @bar
+// CHECK-NEXT:     spv.specConstantComposite @foo_1 (@bar, @bar)
+// CHECK-NEXT:     spv.func @baz
+// CHECK-NEXT:       spv._reference_of @foo_1
+// CHECK-NEXT:       spv.CompositeExtract
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.specConstant @bar = -5 : i32
+  spv.specConstantComposite @foo (@bar, @bar) : !spv.array<2 x i32>
+
+  spv.func @baz() -> i32 "None" {
+    %0 = spv._reference_of @foo : !spv.array<2 x i32>
+    %1 = spv.CompositeExtract %0[0 : i32] : !spv.array<2 x i32>
+    spv.ReturnValue %1 : i32
+  }
+}
+}
+
+// -----
+
+// Resolve conflicting specConstantCompositeOp and funcOp and update the spec
+// constant's references.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.specConstant @bar
+// CHECK-NEXT:     spv.specConstantComposite @foo_1 (@bar, @bar)
+// CHECK-NEXT:     spv.func @baz
+// CHECK-NEXT:       spv._reference_of @foo_1
+// CHECK-NEXT:       spv.CompositeExtract
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.specConstant @bar = -5 : i32
+  spv.specConstantComposite @foo (@bar, @bar) : !spv.array<2 x i32>
+
+  spv.func @baz() -> i32 "None" {
+    %0 = spv._reference_of @foo : !spv.array<2 x i32>
+    %1 = spv.CompositeExtract %0[0 : i32] : !spv.array<2 x i32>
+    spv.ReturnValue %1 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+}
+}
+
+// -----
+
+// Resolve conflicting spec constants and funcOps and update the spec constant's
+// references.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.specConstant @bar_1
+// CHECK-NEXT:     spv.specConstantComposite @foo_2 (@bar_1, @bar_1)
+// CHECK-NEXT:     spv.func @baz
+// CHECK-NEXT:       spv._reference_of @foo_2
+// CHECK-NEXT:       spv.CompositeExtract
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @bar
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.specConstant @bar = -5 : i32
+  spv.specConstantComposite @foo (@bar, @bar) : !spv.array<2 x i32>
+
+  spv.func @baz() -> i32 "None" {
+    %0 = spv._reference_of @foo : !spv.array<2 x i32>
+    %1 = spv.CompositeExtract %0[0 : i32] : !spv.array<2 x i32>
+    spv.ReturnValue %1 : i32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.func @foo(%arg0 : i32) -> i32 "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+
+  spv.func @bar(%arg0 : f32) -> f32 "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+}
+}
+
+// -----
+
+// Resolve conflicting globalVariableOps.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.globalVariable @foo
+
+// CHECK-NEXT:     spv.globalVariable @foo_1
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+}
+
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+}
+}
+
+// -----
+
+// Resolve conflicting globalVariableOp and specConstantOp.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.globalVariable @foo_1
+
+// CHECK-NEXT:     spv.specConstant @foo
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+}
+
+spv.module Logical GLSL450 {
+  spv.specConstant @foo = -5 : i32
+}
+}
+
+// -----
+
+// Resolve conflicting specConstantOp and globalVariableOp.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.specConstant @foo
+
+// CHECK-NEXT:     spv.globalVariable @foo_1
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.specConstant @foo = -5 : i32
+}
+
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+}
+}
+
+// -----
+
+// Resolve conflicting globalVariableOp and specConstantCompositeOp.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.globalVariable @foo_1
+
+// CHECK-NEXT:     spv.specConstant @bar
+// CHECK-NEXT:     spv.specConstantComposite @foo (@bar, @bar)
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+}
+
+spv.module Logical GLSL450 {
+  spv.specConstant @bar = -5 : i32
+  spv.specConstantComposite @foo (@bar, @bar) : !spv.array<2 x i32>
+}
+}
+
+// -----
+
+// Resolve conflicting globalVariableOp and specConstantComposite.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.specConstant @bar
+// CHECK-NEXT:     spv.specConstantComposite @foo (@bar, @bar)
+
+// CHECK-NEXT:     spv.globalVariable @foo_1
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.specConstant @bar = -5 : i32
+  spv.specConstantComposite @foo (@bar, @bar) : !spv.array<2 x i32>
+}
+
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+}
+}

mlir/test/lib/Dialect/SPIRV/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRSPIRVTestPasses
   TestAvailability.cpp
   TestEntryPointAbi.cpp
+  TestModuleCombiner.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SPIRV
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 
   LINK_LIBS PUBLIC
   MLIRGPU
   MLIRIR
   MLIRPass
   MLIRSPIRV
   MLIRSupport
   )

mlir/test/lib/Dialect/SPIRV/TestModuleCombiner.cpp

+#include "mlir/Dialect/SPIRV/ModuleCombiner.h"
+
+#include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/SPIRVTypes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Module.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+class TestModuleCombinerPass
+    : public PassWrapper<TestModuleCombinerPass,
+                         OperationPass<mlir::ModuleOp>> {
+public:
+  TestModuleCombinerPass() = default;
+  TestModuleCombinerPass(const TestModuleCombinerPass &) {}
+  void runOnOperation() override;
+};
+} // namespace
+
+void TestModuleCombinerPass::runOnOperation() {
+  auto modules = llvm::to_vector<4>(getOperation().getOps<spirv::ModuleOp>());
+
+  OpBuilder combinedModuleBuilder(modules[0]);
+  spirv::combine(modules, combinedModuleBuilder);
+
+  for (auto module : modules) {
+    module.erase();
+  }
+}
+
+namespace mlir {
+void registerTestSpirvModuleCombinerPass() {
+  PassRegistration<TestModuleCombinerPass> registration(
+      "test-spirv-module-combiner", "Tests SPIR-V module combiner library");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestPrintDefUsePass();
 void registerTestPrintNestingPass();
 void registerTestRecursiveTypesPass();
 void registerTestReducer();
 void registerTestSpirvEntryPointABIPass();
+void registerTestSpirvModuleCombinerPass();
 void registerTestSCFUtilsPass();
 void registerTestTraitsPass();
 void registerTestVectorConversions();
 void registerVectorizerTestPass();
 } // namespace mlir
 
 #ifdef MLIR_INCLUDE_TESTS
 void registerTestPasses() {
#endif
   registerTestOpaqueLoc();
   registerTestPreparationPassWithAllowedMemrefResults();
   registerTestPrintDefUsePass();
   registerTestPrintNestingPass();
   registerTestRecursiveTypesPass();
   registerTestReducer();
   registerTestGpuParallelLoopMappingPass();
   registerTestSpirvEntryPointABIPass();
+  registerTestSpirvModuleCombinerPass();
   registerTestSCFUtilsPass();
   registerTestTraitsPass();
   registerTestVectorConversions();
   registerVectorizerTestPass();
 }
 #endif
 
 int main(int argc, char **argv) {