binary pim code for reduced memory usage

fast pim code emission

src/PIM/Compiler/PimArtifactWriter.cpp

@@ -11,6 +11,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
+#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
@@ -19,18 +20,35 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-OnnxMlirCompilerErrorCodes writeHostCoreJson(StringRef outputDirPath) {
+OnnxMlirCompilerErrorCodes writeHostCoreArtifacts(StringRef outputDirPath) {
   std::error_code errorCode;
-  std::string outputHostCorePath = outputDirPath.str() + "/core_0.json";
-  raw_fd_ostream hostFileStream(outputHostCorePath, errorCode);
+  std::string outputHostCorePath = outputDirPath.str() + "/core_0.pim";
+  raw_fd_ostream hostFileStream(outputHostCorePath, errorCode, sys::fs::OF_None);
   if (errorCode) {
     errs() << "Error while opening host core file `" << outputHostCorePath << "`: " << errorCode.message() << '\n';
     return InvalidOutputFileAccess;
   }
 
-  // The host core json contains two no-op-like instructions to satisfy pimsim-nn.
-  hostFileStream << "[{\"imm\":0,\"op\":\"sldi\",\"rd\":0},{\"imm\":0,\"op\":\"sldi\",\"rd\":0}]";
+  pim_binary::writeHeader(hostFileStream);
+  pim_binary::InstructionRecord noop;
+  noop.opcode = pim_binary::Opcode::sldi;
+  pim_binary::writeInstructionRecord(hostFileStream, noop);
+  pim_binary::writeInstructionRecord(hostFileStream, noop);
+  pim_binary::patchInstructionCount(hostFileStream, 2);
   hostFileStream.close();
+
+  if (pimEmitJson.getValue()) {
+    std::string outputHostJsonPath = outputDirPath.str() + "/core_0.json";
+    raw_fd_ostream hostJsonStream(outputHostJsonPath, errorCode);
+    if (errorCode) {
+      errs() << "Error while opening host core json file `" << outputHostJsonPath << "`: " << errorCode.message()
+             << '\n';
+      return InvalidOutputFileAccess;
+    }
+    // The host core json contains two no-op-like instructions to satisfy pimsim-nn
+    hostJsonStream << "[{\"imm\":0,\"op\":\"sldi\",\"rd\":0},{\"imm\":0,\"op\":\"sldi\",\"rd\":0}]";
+    hostJsonStream.close();
+  }
   return CompilerSuccess;
 }
 

src/PIM/Compiler/PimArtifactWriter.hpp

@@ -12,7 +12,7 @@ namespace onnx_mlir {
 
 class PimAcceleratorMemory;
 
-OnnxMlirCompilerErrorCodes writeHostCoreJson(llvm::StringRef outputDirPath);
+OnnxMlirCompilerErrorCodes writeHostCoreArtifacts(llvm::StringRef outputDirPath);
 OnnxMlirCompilerErrorCodes writeMemoryBinary(mlir::ModuleOp moduleOp,
                                              mlir::func::FuncOp funcOp,
                                              PimAcceleratorMemory& memory,

src/PIM/Compiler/PimBinaryFormat.hpp

@@ -0,0 +1,381 @@
+#pragma once
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Endian.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <array>
+#include <cassert>
+#include <limits>
+
+namespace onnx_mlir::pim_binary {
+
+inline constexpr char kMagic[4] = {'P', 'I', 'M', 'B'};
+inline constexpr uint32_t kVersion = 1;
+inline constexpr uint64_t kCountOffset = 8;
+inline constexpr size_t kHeaderSize = 12;
+inline constexpr size_t kRecordSize = 20;
+
+enum class Opcode : uint32_t {
+  nop = 0,
+  sldi = 1,
+  sld = 2,
+  sadd = 3,
+  ssub = 4,
+  smul = 5,
+  saddi = 6,
+  smuli = 7,
+  setbw = 8,
+  mvmul = 9,
+  vvadd = 10,
+  vvsub = 11,
+  vvmul = 12,
+  vvdmul = 13,
+  vvmax = 14,
+  vvsll = 15,
+  vvsra = 16,
+  vavg = 17,
+  vrelu = 18,
+  vtanh = 19,
+  vsigm = 20,
+  vsoftmax = 21,
+  vmv = 22,
+  vrsu = 23,
+  vrsl = 24,
+  ld = 25,
+  st = 26,
+  lldi = 27,
+  lmv = 28,
+  send = 29,
+  recv = 30,
+  wait = 31,
+  sync = 32,
+};
+
+struct InstructionRecord {
+  Opcode opcode = Opcode::nop;
+  uint8_t rd = 0;
+  uint8_t r1 = 0;
+  int32_t r2OrImm = 0;
+  int32_t generic1 = 0;
+  int32_t generic2 = 0;
+  int32_t generic3 = 0;
+  uint8_t flags = 0;
+};
+
+inline void writeUint32LE(llvm::raw_ostream& os, uint32_t value) {
+  std::array<char, sizeof(uint32_t)> bytes;
+  llvm::support::endian::write32le(bytes.data(), value);
+  os.write(bytes.data(), bytes.size());
+}
+
+inline void writeInt32LE(llvm::raw_ostream& os, int32_t value) {
+  writeUint32LE(os, static_cast<uint32_t>(value));
+}
+
+inline void writeHeader(llvm::raw_ostream& os) {
+  os.write(kMagic, sizeof(kMagic));
+  writeUint32LE(os, kVersion);
+  writeUint32LE(os, 0);
+}
+
+inline void patchInstructionCount(llvm::raw_pwrite_stream& os, uint32_t instructionCount) {
+  std::array<char, sizeof(uint32_t)> bytes;
+  llvm::support::endian::write32le(bytes.data(), instructionCount);
+  os.pwrite(bytes.data(), bytes.size(), kCountOffset);
+}
+
+inline void writeInstructionRecord(llvm::raw_ostream& os, const InstructionRecord& record) {
+  os << static_cast<char>(static_cast<uint8_t>(record.opcode));
+  os << static_cast<char>(record.rd);
+  os << static_cast<char>(record.r1);
+  os << static_cast<char>(record.flags);
+  writeInt32LE(os, record.r2OrImm);
+  writeInt32LE(os, record.generic1);
+  writeInt32LE(os, record.generic2);
+  writeInt32LE(os, record.generic3);
+}
+
+inline int32_t toI32(int64_t value) {
+  assert(value >= std::numeric_limits<int32_t>::min() && value <= std::numeric_limits<int32_t>::max()
+         && "PIM binary field out of int32 range");
+  return static_cast<int32_t>(value);
+}
+
+inline uint8_t toU8(int64_t value) {
+  assert(value >= 0 && value <= std::numeric_limits<uint8_t>::max() && "PIM binary field out of uint8 range");
+  return static_cast<uint8_t>(value);
+}
+
+inline int32_t getOptionalInt(const llvm::json::Object& object, llvm::StringRef key, int32_t defaultValue = 0) {
+  if (std::optional<int64_t> value = object.getInteger(key))
+    return toI32(*value);
+  return defaultValue;
+}
+
+inline Opcode opcodeFromString(llvm::StringRef opName) {
+  if (opName == "nop")
+    return Opcode::nop;
+  if (opName == "sldi")
+    return Opcode::sldi;
+  if (opName == "sld")
+    return Opcode::sld;
+  if (opName == "sadd")
+    return Opcode::sadd;
+  if (opName == "ssub")
+    return Opcode::ssub;
+  if (opName == "smul")
+    return Opcode::smul;
+  if (opName == "saddi")
+    return Opcode::saddi;
+  if (opName == "smuli")
+    return Opcode::smuli;
+  if (opName == "setbw")
+    return Opcode::setbw;
+  if (opName == "mvmul")
+    return Opcode::mvmul;
+  if (opName == "vvadd")
+    return Opcode::vvadd;
+  if (opName == "vvsub")
+    return Opcode::vvsub;
+  if (opName == "vvmul")
+    return Opcode::vvmul;
+  if (opName == "vvdmul")
+    return Opcode::vvdmul;
+  if (opName == "vvmax")
+    return Opcode::vvmax;
+  if (opName == "vvsll")
+    return Opcode::vvsll;
+  if (opName == "vvsra")
+    return Opcode::vvsra;
+  if (opName == "vavg")
+    return Opcode::vavg;
+  if (opName == "vrelu")
+    return Opcode::vrelu;
+  if (opName == "vtanh")
+    return Opcode::vtanh;
+  if (opName == "vsigm")
+    return Opcode::vsigm;
+  if (opName == "vsoftmax")
+    return Opcode::vsoftmax;
+  if (opName == "vmv")
+    return Opcode::vmv;
+  if (opName == "vrsu")
+    return Opcode::vrsu;
+  if (opName == "vrsl")
+    return Opcode::vrsl;
+  if (opName == "ld")
+    return Opcode::ld;
+  if (opName == "st")
+    return Opcode::st;
+  if (opName == "lldi")
+    return Opcode::lldi;
+  if (opName == "lmv")
+    return Opcode::lmv;
+  if (opName == "send")
+    return Opcode::send;
+  if (opName == "recv")
+    return Opcode::recv;
+  if (opName == "wait")
+    return Opcode::wait;
+  if (opName == "sync")
+    return Opcode::sync;
+  llvm_unreachable("Unsupported PIM binary opcode");
+}
+
+inline llvm::StringRef opcodeToString(Opcode opcode) {
+  switch (opcode) {
+  case Opcode::nop: return "nop";
+  case Opcode::sldi: return "sldi";
+  case Opcode::sld: return "sld";
+  case Opcode::sadd: return "sadd";
+  case Opcode::ssub: return "ssub";
+  case Opcode::smul: return "smul";
+  case Opcode::saddi: return "saddi";
+  case Opcode::smuli: return "smuli";
+  case Opcode::setbw: return "setbw";
+  case Opcode::mvmul: return "mvmul";
+  case Opcode::vvadd: return "vvadd";
+  case Opcode::vvsub: return "vvsub";
+  case Opcode::vvmul: return "vvmul";
+  case Opcode::vvdmul: return "vvdmul";
+  case Opcode::vvmax: return "vvmax";
+  case Opcode::vvsll: return "vvsll";
+  case Opcode::vvsra: return "vvsra";
+  case Opcode::vavg: return "vavg";
+  case Opcode::vrelu: return "vrelu";
+  case Opcode::vtanh: return "vtanh";
+  case Opcode::vsigm: return "vsigm";
+  case Opcode::vsoftmax: return "vsoftmax";
+  case Opcode::vmv: return "vmv";
+  case Opcode::vrsu: return "vrsu";
+  case Opcode::vrsl: return "vrsl";
+  case Opcode::ld: return "ld";
+  case Opcode::st: return "st";
+  case Opcode::lldi: return "lldi";
+  case Opcode::lmv: return "lmv";
+  case Opcode::send: return "send";
+  case Opcode::recv: return "recv";
+  case Opcode::wait: return "wait";
+  case Opcode::sync: return "sync";
+  }
+  llvm_unreachable("Unsupported PIM binary opcode");
+}
+
+inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruction) {
+  InstructionRecord record;
+  std::optional<llvm::StringRef> opName = instruction.getString("op");
+  assert(opName && "Missing op field in PIM instruction");
+  record.opcode = opcodeFromString(*opName);
+  record.rd = toU8(getOptionalInt(instruction, "rd"));
+  record.r1 = toU8(getOptionalInt(instruction, "rs1"));
+
+  switch (record.opcode) {
+  case Opcode::sldi:
+  case Opcode::saddi:
+  case Opcode::smuli:
+  case Opcode::lldi:
+    record.r2OrImm = getOptionalInt(instruction, "imm");
+    break;
+  case Opcode::mvmul:
+    record.r2OrImm = getOptionalInt(instruction, "mbiw");
+    record.generic1 = getOptionalInt(instruction, "relu");
+    record.generic2 = getOptionalInt(instruction, "group");
+    break;
+  case Opcode::setbw:
+    record.generic1 = getOptionalInt(instruction, "ibiw");
+    record.generic2 = getOptionalInt(instruction, "obiw");
+    break;
+  case Opcode::send:
+  case Opcode::recv:
+    record.r2OrImm = getOptionalInt(instruction, "core");
+    record.generic3 = getOptionalInt(instruction, "size");
+    break;
+  default:
+    record.r2OrImm = getOptionalInt(instruction, "rs2");
+    break;
+  }
+
+  if (record.opcode != Opcode::mvmul && record.opcode != Opcode::setbw) {
+    if (auto* offsetValue = instruction.getObject("offset")) {
+      record.generic1 = getOptionalInt(*offsetValue, "offset_select");
+      record.generic2 = getOptionalInt(*offsetValue, "offset_value");
+    }
+  }
+
+  if (instruction.get("len"))
+    record.generic3 = getOptionalInt(instruction, "len");
+  else if (instruction.get("size") && record.opcode != Opcode::send && record.opcode != Opcode::recv)
+    record.generic3 = getOptionalInt(instruction, "size");
+
+  return record;
+}
+
+inline llvm::json::Object makeInstructionJson(const InstructionRecord& record) {
+  llvm::json::Object instruction;
+  instruction["op"] = opcodeToString(record.opcode).str();
+
+  auto addOffset = [&](int32_t offsetSelect, int32_t offsetValue) {
+    llvm::json::Object offset;
+    offset["offset_select"] = offsetSelect;
+    offset["offset_value"] = offsetValue;
+    instruction["offset"] = std::move(offset);
+  };
+
+  switch (record.opcode) {
+  case Opcode::sldi:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["imm"] = record.r2OrImm;
+    break;
+  case Opcode::sld:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["rs1"] = static_cast<int64_t>(record.r1);
+    addOffset(record.generic1, record.generic2);
+    break;
+  case Opcode::sadd:
+  case Opcode::ssub:
+  case Opcode::smul:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["rs1"] = static_cast<int64_t>(record.r1);
+    instruction["rs2"] = record.r2OrImm;
+    break;
+  case Opcode::saddi:
+  case Opcode::smuli:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["rs1"] = static_cast<int64_t>(record.r1);
+    instruction["imm"] = record.r2OrImm;
+    break;
+  case Opcode::setbw:
+    instruction["ibiw"] = record.generic1;
+    instruction["obiw"] = record.generic2;
+    break;
+  case Opcode::mvmul:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["rs1"] = static_cast<int64_t>(record.r1);
+    instruction["mbiw"] = record.r2OrImm;
+    instruction["relu"] = record.generic1;
+    instruction["group"] = record.generic2;
+    break;
+  case Opcode::vvadd:
+  case Opcode::vvsub:
+  case Opcode::vvmul:
+  case Opcode::vvdmul:
+  case Opcode::vvmax:
+  case Opcode::vvsll:
+  case Opcode::vvsra:
+  case Opcode::vavg:
+  case Opcode::vmv:
+  case Opcode::vrsu:
+  case Opcode::vrsl:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["rs1"] = static_cast<int64_t>(record.r1);
+    instruction["rs2"] = record.r2OrImm;
+    addOffset(record.generic1, record.generic2);
+    instruction["len"] = record.generic3;
+    break;
+  case Opcode::vrelu:
+  case Opcode::vtanh:
+  case Opcode::vsigm:
+  case Opcode::vsoftmax:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["rs1"] = static_cast<int64_t>(record.r1);
+    addOffset(record.generic1, record.generic2);
+    instruction["len"] = record.generic3;
+    break;
+  case Opcode::ld:
+  case Opcode::st:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["rs1"] = static_cast<int64_t>(record.r1);
+    addOffset(record.generic1, record.generic2);
+    instruction["size"] = record.generic3;
+    break;
+  case Opcode::lldi:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["imm"] = record.r2OrImm;
+    addOffset(record.generic1, record.generic2);
+    instruction["len"] = record.generic3;
+    break;
+  case Opcode::lmv:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["rs1"] = static_cast<int64_t>(record.r1);
+    addOffset(record.generic1, record.generic2);
+    instruction["len"] = record.generic3;
+    break;
+  case Opcode::send:
+  case Opcode::recv:
+    instruction["rd"] = static_cast<int64_t>(record.rd);
+    instruction["core"] = record.r2OrImm;
+    addOffset(record.generic1, record.generic2);
+    instruction["size"] = record.generic3;
+    break;
+  case Opcode::wait:
+  case Opcode::sync:
+  case Opcode::nop:
+    break;
+  }
+
+  return instruction;
+}
+
+} // namespace onnx_mlir::pim_binary

src/PIM/Compiler/PimCodeGen.cpp

@@ -30,6 +30,7 @@
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
 #include "src/Accelerators/PIM/Compiler/PimBatchEmission.hpp"
+#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
@@ -116,25 +117,29 @@ void PimMemory::allocateCore(Operation* op) {
 
 static void printHostMemoryReportRow(raw_ostream& os, const MemoryReportRow& row) {
   llvm::SmallVector<ReportField, 2> fields = {
-    {"Number of globals", std::to_string(row.numGlobal)},
-    {"Global memory", formatReportMemory(row.sizeGlobal)}};
+    {"Number of globals", std::to_string(row.numGlobal)     },
+    {"Global memory",     formatReportMemory(row.sizeGlobal)}
+  };
   printReportFlatFields(os, fields);
 }
 
 static void printCoreMemoryReportRow(raw_ostream& os, const MemoryReportEntry& entry) {
   llvm::SmallVector<ReportField, 2> fields = {
-    {"Number of allocas", std::to_string(entry.row.numAlloca)},
-    {"Allocated memory", formatReportMemory(entry.row.sizeAlloca)}};
+    {"Number of allocas", std::to_string(entry.row.numAlloca)     },
+    {"Allocated memory",  formatReportMemory(entry.row.sizeAlloca)}
+  };
   printReportFlatFields(os, fields);
 }
 
 static void printBatchMemoryReportRow(raw_ostream& os, const MemoryReportEntry& entry) {
   llvm::SmallVector<ReportField, 2> perCoreFields = {
-    {"Number of allocas", std::to_string(entry.row.numAlloca)},
-    {"Allocated memory", formatReportMemory(entry.row.sizeAlloca)}};
+    {"Number of allocas", std::to_string(entry.row.numAlloca)     },
+    {"Allocated memory",  formatReportMemory(entry.row.sizeAlloca)}
+  };
   llvm::SmallVector<ReportField, 2> totalFields = {
-    {"Number of allocas", std::to_string(entry.totalAllocaCount)},
-    {"Batch memory", formatReportMemory(entry.totalAllocaBytes)}};
+    {"Number of allocas", std::to_string(entry.totalAllocaCount)    },
+    {"Batch memory",      formatReportMemory(entry.totalAllocaBytes)}
+  };
   printReportPerCoreAndTotalFields(os, perCoreFields, totalFields);
 }
 
