From 0c52877e1df2cf09cd1cfe3baaebcb1940f69495 Mon Sep 17 00:00:00 2001
From: Youngsung Kim <youngsung.kim.act2@gmail.com>
Date: Tue, 7 Jul 2026 07:48:26 -0700
Subject: [PATCH] Reduce the number of PIO sync during IOStream write

---
 components/omega/src/base/IO.cpp        |  7 ------
 components/omega/src/infra/IOStream.cpp | 31 +++++++++++++++++--------
 components/omega/src/infra/IOStream.h   |  6 +++--
 3 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/components/omega/src/base/IO.cpp b/components/omega/src/base/IO.cpp
index e650eabeffba..cfeaecff25aa 100644
--- a/components/omega/src/base/IO.cpp
+++ b/components/omega/src/base/IO.cpp
@@ -849,13 +849,6 @@ void writeArray(void *Array,     // [in] array to be written
    if (Err != PIO_NOERR)
       ABORT_ERROR("IO::writeArray: Error in PIO writing distributed array");
 
-   // Make sure write is complete before returning
-   // We may be able to remove this for efficiency later but it was
-   // needed during testing
-   Err = PIOc_sync(FileID);
-   if (Err != PIO_NOERR)
-      LOG_WARN("IO::writeArray: Error in PIO sychronizing file after write");
-
    return;
 
 } // end writeArray
diff --git a/components/omega/src/infra/IOStream.cpp b/components/omega/src/infra/IOStream.cpp
index a30e5da3e64f..e5ac27d7b180 100644
--- a/components/omega/src/infra/IOStream.cpp
+++ b/components/omega/src/infra/IOStream.cpp
@@ -873,9 +873,9 @@ void IOStream::computeDecomp(
                   I4 IGlob = (DimOffsets[4])(I);
                   if (IGlob < 0)
                      continue;
-                  int Add = I * StrideLoc[4] + J * StrideLoc[3] +
-                            K * StrideLoc[2] + M * StrideLoc[1] +
-                            N * StrideLoc[0];
+                  int Add     = I * StrideLoc[4] + J * StrideLoc[3] +
+                                K * StrideLoc[2] + M * StrideLoc[1] +
+                                N * StrideLoc[0];
                   Offset[Add] = IGlob * StrideGlob[4] + JGlob * StrideGlob[3] +
                                 KGlob * StrideGlob[2] + MGlob * StrideGlob[1] +
                                 NGlob * StrideGlob[0];
@@ -963,7 +963,7 @@ void IOStream::writeFieldMeta(
 //------------------------------------------------------------------------------
 // Write a field's data array, performing any manipulations to reduce
 // precision or move data between host and device
-void IOStream::writeFieldData(
+int IOStream::writeFieldData(
     std::shared_ptr<Field> FieldPtr,      // [in] field to write
     int FileID,                           // [in] id assigned to open file
     int FieldID,                          // [in] id assigned to the field
@@ -1615,14 +1615,13 @@ void IOStream::writeFieldData(
       IO::writeArray(DataPtr, LocSize, FillValPtr, FileID, MyDecompID, FieldID,
                      FldFrame);
 
-      // Clean up the decomp
-      IO::destroyDecomp(MyDecompID);
+      return MyDecompID;
 
    } else {
       IO::writeNDVar(DataPtr, FileID, FieldID, FldFrame, &DimLengths);
    }
 
-   return;
+   return -1;
 
 } // end writeFieldData
 
@@ -2665,7 +2664,10 @@ void IOStream::writeStream(
    IO::endDefinePhase(OutFileID);
    DefineMode = false;
 
-   // Now write data arrays for all fields in contents
+   // Now write data arrays for all fields in contents. Each distributed
+   // field's PIO decomposition is not destroyed immediately since its write
+   // may still be buffered by PIO and not yet flushed to disk
+   std::vector<int> DecompIDsToFree;
    for (auto IFld = Contents.begin(); IFld != Contents.end(); ++IFld) {
 
       // Retrieve the field pointer and FieldID
@@ -2674,12 +2676,21 @@ void IOStream::writeStream(
       int FieldID                      = FieldIDs[FieldName];
 
       // Extract and write the data array
-      this->writeFieldData(ThisField, OutFileID, FieldID, AllDimIDs);
+      int DecompID =
+          this->writeFieldData(ThisField, OutFileID, FieldID, AllDimIDs);
+      if (DecompID >= 0)
+         DecompIDsToFree.push_back(DecompID);
    }
 
-   // Close output file
+   // Close output file - this syncs the file, guaranteeing all buffered
+   // writes above have been flushed before we free their decompositions
    IO::closeFile(OutFileID);
 
+   // Now that the file has been synced and closed, it is safe to destroy
+   // the decompositions used by the writes above
+   for (auto &DecompID : DecompIDsToFree)
+      IO::destroyDecomp(DecompID);
+
    // If using pointer files for this stream, write the filename to the pointer
    // after the file is successfully written
    if (UsePointer) {
diff --git a/components/omega/src/infra/IOStream.h b/components/omega/src/infra/IOStream.h
index ba2556867154..7ace3a2dd7ff 100644
--- a/components/omega/src/infra/IOStream.h
+++ b/components/omega/src/infra/IOStream.h
@@ -215,8 +215,10 @@ class IOStream {
    );
 
    /// Write a field's data array, performing any manipulations to reduce
-   /// precision or move data between host and device
-   void
+   /// precision or move data between host and device. Returns the PIO
+   /// decomposition ID used for a distributed field, so the caller can
+   /// destroy it only once the write has been flushed
+   int
    writeFieldData(std::shared_ptr<Field> FieldPtr, ///< [in] field to write
                   int FileID,  ///< [in] id assigned to open file
                   int FieldID, ///< [in] id assigned to the field