Write a constant value to a file without buffering (#2874)

* Write a constant value to single file without buffering to remove spikes in memory consumption.

This PR solves an issue of memory consumption reported in #2821
We found that there is a spike in memory consumption when writing a constant into a file (model.constants.bin). This is because all constants into a buffer once and write it to a file at once. This PR changes to write the constant to the file without the buffering. This removes the spike in memory consumption.

---------

Signed-off-by: Haruki Imai <[email protected]>
Co-authored-by: Tung D. Le <[email protected]>

src/Conversion/KrnlToLLVM/ConvertKrnlToLLVM.cpp

@@ -52,6 +52,7 @@
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Endian.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 
 #include "onnx/onnx_pb.h"
@@ -515,12 +516,13 @@ bool extractConstantsToFile(ModuleOp &module, std::string filepath,
     return (leftAlign < rightAlign);
   });
 
-  // Pack all constants into a single buffer in order to save to file.
+  // Store each constant into single file.
   // Constants with the highest alignment will be packed first in the file.
   // The file will be mmaped later at runtime and aligned at the page boundary,
-  // So every constants must be correctly aligned in the packed constant. Pads
-  // are added if necessary.
-  std::vector<char> packedConst;
+  // So every constants must be correctly aligned. Pads are added if necessary.
+  llvm::sys::fs::remove(filepath);
+  std::ofstream outfile(filepath, std::ios::app | std::ios::binary);
+  uint64_t totalConstSize = 0;
   for (int64_t i = globalOfInterest.size() - 1; i >= 0; --i) {
     KrnlGlobalOp op = globalOfInterest[i];
     ArrayRef<char> rawData = getRawData(op);
@@ -531,27 +533,25 @@ bool extractConstantsToFile(ModuleOp &module, std::string filepath,
       alignment = op.getAlignment().value();
 
     // Padding if necessary.
-    if ((alignment > 0) && (packedConst.size() % alignment != 0)) {
+    if ((alignment > 0) && (totalConstSize % alignment != 0)) {
       uint64_t padSize =
-          ((uint64_t)(packedConst.size() / alignment) + 1) * alignment -
-          packedConst.size();
+          ((uint64_t)(totalConstSize / alignment) + 1) * alignment -
+          totalConstSize;
       SmallVector<char> pads(padSize, (char)0);
-      packedConst.insert(packedConst.end(), pads.begin(), pads.end());
+      outfile.write(pads.data(), pads.size());
+      totalConstSize += pads.size();
     }
 
-    op.setOffsetAttr(b.getI64IntegerAttr(packedConst.size()));
+    op.setOffsetAttr(b.getI64IntegerAttr(totalConstSize));
     op.removeValueAttr();
-    packedConst.insert(packedConst.end(), rawData.begin(), rawData.end());
+    outfile.write(rawData.data(), rawData.size());
+    totalConstSize += rawData.size();
   }
 
   // No constant statisfying thresholds, do not store constants to file.
-  if (packedConst.empty())
+  if (totalConstSize == 0)
     return false;
 
-  // Save to file.
-  std::ofstream outfile(filepath, std::ofstream::binary);
-  outfile.write(packedConst.data(), packedConst.size());
-
   // Create a global op to store the filename in the IR.
   OpBuilder::InsertionGuard guard(b);
   b.setInsertionPointToStart(module.getBody());
@@ -564,7 +564,8 @@ bool extractConstantsToFile(ModuleOp &module, std::string filepath,
   create.llvm.globalOp(llvmI64Ty,
       /*isConstant=*/true, LLVM::Linkage::Internal,
       EXTERNAL_CONSTANT_PREFIX + "filesize",
-      b.getI64IntegerAttr(packedConst.size()));
+      b.getI64IntegerAttr(totalConstSize));
+
   // Create a global to store isLE.
   bool isLE = llvm::endianness::native == llvm::endianness::little;
   create.llvm.globalOp(llvmI8Ty,