@@ -215,12 +220,8 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value, const StaticValu
 void PimAcceleratorMemory::reportHost() { hostReportRow = hostMem.getReportRow(); }
 
 void PimAcceleratorMemory::recordCoreReport(size_t coreId, const MemoryReportRow& row) {
-  reportEntries.push_back({MemoryReportEntry::Kind::Core,
-                           coreId,
-                           {static_cast<int32_t>(coreId)},
-                           row,
-                           row.numAlloca,
-                           row.sizeAlloca});
+  reportEntries.push_back(
+    {MemoryReportEntry::Kind::Core, coreId, {static_cast<int32_t>(coreId)}, row, row.numAlloca, row.sizeAlloca});
 }
 
 void PimAcceleratorMemory::recordBatchReport(uint64_t batchId,
@@ -250,7 +251,8 @@ void PimAcceleratorMemory::flushReport() {
 
   llvm::SmallVector<ReportField, 2> totalFields = {
     {"Global memory", formatReportMemory(totalGlobalMemory)},
-    {"Cores memory", formatReportMemory(totalCoresMemory)}};
+    {"Cores memory",  formatReportMemory(totalCoresMemory) }
+  };
   printReportTotalsBlock(os, totalFields);
 
   if (hostReportRow.has_value()) {
@@ -312,36 +314,25 @@ void PimAcceleratorMemory::clean(mlir::Operation* op) {
   }
 }
 
-json::Object PimCodeGen::createEmptyOffset() {
-  json::Object offset;
-  offset["offset_select"] = 0;
-  offset["offset_value"] = 0;
-  return offset;
-}
-
 size_t PimCodeGen::remapCoreId(size_t coreId) const {
   auto it = emittedCoreIds.find(coreId);
   assert(it != emittedCoreIds.end() && "Missing emitted core id remapping");
   return it->second;
 }
 
-static json::Object createRs1OnlyOffset() {
-  json::Object offset;
-  offset["offset_select"] = 1;
-  offset["offset_value"] = 0;
-  return offset;
-}
-
-void PimCodeGen::emitInstruction(json::Object instruction) const {
-  coreFileStream << json::Value(std::move(instruction)) << ',';
+void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instruction) const {
+  pim_binary::writeInstructionRecord(coreBinaryStream, instruction);
+  ++emittedInstructionCount;
+  if (coreJsonStream)
+    *coreJsonStream << json::Value(pim_binary::makeInstructionJson(instruction)) << ',';
 }
 
 void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const {
-  json::Object json;
-  json["op"] = "sldi";
-  json["rd"] = registerNumber;
-  json["imm"] = immediate;
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::sldi;
+  instruction.rd = static_cast<uint8_t>(registerNumber);
+  instruction.r2OrImm = static_cast<int32_t>(immediate);
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const {
@@ -369,38 +360,41 @@ void PimCodeGen::emitMemCopyOp(StringRef opName,
                                StringRef sizeFieldName) const {
   setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset);
 
-  json::Object json;
-  json["op"] = opName;
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json[sizeFieldName] = size;
-  json["offset"] = createEmptyOffset();
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::opcodeFromString(opName);
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.generic1 = 0;
+  instruction.generic2 = 0;
+  instruction.generic3 = static_cast<int32_t>(size);
+  (void)sizeFieldName;
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::emitCommunicationOp(StringRef opName, size_t bufferAddr, size_t coreId, size_t size) const {
   setupRd(bufferAddr, 0);
 
-  json::Object json;
-  json["op"] = opName;
-  json["rd"] = 0;
-  json["core"] = remapCoreId(coreId);
-  json["size"] = size;
-  json["offset"] = createEmptyOffset();
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::opcodeFromString(opName);
+  instruction.rd = 0;
+  instruction.r2OrImm = static_cast<int32_t>(remapCoreId(coreId));
+  instruction.generic1 = 0;
+  instruction.generic2 = 0;
+  instruction.generic3 = static_cast<int32_t>(size);
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset) const {
   setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset);
 
-  json::Object json;
-  json["op"] = "mvmul";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["group"] = groupId;
-  json["relu"] = 0;
-  json["mbiw"] = 8;
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::mvmul;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 8;
+  instruction.generic1 = 0;
+  instruction.generic2 = static_cast<int32_t>(groupId);
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const {
@@ -508,14 +502,13 @@ void PimCodeGen::codeGenVVAddOp(pim::PimVVAddOp vvaddOp, const StaticValueKnowle
   auto rhsAddr = addressOf(vvaddOp.getRhs(), knowledge);
   setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0);
 
-  json::Object json;
-  json["op"] = "vvadd";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["rs2"] = 2;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vvaddOp.getLhs());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vvadd;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 2;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vvaddOp.getLhs()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVVSubOp(pim::PimVVSubOp vvsubOp, const StaticValueKnowledge& knowledge) const {
@@ -524,14 +517,13 @@ void PimCodeGen::codeGenVVSubOp(pim::PimVVSubOp vvsubOp, const StaticValueKnowle
   auto rhsAddr = addressOf(vvsubOp.getRhs(), knowledge);
   setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0);
 
-  json::Object json;
-  json["op"] = "vvsub";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["rs2"] = 2;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vvsubOp.getLhs());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vvsub;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 2;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vvsubOp.getLhs()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVVMulOp(pim::PimVVMulOp vvmulOp, const StaticValueKnowledge& knowledge) const {
@@ -540,14 +532,13 @@ void PimCodeGen::codeGenVVMulOp(pim::PimVVMulOp vvmulOp, const StaticValueKnowle
   auto rhsAddr = addressOf(vvmulOp.getRhs(), knowledge);
   setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0);
 
-  json::Object json;
-  json["op"] = "vvmul";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["rs2"] = 2;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vvmulOp.getLhs());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vvmul;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 2;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vvmulOp.getLhs()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVVMaxOp(pim::PimVVMaxOp vvmaxOp, const StaticValueKnowledge& knowledge) const {
@@ -556,14 +547,13 @@ void PimCodeGen::codeGenVVMaxOp(pim::PimVVMaxOp vvmaxOp, const StaticValueKnowle
   auto rhsAddr = addressOf(vvmaxOp.getRhs(), knowledge);
   setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0);
 
-  json::Object json;
-  json["op"] = "vvmax";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["rs2"] = 2;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vvmaxOp.getLhs());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vvmax;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 2;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vvmaxOp.getLhs()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVVDMulOp(pim::PimVVDMulOp vvdmulOp, const StaticValueKnowledge& knowledge) const {
@@ -572,14 +562,13 @@ void PimCodeGen::codeGenVVDMulOp(pim::PimVVDMulOp vvdmulOp, const StaticValueKno
   auto rhsAddr = addressOf(vvdmulOp.getRhs(), knowledge);
   setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0);
 
-  json::Object json;
-  json["op"] = "vvdmul";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["rs2"] = 2;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vvdmulOp.getLhs());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vvdmul;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 2;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vvdmulOp.getLhs()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVAvgOp(pim::PimVAvgOp vavgOp, const StaticValueKnowledge& knowledge) const {
@@ -587,14 +576,14 @@ void PimCodeGen::codeGenVAvgOp(pim::PimVAvgOp vavgOp, const StaticValueKnowledge
   auto inputAddr = addressOf(vavgOp.getInput(), knowledge);
   setupRdRs1(outputBufferAddr, 0, inputAddr, 0);
 
-  json::Object json;
-  json["op"] = "vavg";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["rs2"] = 1;
-  json["offset"] = createRs1OnlyOffset();
-  json["len"] = getValueSizeInBytes(vavgOp.getInput());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vavg;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 1;
+  instruction.generic1 = 1;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vavgOp.getInput()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVReluOp(pim::PimVReluOp vreluOp, const StaticValueKnowledge& knowledge) const {
@@ -602,13 +591,12 @@ void PimCodeGen::codeGenVReluOp(pim::PimVReluOp vreluOp, const StaticValueKnowle
   auto inputAddr = addressOf(vreluOp.getInput(), knowledge);
   setupRdRs1(outputBufferAddr, 0, inputAddr, 0);
 
-  json::Object json;
-  json["op"] = "vrelu";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vreluOp.getInput());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vrelu;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vreluOp.getInput()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowledge& knowledge) const {
@@ -616,13 +604,12 @@ void PimCodeGen::codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowle
   auto inputAddr = addressOf(vtanhOp.getInput(), knowledge);
   setupRdRs1(outputBufferAddr, 0, inputAddr, 0);
 
-  json::Object json;
-  json["op"] = "vtanh";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vtanhOp.getInput());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vtanh;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vtanhOp.getInput()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowledge& knowledge) const {
@@ -630,13 +617,12 @@ void PimCodeGen::codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowle
   auto inputAddr = addressOf(vsigmOp.getInput(), knowledge);
   setupRdRs1(outputBufferAddr, 0, inputAddr, 0);
 
-  json::Object json;
-  json["op"] = "vsigm";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vsigmOp.getInput());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vsigm;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vsigmOp.getInput()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticValueKnowledge& knowledge) const {
@@ -644,13 +630,12 @@ void PimCodeGen::codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticVa
   auto inputAddr = addressOf(vsoftmaxOp.getInput(), knowledge);
   setupRdRs1(outputBufferAddr, 0, inputAddr, 0);
 
-  json::Object json;
-  json["op"] = "vsoftmax";
-  json["rd"] = 0;
-  json["rs1"] = 1;
-  json["offset"] = createEmptyOffset();
-  json["len"] = getValueSizeInBytes(vsoftmaxOp.getInput());
-  emitInstruction(std::move(json));
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vsoftmax;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.generic3 = static_cast<int32_t>(getValueSizeInBytes(vsoftmaxOp.getInput()));
+  emitInstruction(instruction);
 }
 
 void PimCodeGen::codeGetGlobalOp(memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const {}
@@ -682,6 +667,30 @@ void PimCodeGen::codeGenTransposeOp(pim::PimTransposeOp transposeOp, const Stati
     dstStrides[i] = dstStrides[i + 1] * dstShape[i + 1];
   }
 
+  bool storagePreserving = true;
+  for (size_t srcFlat = 0; srcFlat < totalElements; srcFlat++) {
+    SmallVector<size_t> srcIdx(rank);
+    size_t remaining = srcFlat;
+    for (size_t d = 0; d < rank; d++) {
+      srcIdx[d] = remaining / srcStrides[d];
+      remaining %= srcStrides[d];
+    }
+
+    size_t dstFlat = 0;
+    for (size_t d = 0; d < rank; d++)
+      dstFlat += srcIdx[perm[d]] * dstStrides[d];
+
+    if (dstFlat != srcFlat) {
+      storagePreserving = false;
+      break;
+    }
+  }
+
+  if (storagePreserving) {
+    emitMemCopyOp("lmv", dstAddr, 0, srcAddr, 0, totalElements * elementSize, "len");
+    return;
+  }
+
   // Emit element-by-element copy with transposed addressing
   for (size_t srcFlat = 0; srcFlat < totalElements; srcFlat++) {
     // Decompose flat source index into multi-dimensional index
@@ -747,9 +756,25 @@ static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
   return coreLikeOps;
 }
 
+static SmallDenseMap<memref::GlobalOp, MemEntry, 16>
+collectMaterializedHostGlobals(ModuleOp moduleOp, func::FuncOp funcOp, const PimAcceleratorMemory& memory) {
+  SmallDenseMap<memref::GlobalOp, MemEntry, 16> materializedHostGlobals;
+  funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
+    if (hasWeightAlways(getGlobalOp))
+      return;
+    auto targetGlobal = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
+    if (!targetGlobal || materializedHostGlobals.contains(targetGlobal))
+      return;
+    auto it = memory.memEntriesMap.find(getGlobalOp.getResult());
+    if (it != memory.memEntriesMap.end())
+      materializedHostGlobals[targetGlobal] = it->second;
+  });
+  return materializedHostGlobals;
+}
+
 static void aliasMaterializedHostGlobals(ModuleOp moduleOp,
-                                         func::FuncOp funcOp,
                                          pim::PimCoreOp coreOp,
+                                         const SmallDenseMap<memref::GlobalOp, MemEntry, 16>& materializedHostGlobals,
                                          PimAcceleratorMemory& memory) {
   coreOp.walk([&](memref::GetGlobalOp getGlobalOp) {
     if (hasWeightAlways(getGlobalOp) || memory.memEntriesMap.contains(getGlobalOp.getResult()))
@@ -759,16 +784,9 @@ static void aliasMaterializedHostGlobals(ModuleOp moduleOp,
     if (!targetGlobal)
       return;
 
-    mlir::Value aliasedValue;
-    funcOp.walk([&](memref::GetGlobalOp candidate) {
-      if (aliasedValue || candidate == getGlobalOp || !memory.memEntriesMap.contains(candidate.getResult()))
-        return;
-      if (lookupGlobalForGetGlobal(moduleOp, candidate) == targetGlobal)
-        aliasedValue = candidate.getResult();
-    });
-
-    if (aliasedValue)
-      memory.memEntriesMap[getGlobalOp.getResult()] = memory.memEntriesMap[aliasedValue];
+    auto it = materializedHostGlobals.find(targetGlobal);
+    if (it != materializedHostGlobals.end())
+      memory.memEntriesMap[getGlobalOp.getResult()] = it->second;
   });
 }
 
@@ -837,7 +855,7 @@ static int64_t codeGenCoreOps(Block& block, PimCodeGen& coreCodeGen) {
   return failed(result) ? -1 : static_cast<int64_t>(processedOperations);
 }
 
-OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::string& outputDirPath) {
+OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::string& outputDirPath) {
   if (!outputDirPath.empty()) {
     if (auto error = sys::fs::create_directory(outputDirPath)) {
       errs() << "Error creating output directory: " << outputDirPath << ": " << error.message() << '\n';
@@ -857,7 +875,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
   if (auto err = writeMemoryBinary(moduleOp, funcOp, memory, outputDirPath))
     return err;
 
-  if (auto err = writeHostCoreJson(outputDirPath))
+  if (auto err = writeHostCoreArtifacts(outputDirPath))
     return err;
 
   // For each core, specify the number of crossbar per array group.
@@ -870,6 +888,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
   auto mapCoreWeightToFileName = createAndPopulateWeightFolder(funcOp, outputDirPath);
 
   SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
+  SmallDenseMap<memref::GlobalOp, MemEntry, 16> materializedHostGlobals =
+    collectMaterializedHostGlobals(moduleOp, funcOp, memory);
   llvm::DenseMap<size_t, size_t> emittedCoreIds;
   size_t nextEmittedCoreId = 1;
 
@@ -899,16 +919,30 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
       maxCoreId = std::max(maxCoreId, coreId);
 
       std::error_code errorCode;
-      auto outputCorePath = outputDirPath + "/core_" + std::to_string(coreId) + ".json";
-      raw_fd_ostream coreFileStream(outputCorePath, errorCode);
+      auto outputCorePath = outputDirPath + "/core_" + std::to_string(coreId) + ".pim";
+      raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
       if (errorCode) {
         errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
         return InvalidOutputFileAccess;
       }
-      coreFileStream << '[';
 
-      PimCodeGen coreCodeGen(memory, coreFileStream, emittedCoreIds);
-      aliasMaterializedHostGlobals(moduleOp, funcOp, coreOp, memory);
+      std::unique_ptr<raw_fd_ostream> coreJsonStream;
+      if (pimEmitJson.getValue()) {
+        std::string outputCoreJsonPath = outputDirPath + "/core_" + std::to_string(coreId) + ".json";
+        errorCode = std::error_code();
+        coreJsonStream = std::make_unique<raw_fd_ostream>(outputCoreJsonPath, errorCode);
+        if (errorCode) {
+          errs() << "Error while opening core json file `" << outputCoreJsonPath << "`: " << errorCode.message()
+                 << '\n';
+          return InvalidOutputFileAccess;
+        }
+        *coreJsonStream << '[';
+      }
+
+      pim_binary::writeHeader(coreBinaryStream);
+
+      PimCodeGen coreCodeGen(memory, coreBinaryStream, coreJsonStream.get(), emittedCoreIds);
+      aliasMaterializedHostGlobals(moduleOp, coreOp, materializedHostGlobals, memory);
       auto& deviceMemory = memory.getOrCreateDeviceMem(coreId);
       deviceMemory.allocateCore(coreOp);
 
@@ -920,9 +954,14 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
       if (reportRow)
         *reportRow = deviceMemory.getReportRow();
 
-      coreFileStream.seek(coreFileStream.tell() - 1);
-      coreFileStream << ']';
-      coreFileStream.close();
+      pim_binary::patchInstructionCount(coreBinaryStream, coreCodeGen.getEmittedInstructionCount());
+      coreBinaryStream.close();
+
+      if (coreJsonStream) {
+        coreJsonStream->seek(coreJsonStream->tell() - 1);
+        *coreJsonStream << ']';
+        coreJsonStream->close();
+      }
 
       auto coreWeightsDirPath = outputDirPath + "/core_" + std::to_string(coreId);
       if (auto error = sys::fs::create_directory(coreWeightsDirPath)) {

src/PIM/Compiler/PimCodeGen.hpp

@@ -13,6 +13,7 @@
 #include "onnx-mlir/Compiler/OMCompilerTypes.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
+#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 namespace onnx_mlir {
@@ -104,16 +105,17 @@ public:
 
 class PimCodeGen {
   PimAcceleratorMemory& memory;
-  llvm::raw_fd_ostream& coreFileStream;
+  llvm::raw_fd_ostream& coreBinaryStream;
+  llvm::raw_fd_ostream* coreJsonStream;
   const llvm::DenseMap<size_t, size_t>& emittedCoreIds;
+  mutable uint32_t emittedInstructionCount = 0;
 
   size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
     return memory.getValueAddress(value, knowledge);
   }
   size_t remapCoreId(size_t coreId) const;
 
-  static llvm::json::Object createEmptyOffset();
-  void emitInstruction(llvm::json::Object instruction) const;
+  void emitInstruction(const pim_binary::InstructionRecord& instruction) const;
 
   void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const;
   void setupRd(size_t rdAddress, size_t rdOffset) const;
@@ -133,9 +135,12 @@ class PimCodeGen {
 
 public:
   PimCodeGen(PimAcceleratorMemory& memory,
-             llvm::raw_fd_ostream& coreJson,
+             llvm::raw_fd_ostream& coreBinary,
+             llvm::raw_fd_ostream* coreJson,
              const llvm::DenseMap<size_t, size_t>& emittedCoreIds)
-  : memory(memory), coreFileStream(coreJson), emittedCoreIds(emittedCoreIds) {}
+  : memory(memory), coreBinaryStream(coreBinary), coreJsonStream(coreJson), emittedCoreIds(emittedCoreIds) {}
+
+  uint32_t getEmittedInstructionCount() const { return emittedInstructionCount; }
 
   void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const;
   void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const;
@@ -164,6 +169,6 @@ public:
   void codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const;
 };
 
-OnnxMlirCompilerErrorCodes compileToPimJson(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);
+OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -24,6 +24,11 @@ llvm::cl::opt<bool> useExperimentalConvImpl("use-experimental-conv-impl",
                                             llvm::cl::init(false),
                                             llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
+                                llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"),
+                                llvm::cl::init(false),
+                                llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<size_t>
   crossbarSize("crossbar-size", llvm::cl::desc("Width and heigth of a single crossbar"), llvm::cl::init(2));
 

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -25,6 +25,7 @@ extern llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget;
 
 extern llvm::cl::opt<bool> pimOnlyCodegen;
 extern llvm::cl::opt<bool> useExperimentalConvImpl;
+extern llvm::cl::opt<bool> pimEmitJson;
 
 extern llvm::cl::opt<size_t> crossbarSize;
 extern llvm::cl::opt<size_t> crossbarCountInCore;

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -52,9 +52,9 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(createPimMaterializeHostConstantsPass());
     pm.addPass(createPimVerificationPass());
     pm.addPass(createMessagePass("Pim verified"));
-    pm.addPass(createEmitPimJsonPass());
+    pm.addPass(createEmitPimCodePass());
     // pm.addPass(createCountInstructionPass());
-    pm.addPass(createMessagePass("Pim json code emitted"));
+    pm.addPass(createMessagePass("Pim code emitted"));
   }
 }
 

src/PIM/Pass/CMakeLists.txt

@@ -7,7 +7,7 @@ add_pim_library(OMPimPasses
   PimCodegen/HostConstantFolding/Patterns/Subview.cpp
   PimCodegen/MaterializeHostConstantsPass.cpp
   PimCodegen/VerificationPass.cpp
-  PimCodegen/EmitPimJsonPass.cpp
+  PimCodegen/EmitPimCodePass.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

src/PIM/Pass/PIMPasses.h

@@ -25,7 +25,7 @@ std::unique_ptr<mlir::Pass> createPimMaterializeHostConstantsPass();
 
 std::unique_ptr<mlir::Pass> createPimVerificationPass();
 
-std::unique_ptr<mlir::Pass> createEmitPimJsonPass();
+std::unique_ptr<mlir::Pass> createEmitPimCodePass();
 
 std::unique_ptr<mlir::Pass> createMessagePass(std::string message);
 

src/PIM/Pass/PimCodegen/EmitPimCodePass.cpp

@@ -0,0 +1,36 @@
+#include "mlir/Pass/Pass.h"
+
+#include "Common/PimCommon.hpp"
+#include "Compiler/PimCodeGen.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+namespace {
+
+struct EmitPimCodePass : PassWrapper<EmitPimCodePass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(EmitPimCodePass);
+  StringRef getArgument() const override { return "emit-pim-code-pass"; }
+  StringRef getDescription() const override { return "Emit PIM simulator code artifacts"; }
+
+  EmitPimCodePass() {}
+  EmitPimCodePass(const EmitPimCodePass& pass) {}
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+
+    std::string pimDir = getOutputDir() + "/pim";
+    createDirectory(pimDir);
+
+    int compiler_error_code = compileToPimCode(moduleOp, pimDir);
+    if (compiler_error_code != CompilerSuccess)
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createEmitPimCodePass() { return std::make_unique<EmitPimCodePass>(); }
+
+} // namespace onnx_mlir

src/PIM/Pass/PimCodegen/EmitPimJsonPass.cpp

@@ -1,36 +0,0 @@
-#include "mlir/Pass/Pass.h"
-
-#include "Common/PimCommon.hpp"
-#include "Compiler/PimCodeGen.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-
-namespace {
-
-struct EmitPimJsonPass : PassWrapper<EmitPimJsonPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(EmitPimJsonPass);
-  StringRef getArgument() const override { return "emit-pim-json-pass"; }
-  StringRef getDescription() const override { return "Emit json code for the pim simulators"; }
-
-  EmitPimJsonPass() {}
-  EmitPimJsonPass(const EmitPimJsonPass& pass) {}
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-
-    std::string pimDir = getOutputDir() + "/pim";
-    createDirectory(pimDir);
-
-    int compiler_error_code = compileToPimJson(moduleOp, pimDir);
-    if (compiler_error_code != CompilerSuccess)
-      signalPassFailure();
-  }
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createEmitPimJsonPass() { return std::make_unique<EmitPimJsonPass>(); }
-
-} // namespace onnx_mlir

src/PIM/PimAccelerator.cpp

@@ -80,7 +80,7 @@ void PimAccelerator::registerPasses(int optLevel) const {
   registerPass(createPimHostConstantFoldingPass);
   registerPass(createPimMaterializeHostConstantsPass);
   registerPass(createPimVerificationPass);
-  registerPass(createEmitPimJsonPass);
+  registerPass(createEmitPimCodePass);
 }
 
 void PimAccelerator::configurePasses() const {