[FLANG][MLIR] Update all module symbol references after changing FuncOp symbol during secondary name mangling

[agozillon]

Summary

This might be a bug, the Symbol References that refer to the current FuncOp symbol were not updated after modification to the FuncOp symbol. This results in not being able to find functions in the module later in the lowering via operations containing SymbolRef's as they have stale symbols.

This showed as a problem in OpenMPToLLVMIRTranslation.cpp where the Flang compiler has lowered from FIR + friends to LLVM Dialect + OpenMP. When we have an OpenMP operation that was populated during the PFT lowering to MLIR; that has a SymbolRef referring to a FuncOp symbol, the symbols would not match later on when we try to lower it further during the OpenMPToLLVMIRTranslation phase where we lower to LLVM IR.

The options -emit-mlir and -emit-fir seem to occur before the rewriting that occurs inside of the ExternalNameConversion.cpp file, so the problem is a little opaque when using just -emit-mlir, however, comparing -emit-llvm and -emit-mlir indicates different function symbol mangling. This occurs during some stage of MLIR lowering to other MLIR dialects, so it may pose issues to other MLIR dialects that get intermixed with Flang.

This SymbolRef mismatch could be an issue for other operations that are rewritten inside of ExternalNameConversion.cpp as they also do not seem to update their SymbolRefs. There are rewrite rules for CallOp, GlobalOp and AddrOfOp within ExternalNameConversion.cpp that may be impacted by similar issues to FuncOp. I currently do not have test cases for these at the moment, so I have left them be for now, but I am happy to modify them if someone with more knowledge believes it's correct to replicate the behavior in this patch across them all (if this is indeed a bug and the desired fix for it).

[clementval]

Please add a test.

[agozillon]

In D142918#4091445, @clementval wrote:

Please add a test.

More than happy to, I wasn't too sure where to add a test for this originally which is why I did not add one, I apologies! Is flang/test/Fir/external-mangling.fir perhaps the best place or is there somewhere more relevant or additional locations? And do you think an apt test would be a FuncOp containing an operation that contains a symbol pointing to the FuncOp and testing if it rewrites the contained symbol after the --external-name-interop pass (if that is the correct pass in this case)?

Thank you very much for your review feedback and help.

[clementval]

In D142918#4093472, @agozillon wrote:

In D142918#4091445, @clementval wrote:

Please add a test.

More than happy to, I wasn't too sure where to add a test for this originally which is why I did not add one, I apologies! Is flang/test/Fir/external-mangling.fir perhaps the best place or is there somewhere more relevant or additional locations? And do you think an apt test would be a FuncOp containing an operation that contains a symbol pointing to the FuncOp and testing if it rewrites the contained symbol after the --external-name-interop pass (if that is the correct pass in this case)?

Thank you very much for your review feedback and help.

No worries. Yes, flang/test/Fir/external-mangling.fir is the right place. A simple funcOp and a fir.call that point to this funcOp should be enough to test this.

[agozillon]

Adding a test that will show regression of the problem. I have used func.call rather than fir.call as these seem to be unaffected but the func.call's appear to display the issue.

[agozillon]

In D142918#4093479, @clementval wrote:

In D142918#4093472, @agozillon wrote:

In D142918#4091445, @clementval wrote:

Please add a test.

More than happy to, I wasn't too sure where to add a test for this originally which is why I did not add one, I apologies! Is flang/test/Fir/external-mangling.fir perhaps the best place or is there somewhere more relevant or additional locations? And do you think an apt test would be a FuncOp containing an operation that contains a symbol pointing to the FuncOp and testing if it rewrites the contained symbol after the --external-name-interop pass (if that is the correct pass in this case)?

Thank you very much for your review feedback and help.

No worries. Yes, flang/test/Fir/external-mangling.fir is the right place. A simple funcOp and a fir.call that point to this funcOp should be enough to test this.

Thank you, I've added a test now, in this case I've used func.call, rather than fir.call, it seems to affect the former but not the latter after some testing with and without the change.

[clementval]

In D142918#4094454, @agozillon wrote:

Thank you, I've added a test now, in this case I've used func.call, rather than fir.call, it seems to affect the former but not the latter after some testing with and without the change.

F18 never uses func.call and I guess the problem doesn't show up with fir.call because we are also modifying the symbol in fir.call operations. Do you have an example of the initial issue you had in the first place?

[agozillon]

In D142918#4094488, @clementval wrote:

In D142918#4094454, @agozillon wrote:

Thank you, I've added a test now, in this case I've used func.call, rather than fir.call, it seems to affect the former but not the latter after some testing with and without the change.

F18 never uses func.call and I guess the problem doesn't show up with fir.call because we are also modifying the symbol in fir.call operations. Do you have an example of the initial issue you had in the first place?

I found this while working with an OpenMP operation that I'm working on: https://reviews.llvm.org/D142910 which will interact with FIR, it holds symbols for the FuncOps and these are given to the operation when it's created during initial creation of MLIR from the F18 parse tree. When this OpenMP operation reaches the OpenMPToLLVMIRTranslation.cpp segment of the compilation process and begins to lower the OpenMP dialect to LLVM IR the symbol in the OpenMP operation no longer matches either the MLIR LLVMFuncOp or the LLVM IR function symbol. This appears to happen from the symbol modification inside of ExternalNameConversion, and unfortunately when you --emit-mlir or --emit-fir the problem doesn't show itself, so I suppose the rewriting occurs somewhere after the emission.

