ilgeco
239 lines
8.7 KiB
C++
239 lines
8.7 KiB
C++
#pragma once
|
|
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/ADT/StringRef.h"
|
|
#include "llvm/Support/Endian.h"
|
|
#include "llvm/Support/JSON.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
|
|
#include <array>
|
|
|
|
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
|
|
|
|
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,
|
|
};
|
|
|
|
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;
|
|
uint8_t r1 = 0;
|
|
int32_t r2OrImm = 0;
|
|
int32_t generic1 = 0;
|
|
int32_t generic2 = 0;
|
|
int32_t generic3 = 0;
|
|
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 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 int32_t toI32(int64_t value) { return onnx_mlir::pim::checkedI32OrCrash(value, "binary field"); }
|
|
|
|
inline uint8_t toU8(int64_t value) {
|
|
return onnx_mlir::pim::checkedU8OrCrash(static_cast<uint64_t>(value), "binary field");
|
|
}
|
|
|
|
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;
|
|
}
|
|
|
|
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, "", "event_register", "wait_value", ""}, // wait
|
|
{false, false, false, "core", "event_register", "", ""}, // sync
|
|
}};
|
|
static_assert(kInstructionJsonFormats.size() == kOpcodeCount);
|
|
|
|
inline Opcode opcodeFromString(llvm::StringRef opName) {
|
|
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) {
|
|
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) {
|
|
InstructionRecord record;
|
|
std::optional<llvm::StringRef> opName = instruction.getString("op");
|
|
assert(opName && "Missing op field in PIM instruction");
|
|
record.opcode = opcodeFromString(*opName);
|
|
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 (!format.generic3.empty())
|
|
record.generic3 = getOptionalInt(instruction, format.generic3);
|
|
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);
|
|
};
|
|
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);
|
|
if (!format.generic3.empty())
|
|
instruction[format.generic3] = record.generic3;
|
|
return instruction;
|
|
}
|
|
|
|
} // namespace onnx_mlir::pim_binary
|