[mlir][func] Fix multiple bugs in `DuplicateFunctionElimination` #109571

CoTinker · 2024-09-22T07:59:36Z

This PR fixes multiple bugs in DuplicateFunctionElimination.

Prevents elimination of function declarations.
Updates all symbol uses to reference unique function representatives.

Fixes #93483.

llvmbot · 2024-09-22T08:00:12Z

@llvm/pr-subscribers-mlir-func

Author: Longsheng Mou (CoTinker)

Changes

This PR fixes multiple bugs in DuplicateFunctionElimination.

Prevents elimination of function declarations.
Updates constant ops to reference unique function representatives.
Simplifies DenseMap by using StringRef as the key instead of StringAttr.

Fixes #93483.

Full diff: https://github.com/llvm/llvm-project/pull/109571.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp (+12-3)
(modified) mlir/test/Dialect/Func/duplicate-function-elimination.mlir (+48)

diff --git a/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp b/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
index d41d6c3e8972f9..5e23207eabf9c4 100644
--- a/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
+++ b/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
@@ -54,6 +54,10 @@ struct DuplicateFuncOpEquivalenceInfo
     if (lhs == getTombstoneKey() || lhs == getEmptyKey() ||
         rhs == getTombstoneKey() || rhs == getEmptyKey())
       return false;
+
+    if (lhs.isDeclaration() || rhs.isDeclaration())
+      return false;
+
     // Check discardable attributes equivalence
     if (lhs->getDiscardableAttrDictionary() !=
         rhs->getDiscardableAttrDictionary())
@@ -87,11 +91,11 @@ struct DuplicateFunctionEliminationPass
 
     // Find unique representant per equivalent func ops.
     DenseSet<func::FuncOp, DuplicateFuncOpEquivalenceInfo> uniqueFuncOps;
