fixes

slightly faster codegen

src/PIM/Compiler/PimBinaryFormat.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/JSON.h"
@@ -53,6 +54,14 @@ enum class Opcode : uint32_t {
   sync = 32,
 };
 
+inline constexpr size_t kOpcodeCount = static_cast<size_t>(Opcode::sync) + 1;
+inline constexpr std::array<llvm::StringLiteral, kOpcodeCount> kOpcodeNames = {
+  "nop",   "sldi", "sld",   "sadd",  "ssub", "smul", "saddi", "smuli", "setbw", "mvmul", "vvadd",
+  "vvsub", "vvmul", "vvdmul", "vvmax", "vvsll", "vvsra", "vavg", "vrelu", "vtanh", "vsigm", "vsoftmax",
+  "vmv",   "vrsu", "vrsl",  "ld",    "st",   "lldi", "lmv",   "send",  "recv",  "wait",  "sync",
+};
+static_assert(kOpcodeNames.size() == kOpcodeCount);
+
 struct InstructionRecord {
   Opcode opcode = Opcode::nop;
   uint8_t rd = 0;
@@ -64,14 +73,29 @@ struct InstructionRecord {
   uint8_t flags = 0;
 };
 
+using EncodedInstruction = std::array<char, kRecordSize>;
+static_assert(kRecordSize == 20);
+static_assert(sizeof(EncodedInstruction) == kRecordSize);
+
+inline EncodedInstruction encodeInstructionRecord(const InstructionRecord& record) {
+  EncodedInstruction encoded = {};
+  encoded[0] = static_cast<char>(static_cast<uint8_t>(record.opcode));
+  encoded[1] = static_cast<char>(record.rd);
+  encoded[2] = static_cast<char>(record.r1);
+  encoded[3] = static_cast<char>(record.flags);
+  llvm::support::endian::write32le(encoded.data() + 4, static_cast<uint32_t>(record.r2OrImm));
+  llvm::support::endian::write32le(encoded.data() + 8, static_cast<uint32_t>(record.generic1));
+  llvm::support::endian::write32le(encoded.data() + 12, static_cast<uint32_t>(record.generic2));
+  llvm::support::endian::write32le(encoded.data() + 16, static_cast<uint32_t>(record.generic3));
+  return encoded;
+}
+
 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);
@@ -84,17 +108,6 @@ inline void patchInstructionCount(llvm::raw_pwrite_stream& os, uint32_t instruct
   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) { return onnx_mlir::pim::checkedI32OrCrash(value, "binary field"); }
 
 inline uint8_t toU8(int64_t value) {
@@ -107,113 +120,64 @@ inline int32_t getOptionalInt(const llvm::json::Object& object, llvm::StringRef
   return defaultValue;
 }
 
+struct InstructionJsonFormat {
+  bool rd;
+  bool r1;
+  bool offset;
+  llvm::StringLiteral r2;
+  llvm::StringLiteral generic1;
+  llvm::StringLiteral generic2;
+  llvm::StringLiteral generic3;
+};
+
+inline constexpr std::array<InstructionJsonFormat, kOpcodeCount> kInstructionJsonFormats = {{
+  {false, false, false, "",     "",     "",      ""    }, // nop
+  {true,  false, false, "imm",  "",     "",      ""    }, // sldi
+  {true,  true,  true,  "",     "",     "",      ""    }, // sld
+  {true,  true,  false, "rs2",  "",     "",      ""    }, // sadd
+  {true,  true,  false, "rs2",  "",     "",      ""    }, // ssub
+  {true,  true,  false, "rs2",  "",     "",      ""    }, // smul
+  {true,  true,  false, "imm",  "",     "",      ""    }, // saddi
+  {true,  true,  false, "imm",  "",     "",      ""    }, // smuli
+  {false, false, false, "",     "ibiw", "obiw",  ""    }, // setbw
+  {true,  true,  false, "mbiw", "relu", "group", ""    }, // mvmul
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvadd
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsub
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvmul
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvdmul
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvmax
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsll
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsra
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vavg
+  {true,  true,  true,  "",     "",     "",      "len" }, // vrelu
+  {true,  true,  true,  "",     "",     "",      "len" }, // vtanh
+  {true,  true,  true,  "",     "",     "",      "len" }, // vsigm
+  {true,  true,  true,  "",     "",     "",      "len" }, // vsoftmax
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vmv
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vrsu
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vrsl
+  {true,  true,  true,  "",     "",     "",      "size"}, // ld
+  {true,  true,  true,  "",     "",     "",      "size"}, // st
+  {true,  false, true,  "imm",  "",     "",      "len" }, // lldi
+  {true,  true,  true,  "",     "",     "",      "len" }, // lmv
+  {true,  false, true,  "core", "",     "",      "size"}, // send
+  {true,  false, true,  "core", "",     "",      "size"}, // recv
+  {false, false, false, "",     "",     "",      ""    }, // wait
+  {false, false, false, "",     "",     "",      ""    }, // sync
+}};
+static_assert(kInstructionJsonFormats.size() == kOpcodeCount);
+
 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;
+  for (auto [index, name] : llvm::enumerate(kOpcodeNames))
+    if (opName == name)
+      return static_cast<Opcode>(index);
   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");
+  size_t index = static_cast<size_t>(opcode);
+  assert(index < kOpcodeNames.size() && "Unsupported PIM binary opcode");
+  return kOpcodeNames[index];
 }
 
 inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruction) {
@@ -221,43 +185,25 @@ inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruc
   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) {
+  const auto& format = kInstructionJsonFormats[static_cast<size_t>(record.opcode)];
+  if (format.rd)
+    record.rd = toU8(getOptionalInt(instruction, "rd"));
+  if (format.r1)
+    record.r1 = toU8(getOptionalInt(instruction, "rs1"));
+  if (!format.r2.empty())
+    record.r2OrImm = getOptionalInt(instruction, format.r2);
+  if (!format.generic1.empty())
+    record.generic1 = getOptionalInt(instruction, format.generic1);
+  if (!format.generic2.empty())
+    record.generic2 = getOptionalInt(instruction, format.generic2);
+  if (format.offset) {
     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");
-
+  if (!format.generic3.empty())
+    record.generic3 = getOptionalInt(instruction, format.generic3);
   return record;
 }
 
@@ -271,98 +217,21 @@ inline llvm::json::Object makeInstructionJson(const InstructionRecord& record) {
     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:
+  const auto& format = kInstructionJsonFormats[static_cast<size_t>(record.opcode)];
+  if (format.rd)
     instruction["rd"] = static_cast<int64_t>(record.rd);
+  if (format.r1)
     instruction["rs1"] = static_cast<int64_t>(record.r1);
+  if (!format.r2.empty())
+    instruction[format.r2] = record.r2OrImm;
+  if (!format.generic1.empty())
+    instruction[format.generic1] = record.generic1;
+  if (!format.generic2.empty())
+    instruction[format.generic2] = record.generic2;
+  if (format.offset)
     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;
-  }
-
+  if (!format.generic3.empty())
+    instruction[format.generic3] = record.generic3;
   return instruction;
 }