[flang-commits] [flang] [flang] Ensure lowering diagnostic handler does not outlive lowering (PR #151608)

[Razvan Lupusoru via flang-commits]

https://github.com/razvanlupusoru updated https://github.com/llvm/llvm-project/pull/151608

>From 0294f249d6874f63d8f8b49253a607ab6170d6e6 Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Thu, 31 Jul 2025 15:18:38 -0700
Subject: [PATCH 1/2] [flang] Ensure lowering diagnostic handler does not
 outlive lowering

When the LoweringBridge is created, it registers an MLIR Diagnostics
handler with the MLIRContext. However, it never deregisters it once
lowering is finished.

This fixes this particular scenario. It also makes it so that the
Diagnostics handler is optional.
---
 flang/include/flang/Lower/Bridge.h            |  2 +
 flang/include/flang/Lower/LoweringOptions.def |  5 ++
 flang/lib/Lower/Bridge.cpp                    | 51 +++++++++++--------
 3 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/flang/include/flang/Lower/Bridge.h b/flang/include/flang/Lower/Bridge.h
index a8c2bcfda31c1..3710e74bf38d9 100644
--- a/flang/include/flang/Lower/Bridge.h
+++ b/flang/include/flang/Lower/Bridge.h
@@ -76,6 +76,7 @@ class LoweringBridge {
                           loweringOptions, envDefaults, languageFeatures,
                           targetMachine, targetOptions, codeGenOptions);
   }
+  ~LoweringBridge();
 
   //===--------------------------------------------------------------------===//
   // Getters
@@ -174,6 +175,7 @@ class LoweringBridge {
   const std::vector<Fortran::lower::EnvironmentDefault> &envDefaults;
   const Fortran::common::LanguageFeatureControl &languageFeatures;
   std::set<std::string> tempNames;
+  std::optional<mlir::DiagnosticEngine::HandlerID> diagHandlerID;
 };
 
 } // namespace lower
diff --git a/flang/include/flang/Lower/LoweringOptions.def b/flang/include/flang/Lower/LoweringOptions.def
index 8135704971aa4..39f197d8d35c8 100644
--- a/flang/include/flang/Lower/LoweringOptions.def
+++ b/flang/include/flang/Lower/LoweringOptions.def
@@ -74,5 +74,10 @@ ENUM_LOWERINGOPT(SkipExternalRttiDefinition, unsigned, 1, 0)
 /// If false, lower to the complex dialect of MLIR.
 /// On by default.
 ENUM_LOWERINGOPT(ComplexDivisionToRuntime, unsigned, 1, 1)
+
+/// When true, it registers MLIRDiagnosticsHandler for the duration
+/// of the lowering pipeline.
+ENUM_LOWERINGOPT(RegisterMLIRDiagnosticsHandler, unsigned, 1, 1)
+
 #undef LOWERINGOPT
 #undef ENUM_LOWERINGOPT
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 1adfb96ab9a98..e2ba20f968fd1 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -6707,27 +6707,30 @@ Fortran::lower::LoweringBridge::LoweringBridge(
       loweringOptions{loweringOptions}, envDefaults{envDefaults},
       languageFeatures{languageFeatures} {
   // Register the diagnostic handler.
-  context.getDiagEngine().registerHandler([](mlir::Diagnostic &diag) {
-    llvm::raw_ostream &os = llvm::errs();
-    switch (diag.getSeverity()) {
-    case mlir::DiagnosticSeverity::Error:
-      os << "error: ";
-      break;
-    case mlir::DiagnosticSeverity::Remark:
-      os << "info: ";
-      break;
-    case mlir::DiagnosticSeverity::Warning:
-      os << "warning: ";
-      break;
-    default:
-      break;
-    }
-    if (!mlir::isa<mlir::UnknownLoc>(diag.getLocation()))
-      os << diag.getLocation() << ": ";
-    os << diag << '\n';
-    os.flush();
-    return mlir::success();
-  });
+  if (loweringOptions.getRegisterMLIRDiagnosticsHandler()) {
+    diagHandlerID =
+        context.getDiagEngine().registerHandler([](mlir::Diagnostic &diag) {
+          llvm::raw_ostream &os = llvm::errs();
+          switch (diag.getSeverity()) {
+          case mlir::DiagnosticSeverity::Error:
+            os << "error: ";
+            break;
+          case mlir::DiagnosticSeverity::Remark:
+            os << "info: ";
+            break;
+          case mlir::DiagnosticSeverity::Warning:
+            os << "warning: ";
+            break;
+          default:
+            break;
+          }
+          if (!mlir::isa<mlir::UnknownLoc>(diag.getLocation()))
+            os << diag.getLocation() << ": ";
+          os << diag << '\n';
+          os.flush();
+          return mlir::success();
+        });
+  }
 
   auto getPathLocation = [&semanticsContext, &context]() -> mlir::Location {
     std::optional<std::string> path;
@@ -6769,6 +6772,12 @@ Fortran::lower::LoweringBridge::LoweringBridge(
     fir::setCommandline(*module, *cgOpts.RecordCommandLine);
 }
 
+Fortran::lower::LoweringBridge::~LoweringBridge() {
+  if (diagHandlerID) {
+    context.getDiagEngine().eraseHandler(*diagHandlerID);
+  }
+}
+
 void Fortran::lower::genCleanUpInRegionIfAny(
     mlir::Location loc, fir::FirOpBuilder &builder, mlir::Region &region,
     Fortran::lower::StatementContext &context) {

>From d82583db82f859438f2ecbef52c9c696aba5445c Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Fri, 1 Aug 2025 08:31:44 -0700
Subject: [PATCH 2/2] Fix braces

---
 flang/lib/Lower/Bridge.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index e2ba20f968fd1..75048c141365e 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -6773,9 +6773,8 @@ Fortran::lower::LoweringBridge::LoweringBridge(
 }
 
 Fortran::lower::LoweringBridge::~LoweringBridge() {
-  if (diagHandlerID) {
+  if (diagHandlerID)
     context.getDiagEngine().eraseHandler(*diagHandlerID);
-  }
 }
 
 void Fortran::lower::genCleanUpInRegionIfAny(