-    DenseMap<StringAttr, func::FuncOp> getRepresentant;
+    DenseMap<StringRef, func::FuncOp> getRepresentant;
     DenseSet<func::FuncOp> toBeErased;
     module.walk([&](func::FuncOp f) {
       auto [repr, inserted] = uniqueFuncOps.insert(f);
-      getRepresentant[f.getSymNameAttr()] = *repr;
+      getRepresentant[f.getSymName()] = *repr;
       if (!inserted) {
         toBeErased.insert(f);
       }
@@ -99,9 +103,14 @@ struct DuplicateFunctionEliminationPass
 
     // Update call ops to call unique func op representants.
     module.walk([&](func::CallOp callOp) {
-      func::FuncOp callee = getRepresentant[callOp.getCalleeAttr().getAttr()];
+      func::FuncOp callee = getRepresentant[callOp.getCallee()];
       callOp.setCallee(callee.getSymName());
     });
+    // Update constant ops to reference unique func op representants.
+    module.walk([&](func::ConstantOp constantOp) {
+      func::FuncOp value = getRepresentant[constantOp.getValue()];
+      constantOp.setValue(value.getSymName());
+    });
 
     // Erase redundant func ops.
     for (auto it : toBeErased) {
diff --git a/mlir/test/Dialect/Func/duplicate-function-elimination.mlir b/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
index 28d059a149bde8..1c6876c1327bc6 100644
--- a/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
+++ b/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
@@ -366,3 +366,51 @@ func.func @user(%p0: i1, %p1: i1, %p2: i1, %p3: i1, %odd: f32, %even: f32)
 // CHECK:     @user
 // CHECK-2:     call @deep_tree
 // CHECK:       call @reverse_deep_tree
+
+// -----
+
+func.func private @func_declaration(i32, i32) -> i32
+func.func private @func_declaration1(i32, i32) -> i32
+
+func.func @user(%arg0: i32, %arg1: i32) -> (i32, i32) {
+  %0 = call @func_declaration(%arg0, %arg1) : (i32, i32) -> i32
+  %1 = call @func_declaration1(%arg0, %arg1) : (i32, i32) -> i32
+  return %0, %1 : i32, i32
+}
+
+// CHECK:  @func_declaration
+// CHECK:  @func_declaration1
+// CHECK:  @user
+// CHECK:    call @func_declaration
+// CHECK:    call @func_declaration1
+
+
+// -----
+
+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 @yet_another_identity(%arg0: tensor<f32>) -> tensor<f32> {
+  return %arg0 : tensor<f32>
+}
+
+func.func @user(%arg0: tensor<f32>) -> tensor<f32> {
+  %f = constant @identity : (tensor<f32>) -> tensor<f32>
+  %0 = call_indirect %f(%arg0) : (tensor<f32>) -> tensor<f32>
+  %f_0 = constant @also_identity : (tensor<f32>) -> tensor<f32>
+  %1 = call_indirect %f_0(%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-2:     constant @identity
+// CHECK:       call @identity

llvmbot · 2024-09-22T08:00:12Z

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

Changes

This PR fixes multiple bugs in DuplicateFunctionElimination.

Prevents elimination of function declarations.
Updates constant ops to reference unique function representatives.
Simplifies DenseMap by using StringRef as the key instead of StringAttr.

Fixes #93483.

Full diff: https://github.com/llvm/llvm-project/pull/109571.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp (+12-3)
(modified) mlir/test/Dialect/Func/duplicate-function-elimination.mlir (+48)

diff --git a/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp b/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
index d41d6c3e8972f9..5e23207eabf9c4 100644
--- a/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
+++ b/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
@@ -54,6 +54,10 @@ struct DuplicateFuncOpEquivalenceInfo
     if (lhs == getTombstoneKey() || lhs == getEmptyKey() ||
         rhs == getTombstoneKey() || rhs == getEmptyKey())
       return false;
+
+    if (lhs.isDeclaration() || rhs.isDeclaration())
+      return false;
+
     // Check discardable attributes equivalence
     if (lhs->getDiscardableAttrDictionary() !=
         rhs->getDiscardableAttrDictionary())
@@ -87,11 +91,11 @@ struct DuplicateFunctionEliminationPass
 
     // Find unique representant per equivalent func ops.
     DenseSet<func::FuncOp, DuplicateFuncOpEquivalenceInfo> uniqueFuncOps;
-    DenseMap<StringAttr, func::FuncOp> getRepresentant;
+    DenseMap<StringRef, func::FuncOp> getRepresentant;
     DenseSet<func::FuncOp> toBeErased;
     module.walk([&](func::FuncOp f) {
       auto [repr, inserted] = uniqueFuncOps.insert(f);
-      getRepresentant[f.getSymNameAttr()] = *repr;
+      getRepresentant[f.getSymName()] = *repr;
       if (!inserted) {
         toBeErased.insert(f);
       }
@@ -99,9 +103,14 @@ struct DuplicateFunctionEliminationPass
 
     // Update call ops to call unique func op representants.
     module.walk([&](func::CallOp callOp) {
-      func::FuncOp callee = getRepresentant[callOp.getCalleeAttr().getAttr()];
+      func::FuncOp callee = getRepresentant[callOp.getCallee()];
       callOp.setCallee(callee.getSymName());
     });
+    // Update constant ops to reference unique func op representants.
+    module.walk([&](func::ConstantOp constantOp) {
+      func::FuncOp value = getRepresentant[constantOp.getValue()];
+      constantOp.setValue(value.getSymName());
+    });
 
     // Erase redundant func ops.
     for (auto it : toBeErased) {
diff --git a/mlir/test/Dialect/Func/duplicate-function-elimination.mlir b/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
index 28d059a149bde8..1c6876c1327bc6 100644
--- a/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
+++ b/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
@@ -366,3 +366,51 @@ func.func @user(%p0: i1, %p1: i1, %p2: i1, %p3: i1, %odd: f32, %even: f32)
 // CHECK:     @user
 // CHECK-2:     call @deep_tree
 // CHECK:       call @reverse_deep_tree
+
+// -----
+
+func.func private @func_declaration(i32, i32) -> i32
+func.func private @func_declaration1(i32, i32) -> i32
+
+func.func @user(%arg0: i32, %arg1: i32) -> (i32, i32) {
+  %0 = call @func_declaration(%arg0, %arg1) : (i32, i32) -> i32
+  %1 = call @func_declaration1(%arg0, %arg1) : (i32, i32) -> i32
+  return %0, %1 : i32, i32
+}
+
+// CHECK:  @func_declaration
+// CHECK:  @func_declaration1
+// CHECK:  @user
+// CHECK:    call @func_declaration
+// CHECK:    call @func_declaration1
+
+
+// -----
+
+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 @yet_another_identity(%arg0: tensor<f32>) -> tensor<f32> {
+  return %arg0 : tensor<f32>
+}
+
+func.func @user(%arg0: tensor<f32>) -> tensor<f32> {
+  %f = constant @identity : (tensor<f32>) -> tensor<f32>
+  %0 = call_indirect %f(%arg0) : (tensor<f32>) -> tensor<f32>
+  %f_0 = constant @also_identity : (tensor<f32>) -> tensor<f32>
+  %1 = call_indirect %f_0(%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-2:     constant @identity
+// CHECK:       call @identity

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

CoTinker · 2024-09-27T08:59:45Z

I've noticed that replaceAllSymbolUsesImpl uses collectSymbolScopes to traverse the region and addReplacement to perform replacements. Perhaps using SymbolTable::replaceAllSymbolUses inside the loop isn't as costly as we think.

llvm-project/mlir/lib/IR/SymbolTable.cpp

Lines 877 to 891 in 8dd817b

    
           /// The implementation of SymbolTable::replaceAllSymbolUses below. 
        
           template <typename SymbolT, typename IRUnitT> 
        
           static LogicalResult 
        
           replaceAllSymbolUsesImpl(SymbolT symbol, StringAttr newSymbol, IRUnitT *limit) { 
        
             // Generate a new attribute to replace the given attribute. 
        
             FlatSymbolRefAttr newLeafAttr = FlatSymbolRefAttr::get(newSymbol); 
        
             for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) { 
        
               SymbolRefAttr oldAttr = scope.symbol; 
        
               SymbolRefAttr newAttr = generateNewRefAttr(scope.symbol, newLeafAttr); 
        
               AttrTypeReplacer replacer; 
        
               replacer.addReplacement( 
        
                   [&](SymbolRefAttr attr) -> std::pair<Attribute, WalkResult> { 
        
                     // Regardless of the match, don't walk nested SymbolRefAttrs, we don't 
        
                     // want to accidentally replace an inner reference. 
        
                     if (attr == oldAttr)

I'm not familiar with this, could you please take a look at it. @River707

CoTinker · 2024-09-27T09:14:32Z

In the following scenario, SymbolTable::replaceAllSymbolUses is also inside a loop.

llvm-project/mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

Lines 214 to 243 in 0df8880

    
           for (auto &op : *combinedModule.getBody()) { 
        
             SymbolOpInterface symbolOp = dyn_cast<SymbolOpInterface>(op); 
        
             if (!symbolOp) 
        
               continue; 
        
             // Do not support ops with operands or results. 
        
             // Global variables, spec constants, and functions won't have 
        
             // operands/results, but just for safety here. 
        
             if (op.getNumOperands() != 0 || op.getNumResults() != 0) 
        
               continue; 
        
             // Deduplicating functions are not supported yet. 
        
             if (isa<FuncOp>(op)) 
        
               continue; 
        
             auto result = hashToSymbolOp.try_emplace(computeHash(symbolOp), symbolOp); 
        
             if (result.second) 
        
               continue; 
        
             SymbolOpInterface replacementSymOp = result.first->second; 
        
             if (failed(SymbolTable::replaceAllSymbolUses( 
        
                     symbolOp, replacementSymOp.getNameAttr(), combinedModule))) { 
        
               symbolOp.emitError("unable to update all symbol uses for ") 
        
                   << symbolOp.getName() << " to " << replacementSymOp.getName(); 
        
               return nullptr; 
        
             } 
        
             eraseList.push_back(symbolOp); 
        
           }

joker-eph · 2024-09-27T09:36:00Z

In the following scenario, SymbolTable::replaceAllSymbolUses is also inside a loop.

I suspect It's just as bad

christopherbate · 2024-10-18T18:29:32Z

@CoTinker I commented above with a suggestion for how to fix.

CoTinker · 2024-10-19T02:25:30Z

@CoTinker I commented above with a suggestion for how to fix.

Thanks, I'll refer to it.

This PR fixes multiple bugs in `DuplicateFunctionElimination`. - Prevents elimination of function declarations. - Updates all symbol uses to reference unique function representatives.

llvmbot added mlir mlir:func labels Sep 22, 2024

CoTinker force-pushed the duplicate branch 2 times, most recently from 780b3d5 to 69f301a Compare September 22, 2024 08:54

llvm deleted a comment from github-actions bot Sep 22, 2024

CoTinker requested review from joker-eph, frgossen and christopherbate September 22, 2024 09:03

joker-eph reviewed Sep 22, 2024

View reviewed changes

mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp Outdated Show resolved Hide resolved

CoTinker requested a review from River707 October 17, 2024 06:52

[mlir][func] Fix multiple bugs in DuplicateFunctionElimination

ef2ce67

This PR fixes multiple bugs in `DuplicateFunctionElimination`. - Prevents elimination of function declarations. - Updates all symbol uses to reference unique function representatives.

CoTinker force-pushed the duplicate branch from 69f301a to ef2ce67 Compare October 19, 2024 11:51

CoTinker requested a review from joker-eph October 19, 2024 11:52

joker-eph approved these changes Oct 19, 2024

View reviewed changes

CoTinker merged commit 2ce655c into llvm:main Oct 22, 2024
8 checks passed

CoTinker deleted the duplicate branch October 22, 2024 01:19

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][func] Fix multiple bugs in `DuplicateFunctionElimination` #109571

[mlir][func] Fix multiple bugs in `DuplicateFunctionElimination` #109571

Uh oh!

CoTinker commented Sep 22, 2024 •

edited

Loading

Uh oh!

llvmbot commented Sep 22, 2024

Uh oh!

llvmbot commented Sep 22, 2024

Uh oh!

Uh oh!

CoTinker commented Sep 27, 2024

Uh oh!

CoTinker commented Sep 27, 2024

Uh oh!

joker-eph commented Sep 27, 2024

Uh oh!

christopherbate commented Oct 18, 2024

Uh oh!

CoTinker commented Oct 19, 2024

Uh oh!

Uh oh!

Uh oh!

[mlir][func] Fix multiple bugs in DuplicateFunctionElimination #109571

[mlir][func] Fix multiple bugs in DuplicateFunctionElimination #109571

Uh oh!

Conversation

CoTinker commented Sep 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 22, 2024

Uh oh!

llvmbot commented Sep 22, 2024

Uh oh!

Uh oh!

CoTinker commented Sep 27, 2024

Uh oh!

CoTinker commented Sep 27, 2024

Uh oh!

joker-eph commented Sep 27, 2024

Uh oh!

christopherbate commented Oct 18, 2024

Uh oh!

CoTinker commented Oct 19, 2024

Uh oh!

Uh oh!

Uh oh!

[mlir][func] Fix multiple bugs in `DuplicateFunctionElimination` #109571

[mlir][func] Fix multiple bugs in `DuplicateFunctionElimination` #109571

CoTinker commented Sep 22, 2024 •

edited

Loading