[MLIR][SPIRV] ModuleCombiner: deduplicate identical functions.

[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 functions that are identical to each
other.

[ergawy]

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

159

This is a partial copy of BlockEquivalenceData from here: https://github.com/llvm/llvm-project/blob/474f7639e3494d9605f4444b087f48e710fbb0d4/mlir/lib/Transforms/Utils/RegionUtils.cpp#L374. I only copied the bare minimum needed for what I want to do in this patch.

However, it would be great if the original (not this stolen adaptation :) ) BlockEquivalenceData be put in some util header file accessible outside RegionUtils.cpp.

[ergawy]

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

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

 ///
 /// 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.
+///   - 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.

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

 
 #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/Hashing.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) {
WalkResult result =
         }
 
         return WalkResult::advance();
       });
 
   return result == WalkResult::interrupt() ? failure() : success();
 }
 
+struct BlockEquivalenceData {
+  BlockEquivalenceData(Block *block);
+
+  llvm::hash_code hash;
+};
+
+BlockEquivalenceData::BlockEquivalenceData(Block *block) : hash(0) {
+  for (Operation &op : *block) {
+    hash = llvm::hash_combine(
+        hash, OperationEquivalence::computeHash(
+                  &op, OperationEquivalence::Flags::IgnoreOperands));
+  }
+}
+
+static LogicalResult deduplicateFunctions(spirv::ModuleOp combinedModule) {
+  DenseMap<llvm::hash_code, spirv::FuncOp> hashToFunctionMap;
+
+  WalkResult result = combinedModule.walk([&](spirv::FuncOp func) {
+    llvm::hash_code funcHash(0);
+    for (auto attr : func.getOperation()->getAttrs()) {
+      if (attr.first == SymbolTable::getSymbolAttrName())
+        continue;
+      funcHash = llvm::hash_combine(funcHash, attr);
+    }
+
+    for (auto &blk : func) {
+      funcHash = llvm::hash_combine(funcHash, BlockEquivalenceData(&blk).hash);
+    }
+
+    spirv::FuncOp replacementSymOp =
+        emplaceOrGetReplacementSymbolName(funcHash, func, hashToFunctionMap);
+
+    if (!replacementSymOp)
+      return WalkResult::advance();
+
+    if (failed(SymbolTable::replaceAllSymbolUses(
+            func, replacementSymOp.getName(), combinedModule)))
+      return WalkResult(func.emitError("unable to update all symbol uses for ")
+                        << func.sym_name() << " to "
+                        << replacementSymOp.getName());
+
+    func.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 module : modules) {
     }
 
     // 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)))
+      failed(deduplicateSpecConstants(combinedModule)) ||
+      failed(deduplicateFunctions(combinedModule)))
     return nullptr;
 
   return combinedModule;
 }
 
 } // namespace spirv
 } // namespace mlir

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

 
 // 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.FAdd
 // CHECK-NEXT:       spv.ReturnValue
 // CHECK-NEXT:     }
 
 // CHECK-NEXT:     spv.func @foo_2
+// CHECK-NEXT:       spv.ISub
 // 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
+    %0 = spv.FAdd %arg0, %arg0 : f32
+    spv.ReturnValue %0 : f32
   }
 }
 
 spv.module Logical GLSL450 {
   spv.func @foo(%arg0 : i32) -> i32 "None" {
-    spv.ReturnValue %arg0 : i32
+    %0 = spv.ISub %arg0, %arg0 : i32
+    spv.ReturnValue %0 : i32
   }
 }
 }
 
 // -----
 
 // Test properly updating references to a renamed funcOp.
 

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

 // 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.FAdd
 // 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
+    %2 = spv.FAdd %1, %1 : f32
+    spv.ReturnValue %2 : f32
   }
 }
 }
 
 // -----
 
 // Deduplicate 2 global variables with the same descriptor set and binding but different types.
 

 
 // 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.FAdd
 // 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
+    %1 = spv.FAdd %0, %0 : f32
+    spv.ReturnValue %1 : f32
+  }
+}
+}
+
+// -----
+
+// CHECK:      module {
+// CHECK-NEXT:   spv.module Logical GLSL450 {
+// CHECK-NEXT:     spv.specConstant @bar spec_id(5)
+
+// CHECK-NEXT:     spv.func @foo(%arg0: f32)
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @foo_different_body(%arg0: f32)
+// CHECK-NEXT:       spv._reference_of
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @baz(%arg0: i32)
+// CHECK-NEXT:       spv.ReturnValue
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @baz_no_return(%arg0: i32)
+// CHECK-NEXT:       spv.Return
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @baz_no_return_different_control
+// CHECK-NEXT:       spv.Return
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @baz_no_return_another_control
+// CHECK-NEXT:       spv.Return
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @kernel
+// CHECK-NEXT:       spv.Return
+// CHECK-NEXT:     }
+
+// CHECK-NEXT:     spv.func @kernel_different_attr
+// CHECK-NEXT:       spv.Return
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT:   }
+
+module {
+spv.module Logical GLSL450 {
+  spv.specConstant @bar spec_id(5) = 1. : f32
+
+  spv.func @foo(%arg0: f32) -> (f32) "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+
+  spv.func @foo_duplicate(%arg0: f32) -> (f32) "None" {
+    spv.ReturnValue %arg0 : f32
+  }
+
+  spv.func @foo_different_body(%arg0: f32) -> (f32) "None" {
+    %0 = spv._reference_of @bar : f32
+    spv.ReturnValue %arg0 : f32
+  }
+
+  spv.func @baz(%arg0: i32) -> (i32) "None" {
+    spv.ReturnValue %arg0 : i32
+  }
+
+  spv.func @baz_no_return(%arg0: i32) "None" {
+    spv.Return
+  }
+
+  spv.func @baz_no_return_duplicate(%arg0: i32) -> () "None" {
+    spv.Return
+  }
+
+  spv.func @baz_no_return_different_control(%arg0: i32) -> () "Inline" {
+    spv.Return
+  }
+
+  spv.func @baz_no_return_another_control(%arg0: i32) -> () "Inline|Pure" {
+    spv.Return
+  }
+
+  spv.func @kernel(
+    %arg0: f32,
+    %arg1: !spv.ptr<!spv.struct<(!spv.array<12 x f32>)>, CrossWorkgroup>) "None"
+  attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}} {
+    spv.Return
+  }
+
+  spv.func @kernel_different_attr(
+    %arg0: f32,
+    %arg1: !spv.ptr<!spv.struct<(!spv.array<12 x f32>)>, CrossWorkgroup>) "None"
+  attributes {spv.entry_point_abi = {local_size = dense<[64, 1, 1]> : vector<3xi32>}} {
+    spv.Return
   }
 }
 }