[MLIR][SPIRV] ModuleCombiner: deduplicate global vars and spec consts.

[ergawy]

Summary

This commit extends the functionality of the SPIR-V module combiner
library by adding new deduplication capabilities. In particular,
implementation of deduplication of global variables and specialization
constants is introduced. For global variables, 2 variables are
considered duplicate if they either have the same descriptor set +
binding or the same built_in attribute. For specialization constants, 2
spec constants are considered duplicate if they have the same spec_id
attribute.

[ergawy]

• Remove unneeded comment.

[ergawy]

I am not familiar with the implicit conversion rules within the spv dialect (or MLIR in general). But, I was expecting that replacing all uses of bar with foo in the following to snippet to fail:

module {
spv.module Logical GLSL450 {
  spv.specConstant @foo spec_id(5) = 1 : i1

  spv.func @use_foo() -> (i1) "None" {
    %0 = spv._reference_of @foo : i1
    spv.ReturnValue %0 : i1
  }
}

spv.module Logical GLSL450 {
  spv.specConstant @bar spec_id(5) = 1. : f32

  spv.func @use_bar() -> (f32) "None" {
    // expected-error @+1 {{floating point value not valid for specified type}}
    %0 = spv._reference_of @bar : f32
    spv.ReturnValue %0 : f32
  }
}
}

@antiagainst , just FYI in case we need more elaborate testing in this area.

[ergawy]

General comment regarding this patch as well as this: https://reviews.llvm.org/D90932.

I think both can be greatly simplified and merged in a more compact patch if I introduce specialized functions to compute hash codes for each of: global vars, spec consts, and functions. Then we would need a single map that maps the hash code to the symbol op and deduplicate based on the hash code in a single walk through the module instead of separating things as I do now.

I will try to do that.

[ergawy]

Abandoned in favor of: https://reviews.llvm.org/D90951

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

 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)
+///   (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
 ///   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.
+/// For deduplication, the following 3 cases are taken into consideration:
+///
+///   - 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
+///   replace one of them using the other.
+///   - TODO: If 2 spv.func's are identical replace one of them using the other.
+///
+/// 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.
+//                               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

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/SPIRVDialect.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 LogicalResult updateSymbolAndAllUses(SymbolOpInterface op,
   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();
 }
 
+template <typename KeyTy, typename SymbolOpTy>
+static SymbolOpTy emplaceOrGetReplacementSymbolName(
+    KeyTy key, SymbolOpTy symbolOp,
+    DenseMap<KeyTy, SymbolOpTy> &deduplicationMap) {
+  auto result = deduplicationMap.try_emplace(key, symbolOp);
+
+  if (result.second)
+    return SymbolOpTy();
+
+  return result.first->second;
+}
+
+static LogicalResult
+deduplicateGlobalVariables(spirv::ModuleOp combinedModule) {
+  DenseMap<std::pair<int64_t, int64_t>, spirv::GlobalVariableOp>
+      descriptorToGlobalVarOpMap;
+  DenseMap<StringRef, spirv::GlobalVariableOp> builtInToGlobalVarOpMap;
+
+  WalkResult result =
+      combinedModule.walk([&](spirv::GlobalVariableOp globalVarOp) {
+        spirv::GlobalVariableOp replacementSymOp;
+
+        IntegerAttr descriptorSet = globalVarOp.getAttrOfType<IntegerAttr>(
+            spirv::SPIRVDialect::getAttributeName(
+                spirv::Decoration::DescriptorSet));
+        IntegerAttr binding = globalVarOp.getAttrOfType<IntegerAttr>(
+            spirv::SPIRVDialect::getAttributeName(spirv::Decoration::Binding));
+
+        if (descriptorSet) {
+          replacementSymOp = emplaceOrGetReplacementSymbolName(
+              {descriptorSet.getInt(), binding.getInt()}, globalVarOp,
+              descriptorToGlobalVarOpMap);
+        } else {
+          StringAttr builtIn = globalVarOp.getAttrOfType<StringAttr>(
+              spirv::SPIRVDialect::getAttributeName(
+                  spirv::Decoration::BuiltIn));
+
+          if (builtIn) {
+            replacementSymOp = emplaceOrGetReplacementSymbolName(
+
+                builtIn.getValue(), globalVarOp, builtInToGlobalVarOpMap);
+          }
+        }
+
+        if (!replacementSymOp)
+          return WalkResult::advance();
+
+        // There is already a global variable with either the same descriptor
+        // set + binding or the same built_in attribute. Deduplicate the current
+        // global variable.
+        if (failed(SymbolTable::replaceAllSymbolUses(
+                globalVarOp, replacementSymOp.getName(), combinedModule)))
+          return WalkResult(
+              globalVarOp.emitError("unable to update all symbol uses for ")
+              << globalVarOp.getName() << " to " << replacementSymOp.getName());
+
+        globalVarOp.erase();
+        return WalkResult::advance();
+      });
+
+  return result == WalkResult::interrupt() ? failure() : success();
+}
+
+static LogicalResult deduplicateSpecConstants(spirv::ModuleOp combinedModule) {
+  DenseMap<int64_t, spirv::SpecConstantOp> specIdToSpecConstCompositeMap;
+
+  WalkResult result =
+      combinedModule.walk([&](spirv::SpecConstantOp specConstOp) {
+        IntegerAttr specId = specConstOp.getAttrOfType<IntegerAttr>(
+            spirv::SPIRVDialect::getAttributeName(spirv::Decoration::SpecId));
+
+        if (specId) {
+          spirv::SpecConstantOp replacementSymOp =
+              emplaceOrGetReplacementSymbolName(specId.getInt(), specConstOp,
+                                                specIdToSpecConstCompositeMap);
+
+          if (!replacementSymOp)
+            return WalkResult::advance();
+
+          if (failed(SymbolTable::replaceAllSymbolUses(
+                  specConstOp, replacementSymOp.getName(), combinedModule)))
+            return WalkResult(
+                specConstOp.emitError("unable to update all symbol uses for ")
+                << specConstOp.sym_name() << " to "
+                << replacementSymOp.getName());
+
+          specConstOp.erase();
+        }
+
+        return WalkResult::advance();
+      });
+
+  return result == WalkResult::interrupt() ? failure() : success();
+}
+
 namespace mlir {
 namespace spirv {
 
 // TODO Properly test symbol rename listener mechanism.
 
 OwningSPIRVModuleRef
 combine(llvm::MutableArrayRef<spirv::ModuleOp> modules,
         OpBuilder &combinedModuleBuilder,
for (auto &op : moduleClone.getBlock().without_terminator()) {
       }
     }
 
     // Clone all the module's ops to the combined module.
     for (auto &op : moduleClone.getBlock().without_terminator())
       combinedModuleBuilder.insert(op.clone());
   }
 
+  if (failed(deduplicateGlobalVariables(combinedModule)) ||
+      failed(deduplicateSpecConstants(combinedModule)))
+    return nullptr;
+
   return combinedModule;
 }
 
 } // namespace spirv
 } // namespace mlir

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

 }
 
 // -----
 
 // Resolve conflicting globalVariableOps.
 
 // CHECK:      module {
 // CHECK-NEXT:   spv.module Logical GLSL450 {
