[flang][driver][openmp] Write MLIR for -save-temps

This patch adds support for producing MLIR files when using -save-temps on
flang. One MLIR file will be produced before lowering and optimization passes,
containing the operations produced by the PFT-to-MLIR lowering bridge, and
another at the end of the process, just before LLVM IR generation.

This is accomplished by forwarding the -save-temps flag from the driver to the
frontend, and modifying it to output MLIR files accordingly.

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

clang/include/clang/Driver/Options.td

@@ -4270,9 +4270,9 @@ def no_offload_add_rpath: Flag<["--"], "no-offload-add-rpath">, Flags<[NoArgumen
   Alias<frtlib_add_rpath>;
 def r : Flag<["-"], "r">, Flags<[LinkerInput,NoArgumentUnused]>,
         Group<Link_Group>;
-def save_temps_EQ : Joined<["-", "--"], "save-temps=">, Flags<[CC1Option, FlangOption, NoXarchOption]>,
+def save_temps_EQ : Joined<["-", "--"], "save-temps=">, Flags<[CC1Option, FlangOption, FC1Option, NoXarchOption]>,
   HelpText<"Save intermediate compilation results.">;
-def save_temps : Flag<["-", "--"], "save-temps">, Flags<[FlangOption, NoXarchOption]>,
+def save_temps : Flag<["-", "--"], "save-temps">, Flags<[FlangOption, FC1Option, NoXarchOption]>,
   Alias<save_temps_EQ>, AliasArgs<["cwd"]>,
   HelpText<"Save intermediate compilation results">;
 def save_stats_EQ : Joined<["-", "--"], "save-stats=">, Flags<[NoXarchOption]>,

clang/lib/Driver/ToolChains/Flang.cpp

@@ -405,6 +405,9 @@ void Flang::ConstructJob(Compilation &C, const JobAction &JA,
 
   assert(Input.isFilename() && "Invalid input.");
 
+  if (Args.getLastArg(options::OPT_save_temps_EQ))
+    Args.AddLastArg(CmdArgs, options::OPT_save_temps_EQ);
+
   addDashXForInput(Args, Input, CmdArgs);
 
   CmdArgs.push_back(Input.getFilename());

flang/include/flang/Frontend/CodeGenOptions.h

@@ -21,6 +21,7 @@
 #include "llvm/Transforms/Instrumentation/AddressSanitizerOptions.h"
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -53,6 +54,9 @@ class CodeGenOptions : public CodeGenOptionsBase {
   /// are offloading binaries containing device images and metadata.
   std::vector<std::string> OffloadObjects;
 
+  /// The directory where temp files are stored if specified by -save-temps
+  std::optional<std::string> SaveTempsDir;
+
   // Define accessors/mutators for code generation options of enumeration type.
 #define CODEGENOPT(Name, Bits, Default)
 #define ENUM_CODEGENOPT(Name, Type, Bits, Default)                             \

flang/lib/Frontend/CompilerInvocation.cpp

@@ -150,6 +150,9 @@ static void parseCodeGenArgs(Fortran::frontend::CodeGenOptions &opts,
       opts.PrepareForThinLTO = true;
   }
 
+  if (auto *a = args.getLastArg(clang::driver::options::OPT_save_temps_EQ))
+    opts.SaveTempsDir = a->getValue();
+
   // -mrelocation-model option.
   if (const llvm::opt::Arg *a =
           args.getLastArg(clang::driver::options::OPT_mrelocation_model)) {

flang/lib/Frontend/FrontendActions.cpp

@@ -13,6 +13,7 @@
 #include "flang/Frontend/FrontendActions.h"
 #include "flang/Common/default-kinds.h"
 #include "flang/Frontend/CompilerInstance.h"
+#include "flang/Frontend/CompilerInvocation.h"
 #include "flang/Frontend/FrontendOptions.h"
 #include "flang/Frontend/PreprocessorOptions.h"
 #include "flang/Lower/Bridge.h"
@@ -39,6 +40,7 @@
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/DiagnosticFrontend.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
@@ -52,10 +54,14 @@
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <memory>
+#include <system_error>
 
 using namespace Fortran::frontend;
 
@@ -64,6 +70,41 @@ using namespace Fortran::frontend;
   llvm::PassPluginLibraryInfo get##Ext##PluginInfo();
 #include "llvm/Support/Extension.def"
 
+/// Save the given \c mlirModule to a temporary .mlir file, in a location
+/// decided by the -save-temps flag. No files are produced if the flag is not
+/// specified.
+static bool saveMLIRTempFile(const CompilerInvocation &ci,
+                             mlir::ModuleOp mlirModule,
+                             llvm::StringRef inputFile,
+                             llvm::StringRef outputTag) {
+  if (!ci.getCodeGenOpts().SaveTempsDir.has_value())
+    return true;
+
+  const llvm::StringRef compilerOutFile = ci.getFrontendOpts().outputFile;
+  const llvm::StringRef saveTempsDir = ci.getCodeGenOpts().SaveTempsDir.value();
+  auto dir = llvm::StringSwitch<llvm::StringRef>(saveTempsDir)
+                 .Case("cwd", "")
+                 .Case("obj", llvm::sys::path::parent_path(compilerOutFile))
+                 .Default(saveTempsDir);
+
+  // Build path from the compiler output file name, triple, cpu and OpenMP
+  // information
+  llvm::SmallString<256> path(dir);
+  llvm::sys::path::append(path, llvm::sys::path::stem(inputFile) + "-" +
+                                    outputTag + ".mlir");
+
+  std::error_code ec;
+  llvm::ToolOutputFile out(path, ec, llvm::sys::fs::OF_Text);
+  if (ec)
+    return false;
+
+  mlirModule->print(out.os());
+  out.os().close();
+  out.keep();
+
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 // Custom BeginSourceFileAction
 //===----------------------------------------------------------------------===//
@@ -208,6 +249,16 @@ bool CodeGenAction::beginSourceFileAction() {
     return false;
   }
 
+  // Print initial full MLIR module, before lowering or transformations, if
+  // -save-temps has been specified.
+  if (!saveMLIRTempFile(ci.getInvocation(), *mlirModule, getCurrentFile(),
+                        "fir")) {
+    unsigned diagID = ci.getDiagnostics().getCustomDiagID(
+        clang::DiagnosticsEngine::Error, "Saving MLIR temp file failed");
+    ci.getDiagnostics().Report(diagID);
+    return false;
+  }
+
   return true;
 }
 
@@ -565,6 +616,16 @@ void CodeGenAction::generateLLVMIR() {
     ci.getDiagnostics().Report(diagID);
   }
 
+  // Print final MLIR module, just before translation into LLVM IR, if
+  // -save-temps has been specified.
+  if (!saveMLIRTempFile(ci.getInvocation(), *mlirModule, getCurrentFile(),
+                        "llvmir")) {
+    unsigned diagID = ci.getDiagnostics().getCustomDiagID(
+        clang::DiagnosticsEngine::Error, "Saving MLIR temp file failed");
+    ci.getDiagnostics().Report(diagID);
+    return;
+  }
+
   // Translate to LLVM IR
   std::optional<llvm::StringRef> moduleName = mlirModule->getName();
   llvmModule = mlir::translateModuleToLLVMIR(

flang/test/Driver/driver-help.f90

@@ -170,6 +170,8 @@
 ! HELP-FC1-NEXT: -pic-level <value>      Value for __PIC__
 ! HELP-FC1-NEXT: -plugin <name>         Use the named plugin action instead of the default action (use "help" to list available options)
 ! HELP-FC1-NEXT: -P                     Disable linemarker output in -E mode
+! HELP-FC1-NEXT: -save-temps=<value>    Save intermediate compilation results.
+! HELP-FC1-NEXT: -save-temps            Save intermediate compilation results
 ! HELP-FC1-NEXT: -std=<value>           Language standard to compile for
 ! HELP-FC1-NEXT: -S                     Only run preprocess and compilation steps
 ! HELP-FC1-NEXT: -target-cpu <value>    Target a specific cpu type

flang/test/Driver/frontend-forwarding.f90

@@ -15,7 +15,8 @@
 ! RUN:     -fassociative-math \
 ! RUN:     -freciprocal-math \
 ! RUN:     -fpass-plugin=Bye%pluginext \
-! RUN:     -mllvm -print-before-all\
+! RUN:     -mllvm -print-before-all \
+! RUN:     -save-temps=obj \
 ! RUN:     -P \
 ! RUN:   | FileCheck %s
 
@@ -34,3 +35,4 @@
 ! CHECK: "-fconvert=little-endian"
 ! CHECK: "-fpass-plugin=Bye
 ! CHECK: "-mllvm" "-print-before-all"
+! CHECK: "-save-temps=obj"

flang/test/Driver/save-mlir-temps.f90

@@ -0,0 +1,58 @@
+! Tests for the `-save-temps` flag. Instead of checking the commands generated
+! by the driver with `-###` (like the save-temps.f90 test does), here we check
+! that the MLIR files are actually produced in the specified location because
+! the driver does not generate specific passes for MLIR. Instead, they are
+! generated during code generation as additional outputs.
+
+! As `flang` does not implement `-fc1as` (i.e. a driver for the integrated
+! assembler), we need to use `-fno-integrated-as` here.
+
+! UNSUPPORTED: system-windows
+
+!--------------------------
+! Invalid output directory
+!--------------------------
+! RUN: not %flang_fc1 -emit-llvm-bc -save-temps=#invalid-dir -o - %s 2>&1 | FileCheck %s -check-prefix=MLIR-ERROR
+! MLIR-ERROR: error: Saving MLIR temp file failed
+
+!--------------------------
+! Save to cwd
+!--------------------------
+! RUN: rm -rf %t && mkdir -p %t
+! RUN: pushd %t && %flang -c -fno-integrated-as -save-temps=cwd -o out.o %s 2>&1
+! RUN: FileCheck %s -input-file=save-mlir-temps-fir.mlir -check-prefix=MLIR-FIR
+! RUN: FileCheck %s -input-file=save-mlir-temps-llvmir.mlir -check-prefix=MLIR-LLVMIR
+! RUN: popd
+
+! RUN: rm -rf %t && mkdir -p %t
+! RUN: pushd %t && %flang -c -fno-integrated-as -save-temps -o out.o %s 2>&1
+! RUN: FileCheck %s -input-file=save-mlir-temps-fir.mlir -check-prefix=MLIR-FIR
+! RUN: FileCheck %s -input-file=save-mlir-temps-llvmir.mlir -check-prefix=MLIR-LLVMIR
+! RUN: popd
+
+!--------------------------
+! Save to output directory
+!--------------------------
+! RUN: rm -rf %t && mkdir -p %t
+! RUN: %flang -c -fno-integrated-as -save-temps=obj -o %t/out.o %s 2>&1
+! RUN: FileCheck %s -input-file=%t/save-mlir-temps-fir.mlir -check-prefix=MLIR-FIR
+! RUN: FileCheck %s -input-file=%t/save-mlir-temps-llvmir.mlir -check-prefix=MLIR-LLVMIR
+
+!--------------------------
+! Save to specific directory
+!--------------------------
+! RUN: rm -rf %t && mkdir -p %t
+! RUN: %flang -c -fno-integrated-as -save-temps=%t -o %t/out.o %s 2>&1
+! RUN: FileCheck %s -input-file=%t/save-mlir-temps-fir.mlir -check-prefix=MLIR-FIR
+! RUN: FileCheck %s -input-file=%t/save-mlir-temps-llvmir.mlir -check-prefix=MLIR-LLVMIR
+
+!--------------------------
+! Content to check from the MLIR outputs
+!--------------------------
+! MLIR-FIR-NOT: llvm.func
+! MLIR-FIR: func.func @{{.*}}main() {
+
+! MLIR-FIR-NOT: func.func
+! MLIR-LLVMIR: llvm.func @{{.*}}main() {
+
+end program