fixes

slightly faster codegen

src/PIM/Compiler/PimArtifactWriter.cpp

@@ -7,7 +7,7 @@
 #include <algorithm>
 #include <cassert>
 #include <cstring>
-#include <vector>
+#include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
@@ -20,6 +20,65 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
+PimInstructionWriter::PimInstructionWriter(raw_pwrite_stream& stream)
+: stream(stream), buffer(kBufferSize) {
+  pim_binary::writeHeader(stream);
+}
+
+PimInstructionWriter::PimInstructionWriter(raw_fd_ostream& stream)
+: PimInstructionWriter(static_cast<raw_pwrite_stream&>(stream)) {
+  fileStream = &stream;
+}
+
+bool PimInstructionWriter::consumeStreamError() {
+  if (!fileStream || !fileStream->has_error())
+    return false;
+  fileStream->clear_error();
+  hasFailure = true;
+  return true;
+}
+
+LogicalResult PimInstructionWriter::flushBuffer() {
+  if (hasFailure)
+    return failure();
+  if (bufferedBytes != 0) {
+    stream.write(buffer.data(), bufferedBytes);
+    bufferedBytes = 0;
+  }
+  return consumeStreamError() ? failure() : success();
+}
+
+LogicalResult PimInstructionWriter::append(const pim_binary::InstructionRecord& record) {
+  if (hasFailure || finalized || count == std::numeric_limits<uint32_t>::max()) {
+    hasFailure = true;
+    return failure();
+  }
+
+  if (bufferedBytes + pim_binary::kRecordSize > buffer.size() && failed(flushBuffer()))
+    return failure();
+
+  pim_binary::EncodedInstruction encoded = pim_binary::encodeInstructionRecord(record);
+  std::memcpy(buffer.data() + bufferedBytes, encoded.data(), encoded.size());
+  bufferedBytes += encoded.size();
+  ++count;
+  return success();
+}
+
+LogicalResult PimInstructionWriter::finalize() {
+  if (finalized)
+    return failure();
+  finalized = true;
+  if (hasFailure || failed(flushBuffer()))
+    return failure();
+
+  stream.flush();
+  if (consumeStreamError())
+    return failure();
+  pim_binary::patchInstructionCount(stream, count);
+  stream.flush();
+  return consumeStreamError() ? failure() : success();
+}
+
 OnnxMlirCompilerErrorCodes
 writeMemoryBinary(ModuleOp moduleOp, func::FuncOp funcOp, PimAcceleratorMemory& memory, StringRef outputDirPath) {
   auto memoryFilePath = (outputDirPath + "/memory.bin").str();