This is an external pull request but might give some further context, essentially just trying to find a function in the module from a symbol given during initial code generation: https://github.com/ROCm-Developer-Tools/llvm-project/pull/185/files#diff-2cbb5651f4570d81d55ac4198deda0f6f7341b2503479752ef2295da3774c586R371

I'm fairly new to Flang/MLIR/OpenMP so I am possibly incorrect and this is not the correct fix to the problem I'm seeing and I am doing something incorrectly or missing some part of the picture that I need to add to properly update the OpenMP operation.

[clementval]

It looks like the right thing to do but I think we can remove the MangleNameOnCallOp because it should be updated by the op.replaceAllSymbolUses. That's why you don't see the issue with fir.call right now but if you remove the conversion pattern it should show up.

[agozillon]

In D142918#4095817, @clementval wrote:

It looks like the right thing to do but I think we can remove the MangleNameOnCallOp because it should be updated by the op.replaceAllSymbolUses. That's why you don't see the issue with fir.call right now but if you remove the conversion pattern it should show up.

Thank you, I'll check that out and update the patch, there are fir::GlobalOp and fir::AddrOfOp rewrites as well that may run into a similar issue, do you believe it'd be sensible to remove these and rely on FuncOp to rewrite them with the symbol update from replaceAllSymbolUses or add a replaceAllSymbolUses within both? or alternatively just leave them as is!

[clementval]

In D142918#4096057, @agozillon wrote:

In D142918#4095817, @clementval wrote:

It looks like the right thing to do but I think we can remove the MangleNameOnCallOp because it should be updated by the op.replaceAllSymbolUses. That's why you don't see the issue with fir.call right now but if you remove the conversion pattern it should show up.

Thank you, I'll check that out and update the patch, there are fir::GlobalOp and fir::AddrOfOp rewrites as well that may run into a similar issue, do you believe it'd be sensible to remove these and rely on FuncOp to rewrite them with the symbol update from replaceAllSymbolUses or add a replaceAllSymbolUses within both? or alternatively just leave them as is!

You can leave fir::GlobalOp and fir::AddrOfOp for now. At least GlobalOp conversion pattern will have to stay because it is not related to funcOp in any way .

[agozillon]

• [FLANG][MLIR] remove MangleNameOnCallOp

[agozillon]

In D142918#4096251, @clementval wrote:

In D142918#4096057, @agozillon wrote:

In D142918#4095817, @clementval wrote:

It looks like the right thing to do but I think we can remove the MangleNameOnCallOp because it should be updated by the op.replaceAllSymbolUses. That's why you don't see the issue with fir.call right now but if you remove the conversion pattern it should show up.

Thank you, I'll check that out and update the patch, there are fir::GlobalOp and fir::AddrOfOp rewrites as well that may run into a similar issue, do you believe it'd be sensible to remove these and rely on FuncOp to rewrite them with the symbol update from replaceAllSymbolUses or add a replaceAllSymbolUses within both? or alternatively just leave them as is!

You can leave fir::GlobalOp and fir::AddrOfOp for now. At least GlobalOp conversion pattern will have to stay because it is not related to funcOp in any way .

Sounds good! I've removed MangleNameOnCallOp and updated the tests, it continues to works as expected.

[clementval]

Just couple of nits

flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp

62

We use fully qualified names

66

[agozillon]

• [FLANG][MLIR] remove MangleNameOnCallOp

[agozillon]

Sorry, done, I've added the fully qualified names.

[clementval]

LGTM. Thanks

[agozillon]

In D142918#4096668, @clementval wrote:

LGTM. Thanks

Thank you very much for your help and the review! I believe I have lost push rights since it's been a long while since I submitted any changes, so I will have to ask for access again so it'll take a day or two for this to land.

flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // Rewrite patterns
 //===----------------------------------------------------------------------===//
 
 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;
 
   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;
   }
 };
 
 struct MangleNameForCommonBlock : public mlir::OpRewritePattern<fir::GlobalOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
 
   mlir::LogicalResult

 };
 } // namespace
 
 void ExternalNameConversionPass::runOnOperation() {
   auto op = getOperation();
   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());
   });
 
   target.addDynamicallyLegalOp<fir::GlobalOp>([](fir::GlobalOp op) {
     return !fir::NameUniquer::needExternalNameMangling(
         op.getSymref().getRootReference().getValue());
   });

flang/test/Fir/external-mangling.fir

 // LLVMIR: %{{.*}} = llvm.mlir.addressof @__BLNK__ : !llvm.ptr<array<4 x i8>>
 // LLVMIR: llvm.call @bar_(%{{.*}}) : (!llvm.ptr<i32>) -> ()
 // LLVMIR: llvm.call @bar2_(%{{.*}}) : (!llvm.ptr<f32>) -> ()
 
 // LLVMIR: llvm.mlir.global common @a_(dense<0> : vector<4xi8>) {{.*}} : !llvm.array<4 x i8>
 // 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_() : () -> ()