[flang-commits] [flang] [flang] Fix seg fault `CodeGenAction::executeAction()` (PR #78269)

[Kareem Ergawy via flang-commits]

https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/78269

>From ae81413d98acc20141157bce60f14ec404e58859 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Tue, 16 Jan 2024 06:56:23 -0600
Subject: [PATCH] [flang] Fix seg fault `CodeGenAction::executeAction()`

If `generateLLVMIR()` fails, we still continue using the module we
failed to generate which causes a seg fault if LLVM code-gen failed for
some reason or another. This commit fixes this issue.
---
 flang/lib/Frontend/FrontendActions.cpp        |   5 +
 flang/unittests/Frontend/CMakeLists.txt       |   1 +
 .../unittests/Frontend/CodeGenActionTest.cpp  | 109 ++++++++++++++++++
 3 files changed, 115 insertions(+)
 create mode 100644 flang/unittests/Frontend/CodeGenActionTest.cpp

diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
index 74e3992d5ab62ba..65c4df7388f97b2 100644
--- a/flang/lib/Frontend/FrontendActions.cpp
+++ b/flang/lib/Frontend/FrontendActions.cpp
@@ -1202,6 +1202,11 @@ void CodeGenAction::executeAction() {
   if (!llvmModule)
     generateLLVMIR();
 
+  // If generating the LLVM module failed, abort! No need for further error
+  // reporting since generateLLVMIR() does this already.
+  if (!llvmModule)
+    return;
+
   // Set the triple based on the targetmachine (this comes compiler invocation
   // and the command-line target option if specified, or the default if not
   // given on the command-line).
diff --git a/flang/unittests/Frontend/CMakeLists.txt b/flang/unittests/Frontend/CMakeLists.txt
index 79a394f161ed1ec..3bcc37bed7f6d11 100644
--- a/flang/unittests/Frontend/CMakeLists.txt
+++ b/flang/unittests/Frontend/CMakeLists.txt
@@ -4,6 +4,7 @@ set(LLVM_LINK_COMPONENTS
 )
 
 add_flang_unittest(FlangFrontendTests
+  CodeGenActionTest.cpp
   CompilerInstanceTest.cpp
   FrontendActionTest.cpp
 )
diff --git a/flang/unittests/Frontend/CodeGenActionTest.cpp b/flang/unittests/Frontend/CodeGenActionTest.cpp
new file mode 100644
index 000000000000000..9d798c7678ad158
--- /dev/null
+++ b/flang/unittests/Frontend/CodeGenActionTest.cpp
@@ -0,0 +1,109 @@
+//===- unittests/Frontend/CodeGenActionTest.cpp --- FrontendAction tests --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Unit tests for CodeGenAction.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Builders.h"
+#include "flang/Frontend/CompilerInstance.h"
+#include "flang/Frontend/FrontendActions.h"
+#include "flang/Frontend/TextDiagnosticPrinter.h"
+
+#include "gtest/gtest.h"
+
+#include <memory>
+
+using namespace Fortran::frontend;
+
+namespace test {
+class DummyDialect : public ::mlir::Dialect {
+  explicit DummyDialect(::mlir::MLIRContext *context)
+      : ::mlir::Dialect(getDialectNamespace(), context,
+            ::mlir::TypeID::get<DummyDialect>()) {
+    initialize();
+  }
+
+  void initialize();
+  friend class ::mlir::MLIRContext;
+
+public:
+  ~DummyDialect() override = default;
+  static constexpr ::llvm::StringLiteral getDialectNamespace() {
+    return ::llvm::StringLiteral("dummy");
+  }
+};
+
+namespace dummy {
+class FakeOp : public ::mlir::Op<FakeOp> {
+public:
+  using Op::Op;
+
+  static llvm::StringRef getOperationName() { return "dummy.fake"; }
+
+  static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { return {}; }
+
+  static void build(
+      ::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState) {}
+};
+} // namespace dummy
+} // namespace test
+
+MLIR_DECLARE_EXPLICIT_TYPE_ID(::test::DummyDialect)
+MLIR_DEFINE_EXPLICIT_TYPE_ID(::test::DummyDialect)
+
+namespace test {
+
+void DummyDialect::initialize() { addOperations<::test::dummy::FakeOp>(); }
+} // namespace test
+
+// A test CodeGenAction to verify that we gracefully handle failure to convert
+// from MLIR to LLVM IR.
+class LLVMConversionFailureCodeGenAction : public CodeGenAction {
+public:
+  LLVMConversionFailureCodeGenAction()
+      : CodeGenAction(BackendActionTy::Backend_EmitLL) {
+    mlirCtx = std::make_unique<mlir::MLIRContext>();
+    mlirCtx->loadDialect<test::DummyDialect>();
+
+    mlir::Location loc(mlir::UnknownLoc::get(mlirCtx.get()));
+    mlirModule =
+        std::make_unique<mlir::ModuleOp>(mlir::ModuleOp::create(loc, "mod"));
+
+    mlir::OpBuilder builder(mlirCtx.get());
+    builder.setInsertionPointToStart(&mlirModule->getRegion().front());
+    // Create a fake op to trip conversion to LLVM.
+    builder.create<test::dummy::FakeOp>(loc);
+
+    llvmCtx = std::make_unique<llvm::LLVMContext>();
+  }
+};
+
+TEST(CodeGenAction, GracefullyHandleLLVMConversionFailure) {
+  std::string diagnosticOutput;
+  llvm::raw_string_ostream diagnosticsOS(diagnosticOutput);
+  auto diagPrinter = std::make_unique<Fortran::frontend::TextDiagnosticPrinter>(
+      diagnosticsOS, new clang::DiagnosticOptions());
+
+  CompilerInstance ci;
+  ci.createDiagnostics(diagPrinter.get(), /*ShouldOwnClient=*/false);
+  ci.setInvocation(std::make_shared<CompilerInvocation>());
+  ci.setOutputStream(std::make_unique<llvm::raw_null_ostream>());
+  ci.getInvocation().getCodeGenOpts().OptimizationLevel = 0;
+
+  FrontendInputFile file("/dev/null", InputKind());
+
+  LLVMConversionFailureCodeGenAction action;
+  action.setInstance(&ci);
+  action.setCurrentInput(file);
+
+  consumeError(action.execute());
+  ASSERT_EQ(diagnosticsOS.str(),
+      "error: Lowering to LLVM IR failed\n"
+      "error: failed to create the LLVM module\n");
+}