[flang] Ensure lowering diagnostic handler does not outlive lowering

[razvanlupusoru]

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.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Razvan Lupusoru (razvanlupusoru)

Changes

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.

---

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

3 Files Affected:

• (modified) flang/include/flang/Lower/Bridge.h (+2)
• (modified) flang/include/flang/Lower/LoweringOptions.def (+5)
• (modified) flang/lib/Lower/Bridge.cpp (+30-21)

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) {

[clementval]

Braces

[clementval]

What's the use case when we don't need to register it?

[razvanlupusoru]

When an appropriate handler is already registered with the MLIRContext. Lowering should not need to register its own if one can already handle all of the MLIR diagnostics.

[clementval]

Ok that makes sense. Thanks for the explanation!

[clementval]

LGTM