[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

2

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

+//===- ModuleCombiner.h - MLIR SPIR-V Module Combiner -----------*- C++ -*-===//
+//
+// 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 declares the entry point to the SPIR-V module combiner library.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_MODULECOMBINER_H_
+#define MLIR_DIALECT_SPIRV_MODULECOMBINER_H_
+
+#include "mlir/Dialect/SPIRV/SPIRVModule.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+class OpBuilder;
+
+namespace spirv {
+class ModuleOp;
+
+/// To combine a number of MLIR 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. (TODO)
+///
+/// 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
+///   functions.
+///   - 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.
+///
+/// In all cases, the references to the updated symbol 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.
+OwningSPIRVModuleRef
+combine(llvm::MutableArrayRef<ModuleOp> modules,
+        OpBuilder &combinedModuleBuilder,
+        llvm::function_ref<void(ModuleOp, StringRef, StringRef)>
+            symbRenameListener);
+} // namespace spirv
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPIRV_MODULECOMBINER_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

+//===- ModuleCombiner.cpp - MLIR SPIR-V Module Combiner ---------*- C++ -*-===//
+//
+// 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 implements the the SPIR-V module combiner library.
+//
+//===----------------------------------------------------------------------===//
+
+#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/ArrayRef.h"
+#include "llvm/ADT/StringExtras.h"
+
+using namespace mlir;
+
+static constexpr unsigned maxFreeID = 1 << 20;
+
+static SmallString<64> renameSymbol(StringRef oldSymName, unsigned &lastUsedID,
+                                    spirv::ModuleOp combinedModule) {
+  SmallString<64> newSymName(oldSymName);
+  newSymName.push_back('_');
+
+  while (lastUsedID < maxFreeID) {
+    std::string possible = (newSymName + llvm::utostr(++lastUsedID)).str();
+
+    if (!SymbolTable::lookupSymbolIn(combinedModule, possible)) {
+      newSymName += llvm::utostr(lastUsedID);
+      break;
+    }
+  }
+
+  return newSymName;
+}
+
+/// Check if a symbol with the same name as op already exists in source. If so,
+/// rename op and update all its references in target.
+static LogicalResult updateSymbolAndAllUses(SymbolOpInterface op,
+                                            spirv::ModuleOp target,
+                                            spirv::ModuleOp source,
+                                            unsigned &lastUsedID) {
+  if (!SymbolTable::lookupSymbolIn(source, op.getName()))
+    return success();
+
+  StringRef oldSymName = op.getName();
+  SmallString<64> newSymName = renameSymbol(oldSymName, lastUsedID, target);
+
+  if (failed(SymbolTable::replaceAllSymbolUses(op, newSymName, target)))
+    return op.emitError("unable to update all symbol uses for ")
+           << oldSymName << " to " << newSymName;
+
+  SymbolTable::setSymbolName(op, newSymName);
+  return success();
+}
+
+namespace mlir {
+namespace spirv {
+
+// TODO Properly test symbol rename listener mechanism.
+
+OwningSPIRVModuleRef
+combine(llvm::MutableArrayRef<spirv::ModuleOp> modules,
+        OpBuilder &combinedModuleBuilder,
+        llvm::function_ref<void(ModuleOp, StringRef, StringRef)>
+            symRenameListener) {
+  unsigned lastUsedID = 0;
+
+  if (modules.empty())
+    return nullptr;
+
+  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());
+
+  // In some cases, a symbol in the (current state of the) combined module is
+  // renamed in order to maintain the conflicting symbol in the input module
+  // being merged. For example, if the conflict is between a global variable in
+  // the current combined module and a function in the input module, the global
+  // varaible is renamed. In order to notify listeners of the symbol updates in
+  // such cases, we need to keep track of the module from which the renamed
+  // symbol in the combined module originated. This map keeps such information.
+  DenseMap<StringRef, spirv::ModuleOp> symNameToModuleMap;
+
+  for (auto module : modules) {
+    if (module.addressing_model() != addressingModel ||
+        module.memory_model() != memoryModel) {
+      module.emitError(
+          "input modules differ in addressing model and/or memory model");
+      return nullptr;
+    }
+
+    spirv::ModuleOp moduleClone = module.clone();
+
+    // In the combined module, rename all symbols that conflict with symbols
+    // from the current input module. This renmaing applies to all ops except
+    // for spv.funcs. This way, if the conflicting op in the input module is
+    // non-spv.func, we rename that symbol instead and maintain the spv.func in
+    // the combined module name as it is.
+    for (auto &op : combinedModule.getBlock().without_terminator()) {
+      auto symbolOp = dyn_cast<SymbolOpInterface>(op);
+      if (!symbolOp)
+        continue;
+
+      StringRef oldSymName = symbolOp.getName();
+
+      if (!isa<FuncOp>(op) &&
+          failed(updateSymbolAndAllUses(symbolOp, combinedModule, moduleClone,
+                                        lastUsedID)))
+        return nullptr;
+
+      StringRef newSymName = symbolOp.getName();
+
+      if (symRenameListener && oldSymName != newSymName) {
+        spirv::ModuleOp originalModule = symNameToModuleMap.lookup(oldSymName);
+
+        if (!originalModule) {
+          module.emitError("unable to find original ModuleOp for symbol ")
+              << oldSymName;
+          return nullptr;
+        }
+
+        symRenameListener(originalModule, oldSymName, newSymName);
+
+        // Since the symbol name is updated, there is no need to maintain the
+        // entry that assocaites the old symbol name with the original module.
+        symNameToModuleMap.erase(oldSymName);
+        // Instead, add a new entry to map the new symbol name to the original
+        // module in case it gets renamed again later.
+        symNameToModuleMap[newSymName] = originalModule;
+      }
+    }
+
+    // In the current input module, rename all symbols that conflict with
+    // symbols from the combined module. This includes renaming spv.funcs.
+    for (auto &op : moduleClone.getBlock().without_terminator()) {
+      auto symbolOp = dyn_cast<SymbolOpInterface>(op);
+      if (!symbolOp)
+        continue;
+
+      StringRef oldSymName = symbolOp.getName();
+
+      if (failed(updateSymbolAndAllUses(symbolOp, moduleClone, combinedModule,
+                                        lastUsedID)))
+        return nullptr;
+
+      StringRef newSymName = symbolOp.getName();
+
+      if (symRenameListener && oldSymName != newSymName) {
+        symRenameListener(module, oldSymName, newSymName);
+
+        // Insert the module associated with the symbol name.
+        auto emplaceResult = symNameToModuleMap.try_emplace(newSymName, module);
+
+        // If an entry with the same symbol name is already present, this must
+        // be a problem with the implementation, specially clean-up of the map
+        // while iterating over the combined module above.
+        if (!emplaceResult.second) {
+          module.emitError("did not expect to find an entry for symbol ")
+              << newSymName;
+          return nullptr;
+        }
+      }
+    }
+
+    // Clone all the module's ops to the combined module.
+    for (auto &op : moduleClone.getBlock().without_terminator())
+      combinedModuleBuilder.insert(op.clone());
+  }
+
+  return combinedModule;
+}
+
+} // 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 {
+spv.module Physical64 GLSL450 {
+}
+
+// expected-error @+1 {{input modules differ in addressing model and/or memory model}}
+spv.module Logical GLSL450 {
+}
+}
+
+// -----
+
+module {
+spv.module Logical Simple {
+}
+
+// expected-error @+1 {{input modules differ in addressing model and/or memory model}}
+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_1
+
+// CHECK-NEXT:     spv.globalVariable @foo
+// 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_1
+
+// CHECK-NEXT:     spv.globalVariable @foo
+// 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_1 (@bar, @bar)
+
+// CHECK-NEXT:     spv.globalVariable @foo
+// 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

+//===- TestModuleCombiner.cpp - Pass to test SPIR-V module combiner lib ---===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#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;
+
+private:
+  mlir::spirv::OwningSPIRVModuleRef combinedModule;
+};
+} // namespace
+
+void TestModuleCombinerPass::runOnOperation() {
+  auto modules = llvm::to_vector<4>(getOperation().getOps<spirv::ModuleOp>());
+
+  OpBuilder combinedModuleBuilder(modules[0]);
+  combinedModule = spirv::combine(modules, combinedModuleBuilder, nullptr);
+
+  for (spirv::ModuleOp 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) {