[flang-commits] [flang] f86209f - [FLANG][MLIR] Update all module symbol references after changing FuncOp symbol during external name mangling

[Andrew Gozillon via flang-commits]

Author: Andrew Gozillon
Date: 2023-02-02T04:59:32-06:00
New Revision: f86209fc80d8a646ed5802d58f4ed6934d980d4e

URL: https://github.com/llvm/llvm-project/commit/f86209fc80d8a646ed5802d58f4ed6934d980d4e
DIFF: https://github.com/llvm/llvm-project/commit/f86209fc80d8a646ed5802d58f4ed6934d980d4e.diff

LOG: [FLANG][MLIR] Update all module symbol references after changing FuncOp symbol during external name mangling

This fixes an issue where the symbols for operations that were not directly
handled by the rewriting in ExternalNameConversion.cpp were not updated
accurately when a FuncOp symbol was modified. Resulting in a name
mismatch between the FuncOp and the operation holding a symbol to
the FuncOp.

This fix works by updating all of the symbols relating to a FuncOp in a
module, this did not show up as an issue previously as fir::CallOps were
getting specific handling and only fir::CallOps were being tested. So
as the more larger case is now being handled the specific handling for
fir::CallOps has been removed (but is still handled by the fix).

Reviewers:
clementval

Differential Revision: https://reviews.llvm.org/D142918

Added: 
    

Modified: 
    flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
    flang/test/Fir/external-mangling.fir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
index d9b24c909582..0d3df82fc583 100644
--- a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
@@ -45,27 +45,6 @@ mangleExternalName(const std::pair<fir::NameUniquer::NameKind,
 
 namespace {
 
-class MangleNameOnCallOp : public mlir::OpRewritePattern<fir::CallOp> {
-public:
-  using OpRewritePattern::OpRewritePattern;
-
-  mlir::LogicalResult
-  matchAndRewrite(fir::CallOp op,
-                  mlir::PatternRewriter &rewriter) const override {
-    rewriter.startRootUpdate(op);
-    auto callee = op.getCallee();
-    if (callee) {
-      auto result =
-          fir::NameUniquer::deconstruct(callee->getRootReference().getValue());
-      if (fir::NameUniquer::isExternalFacingUniquedName(result))
-        op.setCalleeAttr(
-            SymbolRefAttr::get(op.getContext(), mangleExternalName(result)));
-    }
-    rewriter.finalizeRootUpdate(op);
-    return success();
-  }
-};
-
 struct MangleNameOnFuncOp : public mlir::OpRewritePattern<mlir::func::FuncOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
@@ -73,12 +52,22 @@ struct MangleNameOnFuncOp : public mlir::OpRewritePattern<mlir::func::FuncOp> {
   mlir::LogicalResult
   matchAndRewrite(mlir::func::FuncOp op,
                   mlir::PatternRewriter &rewriter) const override {
+    mlir::LogicalResult ret = success();
     rewriter.startRootUpdate(op);
     auto result = fir::NameUniquer::deconstruct(op.getSymName());
-    if (fir::NameUniquer::isExternalFacingUniquedName(result))
-      op.setSymNameAttr(rewriter.getStringAttr(mangleExternalName(result)));
+    if (fir::NameUniquer::isExternalFacingUniquedName(result)) {
+      auto newSymbol = rewriter.getStringAttr(mangleExternalName(result));
+
+      // Try to update all SymbolRef's in the module that match the current op
+      if (mlir::ModuleOp mod = op->getParentOfType<mlir::ModuleOp>())
+        ret = op.replaceAllSymbolUses(newSymbol, mod);
+
+      op.setSymNameAttr(newSymbol);
+      mlir::SymbolTable::setSymbolName(op, newSymbol);
+    }
+
     rewriter.finalizeRootUpdate(op);
-    return success();
+    return ret;
   }
 };
 
@@ -134,20 +123,13 @@ void ExternalNameConversionPass::runOnOperation() {
   auto *context = &getContext();
 
   mlir::RewritePatternSet patterns(context);
-  patterns.insert<MangleNameOnCallOp, MangleNameOnCallOp, MangleNameOnFuncOp,
-                  MangleNameForCommonBlock, MangleNameOnAddrOfOp>(context);
+  patterns.insert<MangleNameOnFuncOp, MangleNameForCommonBlock,
+                  MangleNameOnAddrOfOp>(context);
 
   ConversionTarget target(*context);
   target.addLegalDialect<fir::FIROpsDialect, LLVM::LLVMDialect,
                          acc::OpenACCDialect, omp::OpenMPDialect>();
 
-  target.addDynamicallyLegalOp<fir::CallOp>([](fir::CallOp op) {
-    if (op.getCallee())
-      return !fir::NameUniquer::needExternalNameMangling(
-          op.getCallee()->getRootReference().getValue());
-    return true;
-  });
-
   target.addDynamicallyLegalOp<mlir::func::FuncOp>([](mlir::func::FuncOp op) {
     return !fir::NameUniquer::needExternalNameMangling(op.getSymName());
   });

diff  --git a/flang/test/Fir/external-mangling.fir b/flang/test/Fir/external-mangling.fir
index 794132847c6d..ce728c1b706c 100644
--- a/flang/test/Fir/external-mangling.fir
+++ b/flang/test/Fir/external-mangling.fir
@@ -39,3 +39,23 @@ func.func private @_QPbar2(!fir.ref<f32>)
 // LLVMIR: llvm.mlir.global common @__BLNK__(dense<0> : vector<4xi8>) {{.*}}  : !llvm.array<4 x i8>
 // LLVMIR: llvm.func @bar_(!llvm.ptr<i32>) attributes {sym_visibility = "private"}
 // LLVMIR: llvm.func @bar2_(!llvm.ptr<f32>) attributes {sym_visibility = "private"}
+
+// ----- 
+
+func.func @_QPcallee() {
+  fir.call @_QPcallee() : () -> ()
+  return
+}
+
+func.func @_QPcaller() {
+  fir.call @_QPcallee() : () -> ()
+  return 
+}
+
+// CHECK: func @callee_
+// CHECK: fir.call @callee_
+// CHECK: func @caller_
+// CHECK: fir.call @callee_
+
+// LLVMIR: llvm.call @callee_() : () -> ()
+// LLVMIR: llvm.call @callee_() : () -> ()