[llvm-branch-commits] [flang][runtime] Added self-printing for InternalUnit. (PR #85181)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-runtime

Author: Slava Zakharin (vzakhari)

<details>
<summary>Changes</summary>

An InternalUnit might be constructed to allocate its own "output"
buffer of a predefined size. The buffer is then used for collecting
all the output, and it printed by std::printf at the end of the statement.

This is a suggested way for supporting 'PRINT *, ...' in the device code.
It might be not ideal, because the output is not formatted the same way
as the UNIT=5 output is formatted by default.


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


3 Files Affected:

- (modified) flang/runtime/internal-unit.cpp (+34-3) 
- (modified) flang/runtime/internal-unit.h (+12-1) 
- (modified) flang/runtime/io-stmt.cpp (+5-1) 


``````````diff
diff --git a/flang/runtime/internal-unit.cpp b/flang/runtime/internal-unit.cpp
index 66140e00588723..39fefd8c58cafb 100644
--- a/flang/runtime/internal-unit.cpp
+++ b/flang/runtime/internal-unit.cpp
@@ -15,13 +15,24 @@
 namespace Fortran::runtime::io {
 
 template <Direction DIR>
-InternalDescriptorUnit<DIR>::InternalDescriptorUnit(
-    Scalar scalar, std::size_t length, int kind) {
+InternalDescriptorUnit<DIR>::InternalDescriptorUnit(Scalar scalar,
+    std::size_t length, int kind, const Terminator &terminator,
+    bool allocateOwnOutput) {
   internalIoCharKind = kind;
   recordLength = length;
   endfileRecordNumber = 2;
   void *pointer{reinterpret_cast<void *>(const_cast<char *>(scalar))};
-  descriptor().Establish(TypeCode{TypeCategory::Character, kind}, length * kind,
+  std::size_t bufferSize{length * kind};
+  if (allocateOwnOutput) {
+    RUNTIME_CHECK(terminator, DIR == Direction::Output);
+    pointer = AllocateMemoryOrCrash(terminator, ownBufferSizeInBytes);
+    usesOwnBuffer = true;
+
+    // Reserve one character of the kind for '\0' at the end.
+    bufferSize = (ownBufferSizeInBytes / kind) * kind - kind;
+    recordLength = bufferSize / kind;
+  }
+  descriptor().Establish(TypeCode{TypeCategory::Character, kind}, bufferSize,
       pointer, 0, nullptr, CFI_attribute_pointer);
 }
 
@@ -41,6 +52,26 @@ InternalDescriptorUnit<DIR>::InternalDescriptorUnit(
   endfileRecordNumber = d.Elements() + 1;
 }
 
+template <Direction DIR> void InternalDescriptorUnit<DIR>::EndIoStatement() {
+  if constexpr (DIR == Direction::Output) {
+    if (usesOwnBuffer) {
+      // Null terminate the buffer that contains just a single record.
+      Terminator terminator{__FILE__, __LINE__};
+      RUNTIME_CHECK(terminator,
+          furthestPositionInRecord <
+              static_cast<std::int64_t>(ownBufferSizeInBytes));
+      *reinterpret_cast<char *>(CurrentRecord() + furthestPositionInRecord) =
+          '\0';
+
+      // Print the buffer and deallocate memory.
+      // FIXME: this output does not match the regular unit 5 output.
+      std::printf("%s\n", descriptor().OffsetElement());
+      FreeMemory(descriptor().OffsetElement());
+      descriptor().set_base_addr(nullptr);
+    }
+  }
+}
+
 template <Direction DIR>
 bool InternalDescriptorUnit<DIR>::Emit(
     const char *data, std::size_t bytes, IoErrorHandler &handler) {
diff --git a/flang/runtime/internal-unit.h b/flang/runtime/internal-unit.h
index b536ffb831d550..fdf6e498eabf7e 100644
--- a/flang/runtime/internal-unit.h
+++ b/flang/runtime/internal-unit.h
@@ -12,6 +12,7 @@
 #define FORTRAN_RUNTIME_IO_INTERNAL_UNIT_H_
 
 #include "connection.h"
+#include "flang/Runtime/api-attrs.h"
 #include "flang/Runtime/descriptor.h"
 #include <cinttypes>
 #include <type_traits>
@@ -21,13 +22,19 @@ namespace Fortran::runtime::io {
 class IoErrorHandler;
 
 // Points to (but does not own) a CHARACTER scalar or array for internal I/O.
+// The internal unit does not own the scalar buffer unless it is constructed
+// with allocateOwnOutput set to true: in this case, it owns the buffer
+// and also prints it to stdout at the end of the statement.
+// This is used to support output on offload devices.
 // Does not buffer.
 template <Direction DIR> class InternalDescriptorUnit : public ConnectionState {
 public:
   using Scalar =
       std::conditional_t<DIR == Direction::Input, const char *, char *>;
-  InternalDescriptorUnit(Scalar, std::size_t chars, int kind);
+  InternalDescriptorUnit(Scalar, std::size_t chars, int kind,
+      const Terminator &terminator, bool allocateOwnOutput = false);
   InternalDescriptorUnit(const Descriptor &, const Terminator &);
+  void EndIoStatement();
 
   bool Emit(const char *, std::size_t, IoErrorHandler &);
   std::size_t GetNextInputBytes(const char *&, IoErrorHandler &);
@@ -48,6 +55,10 @@ template <Direction DIR> class InternalDescriptorUnit : public ConnectionState {
   void BlankFillOutputRecord();
 
   StaticDescriptor<maxRank, true /*addendum*/> staticDescriptor_;
+  RT_OFFLOAD_VAR_GROUP_BEGIN
+  static constexpr std::size_t ownBufferSizeInBytes{1024};
+  RT_OFFLOAD_VAR_GROUP_END
+  bool usesOwnBuffer{false};
 };
 
 extern template class InternalDescriptorUnit<Direction::Output>;
diff --git a/flang/runtime/io-stmt.cpp b/flang/runtime/io-stmt.cpp
index efefbc5e1a1c08..417225f79006ff 100644
--- a/flang/runtime/io-stmt.cpp
+++ b/flang/runtime/io-stmt.cpp
@@ -82,7 +82,8 @@ void IoStatementBase::BadInquiryKeywordHashCrash(InquiryKeywordHash inquiry) {
 template <Direction DIR>
 InternalIoStatementState<DIR>::InternalIoStatementState(
     Buffer scalar, std::size_t length, const char *sourceFile, int sourceLine)
-    : IoStatementBase{sourceFile, sourceLine}, unit_{scalar, length, 1} {}
+    : IoStatementBase{sourceFile, sourceLine},
+      unit_{scalar, length, /*kind=*/1, *this} {}
 
 template <Direction DIR>
 InternalIoStatementState<DIR>::InternalIoStatementState(
@@ -119,6 +120,9 @@ template <Direction DIR> void InternalIoStatementState<DIR>::BackspaceRecord() {
 }
 
 template <Direction DIR> int InternalIoStatementState<DIR>::EndIoStatement() {
+  if constexpr (DIR == Direction::Output) {
+    unit_.EndIoStatement();
+  }
   auto result{IoStatementBase::EndIoStatement()};
   if (free_) {
     FreeMemory(this);

``````````

</details>


https://github.com/llvm/llvm-project/pull/85181