nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity
Files
45578ef4c4c2c7df3cc8d139be8a7743483b467d
Raptor/src/PIM/Compiler/PimCodeGen.cpp
T
NiccoloN f47fcebb83
Validate Operations / validate-operations (push) Has been cancelled
Details
unify memory opimization passes 
better reports
cleanups
2026-07-21 17:26:01 +02:00

1461 lines
61 KiB
C++
Raw Blame History

#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Threading.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/Verifier.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/raw_ostream.h"
#include <absl/types/compare.h>
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <fstream>
#include <limits>
#include <memory>
#include <numeric>
#include <string>
#include <utility>
#include "Common/IR/CompactAsmUtils.hpp"
#include "Common/PimCommon.hpp"
#include "Common/Support/CheckedArithmetic.hpp"
#include "Common/Support/Diagnostics.hpp"
#include "Common/Support/ReportUtils.hpp"
#include "Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.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"
#include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
#include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
using namespace llvm;
using namespace mlir;
using namespace onnx_mlir;
using namespace onnx_mlir::compact_asm;
namespace {
static std::optional<unsigned> getLaneForMemoryValue(mlir::Value value, std::optional<unsigned> lane) {
if (!lane)
return std::nullopt;
auto allocOp = value.getDefiningOp<memref::AllocOp>();
if (!allocOp || !allocOp->getParentOfType<pim::PimCoreBatchOp>())
return std::nullopt;
return lane;
}
static mlir::Value resolveCachedAlias(mlir::Value value, const StaticValueKnowledge& knowledge) {
auto iter = knowledge.aliases.find(value);
while (iter != knowledge.aliases.end()) {
value = iter->second;
iter = knowledge.aliases.find(value);
}
return value;
}
static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional<unsigned> lane = std::nullopt) {
return {value, getLaneForMemoryValue(value, lane)};
}
static bool isInsidePimCoreLikeOp(memref::AllocOp allocOp) {
return allocOp->getParentOfType<pim::PimCoreOp>() || allocOp->getParentOfType<pim::PimCoreBatchOp>();
}
static MemoryReportKind classifyMemoryReportKind(mlir::Value value) {
if (isa<mlir::BlockArgument>(value))
return MemoryReportKind::Input;
if (auto* op = value.getDefiningOp()) {
if (isa<memref::AllocOp>(op))
return MemoryReportKind::Alloca;
if (isa<memref::GetGlobalOp>(op))
return MemoryReportKind::Global;
}
return MemoryReportKind::None;
}
static int32_t getVectorByteSizeOrCrash(ShapedType type) {
auto byteSize = pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "vector byte size");
if (failed(byteSize))
llvm_unreachable("Failed to compute checked vector byte size");
return pim::checkedI32OrCrash(*byteSize, "vector byte size");
}
static Operation* getDiagnosticAnchor(mlir::Value value) {
if (Operation* definingOp = value.getDefiningOp())
return definingOp;
if (auto blockArg = dyn_cast<BlockArgument>(value))
return blockArg.getOwner()->getParentOp();
return nullptr;
}
// PIM instruction immediates are serialized as signed int32_t fields today
// (`sldi` goes through checkedI32OrCrash), so local addresses must stay within
// the non-negative int32_t range.
static FailureOr<size_t> checkedAlignTo(size_t value, size_t alignment, Operation* anchor, StringRef fieldName) {
if (alignment == 0)
return value;
size_t remainder = value % alignment;
if (remainder == 0)
return value;
return pim::checkedAdd(value, alignment - remainder, anchor, fieldName);
}
static void printMemoryOverflowDiagnostic(const MemoryValueKey& key,
size_t requestedSize,
size_t currentFirstAvailableAddress,
size_t alignedEndAddress) {
llvm::errs() << "PIM local memory allocation overflow\n";
llvm::errs() << "Requested allocation size: " << requestedSize << " bytes\n";
llvm::errs() << "Current firstAvailableAddress: " << currentFirstAvailableAddress << "\n";
llvm::errs() << "Aligned end address: " << alignedEndAddress << "\n";
llvm::errs() << "Address limit: " << kPimLocalMemoryAddressLimit << " (signed int32_t immediate range)\n";
if (key.lane)
llvm::errs() << "Lane: " << *key.lane << "\n";
llvm::errs() << "Value: ";
key.value.print(llvm::errs());
llvm::errs() << "\n";
llvm::errs() << "Value type: " << key.value.getType() << "\n";
if (Operation* definingOp = key.value.getDefiningOp()) {
llvm::errs() << "Defining op:\n";
definingOp->print(llvm::errs());
llvm::errs() << "\n";
}
}
static std::fstream openMemoryReport(bool enabled) {
if (std::string outputRoot = getOutputDir(); !outputRoot.empty()) {
for (std::string name : {std::string("memory_") + "coalescing_report.txt",
std::string("pim_memory_") + "liveness_report.txt",
std::string("pim_memory_") + "liveness_report.json",
std::string("pim_memory_") + "liveness_timeline.dot"})
sys::fs::remove(outputRoot + "/reports/" + name);
if (!enabled)
sys::fs::remove(outputRoot + "/reports/memory_report.txt");
}
return enabled ? openReportFile("memory_report") : std::fstream();
}
} // namespace
size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) {
size_t address = firstAvailableAddress;
Operation* anchor = getDiagnosticAnchor(key.value);
auto checkedEnd = pim::checkedAdd(address, size, anchor, "local memory end");
FailureOr<size_t> checkedAlignedEnd = failure();
if (succeeded(checkedEnd))
checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit
|| failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimLocalMemoryAddressLimit) {
printMemoryOverflowDiagnostic(
key,
size,
firstAvailableAddress,
succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit);
llvm_unreachable("PIM local memory allocation overflow");
}
firstAvailableAddress = *checkedAlignedEnd;
return address;
}
MemEntry* PimMemory::gatherMemEntry(mlir::Value value, std::optional<unsigned> lane) {
auto type = cast<ShapedType>(value.getType());
assert("Only static shape is supported" && type.hasStaticShape());
auto checkedAllocSize =
pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "memory allocation byte size");
if (failed(checkedAllocSize))
llvm_unreachable("Failed to compute checked allocation byte size");
PendingMemEntry pending;
pending.memEntry = {0, *checkedAllocSize};
pending.key = getMemoryValueKey(value, lane);
pending.reportKind = classifyMemoryReportKind(value);
return &memEntries.emplace_back(std::move(pending)).memEntry;
}
void PimMemory::allocateGatheredMemory() {
llvm::sort(memEntries, [](const PendingMemEntry& lhs, const PendingMemEntry& rhs) {
return lhs.memEntry.size > rhs.memEntry.size;
});
for (PendingMemEntry& pending : memEntries)
allocateMemoryForValue(pending.key, pending.memEntry, pending.reportKind);
memEntries.clear();
}
void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind) {
memEntry.address = allocateAddress(memEntry.size, key);
ownedMemEntriesMap[key] = memEntry;
globalMemEntriesMap[key] = memEntry;
switch (reportKind) {
case MemoryReportKind::Alloca:
case MemoryReportKind::Global:
case MemoryReportKind::Input:
++reportRow.hostObjectCount;
break;
case MemoryReportKind::None: break;
}
}
void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
SmallVector<mlir::Value> args;
for (mlir::Value arg : funcOp.getArguments()) {
gatherMemEntry(arg);
args.push_back(arg);
}
funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
if (!hasWeightAlways(getGlobalOp)) {
auto globalMemrefOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
if (globalMemrefOp.getName().starts_with("arg")) {
StringRef indexStr = globalMemrefOp.getName().substr(4);
int index = 0;
llvm::to_integer(indexStr, index, 10);
globalAliases.push_back({getGlobalOp.getResult(), args[index]});
}
auto [iter, inserted] = globalConstants.try_emplace(globalMemrefOp, getGlobalOp.getResult());
if (inserted)
gatherMemEntry(getGlobalOp.getResult());
else
globalAliases.push_back({getGlobalOp.getResult(), iter->second});
}
});
funcOp.walk([&](memref::AllocOp allocOp) {
if (!isInsidePimCoreLikeOp(allocOp))
gatherMemEntry(allocOp.getResult());
});
allocateGatheredMemory();
for (auto [alias, original] : globalAliases)
globalMemEntriesMap[getMemoryValueKey(alias)] = getMemEntry(getMemoryValueKey(original));
}
void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane) {
if (!localArenaSize) {
localArenaSize = plan.arenaSize;
reportRow.logicalLocalAllocationCount = plan.logicalAllocationCount;
reportRow.logicalLocalBytes = plan.logicalBytes;
}
else if (*localArenaSize != plan.arenaSize
|| reportRow.logicalLocalAllocationCount != plan.logicalAllocationCount
|| reportRow.logicalLocalBytes != plan.logicalBytes)
llvm_unreachable("inconsistent PIM local-memory plan across core-batch lanes");
for (const CompiledLocalMemoryEntry& entry : plan.entries) {
MemoryValueKey key = getMemoryValueKey(entry.value, lane);
ownedMemEntriesMap[key] = entry.memory;
globalMemEntriesMap[key] = entry.memory;
}
}
MemoryReportRow PimMemory::getReportRow() const {
MemoryReportRow row = reportRow;
row.hostBytes = firstAvailableAddress;
row.physicalLocalBytes = localArenaSize.value_or(0);
return row;
}
MemEntry PimMemory::getMemEntry(const MemoryValueKey& key) const {
auto iter = globalMemEntriesMap.find(key);
assert("Missing memEntry for value" && iter != globalMemEntriesMap.end());
return iter->second;
}
PimMemory& PimAcceleratorMemory::getOrCreateDeviceMem(size_t id) {
return deviceMem.try_emplace(id, memEntriesMap).first->second;
}
size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
const StaticValueKnowledge& knowledge,
std::optional<unsigned> lane) const {
value = resolveCachedAlias(value, knowledge);
FailureOr<ResolvedContiguousAddress> resolvedAddress = resolveContiguousAddress(value, knowledge);
if (failed(resolvedAddress)) {
auto compiledIt = compiledAddressExprs.find(value);
if (compiledIt == compiledAddressExprs.end()) {
auto compiledExpr = compileContiguousAddressExpr(value);
if (failed(compiledExpr)) {
errs() << "Failed to compile contiguous address for value: ";
value.print(errs());
errs() << " : " << value.getType();
errs() << "\n";
llvm_unreachable("Failed to compile contiguous address");
}
compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
}
resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
if (failed(resolvedAddress)) {
errs() << "Failed to evaluate contiguous address for value: ";
value.print(errs());
errs() << " : " << value.getType();
errs() << "\n";
if (auto* definingOp = value.getDefiningOp()) {
errs() << "Defining op:\n";
definingOp->print(errs());
errs() << "\n";
}
llvm_unreachable("Failed to resolve contiguous address");
}
}
MemoryValueKey key = getMemoryValueKey(resolvedAddress->base, lane);
auto iter = memEntriesMap.find(key);
if (iter == memEntriesMap.end()) {
errs() << "Missing mem entry for value: ";
resolvedAddress->base.print(errs());
errs() << "\n";
if (key.lane)
errs() << "Lane: " << *key.lane << "\n";
if (auto* definingOp = resolvedAddress->base.getDefiningOp()) {
errs() << "Defining op:\n";
definingOp->print(errs());
errs() << "\n";
}
llvm_unreachable("Missing mem entry");
}
size_t byteOffset = pim::checkedSizeOrCrash(resolvedAddress->byteOffset, "resolved PIM byte offset");
return pim::checkedAddOrCrash(iter->second.address, byteOffset, "resolved PIM address");
}
llvm::FailureOr<int64_t> PimAcceleratorMemory::getIndexValue(mlir::Value value,
const StaticValueKnowledge& knowledge) const {
value = resolveCachedAlias(value, knowledge);
auto compiledIt = compiledIndexExprs.find(value);
if (compiledIt == compiledIndexExprs.end()) {
auto compiledExpr = compileIndexExpr(value);
if (failed(compiledExpr))
return mlir::failure();
compiledIt = compiledIndexExprs.try_emplace(value, *compiledExpr).first;
}
return compiledIt->second.evaluate(knowledge);
}
PimAcceleratorMemory::PimAcceleratorMemory()
: hostMem(memEntriesMap), fileReport(openMemoryReport(pimMemoryReport == PimMemoryReportSummary)) {}
PimAcceleratorMemory::PimAcceleratorMemory(
const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport)
: memEntriesMap(initialMemEntries),
hostMem(memEntriesMap),
fileReport(enableReport ? openMemoryReport(true) : std::fstream()) {}
void PimAcceleratorMemory::reportHost() { hostReportRow = hostMem.getReportRow(); }
void PimAcceleratorMemory::recordCoreReport(size_t coreId, const MemoryReportRow& row) {
reportEntries.push_back({MemoryReportEntry::Kind::Core,
coreId,
{pim::checkedI32OrCrash(coreId, "memory report core id")},
row});
}
void PimAcceleratorMemory::recordBatchReport(uint64_t batchId,
ArrayRef<int32_t> coreIds,
const MemoryReportRow& perCoreRow) {
MemoryReportEntry entry;
entry.kind = MemoryReportEntry::Kind::Batch;
entry.id = batchId;
llvm::append_range(entry.coreIds, coreIds);
entry.row = perCoreRow;
reportEntries.push_back(std::move(entry));
}
void PimAcceleratorMemory::flushReport() {
if (!fileReport.is_open())
return;
llvm::raw_os_ostream os(fileReport);
struct ReportGroup {
MemoryReportRow row;
SmallVector<int32_t, 8> coreIds;
std::optional<uint64_t> batchId;
};
SmallVector<ReportGroup, 32> groups;
for (const MemoryReportEntry& entry : reportEntries) {
auto group = llvm::find_if(groups, [&](const ReportGroup& candidate) {
return candidate.row.logicalLocalBytes == entry.row.logicalLocalBytes
&& candidate.row.physicalLocalBytes == entry.row.physicalLocalBytes;
});
if (group == groups.end()) {
ReportGroup newGroup;
newGroup.row = entry.row;
llvm::append_range(newGroup.coreIds, entry.coreIds);
if (entry.kind == MemoryReportEntry::Kind::Batch)
newGroup.batchId = entry.id;
groups.push_back(std::move(newGroup));
continue;
}
llvm::append_range(group->coreIds, entry.coreIds);
group->batchId.reset();
}
llvm::sort(groups, [](const ReportGroup& lhs, const ReportGroup& rhs) {
if (lhs.row.physicalLocalBytes != rhs.row.physicalLocalBytes)
return lhs.row.physicalLocalBytes > rhs.row.physicalLocalBytes;
return lhs.row.logicalLocalBytes > rhs.row.logicalLocalBytes;
});
uint64_t logicalBytes = 0;
uint64_t physicalBytes = 0;
for (const ReportGroup& group : groups) {
logicalBytes += group.row.logicalLocalBytes * group.coreIds.size();
physicalBytes += group.row.physicalLocalBytes * group.coreIds.size();
}
uint64_t savedBytes = logicalBytes - physicalBytes;
double savedPercent = logicalBytes == 0 ? 0.0 : 100.0 * savedBytes / logicalBytes;
uint64_t largest = groups.empty() ? 0 : groups.front().row.physicalLocalBytes;
auto printLabel = [&](const ReportGroup& group) {
if (group.batchId)
os << "Batch " << *group.batchId << " — cores ";
else if (group.coreIds.size() == 1)
os << "Core ";
else
os << "Cores ";
printCompressedIntegerEntries(os, ArrayRef<int32_t>(group.coreIds));
};
os << "Summary\n";
os << " Host memory: " << formatReportMemory(hostReportRow ? hostReportRow->hostBytes : 0) << "\n";
os << " Weights memory: " << formatReportMemory(totalWeightBytes) << "\n";
os << " Local memory before reuse: " << formatReportMemory(logicalBytes) << "\n";
os << " Local memory after reuse: " << formatReportMemory(physicalBytes) << "\n";
os << " Saved local memory: " << formatReportMemory(savedBytes) << " ("
<< formatv("{0:F1}%", savedPercent) << ")\n";
os << " Largest core local memory: " << formatReportMemory(largest) << "\n";
if (!groups.empty()) {
os << " ";
printLabel(groups.front());
os << "\n";
}
os << "\nLargest core groups\n";
constexpr size_t kGroupLimit = 8;
for (const ReportGroup& group : ArrayRef(groups).take_front(kGroupLimit)) {
printLabel(group);
os << "\n";
uint64_t groupSaved = group.row.logicalLocalBytes - group.row.physicalLocalBytes;
double groupPercent = group.row.logicalLocalBytes == 0
? 0.0
: 100.0 * groupSaved / group.row.logicalLocalBytes;
if (group.coreIds.size() == 1) {
os << " Local memory: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
<< formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent)
<< ")\n";
}
else {
os << " Per core: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
<< formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent)
<< ")\n";
os << " Total after reuse: " << formatReportMemory(group.row.physicalLocalBytes * group.coreIds.size())
<< "\n";
}
}
if (groups.size() > kGroupLimit)
os << " " << groups.size() - kGroupLimit << " smaller groups omitted\n";
os.flush();
fileReport.close();
}
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;
}
void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instruction) const {
if (failed(instructionWriter.append(instruction)))
return;
if (coreJsonStream)
*coreJsonStream << json::Value(pim_binary::makeInstructionJson(instruction)) << ',';
updateScalarRegisterCache(instruction);
}
void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const {
switch (instruction.opcode) {
case pim_binary::Opcode::sldi: scalarRegisterValues[instruction.rd] = instruction.r2OrImm; break;
case pim_binary::Opcode::sld:
case pim_binary::Opcode::sadd:
case pim_binary::Opcode::ssub:
case pim_binary::Opcode::smul:
case pim_binary::Opcode::saddi:
case pim_binary::Opcode::smuli: scalarRegisterValues[instruction.rd].reset(); break;
default: break;
}
}
void PimCodeGen::genSetRegisterImmediate(uint8_t registerNumber, int32_t immediate) const {
if (scalarRegisterValues[registerNumber] == immediate)
return;
pim_binary::InstructionRecord instruction;
instruction.opcode = pim_binary::Opcode::sldi;
instruction.rd = registerNumber;
instruction.r2OrImm = immediate;
emitInstruction(instruction);
}
void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const {
genSetRegisterImmediate(pim::checkedU8OrCrash(registerNumber, "register number"),
pim::checkedI32OrCrash(immediate, "register immediate"));
}
void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const {
genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
}
void PimCodeGen::setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const {
genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
}
void PimCodeGen::setupRdRs1Rs2(
size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const {
genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
genSetRegisterImmediateUnsigned(2, pim::checkedAddOrCrash(rs2Address, rs2Offset, "rs2 address"));
}
void PimCodeGen::emitMemCopyOp(pim_binary::Opcode opcode,
size_t rdAddr,
size_t rdOffset,
size_t rs1Addr,
size_t rs1Offset,
size_t size,
StringRef sizeFieldName) const {
setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset);
pim_binary::InstructionRecord instruction;
instruction.opcode = opcode;
instruction.rd = 0;
instruction.r1 = 1;
instruction.generic1 = 0;
instruction.generic2 = 0;
instruction.generic3 = pim::checkedI32OrCrash(size, sizeFieldName);
emitInstruction(instruction);
}
void PimCodeGen::emitCommunicationOp(pim_binary::Opcode opcode, size_t bufferAddr, size_t coreId, size_t size) const {
setupRd(bufferAddr, 0);
pim_binary::InstructionRecord instruction;
instruction.opcode = opcode;
instruction.rd = 0;
instruction.r2OrImm = pim::checkedI32OrCrash(remapCoreId(coreId), "communication core id");
instruction.generic1 = 0;
instruction.generic2 = 0;
instruction.generic3 = pim::checkedI32OrCrash(size, "communication byte 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);
pim_binary::InstructionRecord instruction;
instruction.opcode = pim_binary::Opcode::mvmul;
instruction.rd = 0;
instruction.r1 = 1;
instruction.r2OrImm = 8;
instruction.generic1 = 0;
instruction.generic2 = pim::checkedI32OrCrash(groupId, "mvm group id");
emitInstruction(instruction);
}
void PimCodeGen::codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const {
auto deviceTargetOffset = indexOf(loadOp.getDeviceTargetOffset(), knowledge);
auto hostSourceOffset = indexOf(loadOp.getHostSourceOffset(), knowledge);
assert(succeeded(deviceTargetOffset) && succeeded(hostSourceOffset)
&& "pim.memcp_hd offsets must be statically resolvable during codegen");
emitMemCopyOp(pim_binary::Opcode::ld,
addressOf(loadOp.getDeviceTarget(), knowledge),
*deviceTargetOffset,
addressOf(loadOp.getHostSource(), knowledge),
*hostSourceOffset,
loadOp.getSize());
}
void PimCodeGen::codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const {
auto hostTargetOffset = indexOf(storeOp.getHostTargetOffset(), knowledge);
auto deviceSourceOffset = indexOf(storeOp.getDeviceSourceOffset(), knowledge);
assert(succeeded(hostTargetOffset) && succeeded(deviceSourceOffset)
&& "pim.memcp_dh offsets must be statically resolvable during codegen");
emitMemCopyOp(pim_binary::Opcode::st,
addressOf(storeOp.getHostTarget(), knowledge),
*hostTargetOffset,
addressOf(storeOp.getDeviceSource(), knowledge),
*deviceSourceOffset,
storeOp.getSize());
}
void PimCodeGen::codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledge& knowledge) const {
auto targetOffset = indexOf(lmvOp.getTargetOffset(), knowledge);
auto sourceOffset = indexOf(lmvOp.getSourceOffset(), knowledge);
assert(succeeded(targetOffset) && succeeded(sourceOffset)
&& "pim.memcp offsets must be statically resolvable during codegen");
emitMemCopyOp(pim_binary::Opcode::lmv,
addressOf(lmvOp.getTarget(), knowledge),
*targetOffset,
addressOf(lmvOp.getSource(), knowledge),
*sourceOffset,
lmvOp.getSize(),
"len");
}
void PimCodeGen::codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const {
auto sourceCoreId = indexOf(receiveOp.getSourceCoreId(), knowledge);
assert(succeeded(sourceCoreId) && "pim.receive source core id must be statically resolvable during codegen");
emitCommunicationOp(
pim_binary::Opcode::recv, addressOf(receiveOp.getOutputBuffer(), knowledge), *sourceCoreId, receiveOp.getSize());
}
void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const {
auto targetCoreId = indexOf(sendOp.getTargetCoreId(), knowledge);
assert(succeeded(targetCoreId) && "pim.send target core id must be statically resolvable during codegen");
emitCommunicationOp(
pim_binary::Opcode::send, addressOf(sendOp.getInput(), knowledge), *targetCoreId, sendOp.getSize());
}
void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const {
auto outputType = cast<ShapedType>(concatOp.getOutputBuffer().getType());
assert(outputType.hasStaticShape() && "concat codegen requires static output shape");
int64_t axis = concatOp.getAxis();
ArrayRef<int64_t> outputShape = outputType.getShape();
size_t elementSize = getElementTypeSizeInBytes(outputType.getElementType());
size_t outputAddr = addressOf(concatOp.getOutputBuffer(), knowledge);
size_t outerCount = 1;
for (int64_t dim = 0; dim < axis; ++dim)
outerCount *= static_cast<size_t>(outputShape[dim]);
size_t innerCount = 1;
for (size_t dim = static_cast<size_t>(axis) + 1; dim < outputShape.size(); ++dim)
innerCount *= static_cast<size_t>(outputShape[dim]);
size_t outputConcatDim = static_cast<size_t>(outputShape[axis]);
size_t concatOffset = 0;
for (mlir::Value input : concatOp.getInputs()) {
auto inputType = cast<ShapedType>(input.getType());
assert(inputType.hasStaticShape() && "concat codegen requires static input shapes");
size_t inputConcatDim = static_cast<size_t>(inputType.getDimSize(axis));
size_t blockSizeInBytes = inputConcatDim * innerCount * elementSize;
size_t inputAddr = addressOf(input, knowledge);
for (size_t outerIndex = 0; outerIndex < outerCount; ++outerIndex) {
size_t dstOffset = (outerIndex * outputConcatDim + concatOffset) * innerCount * elementSize;
size_t srcOffset = outerIndex * inputConcatDim * innerCount * elementSize;
emitMemCopyOp(pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len");
}
concatOffset += inputConcatDim;
}
}
template <typename MVMTy>
void PimCodeGen::codeGenMVMLikeOp(size_t mvmId,
MVMTy mvmLikeOp,
bool transposeMatrix,
const StaticValueKnowledge& knowledge) {
emitMvmOp(mvmId, addressOf(mvmLikeOp.getOutputBuffer(), knowledge), 0, addressOf(mvmLikeOp.getInput(), knowledge), 0);
// TODO: save weights somewhere (if transposeMatrix=true, transpose the weight matrix)
}
void PimCodeGen::emitBinaryVectorOp(pim_binary::Opcode opcode,
mlir::Value output,
mlir::Value lhs,
mlir::Value rhs,
size_t byteSize,
const StaticValueKnowledge& knowledge) const {
setupRdRs1Rs2(addressOf(output, knowledge), 0, addressOf(lhs, knowledge), 0, addressOf(rhs, knowledge), 0);
pim_binary::InstructionRecord instruction;
instruction.opcode = opcode;
instruction.rd = 0;
instruction.r1 = 1;
instruction.r2OrImm = 2;
instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size");
emitInstruction(instruction);
}
void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode,
mlir::Value output,
mlir::Value input,
size_t byteSize,
const StaticValueKnowledge& knowledge,
int32_t r2OrImm,
int32_t generic1) const {
setupRdRs1(addressOf(output, knowledge), 0, addressOf(input, knowledge), 0);
pim_binary::InstructionRecord instruction;
instruction.opcode = opcode;
instruction.rd = 0;
instruction.r1 = 1;
instruction.r2OrImm = r2OrImm;
instruction.generic1 = generic1;
instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size");
emitInstruction(instruction);
}
void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const {
auto srcAddr = addressOf(plan.source, knowledge);
auto dstAddr = addressOf(plan.destination, knowledge);
size_t maxElementOffset = plan.totalBytes == 0 ? 0 : plan.totalBytes - plan.elementBytes;
int32_t maxSourceAddress = pim::checkedI32OrCrash(
pim::checkedAddOrCrash(srcAddr, maxElementOffset, "transpose source address"), "transpose source address");
int32_t maxDestinationAddress =
pim::checkedI32OrCrash(pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"),
"transpose destination address");
(void) maxSourceAddress;
(void) maxDestinationAddress;
pim_binary::InstructionRecord copyInstruction;
copyInstruction.opcode = pim_binary::Opcode::lmv;
copyInstruction.rd = 0;
copyInstruction.r1 = 1;
size_t maxRunElements = static_cast<size_t>(std::numeric_limits<int32_t>::max()) / plan.elementBytes;
auto emitRun = [&](size_t sourceStart, size_t destinationStart, size_t runLength) {
while (runLength != 0) {
size_t chunkElements = std::min(runLength, maxRunElements);
// totalBytes was checked when the plan was compiled, so these bounded
// products cannot overflow.
size_t sourceOffset = sourceStart * plan.elementBytes;
size_t destinationOffset = destinationStart * plan.elementBytes;
size_t byteSize = chunkElements * plan.elementBytes;
assert(sourceOffset <= plan.totalBytes - byteSize && destinationOffset <= plan.totalBytes - byteSize);
genSetRegisterImmediate(0, static_cast<int32_t>(dstAddr + destinationOffset));
genSetRegisterImmediate(1, static_cast<int32_t>(srcAddr + sourceOffset));
copyInstruction.generic3 = static_cast<int32_t>(byteSize);
emitInstruction(copyInstruction);
sourceStart += chunkElements;
destinationStart += chunkElements;
runLength -= chunkElements;
}
};
if (plan.storagePreserving) {
emitRun(0, 0, plan.totalElements);
return;
}
size_t rank = plan.sourceShape.size();
SmallVector<size_t> sourceIndices(rank, 0);
size_t destinationFlat = 0;
size_t runSourceStart = 0;
size_t runDestinationStart = 0;
size_t runLength = 0;
for (size_t sourceFlat = 0; sourceFlat < plan.totalElements; ++sourceFlat) {
if (runLength != 0 && destinationFlat != runDestinationStart + runLength) {
emitRun(runSourceStart, runDestinationStart, runLength);
runSourceStart = sourceFlat;
runDestinationStart = destinationFlat;
runLength = 0;
}
if (runLength == 0) {
runSourceStart = sourceFlat;
runDestinationStart = destinationFlat;
}
++runLength;
if (runLength == maxRunElements) {
emitRun(runSourceStart, runDestinationStart, runLength);
runLength = 0;
}
if (sourceFlat + 1 == plan.totalElements)
break;
for (size_t sourceDim = rank; sourceDim-- > 0;) {
destinationFlat += plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]];
if (++sourceIndices[sourceDim] < plan.sourceShape[sourceDim])
break;
sourceIndices[sourceDim] = 0;
destinationFlat -= plan.destinationRewinds[sourceDim];
}
}
if (runLength != 0)
emitRun(runSourceStart, runDestinationStart, runLength);
}
static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
SmallVector<Operation*> coreLikeOps;
for (Operation& op : funcOp.getBody().front())
if (dyn_cast<pim::PimCoreOp>(&op) || dyn_cast<pim::PimCoreBatchOp>(&op))
coreLikeOps.push_back(&op);
return coreLikeOps;
}
static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLikeOp) {
CompiledCoreMemoryPlan plan;
auto arenaAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemorySizeAttrName);
if (!arenaAttr || arenaAttr.getInt() < 0
|| static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
coreLikeOp->emitError("requires a valid pim.local_memory_size attribute before codegen");
return failure();
}
plan.arenaSize = static_cast<size_t>(arenaAttr.getInt());
bool hasFailure = false;
coreLikeOp->walk([&](memref::AllocOp allocOp) {
if (hasFailure)
return;
auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
if (!addressAttr || addressAttr.getInt() < 0) {
allocOp.emitOpError("requires a non-negative pim.local_memory_address attribute before codegen");
hasFailure = true;
return;
}
auto checkedSize = pim::getCheckedShapedTypeSizeInBytes(
cast<ShapedType>(allocOp.getType()), allocOp, "planned local allocation byte size");
if (failed(checkedSize)) {
hasFailure = true;
return;
}
size_t address = static_cast<size_t>(addressAttr.getInt());
if (address > plan.arenaSize || *checkedSize > plan.arenaSize - address) {
allocOp.emitOpError("has a planned local-memory range outside its core arena");
hasFailure = true;
return;
}
plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
auto logicalBytes = pim::checkedAdd(
static_cast<size_t>(plan.logicalBytes), static_cast<size_t>(*checkedSize), allocOp, "logical local bytes");
if (failed(logicalBytes)) {
hasFailure = true;
return;
}
plan.logicalBytes = *logicalBytes;
++plan.logicalAllocationCount;
});
if (hasFailure)
return failure();
return plan;
}
struct CoreEmissionResult {
static constexpr size_t kMaxStoredCodegenDiagnostics = 8;
struct DiagnosticRecord {
Operation* op = nullptr;
std::string message;
};
OnnxMlirCompilerErrorCodes status = CompilerSuccess;
MemoryReportRow reportRow;
llvm::SmallVector<ResolvedWeightView, 8> usedWeights;
llvm::SmallVector<DiagnosticRecord, kMaxStoredCodegenDiagnostics> diagnostics;
size_t diagnosticCount = 0;
void recordDiagnostic(Operation* op, StringRef message) {
++diagnosticCount;
if (diagnostics.size() < kMaxStoredCodegenDiagnostics)
diagnostics.push_back({op, message.str()});
}
};
static StaticValueKnowledge seedCoreCodegenKnowledge(pim::PimCoreOp coreOp) {
StaticValueKnowledge knowledge;
for (auto [index, weight] : llvm::enumerate(coreOp.getWeights()))
knowledge.aliases[coreOp.getWeightArgument(index)] = weight;
return knowledge;
}
static StaticValueKnowledge seedCoreBatchCodegenKnowledge(pim::PimCoreBatchOp coreBatchOp, unsigned lane) {
StaticValueKnowledge knowledge;
knowledge.indexValues[coreBatchOp.getLaneArgument()] = lane;
for (auto [index, weight] : llvm::enumerate(coreBatchOp.getWeights()))
knowledge.aliases[coreBatchOp.getWeightArgument(index)] = weight;
for (auto [index, input] : llvm::enumerate(coreBatchOp.getInputs()))
knowledge.aliases[coreBatchOp.getInputArgument(index)] = input;
return knowledge;
}
template <typename MapTy>
class ScopedMapBindings {
using KeyTy = typename MapTy::key_type;
using ValueTy = typename MapTy::mapped_type;
MapTy& map;
llvm::SmallVector<std::pair<KeyTy, std::optional<ValueTy>>, 8> savedEntries;
public:
explicit ScopedMapBindings(MapTy& map)
: map(map) {}
void bind(const KeyTy& key, const ValueTy& value) {
auto it = map.find(key);
if (it == map.end())
savedEntries.emplace_back(key, std::nullopt);
else
savedEntries.emplace_back(key, it->second);
map[key] = value;
}
~ScopedMapBindings() {
for (auto it = savedEntries.rbegin(); it != savedEntries.rend(); ++it)
if (it->second)
map[it->first] = *it->second;
else
map.erase(it->first);
}
};
static LogicalResult executeCompiledCorePlan(
const llvm::SmallVectorImpl<CompiledCoreNode>& plan,
PimCodeGen& coreCodeGen,
StaticValueKnowledge& knowledge,
llvm::function_ref<llvm::FailureOr<unsigned>(pim::PimVMMOp, const StaticValueKnowledge&)> resolveWeightSlot) {
for (const CompiledCoreNode& node : plan) {
if (node.kind == CompiledCoreNode::Kind::Loop) {
auto lowerBound = node.lowerBound.evaluate(knowledge);
auto upperBound = node.upperBound.evaluate(knowledge);
auto step = node.step.evaluate(knowledge);
auto forOp = cast<mlir::scf::ForOp>(node.op);
if (failed(lowerBound) || failed(upperBound) || failed(step) || *step <= 0) {
forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
return failure();
}
llvm::SmallVector<mlir::Value> iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end());
for (int64_t inductionValue = *lowerBound; inductionValue < *upperBound; inductionValue += *step) {
ScopedMapBindings<decltype(knowledge.indexValues)> indexBindings(knowledge.indexValues);
ScopedMapBindings<decltype(knowledge.aliases)> aliasBindings(knowledge.aliases);
indexBindings.bind(forOp.getInductionVar(), inductionValue);
for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), iterValues))
aliasBindings.bind(iterArg, iterValue);
if (failed(executeCompiledCorePlan(*node.loopBody, coreCodeGen, knowledge, resolveWeightSlot)))
return failure();
auto yieldOp = cast<mlir::scf::YieldOp>(forOp.getRegion().front().getTerminator());
for (auto [index, yieldedValue] : llvm::enumerate(yieldOp.getOperands()))
iterValues[index] = resolveLoopCarriedAlias(yieldedValue, knowledge);
}
continue;
}
if (node.kind == CompiledCoreNode::Kind::If) {
auto condition = node.condition.evaluate(knowledge);
auto ifOp = cast<mlir::scf::IfOp>(node.op);
if (failed(condition)) {
ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
return failure();
}
const auto& selectedBody = *condition != 0 ? node.thenBody : node.elseBody;
if (selectedBody && failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge, resolveWeightSlot)))
return failure();
continue;
}
if (node.kind == CompiledCoreNode::Kind::IndexSwitch) {
auto selector = node.condition.evaluate(knowledge);
auto switchOp = cast<mlir::scf::IndexSwitchOp>(node.op);
if (failed(selector)) {
switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
return failure();
}
const llvm::SmallVectorImpl<CompiledCoreNode>* selectedBody = node.defaultBody.get();
mlir::Region* selectedRegion = &switchOp.getDefaultRegion();
for (auto [index, caseValue] : llvm::enumerate(node.caseValues))
if (caseValue == *selector) {
selectedBody = node.caseBodies[index].get();
selectedRegion = &switchOp.getCaseRegions()[index];
break;
}
if (failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge, resolveWeightSlot)))
return failure();
auto yield = cast<mlir::scf::YieldOp>(selectedRegion->front().getTerminator());
for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
continue;
}
auto emitBinary = [&](auto op, pim_binary::Opcode opcode) {
coreCodeGen.emitBinaryVectorOp(opcode,
op.getOutputBuffer(),
op.getLhs(),
op.getRhs(),
getVectorByteSizeOrCrash(cast<ShapedType>(op.getLhs().getType())),
knowledge);
};
auto emitUnary = [&](auto op, pim_binary::Opcode opcode, int32_t r2OrImm, int32_t generic1) {
coreCodeGen.emitUnaryVectorOp(opcode,
op.getOutputBuffer(),
op.getInput(),
getVectorByteSizeOrCrash(cast<ShapedType>(op.getInput().getType())),
knowledge,
r2OrImm,
generic1);
};
switch (node.opKind) {
case CompiledCoreOpKind::Load:
coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);
break;
case CompiledCoreOpKind::Store:
coreCodeGen.codeGenStoreOp(cast<pim::PimMemCopyDevToHostOp>(node.op), knowledge);
break;
case CompiledCoreOpKind::Lmv: coreCodeGen.codeGenLmvOp(cast<pim::PimMemCopyOp>(node.op), knowledge); break;
case CompiledCoreOpKind::Receive: coreCodeGen.codeGenReceiveOp(cast<pim::PimReceiveOp>(node.op), knowledge); break;
case CompiledCoreOpKind::Send: coreCodeGen.codeGenSendOp(cast<pim::PimSendOp>(node.op), knowledge); break;
case CompiledCoreOpKind::Concat: coreCodeGen.codeGenConcatOp(cast<pim::PimConcatOp>(node.op), knowledge); break;
case CompiledCoreOpKind::Vmm:
if (auto weightSlot = resolveWeightSlot(cast<pim::PimVMMOp>(node.op), knowledge); succeeded(weightSlot))
coreCodeGen.codeGenMVMLikeOp<pim::PimVMMOp>(*weightSlot, cast<pim::PimVMMOp>(node.op), true, knowledge);
else
return failure();
break;
case CompiledCoreOpKind::Transpose: coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break;
case CompiledCoreOpKind::VVAdd: emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
case CompiledCoreOpKind::VVSub: emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
case CompiledCoreOpKind::VVMul: emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
case CompiledCoreOpKind::VVMax: emitBinary(cast<pim::PimVVMaxOp>(node.op), pim_binary::Opcode::vvmax); break;
case CompiledCoreOpKind::VVDMul: emitBinary(cast<pim::PimVVDMulOp>(node.op), pim_binary::Opcode::vvdmul); break;
case CompiledCoreOpKind::VAvg: emitUnary(cast<pim::PimVAvgOp>(node.op), pim_binary::Opcode::vavg, 1, 1); break;
case CompiledCoreOpKind::VRelu: emitUnary(cast<pim::PimVReluOp>(node.op), pim_binary::Opcode::vrelu, 0, 0); break;
case CompiledCoreOpKind::VTanh: emitUnary(cast<pim::PimVTanhOp>(node.op), pim_binary::Opcode::vtanh, 0, 0); break;
case CompiledCoreOpKind::VSigm: emitUnary(cast<pim::PimVSigmOp>(node.op), pim_binary::Opcode::vsigm, 0, 0); break;
case CompiledCoreOpKind::VSoftmax:
emitUnary(cast<pim::PimVSoftmaxOp>(node.op), pim_binary::Opcode::vsoftmax, 0, 0);
break;
}
}
return success();
}
static LogicalResult executeCompiledCoreProgram(
const CompiledCoreProgram& program,
PimCodeGen& coreCodeGen,
const StaticValueKnowledge& initialKnowledge,
llvm::function_ref<llvm::FailureOr<unsigned>(pim::PimVMMOp, const StaticValueKnowledge&)> resolveWeightSlot) {
StaticValueKnowledge knowledge = initialKnowledge;
return executeCompiledCorePlan(program.nodes, coreCodeGen, knowledge, resolveWeightSlot);
}
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(getMemoryValueKey(getGlobalOp.getResult()));
if (it != memory.memEntriesMap.end())
materializedHostGlobals[targetGlobal] = it->second;
});
return materializedHostGlobals;
}
template <typename CoreLikeOpTy>
static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp,
ModuleOp moduleOp,
const SmallDenseMap<memref::GlobalOp, MemEntry, 16>& materializedHostGlobals,
PimAcceleratorMemory& memory) {
coreLikeOp.walk([&](memref::GetGlobalOp getGlobalOp) {
MemoryValueKey key = getMemoryValueKey(getGlobalOp.getResult());
if (hasWeightAlways(getGlobalOp) || memory.memEntriesMap.contains(key))
return;
auto targetGlobal = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
if (!targetGlobal)
return;
auto it = materializedHostGlobals.find(targetGlobal);
if (it != materializedHostGlobals.end())
memory.memEntriesMap[key] = it->second;
});
}
static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
std::error_code errorCode;
raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
if (errorCode) {
errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
return InvalidOutputFileAccess;
}
PimInstructionWriter instructionWriter(coreBinaryStream);
if (failed(instructionWriter.finalize())) {
coreBinaryStream.close();
errs() << "Error while writing core file `" << outputCorePath << "`\n";
return InvalidOutputFileAccess;
}
coreBinaryStream.close();
if (coreBinaryStream.has_error()) {
coreBinaryStream.clear_error();
errs() << "Error while closing core file `" << outputCorePath << "`\n";
return InvalidOutputFileAccess;
}
if (!pimEmitJson.getValue())
return CompilerSuccess;
std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
errorCode = std::error_code();
raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
if (errorCode) {
errs() << "Error while opening core json file `" << outputCoreJsonPath << "`: " << errorCode.message() << '\n';
return InvalidOutputFileAccess;
}
coreJsonStream << "[]";
coreJsonStream.close();
return CompilerSuccess;
}
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';
return InvalidOutputFileAccess;
}
}
auto entryFunc = getPimEntryFunc(moduleOp);
if (failed(entryFunc))
return CompilerFailure;
auto funcOp = *entryFunc;
PimAcceleratorMemory memory;
memory.hostMem.allocateHost(moduleOp, funcOp);
memory.reportHost();
if (auto err = writeMemoryBinary(moduleOp, funcOp, memory, outputDirPath))
return err;
json::Object xbarsPerArrayGroup;
size_t maxCoreId = 0;
uint64_t nextBatchReportId = 0;
SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
SmallDenseMap<memref::GlobalOp, MemEntry, 16> materializedHostGlobals =
collectMaterializedHostGlobals(moduleOp, funcOp, memory);
llvm::DenseMap<Operation*, std::unique_ptr<CompiledCoreProgram>> compiledPrograms;
llvm::DenseMap<Operation*, std::unique_ptr<CompiledCoreMemoryPlan>> memoryPlans;
for (Operation* op : coreLikeOps) {
auto program = std::make_unique<CompiledCoreProgram>();
if (failed(compileCoreProgram(op, *program)))
return CompilerFailure;
compiledPrograms.try_emplace(op, std::move(program));
auto memoryPlan = compileCoreMemoryPlan(op);
if (failed(memoryPlan))
return CompilerFailure;
memoryPlans.try_emplace(op, std::make_unique<CompiledCoreMemoryPlan>(std::move(*memoryPlan)));
}
auto getCompiledProgram = [&](Operation* op) {
auto it = compiledPrograms.find(op);
assert(it != compiledPrograms.end() && "missing compiled PIM core program");
return it->second.get();
};
auto getMemoryPlan = [&](Operation* op) {
auto it = memoryPlans.find(op);
assert(it != memoryPlans.end() && "missing PIM core memory plan");
return it->second.get();
};
llvm::DenseMap<size_t, size_t> emittedCoreIds;
size_t nextEmittedCoreId = 0;
for (Operation* op : coreLikeOps) {
if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
if (!emittedCoreIds.contains(originalCoreId))
emittedCoreIds[originalCoreId] = nextEmittedCoreId++;
continue;
}
auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
auto batchCoreIds = getBatchCoreIds(coreBatchOp);
for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane) {
size_t originalCoreId = static_cast<size_t>(batchCoreIds[lane]);
if (!emittedCoreIds.contains(originalCoreId))
emittedCoreIds[originalCoreId] = nextEmittedCoreId++;
}
}
SmallVector<CoreEmissionJob> jobs;
SmallVector<SmallVector<size_t>> batchJobIndices;
for (Operation* op : coreLikeOps) {
if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
CoreEmissionJob job;
job.coreLikeOp = coreOp;
job.program = getCompiledProgram(op);
job.memoryPlan = getMemoryPlan(op);
job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
jobs.push_back(std::move(job));
continue;
}
auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
auto batchCoreIds = getBatchCoreIds(coreBatchOp);
llvm::DenseMap<size_t, SmallVector<unsigned>> lanesByCoreId;
for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
lanesByCoreId[static_cast<size_t>(batchCoreIds[lane])].push_back(lane);
SmallVector<size_t> jobIndices;
SmallVector<size_t> orderedOriginalCoreIds = llvm::to_vector(lanesByCoreId.keys());
llvm::sort(orderedOriginalCoreIds,
[&](size_t lhs, size_t rhs) { return emittedCoreIds.lookup(lhs) < emittedCoreIds.lookup(rhs); });
for (size_t originalCoreId : orderedOriginalCoreIds) {
CoreEmissionJob job;
job.coreLikeOp = coreBatchOp;
job.program = getCompiledProgram(op);
job.memoryPlan = getMemoryPlan(op);
job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
job.lanes = lanesByCoreId.lookup(originalCoreId);
job.batchReportId = nextBatchReportId;
jobIndices.push_back(jobs.size());
jobs.push_back(std::move(job));
}
batchJobIndices.push_back(std::move(jobIndices));
++nextBatchReportId;
}
auto linkCoreWeights =
[&](size_t coreId, ArrayRef<std::string> weightFiles, json::Array& xbarsPerGroup) -> OnnxMlirCompilerErrorCodes {
auto coreWeightsDirPath = outputDirPath + "/core_" + std::to_string(coreId);
if (auto error = sys::fs::create_directory(coreWeightsDirPath); error && error != std::errc::file_exists) {
errs() << "Error creating core directory: " << coreWeightsDirPath << ": " << error.message() << '\n';
return InvalidOutputFileAccess;
}
for (auto [slot, fileName] : llvm::enumerate(weightFiles)) {
xbarsPerGroup.push_back(1);
std::string sourcePath = outputDirPath + "/weights/" + fileName;
std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin";
sys::fs::remove(targetPath);
if (auto error = sys::fs::create_link(sourcePath, targetPath)) {
errs() << "Error creating link file: " << sourcePath << " to " << targetPath << "\nError:" << error.message()
<< '\n';
return InvalidOutputFileAccess;
}
}
return CompilerSuccess;
};
auto emitJob = [&](const CoreEmissionJob& job) -> CoreEmissionResult {
CoreEmissionResult result;
PimAcceleratorMemory jobMemory(memory.memEntriesMap, false);
llvm::SmallVector<ResolvedWeightView, 8> usedWeights;
auto resolveWeightSlot = [&](pim::PimVMMOp vmmOp,
const StaticValueKnowledge& knowledge) -> llvm::FailureOr<unsigned> {
auto weightView = onnx_mlir::resolveWeightView(job.coreLikeOp, vmmOp.getWeight(), knowledge);
if (failed(weightView)) {
std::string message;
llvm::raw_string_ostream os(message);
os << "requires a statically resolvable dense global weight view during PIM codegen; weight="
<< vmmOp.getWeight() << " type=" << vmmOp.getWeight().getType();
result.recordDiagnostic(vmmOp, os.str());
return failure();
}
if (weightView->shape.size() != 2) {
std::string message;
llvm::raw_string_ostream os(message);
os << "requires a rank-2 matrix weight view during PIM codegen; resolved shape=[";
llvm::interleaveComma(weightView->shape, os);
os << "] weight=" << vmmOp.getWeight() << " type=" << vmmOp.getWeight().getType();
result.recordDiagnostic(vmmOp, os.str());
return failure();
}
if (auto it = llvm::find(usedWeights, *weightView); it != usedWeights.end())
return static_cast<unsigned>(std::distance(usedWeights.begin(), it));
usedWeights.push_back(*weightView);
return static_cast<unsigned>(usedWeights.size() - 1);
};
std::error_code errorCode;
auto outputCorePath = outputDirPath + "/core_" + std::to_string(job.emittedCoreId) + ".pim";
raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
if (errorCode) {
errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
result.status = InvalidOutputFileAccess;
return result;
}
std::unique_ptr<raw_fd_ostream> coreJsonStream;
if (pimEmitJson.getValue()) {
std::string outputCoreJsonPath = outputDirPath + "/core_" + std::to_string(job.emittedCoreId) + ".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';
result.status = InvalidOutputFileAccess;
return result;
}
*coreJsonStream << '[';
}
PimInstructionWriter instructionWriter(coreBinaryStream);
PimCodeGen coreCodeGen(jobMemory, instructionWriter, coreJsonStream.get(), emittedCoreIds);
auto finalizeInstructions = [&]() {
bool succeeded = mlir::succeeded(instructionWriter.finalize());
coreBinaryStream.close();
if (coreBinaryStream.has_error()) {
coreBinaryStream.clear_error();
succeeded = false;
}
return succeeded;
};
if (auto coreOp = dyn_cast<pim::PimCoreOp>(job.coreLikeOp)) {
aliasMaterializedHostGlobals(coreOp, moduleOp, materializedHostGlobals, jobMemory);
auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
deviceMemory.allocateCore(*job.memoryPlan);
StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp);
if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
(void) finalizeInstructions();
result.status = CompilerFailure;
return result;
}
result.reportRow = deviceMemory.getReportRow();
result.usedWeights = std::move(usedWeights);
}
else {
auto coreBatchOp = cast<pim::PimCoreBatchOp>(job.coreLikeOp);
aliasMaterializedHostGlobals(coreBatchOp, moduleOp, materializedHostGlobals, jobMemory);
auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
for (unsigned lane : job.lanes) {
StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane);
deviceMemory.allocateCore(*job.memoryPlan, lane);
coreCodeGen.setBatchLane(lane);
if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
(void) finalizeInstructions();
result.status = CompilerFailure;
return result;
}
}
result.reportRow = deviceMemory.getReportRow();
result.usedWeights = std::move(usedWeights);
}
if (!finalizeInstructions()) {
result.status = InvalidOutputFileAccess;
return result;
}
if (coreJsonStream) {
coreJsonStream->seek(coreJsonStream->tell() - 1);
*coreJsonStream << ']';
coreJsonStream->close();
}
return result;
};
std::vector<CoreEmissionResult> jobResults(jobs.size());
mlir::parallelFor(
moduleOp.getContext(), 0, jobs.size(), [&](size_t index) { jobResults[index] = emitJob(jobs[index]); });
pim::CappedDiagnosticReporter diagnostics;
Operation* summaryAnchor = nullptr;
for (const CoreEmissionResult& result : jobResults) {
if (!summaryAnchor && !result.diagnostics.empty())
summaryAnchor = result.diagnostics.front().op;
for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics)
diagnostics.report(diagnostic.op, [&](Operation* op) { op->emitError() << diagnostic.message; });
size_t unreportedCount = result.diagnosticCount - result.diagnostics.size();
diagnostics.noteFailures(static_cast<int64_t>(unreportedCount));
}
if (diagnostics.hasFailure())
diagnostics.emitSuppressedSummary(summaryAnchor ? summaryAnchor : moduleOp.getOperation(),
"PIM codegen diagnostic(s)");
for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex)
if (jobResults[jobIndex].status != CompilerSuccess)
return jobResults[jobIndex].status;
if (jobs.empty()) {
if (auto err = emitEmptyCoreArtifacts(outputDirPath, 0))
return err;
xbarsPerArrayGroup["core0"] = json::Array {};
memory.recordCoreReport(0, MemoryReportRow {});
}
llvm::SmallVector<WeightFileRequest, 8> weightRequests;
weightRequests.reserve(jobs.size());
for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
WeightFileRequest request;
request.coreId = jobs[jobIndex].emittedCoreId;
request.weights = jobResults[jobIndex].usedWeights;
weightRequests.push_back(std::move(request));
}
auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath);
memory.setTotalWeightBytes(weightEmission.totalWeightBytes);
auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName;
for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
const CoreEmissionJob& job = jobs[jobIndex];
const CoreEmissionResult& result = jobResults[jobIndex];
json::Array xbarsPerGroup;
if (auto coreOp = dyn_cast<pim::PimCoreOp>(job.coreLikeOp)) {
if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup))
return err;
xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
memory.recordCoreReport(job.emittedCoreId, result.reportRow);
continue;
}
}
for (const SmallVector<size_t>& group : batchJobIndices) {
SmallVector<int32_t> reportedCoreIds;
std::optional<MemoryReportRow> batchPerCoreRow;
for (size_t jobIndex : group) {
const CoreEmissionJob& job = jobs[jobIndex];
const CoreEmissionResult& result = jobResults[jobIndex];
json::Array xbarsPerGroup;
if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup))
return err;
xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
reportedCoreIds.push_back(pim::checkedI32OrCrash(job.emittedCoreId, "batch report core id"));
if (!batchPerCoreRow)
batchPerCoreRow = result.reportRow;
else if (!(*batchPerCoreRow == result.reportRow))
llvm_unreachable("one PIM core batch produced inconsistent per-core memory reports");
}
uint64_t batchReportId = jobs[group.front()].batchReportId.value_or(0);
memory.recordBatchReport(
batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {}));
}
maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1;
memory.flushReport();
return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
}
Reference in New Issue View Git Blame Copy Permalink