[MLIR] Address post-commit comments on function deduplication

[mehdi_amini]

mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp

74

There shouldn't be anything special about "ModuleOp", I suspect you probably want to walk operations that define a symbol table?

76

Aren't there other potential uses of a function? I would think that there are operations taking the address of a function for example. How does the Inliner finds all uses of a symbol?

78

I don't understand what is this check intended to catch? Nested modules?

I would walk all Operations, check if the operation defined a SymbolTable, and return WalkResult::Skip to avoid walking through it.

[frgossen]

Addrwss comments

mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp

76

Right, there may be other uses. For now, we only need call sites in XLA:CPU and I don't see an interface for this. I suggest we add them as needed. Otherwise, I would probably move this pass to the XLA repo for now.

78

Yes, this is for nested modules. Skip did not work as expected for some reason :/

[mehdi_amini]

mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp

76

I don’t quite get your point here: I see this as an actual bug making the pass unsafe to use, you seem to think this is some « nice to have » extra features or something?

95

Without skip the operation will be traversed incorrectly, won’t it?

[frgossen]

Address comments

[rriddle]

mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp

76

Yeah, that's kind of a serious bug right now. That should be handled here, otherwise the pass isn't really usable outside of very narrow situations.

78

Default walk is post-order, i.e. you walk the inner operations first (and skip is meaningless). You need a walk<PreOrder> for that to be useful.

[mehdi_amini]

Ping on this revision?

[frgossen]

Haven't forgotten about this but, unfortunately, there were more urgent things to do.
If you prefer, I can move the pass to XLA for now and move it here when I find the time.

mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp

76

Why is that a bug? The old functions remain untouched as we no longer delete them.
If we wanted to cover all cases here we'd (i) need an interface for all the ops that may reference a function or (ii) implement all the cases here.

95

It would have skipped all the irrelevant ops. With WalkOrder::PreOrder this works and is a bit nicer, I guess.

[mehdi_amini]

In D146456#4249863, @frgossen wrote:

Haven't forgotten about this but, unfortunately, there were more urgent things to do.
If you prefer, I can move the pass to XLA for now and move it here when I find the time.

It's up to you actually, there is just a bit more time pressure to address post-review comment because code becomes a bit more "load bearing" after it lands. I'd think that if you can fix this in the coming week it is fine.

mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp

76

Why is that a bug? The old functions remain untouched as we no longer delete them.

Ah you may very well be right!
That's a nice side-effect of removing the code that delete them :)

If we wanted to cover all cases here we'd (i) need an interface for all the ops that may reference a function or (ii) implement all the cases here.

I thought we had such an interface?

mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp

 //===- DuplicateFunctionElimination.cpp - Duplicate function elimination --===//
 //
 // 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/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/Passes.h"
+#include "mlir/IR/Visitors.h"
+
+using namespace mlir;
 