-// CHECK-NEXT:     spv.globalVariable @foo_1
+// CHECK-NEXT:     spv.globalVariable @foo_1 bind(1, 0)
 
-// CHECK-NEXT:     spv.globalVariable @foo
+// CHECK-NEXT:     spv.globalVariable @foo bind(2, 0)
 // 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>
+  spv.globalVariable @foo bind(2, 0) : !spv.ptr<f32, Input>
+}
+}
+
+// -----
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.globalVariable @foo_1 built_in("GlobalInvocationId")
+
+// CHECK-NEXT:     spv.globalVariable @foo built_in("LocalInvocationId")
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo built_in("GlobalInvocationId") : !spv.ptr<vector<3xi32>, Input>
+}
+
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo built_in("LocalInvocationId") : !spv.ptr<vector<3xi32>, Input>
 }
 }
 
 // -----
 
 // Resolve conflicting globalVariableOp and specConstantOp.
 
 // CHECK:      module {

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

+// RUN: mlir-opt -test-spirv-module-combiner -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// Deduplicate 2 global variables with the same descriptor set and binding.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.globalVariable @foo
+
+// CHECK-NEXT:     spv.func @use_foo
+// CHECK-NEXT:       spv._address_of @foo
+// CHECK-NEXT:       spv.Load
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @use_bar
+// CHECK-NEXT:       spv._address_of @foo
+// CHECK-NEXT:       spv.Load
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<f32, Input>
+
+  spv.func @use_foo() -> 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.globalVariable @bar bind(1, 0) : !spv.ptr<f32, Input>
+
+  spv.func @use_bar() -> f32 "None" {
+    %0 = spv._address_of @bar : !spv.ptr<f32, Input>
+    %1 = spv.Load "Input" %0 : f32
+    spv.ReturnValue %1 : f32
+  }
+}
+}
+
+// -----
+
+// Deduplicate 2 global variables with the same descriptor set and binding but different types.
+
+module {
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo bind(1, 0) : !spv.ptr<i32, Input>
+}
+
+spv.module Logical GLSL450 {
+  spv.globalVariable @bar bind(1, 0) : !spv.ptr<f32, Input>
+
+  spv.func @use_bar() -> f32 "None" {
+    // expected-error @+1 {{result type mismatch with the referenced global variable's type}}
+    %0 = spv._address_of @bar : !spv.ptr<f32, Input>
+    %1 = spv.Load "Input" %0 : f32
+    spv.ReturnValue %1 : f32
+  }
+}
+}
+
+// -----
+
+// Deduplicate 2 global variables with the same built-in attribute.
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.globalVariable @foo built_in("GlobalInvocationId")
+// CHECK-NEXT:     spv.func @use_bar
+// CHECK-NEXT:       spv._address_of @foo
+// CHECK-NEXT:       spv.Load
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.globalVariable @foo built_in("GlobalInvocationId") : !spv.ptr<vector<3xi32>, Input>
+}
+
+spv.module Logical GLSL450 {
+  spv.globalVariable @bar built_in("GlobalInvocationId") : !spv.ptr<vector<3xi32>, Input>
+
+  spv.func @use_bar() -> vector<3xi32> "None" {
+    %0 = spv._address_of @bar : !spv.ptr<vector<3xi32>, Input>
+    %1 = spv.Load "Input" %0 : vector<3xi32>
+    spv.ReturnValue %1 : vector<3xi32>
+  }
+}
+}
+
+// -----
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.specConstant @foo spec_id(5)
+
+// CHECK-NEXT:     spv.func @use_foo()
+// CHECK-NEXT:       %0 = spv._reference_of @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @use_bar()
+// CHECK-NEXT:       %0 = spv._reference_of @foo
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+module {
+spv.module Logical GLSL450 {
+  spv.specConstant @foo spec_id(5) = 1. : f32
+
+  spv.func @use_foo() -> (f32) "None" {
+    %0 = spv._reference_of @foo : f32
+    spv.ReturnValue %0 : f32
+  }
+}
+
+spv.module Logical GLSL450 {
+  spv.specConstant @bar spec_id(5) = 1. : f32
+
+  spv.func @use_bar() -> (f32) "None" {
+    %0 = spv._reference_of @bar : f32
+    spv.ReturnValue %0 : f32
+  }
+}
+}