[flang-rt] Add support for formatted I/O on the GPU (#182580)

Summary:
Expands on the previous support to enable formatted output, characters,
and checking basic iostat. We intentionally do not handle cases where
the descriptor is non-null as this is a non-trivial class that cannot
easily be shepherded across the wire.

flang-rt/lib/runtime/io-api-gpu.cpp

@@ -6,13 +6,12 @@
 //
 //===----------------------------------------------------------------------===//
 
-// Implements the subset of the I/O statement API needed for basic
-// list-directed output (PRINT *) of intrinsic types for the GPU.
+// Implements the subset of the I/O statement API needed for basic list-directed
+// output (PRINT *) of intrinsic types for the GPU.
 //
 // The RPC interface forwards each runtime call from the client to the server
 // using a shared buffer. These calls are buffered on the server, so only the
-// return value from 'BeginExternalListOutput' and 'EndIoStatement' are
-// meaningful.
+// return values from 'Begin' and 'EndIoStatement' are meaningful.
 
 #include "io-api-gpu.h"
 #include "flang/Runtime/io-api.h"
@@ -33,6 +32,21 @@ Cookie IODEF(BeginExternalListOutput)(
       IONAME(BeginExternalListOutput), unitNumber, sourceFile, sourceLine);
 }
 
+Cookie IODEF(BeginExternalFormattedOutput)(const char *format,
+    std::size_t formatLength, const Descriptor *formatDescriptor,
+    ExternalUnit unitNumber, const char *sourceFile, int sourceLine) {
+  return rpc::dispatch<BeginExternalFormattedOutput_Opcode>(client,
+      IONAME(BeginExternalFormattedOutput),
+      rpc::span<const char>{format, formatLength}, formatLength,
+      formatDescriptor, unitNumber, sourceFile, sourceLine);
+}
+
+void IODEF(EnableHandlers)(Cookie cookie, bool hasIoStat, bool hasErr,
+    bool hasEnd, bool hasEor, bool hasIoMsg) {
+  return rpc::dispatch<EnableHandlers_Opcode>(client, IONAME(EnableHandlers),
+      cookie, hasIoStat, hasErr, hasEnd, hasEor, hasIoMsg);
+}
+
 enum Iostat IODEF(EndIoStatement)(Cookie cookie) {
   return rpc::dispatch<EndIoStatement_Opcode>(
       client, IONAME(EndIoStatement), cookie);
@@ -86,8 +100,14 @@ bool IODEF(OutputComplex64)(Cookie cookie, double re, double im) {
 }
 
 bool IODEF(OutputAscii)(Cookie cookie, const char *x, std::size_t length) {
-  return rpc::dispatch<OutputAscii_Opcode>(
-      client, IONAME(OutputAscii), cookie, x, length);
+  return rpc::dispatch<OutputAscii_Opcode>(client, IONAME(OutputAscii), cookie,
+      rpc::span<const char>{x, length}, length);
+}
+
+bool IODEF(OutputCharacter)(
+    Cookie cookie, const char *x, std::size_t length, int kind) {
+  return rpc::dispatch<OutputCharacter_Opcode>(client, IONAME(OutputCharacter),
+      cookie, rpc::span<const char>{x, length * kind}, length, kind);
 }
 
 bool IODEF(OutputLogical)(Cookie cookie, bool truth) {

flang-rt/lib/runtime/io-api-gpu.h

@@ -20,18 +20,21 @@ constexpr std::uint32_t MakeOpcode(std::uint32_t base) {
 // Opcodes shared between the client and server for each function we support.
 enum RPCOpcodes : std::uint32_t {
   BeginExternalListOutput_Opcode = MakeOpcode(0),
-  EndIoStatement_Opcode = MakeOpcode(1),
-  OutputInteger8_Opcode = MakeOpcode(2),
-  OutputInteger16_Opcode = MakeOpcode(3),
-  OutputInteger32_Opcode = MakeOpcode(4),
-  OutputInteger64_Opcode = MakeOpcode(5),
-  OutputInteger128_Opcode = MakeOpcode(6),
-  OutputReal32_Opcode = MakeOpcode(7),
-  OutputReal64_Opcode = MakeOpcode(8),
-  OutputComplex32_Opcode = MakeOpcode(9),
-  OutputComplex64_Opcode = MakeOpcode(10),
-  OutputAscii_Opcode = MakeOpcode(11),
-  OutputLogical_Opcode = MakeOpcode(12),
+  BeginExternalFormattedOutput_Opcode = MakeOpcode(1),
+  EnableHandlers_Opcode = MakeOpcode(2),
+  EndIoStatement_Opcode = MakeOpcode(3),
+  OutputInteger8_Opcode = MakeOpcode(4),
+  OutputInteger16_Opcode = MakeOpcode(5),
+  OutputInteger32_Opcode = MakeOpcode(6),
+  OutputInteger64_Opcode = MakeOpcode(7),
+  OutputInteger128_Opcode = MakeOpcode(8),
+  OutputReal32_Opcode = MakeOpcode(9),
+  OutputReal64_Opcode = MakeOpcode(10),
+  OutputComplex32_Opcode = MakeOpcode(11),
+  OutputComplex64_Opcode = MakeOpcode(12),
+  OutputAscii_Opcode = MakeOpcode(13),
+  OutputCharacter_Opcode = MakeOpcode(14),
+  OutputLogical_Opcode = MakeOpcode(15),
 };
 
 } // namespace Fortran::runtime::io

flang-rt/lib/runtime/io-api-server.cpp

@@ -60,17 +60,23 @@ struct DeferredFunctionBase {
 
   void execute(IOContext &ctx) { execute_(impl_, ctx); }
 
-  static OwningPtr<char> TempString(const char *str) {
+  static OwningPtr<char> TempString(const char *str, std::size_t size) {
     if (!str) {
       return {};
     }
 
-    const auto size = std::strlen(str) + 1;
     OwningPtr<char> temp = SizedNew<char>{Terminator{__FILE__, __LINE__}}(size);
     std::memcpy(temp.get(), str, size);
     return OwningPtr<char>(temp.release());
   }
 
+  static OwningPtr<char> TempString(const char *str) {
+    if (!str) {
+      return {};
+    }
+    return TempString(str, std::strlen(str) + 1);
+  }
+
 private:
   void reset() {
     if (impl_) {
@@ -167,6 +173,36 @@ rpc::Status HandleOpcodesImpl(rpc::Server::Port &port) {
           return reinterpret_cast<Cookie>(ctx);
         });
     break;
+  case BeginExternalFormattedOutput_Opcode:
+    rpc::invoke<NumLanes>(port,
+        [](const char *format, std::size_t formatLength,
+            const Descriptor *formatDescriptor, ExternalUnit unitNumber,
+            const char *sourceFile, int sourceLine) -> Cookie {
+          Terminator terminator{__FILE__, __LINE__};
+          if (formatDescriptor)
+            terminator.Crash("Non-trivial format descriptors are unsupported");
+
+          DeferredContext *ctx =
+              new (AllocateMemoryOrCrash(terminator, sizeof(DeferredContext)))
+                  DeferredContext;
+
+          ctx->commands.emplace_back(
+              MakeDeferred(IONAME(BeginExternalFormattedOutput),
+                  DeferredFunctionBase::TempString(format, formatLength),
+                  formatLength, formatDescriptor, unitNumber,
+                  DeferredFunctionBase::TempString(sourceFile), sourceLine));
+
+          return reinterpret_cast<Cookie>(ctx);
+        });
+    break;
+  case EnableHandlers_Opcode:
+    rpc::invoke<NumLanes>(port,
+        [](Cookie cookie, bool hasIoStat, bool hasErr, bool hasEnd, bool hasEor,
+            bool hasIoMsg) -> void {
+          EnqueueDeferred(IONAME(EnableHandlers), cookie, hasIoStat, hasErr,
+              hasEnd, hasEor, hasIoMsg);
+        });
+    break;
   case EndIoStatement_Opcode:
     rpc::invoke<NumLanes>(port, [](Cookie cookie) -> Iostat {
       DeferredContext *ctx = reinterpret_cast<DeferredContext *>(cookie);
@@ -183,13 +219,6 @@ rpc::Status HandleOpcodesImpl(rpc::Server::Port &port) {
       return result;
     });
     break;
-  case OutputAscii_Opcode:
-    rpc::invoke<NumLanes>(
-        port, [](Cookie cookie, const char *x, std::size_t length) -> bool {
-          return EnqueueDeferred(IONAME(OutputAscii), cookie,
-              DeferredFunctionBase::TempString(x), length);
-        });
-    break;
   case OutputInteger8_Opcode:
     rpc::invoke<NumLanes>(port, [](Cookie cookie, std::int8_t n) -> bool {
       return EnqueueDeferred(IONAME(OutputInteger8), cookie, n);
@@ -238,6 +267,20 @@ rpc::Status HandleOpcodesImpl(rpc::Server::Port &port) {
           return EnqueueDeferred(IONAME(OutputComplex64), cookie, re, im);
         });
     break;
+  case OutputAscii_Opcode:
+    rpc::invoke<NumLanes>(
+        port, [](Cookie cookie, const char *x, std::size_t length) -> bool {
+          return EnqueueDeferred(IONAME(OutputAscii), cookie,
+              DeferredFunctionBase::TempString(x, length), length);
+        });
+    break;
+  case OutputCharacter_Opcode:
+    rpc::invoke<NumLanes>(port,
+        [](Cookie cookie, const char *x, std::size_t length, int kind) -> bool {
+          return EnqueueDeferred(IONAME(OutputCharacter), cookie,
+              DeferredFunctionBase::TempString(x, length * kind), length, kind);
+        });
+    break;
   case OutputLogical_Opcode:
     rpc::invoke<NumLanes>(port, [](Cookie cookie, bool truth) -> bool {
       return EnqueueDeferred(IONAME(OutputLogical), cookie, truth);

offload/test/offloading/fortran/formatted-io.f90

@@ -0,0 +1,67 @@
+! REQUIRES: flang, libc
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+
+program formatted_io_test
+  implicit none
+
+  integer :: i, ios
+  real :: r
+  complex :: c
+  logical :: l
+  character(len=5) :: s
+
+  i = 42
+  r = 3.14
+  c = (1.0, -1.0)
+  l = .true.
+  s = "Hello"
+
+  ! CHECK: i=  42
+  !$omp target
+    write(*, '(A,I4)') "i=", i
+  !$omp end target
+
+  ! CHECK: r= 3.14
+  !$omp target
+    write(*, '(A,F5.2)') "r=", r
+  !$omp end target
+
+  ! CHECK: s=Hello
+  !$omp target map(to : s)
+    write(*, '(A,A)') "s=", s
+  !$omp end target
+
+  ! CHECK: l=T
+  !$omp target
+    write(*, '(A,L1)') "l=", l
+  !$omp end target
+
+  ! CHECK: c=(  1.00, -1.00)
+  !$omp target
+    write(*, '(A,A,F6.2,A,F6.2,A)') "c=", "(", real(c), ",", aimag(c), ")"
+  !$omp end target
+
+  ! CHECK: mixed: Hello   42 3.14 T
+  !$omp target map(to : s)
+    write(*, '(A,A,I5,F5.2,L2)') "mixed: ", s, i, r, l
+  !$omp end target
+
+  ! Test IOSTAT handling
+  ! CHECK: iostat:           0
+  !$omp target
+    write(*, '(A)', iostat=ios) "dummy"
+    write(*, '(A,I12)') "iostat:", ios
+  !$omp end target
+
+  ! Test formatted output from multiple teams.
+  ! CHECK: val=  42
+  ! CHECK: val=  42
+  ! CHECK: val=  42
+  ! CHECK: val=  42
+  !$omp target teams num_teams(4)
+  !$omp parallel num_threads(1)
+    write(*, '(A,I4)') "val=", i
+  !$omp end parallel
+  !$omp end target teams
+
+end program formatted_io_test