-namespace mlir {
 namespace {
 
 #define GEN_PASS_DEF_DUPLICATEFUNCTIONELIMINATIONPASS
 #include "mlir/Dialect/Func/Transforms/Passes.h.inc"
 
 // Define a notion of function equivalence that allows for reuse. Ignore the
 // symbol name for this purpose.
 struct DuplicateFuncOpEquivalenceInfo
-    : public llvm::DenseMapInfo<func::FuncOp> {
+    : public llvm::DenseMapInfo<FunctionOpInterface> {
 
-  static unsigned getHashValue(const func::FuncOp cFunc) {
-    if (!cFunc) {
-      return DenseMapInfo<func::FuncOp>::getHashValue(cFunc);
-    }
-
-    // Aggregate attributes, ignoring the symbol name.
-    llvm::hash_code hash = {};
-    func::FuncOp func = const_cast<func::FuncOp &>(cFunc);
-    StringAttr symNameAttrName = func.getSymNameAttrName();
-    for (NamedAttribute namedAttr : cFunc->getAttrs()) {
-      StringAttr attrName = namedAttr.getName();
-      if (attrName == symNameAttrName)
-        continue;
-      hash = llvm::hash_combine(hash, namedAttr);
-    }
+  static unsigned getHashValue(const FunctionOpInterface cFunc) {
+    if (!cFunc)
+      return DenseMapInfo<FunctionOpInterface>::getHashValue(cFunc);
+
+    // Base hash on function type.
+    auto &func = const_cast<FunctionOpInterface &>(cFunc);
+    llvm::hash_code hash = llvm::hash_combine(func.getFunctionType());
 
     // Also hash the func body.
-    func.getBody().walk([&](Operation *op) {
+    func.getFunctionBody().walk([&](Operation *op) {
       hash = llvm::hash_combine(
           hash, OperationEquivalence::computeHash(
                     op, /*hashOperands=*/OperationEquivalence::ignoreHashValue,
                     /*hashResults=*/OperationEquivalence::ignoreHashValue,
                     OperationEquivalence::IgnoreLocations));
     });
 
     return hash;
   }
 
-  static bool isEqual(const func::FuncOp cLhs, const func::FuncOp cRhs) {
-    if (cLhs == cRhs) {
+  static bool isEqual(const FunctionOpInterface cLhs,
+                      const FunctionOpInterface cRhs) {
+    if (cLhs == cRhs)
       return true;
-    }
     if (cLhs == getTombstoneKey() || cLhs == getEmptyKey() ||
-        cRhs == getTombstoneKey() || cRhs == getEmptyKey()) {
+        cRhs == getTombstoneKey() || cRhs == getEmptyKey())
       return false;
-    }
 
-    // Check attributes equivalence, ignoring the symbol name.
-    if (cLhs->getAttrDictionary().size() != cRhs->getAttrDictionary().size()) {
+    // Check function type equivalence.
+    auto &lhs = const_cast<FunctionOpInterface &>(cLhs);
+    auto &rhs = const_cast<FunctionOpInterface &>(cRhs);
+    if (lhs.getFunctionType() != rhs.getFunctionType())
       return false;
-    }
-    func::FuncOp lhs = const_cast<func::FuncOp &>(cLhs);
-    StringAttr symNameAttrName = lhs.getSymNameAttrName();
-    for (NamedAttribute namedAttr : cLhs->getAttrs()) {
-      StringAttr attrName = namedAttr.getName();
-      if (attrName == symNameAttrName) {
-        continue;
-      }
-      if (namedAttr.getValue() != cRhs->getAttr(attrName)) {
-        return false;
-      }
-    }
 
     // Compare inner workings.
-    func::FuncOp rhs = const_cast<func::FuncOp &>(cRhs);
     return OperationEquivalence::isRegionEquivalentTo(
-        &lhs.getBody(), &rhs.getBody(), OperationEquivalence::IgnoreLocations);
+        &lhs.getFunctionBody(), &rhs.getFunctionBody(),
+        OperationEquivalence::IgnoreLocations);
   }
 };
 
 struct DuplicateFunctionEliminationPass
     : public impl::DuplicateFunctionEliminationPassBase<
           DuplicateFunctionEliminationPass> {
 
   using DuplicateFunctionEliminationPassBase<
       DuplicateFunctionEliminationPass>::DuplicateFunctionEliminationPassBase;
 
   void runOnOperation() override {
-    auto module = getOperation();
 
-    // Find unique representant per equivalent func ops.
-    DenseSet<func::FuncOp, DuplicateFuncOpEquivalenceInfo> uniqueFuncOps;
-    DenseMap<StringAttr, func::FuncOp> getRepresentant;
-    DenseSet<func::FuncOp> toBeErased;
-    module.walk([&](func::FuncOp f) {
+    // Traverse all symbol table ops.
+    getOperation().walk([](Operation *root) {
+      if (!root->hasTrait<OpTrait::SymbolTable>())
+        return;
+
+      // Find unique representant per equivalent func ops. Analyze direct
+      // children of the root op.
+      DenseSet<FunctionOpInterface, DuplicateFuncOpEquivalenceInfo>
+          uniqueFuncOps;
+      DenseMap<StringAttr, FunctionOpInterface> getRepresentant;
+      root->walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
+        // Don't enter inner modules.
+        if (op != root && op->hasTrait<OpTrait::SymbolTable>())
+          return WalkResult::skip();
+
+        if (auto f = llvm::dyn_cast<FunctionOpInterface>(op)) {
           auto [repr, inserted] = uniqueFuncOps.insert(f);
-      getRepresentant[f.getSymNameAttr()] = *repr;
-      if (!inserted) {
-        toBeErased.insert(f);
+          getRepresentant[f.getNameAttr()] = *repr;
         }
-    });
 
-    // Update call ops to call unique func op representants.
-    module.walk([&](func::CallOp callOp) {
-      func::FuncOp callee = getRepresentant[callOp.getCalleeAttr().getAttr()];
-      callOp.setCallee(callee.getSymName());
+        return WalkResult::advance();
       });
 
-    // Erase redundant func ops.
-    for (auto it : toBeErased) {
-      it.erase();
+      // Update symbol users.
+      root->walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
+        // Don't enter inner modules.
+        if (op != root && op->hasTrait<OpTrait::SymbolTable>())
+          return WalkResult::skip();
+
+        // Update call sites.
+        if (func::CallOp callOp = llvm::dyn_cast<func::CallOp>(op)) {
+          FunctionOpInterface callee =
+              getRepresentant[callOp.getCalleeAttr().getAttr()];
+          callOp.setCallee(callee.getNameAttr());
         }
+
+        return WalkResult::advance();
+      });
+    });
   }
 };
 
 } // namespace
 
-std::unique_ptr<Pass> mlir::func::createDuplicateFunctionEliminationPass() {
+std::unique_ptr<Pass> func::createDuplicateFunctionEliminationPass() {
   return std::make_unique<DuplicateFunctionEliminationPass>();
 }
-
-} // namespace mlir

mlir/test/Dialect/Func/duplicate-function-elimination.mlir

 
 func.func @user(%arg0: tensor<f32>) -> tensor<f32> {
   %0 = call @identity(%arg0) : (tensor<f32>) -> tensor<f32>
   %1 = call @also_identity(%0) : (tensor<f32>) -> tensor<f32>
   %2 = call @yet_another_identity(%1) : (tensor<f32>) -> tensor<f32>
   return %2 : tensor<f32>
 }
 
-// CHECK:     @identity
-// CHECK-NOT: @also_identity
-// CHECK-NOT: @yet_another_identity
-// CHECK:     @user
-// CHECK-3:     call @identity
+// CHECK:   @identity
+// CHECK:   @also_identity
+// CHECK:   @yet_another_identity
+// CHECK:   @user
+// CHECK-3:   call @identity
 
 // -----
 
 func.func @add_lr(%arg0: f32, %arg1: f32) -> f32 {
   %0 = arith.addf %arg0, %arg1 : f32
   return %0 : f32
 }
 

 func.func @user(%arg0: f32, %arg1: f32) -> f32 {
   %0 = call @add_lr(%arg0, %arg1) : (f32, f32) -> f32
   %1 = call @also_add_lr(%arg0, %arg1) : (f32, f32) -> f32
   %2 = call @add_rl(%0, %1) : (f32, f32) -> f32
   %3 = call @also_add_rl(%arg0, %2) : (f32, f32) -> f32
  return %3 : f32
 }
 
-// CHECK:     @add_lr
-// CHECK-NOT: @also_add_lr
-// CHECK-NOT: @add_rl
-// CHECK-NOT: @also_add_rl
-// CHECK:     @user
-// CHECK-4:     call @add_lr
+// CHECK:   @add_lr
+// CHECK:   @also_add_lr
+// CHECK:   @add_rl
+// CHECK:   @also_add_rl
+// CHECK:   @user
+// CHECK-4:   call @add_lr
 
 // -----
 
 func.func @ite(%pred: i1, %then: f32, %else: f32) -> f32 {
   %0 = scf.if %pred -> f32 {
     scf.yield %then : f32
   } else {
     scf.yield %else : f32

 
 func.func @user(%pred : i1, %arg0: f32, %arg1: f32) -> f32 {
   %0 = call @also_ite(%pred, %arg0, %arg1) : (i1, f32, f32) -> f32
   %1 = call @ite(%pred, %arg0, %arg1) : (i1, f32, f32) -> f32
   %2 = call @reverse_ite(%pred, %0, %1) : (i1, f32, f32) -> f32
  return %2 : f32
 }
 
-// CHECK:     @ite
-// CHECK-NOT: @also_ite
-// CHECK:     @reverse_ite
-// CHECK:     @user
-// CHECK-2:     call @ite
-// CHECK:       call @reverse_ite
+// CHECK:   @ite
+// CHECK:   @also_ite
+// CHECK:   @reverse_ite
+// CHECK:   @user
+// CHECK-2:   call @ite
+// CHECK:     call @reverse_ite
 
 // -----
 
 func.func @deep_tree(%p0: i1, %p1: i1, %p2: i1, %p3: i1, %even: f32, %odd: f32) 
     -> f32 {
   %0 = scf.if %p0 -> f32 {
     %1 = scf.if %p1 -> f32 {
       %2 = scf.if %p2 -> f32 {
%0 = call @deep_tree(%p0, %p1, %p2, %p3, %odd, %even)
       : (i1, i1, i1, i1, f32, f32) -> f32
   %1 = call @also_deep_tree(%p0, %p1, %p2, %p3, %odd, %even) 
       : (i1, i1, i1, i1, f32, f32) -> f32
   %2 = call @reverse_deep_tree(%p0, %p1, %p2, %p3, %odd, %even) 
       : (i1, i1, i1, i1, f32, f32) -> f32
   return %0, %1, %2 : f32, f32, f32
 }
 
-// CHECK:     @deep_tree
-// CHECK-NOT: @also_deep_tree
-// CHECK:     @reverse_deep_tree
-// CHECK:     @user
-// CHECK-2:     call @deep_tree
-// CHECK:       call @reverse_deep_tree
+// CHECK:   @deep_tree
+// CHECK:   @also_deep_tree
+// CHECK:   @reverse_deep_tree
+// CHECK:   @user
+// CHECK-2:   call @deep_tree
+// CHECK:     call @reverse_deep_tree
+
+// -----
+
+module {
+
+  func.func @inner_identity(%arg0: tensor<f32>) -> tensor<f32> {
+    return %arg0 : tensor<f32>
+  }
+
+  func.func @also_inner_identity(%arg0: tensor<f32>) -> tensor<f32> {
+    return %arg0 : tensor<f32>
+  }
+
+  func.func @inner_user(%arg0: tensor<f32>) -> tensor<f32> {
+    %0 = call @inner_identity(%arg0) : (tensor<f32>) -> tensor<f32>
+    %1 = call @also_inner_identity(%0) : (tensor<f32>) -> tensor<f32>
+    return %1 : tensor<f32>
+  }
+
+}
+
+func.func @identity(%arg0: tensor<f32>) -> tensor<f32> {
+  return %arg0 : tensor<f32>
+}
+
+func.func @also_identity(%arg0: tensor<f32>) -> tensor<f32> {
+  return %arg0 : tensor<f32>
+}
+
+func.func @user(%arg0: tensor<f32>) -> tensor<f32> {
+  %0 = call @identity(%arg0) : (tensor<f32>) -> tensor<f32>
+  %1 = call @also_identity(%0) : (tensor<f32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK: module 
+// CHECK:   module
+// CHECK:     @inner_identity
+// CHECK:     @also_inner_identity
+// CHECK:     @inner_user
+// CHECK:       call @inner_identity
+// CHECK:       call @inner_identity
+// CHECK:   @identity
+// CHECK:   @also_identity
+// CHECK:   @user
+// CHECK:     call @identity
+// CHECK:     call @identity