From a62b60167ddbc23602ea20d65a41b2cd1ce5953a Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler@nvidia.com>
Date: Thu, 18 Nov 2021 09:34:43 -0800
Subject: [PATCH] [flang] Predefine unit 0 connected to stderr

This is a near-universal language extension; external unit 0
is preconnected to the standard error output.

Differential Revision: https://reviews.llvm.org/D114298
---
 flang/docs/Extensions.md         |  2 ++
 flang/module/iso_fortran_env.f90 |  2 +-
 flang/runtime/unit.cpp           | 35 +++++++++++++++++++++++++-------
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/flang/docs/Extensions.md b/flang/docs/Extensions.md
index 02276fbe45e52..3c63d83aacc4f 100644
--- a/flang/docs/Extensions.md
+++ b/flang/docs/Extensions.md
@@ -188,6 +188,8 @@ end
   if an actual argument acceptable to one could not be passed to
   the other & vice versa because exactly one is polymorphic or
   exactly one is unlimited polymorphic).
+* External unit 0 is predefined and connected to the standard error output,
+  and defined as `ERROR_UNIT` in the intrinsic `ISO_FORTRAN_ENV` module.
 
 ### Extensions supported when enabled by options
 
diff --git a/flang/module/iso_fortran_env.f90 b/flang/module/iso_fortran_env.f90
index 3d77485c13e31..e85d2e2a7a36c 100644
--- a/flang/module/iso_fortran_env.f90
+++ b/flang/module/iso_fortran_env.f90
@@ -129,7 +129,7 @@ module iso_fortran_env
   integer, parameter :: current_team = -1, initial_team = -2, parent_team = -3
 
   integer, parameter :: input_unit = 5, output_unit = 6
-  integer, parameter :: error_unit = output_unit
+  integer, parameter :: error_unit = 0
   integer, parameter :: iostat_end = -1, iostat_eor = -2
   integer, parameter :: iostat_inquire_internal_unit = -1
 
diff --git a/flang/runtime/unit.cpp b/flang/runtime/unit.cpp
index 1c915f3b49753..e24914bc5fa68 100644
--- a/flang/runtime/unit.cpp
+++ b/flang/runtime/unit.cpp
@@ -21,19 +21,23 @@ namespace Fortran::runtime::io {
 // should work without a Fortran main program.
 static Lock unitMapLock;
 static UnitMap *unitMap{nullptr};
-static ExternalFileUnit *defaultInput{nullptr};
-static ExternalFileUnit *defaultOutput{nullptr};
+static ExternalFileUnit *defaultInput{nullptr}; // unit 5
+static ExternalFileUnit *defaultOutput{nullptr}; // unit 6
+static ExternalFileUnit *errorOutput{nullptr}; // unit 0 extension
 
 void FlushOutputOnCrash(const Terminator &terminator) {
-  if (!defaultOutput) {
+  if (!defaultOutput && !errorOutput) {
     return;
   }
+  IoErrorHandler handler{terminator};
+  handler.HasIoStat(); // prevent nested crash if flush has error
   CriticalSection critical{unitMapLock};
   if (defaultOutput) {
-    IoErrorHandler handler{terminator};
-    handler.HasIoStat(); // prevent nested crash if flush has error
     defaultOutput->FlushOutput(handler);
   }
+  if (errorOutput) {
+    errorOutput->FlushOutput(handler);
+  }
 }
 
 ExternalFileUnit *ExternalFileUnit::LookUp(int unit) {
@@ -210,6 +214,7 @@ UnitMap &ExternalFileUnit::GetUnitMap() {
   Terminator terminator{__FILE__, __LINE__};
   IoErrorHandler handler{terminator};
   UnitMap *newUnitMap{New<UnitMap>{terminator}().release()};
+
   bool wasExtant{false};
   ExternalFileUnit &out{newUnitMap->LookUpOrCreate(6, terminator, wasExtant)};
   RUNTIME_CHECK(terminator, !wasExtant);
@@ -217,12 +222,21 @@ UnitMap &ExternalFileUnit::GetUnitMap() {
   out.SetDirection(Direction::Output, handler);
   out.isUnformatted = false;
   defaultOutput = &out;
+
   ExternalFileUnit &in{newUnitMap->LookUpOrCreate(5, terminator, wasExtant)};
   RUNTIME_CHECK(terminator, !wasExtant);
   in.Predefine(0);
   in.SetDirection(Direction::Input, handler);
   in.isUnformatted = false;
   defaultInput = &in;
+
+  ExternalFileUnit &error{newUnitMap->LookUpOrCreate(0, terminator, wasExtant)};
+  RUNTIME_CHECK(terminator, !wasExtant);
+  error.Predefine(2);
+  error.SetDirection(Direction::Output, handler);
+  error.isUnformatted = false;
+  errorOutput = &error;
+
   // TODO: Set UTF-8 mode from the environment
   unitMap = newUnitMap;
   return *unitMap;
@@ -235,6 +249,8 @@ void ExternalFileUnit::CloseAll(IoErrorHandler &handler) {
     FreeMemoryAndNullify(unitMap);
   }
   defaultOutput = nullptr;
+  defaultInput = nullptr;
+  errorOutput = nullptr;
 }
 
 void ExternalFileUnit::FlushAll(IoErrorHandler &handler) {
@@ -627,8 +643,13 @@ void ExternalFileUnit::BeginSequentialVariableUnformattedInputRecord(
 
 void ExternalFileUnit::BeginSequentialVariableFormattedInputRecord(
     IoErrorHandler &handler) {
-  if (this == defaultInput && defaultOutput) {
-    defaultOutput->FlushOutput(handler);
+  if (this == defaultInput) {
+    if (defaultOutput) {
+      defaultOutput->FlushOutput(handler);
+    }
+    if (errorOutput) {
+      errorOutput->FlushOutput(handler);
+    }
   }
   std::size_t length{0};
   do {