nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

blazingly fasterer
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source
This commit is contained in:
NiccoloN
2026-07-21 15:43:35 +02:00
parent a893d23a74
commit 3e468b58c8
37 changed files with 1119 additions and 933 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Compiler/CMakeLists.txt
+1 -2
View File
@@ -18,7 +18,6 @@ add_pim_library(OMPimCompilerUtils
PimArtifactWriter.cpp
PimCodeGen.cpp
PimCoreProgram.cpp
PimMemoryLiveness.cpp
PimWeightEmitter.cpp
EXCLUDE_FROM_OM_LIBS
@@ -31,8 +30,8 @@ add_pim_library(OMPimCompilerUtils
OMPimCompilerOptions
OMPimCommon
OMPimBufferization
OMPimMemoryCoalescing
OMPimHostConstantFolding
OMPimLocalMemoryPlanning
OMPimVerification
OMPimPasses
OMONNXToSpatial
src/PIM/Compiler/PimCodeGen.cpp
+160 -157
View File
@@ -15,6 +15,7 @@
#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"
@@ -31,8 +32,8 @@
#include "Common/IR/CompactAsmUtils.hpp"
#include "Common/PimCommon.hpp"
#include "Common/Support/Diagnostics.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"
@@ -43,7 +44,6 @@
#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/PimMemoryLiveness.hpp"
#include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
@@ -110,8 +110,6 @@ static Operation* getDiagnosticAnchor(mlir::Value value) {
// 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 constexpr size_t kPimAddressLimit = static_cast<size_t>(std::numeric_limits<int32_t>::max());
static FailureOr<size_t> checkedAlignTo(size_t value, size_t alignment, Operation* anchor, StringRef fieldName) {
if (alignment == 0)
return value;
@@ -129,7 +127,7 @@ static void printMemoryOverflowDiagnostic(const MemoryValueKey& key,
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: " << kPimAddressLimit << " (signed int32_t immediate range)\n";
llvm::errs() << "Address limit: " << kPimLocalMemoryAddressLimit << " (signed int32_t immediate range)\n";
if (key.lane)
llvm::errs() << "Lane: " << *key.lane << "\n";
llvm::errs() << "Value: ";
@@ -152,10 +150,13 @@ size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) {
FailureOr<size_t> checkedAlignedEnd = failure();
if (succeeded(checkedEnd))
checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
if (address > kPimAddressLimit || failed(checkedEnd) || *checkedEnd > kPimAddressLimit
|| failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimAddressLimit) {
if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit
|| failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimLocalMemoryAddressLimit) {
printMemoryOverflowDiagnostic(
key, size, firstAvailableAddress, succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
key,
size,
firstAvailableAddress,
succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit);
llvm_unreachable("PIM local memory allocation overflow");
}
firstAvailableAddress = *checkedAlignedEnd;
@@ -202,15 +203,6 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
}
}
PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key) {
PhysicalSlotInfo slot;
slot.id = nextPhysicalSlotId++;
slot.address = allocateAddress(slotSize, key);
slot.size = slotSize;
localPhysicalSlots.push_back(slot);
return slot;
}
void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
@@ -249,71 +241,17 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
globalMemEntriesMap[getMemoryValueKey(alias)] = getMemEntry(getMemoryValueKey(original));
}
void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
auto intervals = buildLocalAllocIntervals(op, lane, pimMemoryReport == PimMemoryReportFull);
SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);
SmallVector<size_t> slotOrder(plannedSlots.size());
std::iota(slotOrder.begin(), slotOrder.end(), 0);
llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) {
const PlannedPhysicalSlot& lhs = plannedSlots[lhsIndex];
const PlannedPhysicalSlot& rhs = plannedSlots[rhsIndex];
if (lhs.requiredSize != rhs.requiredSize)
return lhs.requiredSize > rhs.requiredSize;
return lhs.id < rhs.id;
});
SmallVector<bool, 16> usedExistingSlots(localPhysicalSlots.size(), false);
for (size_t slotIndex : slotOrder) {
PlannedPhysicalSlot& slot = plannedSlots[slotIndex];
size_t bestExistingIndex = std::numeric_limits<size_t>::max();
auto bestKey = std::tuple<size_t, size_t, size_t>(
std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
for (size_t existingIndex = 0; existingIndex < localPhysicalSlots.size(); ++existingIndex) {
if (usedExistingSlots[existingIndex])
continue;
const PhysicalSlotInfo& existingSlot = localPhysicalSlots[existingIndex];
if (existingSlot.size < slot.requiredSize)
continue;
auto candidateKey =
std::tuple<size_t, size_t, size_t>(existingSlot.size - slot.requiredSize, existingSlot.size, existingSlot.id);
if (candidateKey < bestKey) {
bestKey = candidateKey;
bestExistingIndex = existingIndex;
}
}
if (bestExistingIndex != std::numeric_limits<size_t>::max()) {
const PhysicalSlotInfo& existingSlot = localPhysicalSlots[bestExistingIndex];
slot.id = existingSlot.id;
slot.address = existingSlot.address;
slot.size = existingSlot.size;
usedExistingSlots[bestExistingIndex] = true;
}
else {
PhysicalSlotInfo newSlot = allocatePhysicalSlot(slot.requiredSize, intervals[slot.intervalIndices.front()].key);
slot.id = newSlot.id;
slot.address = newSlot.address;
slot.size = newSlot.size;
usedExistingSlots.push_back(true);
}
for (size_t intervalIndex : slot.intervalIndices) {
LocalAllocInterval& interval = intervals[intervalIndex];
interval.physicalSlotId = slot.id;
interval.assignedAddress = slot.address;
interval.physicalSlotSize = slot.size;
MemEntry memEntry {slot.address, interval.size};
ownedMemEntriesMap[interval.key] = memEntry;
globalMemEntriesMap[interval.key] = memEntry;
}
void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane) {
if (localPhysicalSlots.empty()) {
llvm::append_range(localPhysicalSlots, plan.slots);
reportRow.logicalAllocaBytes = plan.logicalBytes;
reportRow.fallbackIntervals = plan.fallbackIntervals;
reportRow.nestedSingleUseIntervals = plan.nestedSingleUseIntervals;
}
if (pimMemoryReport != PimMemoryReportNone) {
MemoryPlanArtifacts artifacts =
buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport);
livenessArtifacts += artifacts;
for (const CompiledLocalMemoryEntry& entry : plan.entries) {
MemoryValueKey key = getMemoryValueKey(entry.value, lane);
ownedMemEntriesMap[key] = entry.memory;
globalMemEntriesMap[key] = entry.memory;
}
}
@@ -475,13 +413,41 @@ void PimAcceleratorMemory::flushReport() {
uint64_t totalGlobalMemory = hostReportRow.has_value() ? hostReportRow->sizeGlobal : 0;
uint64_t totalWeightsMemory = totalWeightBytes;
uint64_t totalCoresMemory = 0;
for (const MemoryReportEntry& entry : reportEntries)
uint64_t totalLogicalCoreMemory = 0;
uint64_t totalFallbackIntervals = 0;
uint64_t totalNestedSingleUseIntervals = 0;
uint64_t maximumCorePeak = 0;
std::string worstEntry = "<none>";
for (const MemoryReportEntry& entry : reportEntries) {
totalCoresMemory += entry.totalAllocaBytes;
uint64_t factor = entry.kind == MemoryReportEntry::Kind::Batch ? entry.coreIds.size() : 1;
totalLogicalCoreMemory += entry.row.logicalAllocaBytes * factor;
totalFallbackIntervals += entry.row.fallbackIntervals * factor;
totalNestedSingleUseIntervals += entry.row.nestedSingleUseIntervals * factor;
if (entry.row.sizeAlloca <= maximumCorePeak)
continue;
maximumCorePeak = entry.row.sizeAlloca;
worstEntry = entry.kind == MemoryReportEntry::Kind::Batch ? "Batch " + std::to_string(entry.id)
: "Core " + std::to_string(entry.coreIds.front());
}
uint64_t savedCoreMemory =
totalLogicalCoreMemory >= totalCoresMemory ? totalLogicalCoreMemory - totalCoresMemory : 0;
double savedPercent = totalLogicalCoreMemory == 0
? 0.0
: 100.0 * static_cast<double>(savedCoreMemory) / static_cast<double>(totalLogicalCoreMemory);
llvm::SmallVector<ReportField, 3> totalFields = {
{"Global memory", formatReportMemory(totalGlobalMemory) },
{"Weights memory", formatReportMemory(totalWeightsMemory)},
{"Cores memory", formatReportMemory(totalCoresMemory) }
llvm::SmallVector<ReportField, 10> totalFields = {
{"Global memory", formatReportMemory(totalGlobalMemory) },
{"Weights memory", formatReportMemory(totalWeightsMemory) },
{"Logical core allocations", formatReportMemory(totalLogicalCoreMemory) },
{"Physical core memory", formatReportMemory(totalCoresMemory) },
{"Cores memory", formatReportMemory(totalCoresMemory) },
{"Saved core memory", formatReportMemory(savedCoreMemory) },
{"Saved core memory percent", formatv("{0:F2}%", savedPercent).str() },
{"Maximum core peak", formatReportMemory(maximumCorePeak) },
{"Worst core/batch", worstEntry },
{"Fallback intervals", std::to_string(totalFallbackIntervals) },
{"Nested single-use intervals", std::to_string(totalNestedSingleUseIntervals) }
};
printReportTotalsBlock(os, totalFields);
@@ -552,19 +518,14 @@ void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instructio
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::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;
case pim_binary::Opcode::smuli: scalarRegisterValues[instruction.rd].reset(); break;
default: break;
}
}
@@ -580,8 +541,8 @@ void PimCodeGen::genSetRegisterImmediate(uint8_t registerNumber, int32_t immedia
}
void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const {
genSetRegisterImmediate(
pim::checkedU8OrCrash(registerNumber, "register number"), pim::checkedI32OrCrash(immediate, "register immediate"));
genSetRegisterImmediate(pim::checkedU8OrCrash(registerNumber, "register number"),
pim::checkedI32OrCrash(immediate, "register immediate"));
}
void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const {
@@ -729,8 +690,7 @@ void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKno
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");
emitMemCopyOp(pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len");
}
concatOffset += inputConcatDim;
@@ -788,10 +748,11 @@ void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const Sta
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;
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;
@@ -871,6 +832,60 @@ static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
return coreLikeOps;
}
static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLikeOp) {
CompiledCoreMemoryPlan plan;
auto fallbackAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryFallbackCountAttrName);
auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
coreLikeOp->emitError("requires complete PIM local-memory planning summary attributes before codegen");
return failure();
}
plan.fallbackIntervals = static_cast<uint64_t>(fallbackAttr.getInt());
plan.nestedSingleUseIntervals = static_cast<uint64_t>(nestedSingleUseAttr.getInt());
SmallDenseMap<size_t, size_t, 32> slotIndices;
bool hasFailure = false;
coreLikeOp->walk([&](memref::AllocOp allocOp) {
if (hasFailure)
return;
auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
auto slotAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotAttrName);
auto slotSizeAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotSizeAttrName);
if (!addressAttr || !slotAttr || !slotSizeAttr || addressAttr.getInt() < 0 || slotAttr.getInt() < 0
|| slotSizeAttr.getInt() < 0) {
allocOp.emitOpError("requires a complete non-negative PIM local-memory plan 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());
size_t slotId = static_cast<size_t>(slotAttr.getInt());
size_t slotSize = static_cast<size_t>(slotSizeAttr.getInt());
plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
plan.logicalBytes += *checkedSize;
auto [slotIt, inserted] = slotIndices.try_emplace(slotId, plan.slots.size());
if (inserted) {
plan.slots.push_back({slotId, 0, slotSize});
return;
}
const PhysicalSlotInfo& existing = plan.slots[slotIt->second];
if (existing.size != slotSize) {
allocOp.emitOpError("has a PIM local-memory arena inconsistent with another allocation");
hasFailure = true;
}
});
if (hasFailure)
return failure();
llvm::sort(plan.slots, [](const PhysicalSlotInfo& lhs, const PhysicalSlotInfo& rhs) { return lhs.id < rhs.id; });
return plan;
}
struct CoreEmissionResult {
static constexpr size_t kMaxStoredCodegenDiagnostics = 8;
@@ -882,10 +897,8 @@ struct CoreEmissionResult {
OnnxMlirCompilerErrorCodes status = CompilerSuccess;
MemoryReportRow reportRow;
llvm::SmallVector<ResolvedWeightView, 8> usedWeights;
MemoryPlanArtifacts livenessArtifacts;
llvm::SmallVector<DiagnosticRecord, kMaxStoredCodegenDiagnostics> diagnostics;
size_t diagnosticCount = 0;
void recordDiagnostic(Operation* op, StringRef message) {
++diagnosticCount;
if (diagnostics.size() < kMaxStoredCodegenDiagnostics)
@@ -1012,13 +1025,21 @@ static LogicalResult executeCompiledCorePlan(
}
auto emitBinary = [&](auto op, pim_binary::Opcode opcode) {
coreCodeGen.emitBinaryVectorOp(opcode, op.getOutputBuffer(), op.getLhs(), op.getRhs(),
getVectorByteSizeOrCrash(cast<ShapedType>(op.getLhs().getType())), knowledge);
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(),
coreCodeGen.emitUnaryVectorOp(opcode,
op.getOutputBuffer(),
op.getInput(),
getVectorByteSizeOrCrash(cast<ShapedType>(op.getInput().getType())),
knowledge, r2OrImm, generic1);
knowledge,
r2OrImm,
generic1);
};
switch (node.opKind) {
@@ -1038,20 +1059,19 @@ static LogicalResult executeCompiledCorePlan(
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::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;
emitUnary(cast<pim::PimVSoftmaxOp>(node.op), pim_binary::Opcode::vsoftmax, 0, 0);
break;
}
}
return success();
@@ -1103,8 +1123,7 @@ static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp,
}
static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
std::string outputCorePath =
(outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
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) {
@@ -1128,8 +1147,7 @@ static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath
if (!pimEmitJson.getValue())
return CompilerSuccess;
std::string outputCoreJsonPath =
(outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
errorCode = std::error_code();
raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
if (errorCode) {
@@ -1169,18 +1187,27 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
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;
@@ -1210,6 +1237,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
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;
@@ -1229,6 +1257,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
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;
@@ -1332,17 +1361,16 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
if (auto coreOp = dyn_cast<pim::PimCoreOp>(job.coreLikeOp)) {
aliasMaterializedHostGlobals(coreOp, moduleOp, materializedHostGlobals, jobMemory);
auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
deviceMemory.allocateCore(coreOp);
deviceMemory.allocateCore(*job.memoryPlan);
StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp);
if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
(void)finalizeInstructions();
(void) finalizeInstructions();
result.status = CompilerFailure;
return result;
}
result.reportRow = deviceMemory.getReportRow();
result.usedWeights = std::move(usedWeights);
result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
}
else {
auto coreBatchOp = cast<pim::PimCoreBatchOp>(job.coreLikeOp);
@@ -1352,10 +1380,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
for (unsigned lane : job.lanes) {
StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane);
deviceMemory.allocateCore(coreBatchOp, lane);
deviceMemory.allocateCore(*job.memoryPlan, lane);
coreCodeGen.setBatchLane(lane);
if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
(void)finalizeInstructions();
(void) finalizeInstructions();
result.status = CompilerFailure;
return result;
}
@@ -1363,7 +1391,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
result.reportRow = deviceMemory.getReportRow();
result.usedWeights = std::move(usedWeights);
result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
}
if (!finalizeInstructions()) {
@@ -1388,9 +1415,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
for (const CoreEmissionResult& result : jobResults) {
if (!summaryAnchor && !result.diagnostics.empty())
summaryAnchor = result.diagnostics.front().op;
for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics) {
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));
}
@@ -1420,19 +1446,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath);
memory.setTotalWeightBytes(weightEmission.totalWeightBytes);
auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName;
if (std::string reportsRoot = getOutputDir(); !reportsRoot.empty()) {
std::string reportsDir = reportsRoot + "/reports";
sys::fs::remove(reportsDir + "/pim_memory_liveness_report.txt");
sys::fs::remove(reportsDir + "/pim_memory_liveness_report.json");
sys::fs::remove(reportsDir + "/pim_memory_liveness_timeline.dot");
}
std::fstream livenessReportFile;
std::unique_ptr<llvm::raw_os_ostream> livenessReportOs;
if (pimMemoryReport != PimMemoryReportNone) {
livenessReportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt");
livenessReportOs = std::make_unique<llvm::raw_os_ostream>(livenessReportFile);
}
for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
const CoreEmissionJob& job = jobs[jobIndex];
const CoreEmissionResult& result = jobResults[jobIndex];
@@ -1443,8 +1456,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
return err;
xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
memory.recordCoreReport(job.emittedCoreId, result.reportRow);
if (livenessReportFile.is_open())
*livenessReportOs << "Core " << job.emittedCoreId << ":\n" << result.livenessArtifacts;
continue;
}
}
@@ -1473,18 +1484,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
batchPerCoreRow.value_or(MemoryReportRow {}),
batchRow.numAlloca,
batchRow.sizeAlloca);
if (livenessReportFile.is_open())
for (size_t jobIndex : group)
*livenessReportOs << "Batch " << batchReportId << " core " << jobs[jobIndex].emittedCoreId << ":\n"
<< jobResults[jobIndex].livenessArtifacts;
}
maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1;
if (livenessReportFile.is_open()) {
livenessReportOs->flush();
livenessReportFile.close();
}
memory.flushReport();
return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
}
src/PIM/Compiler/PimCodeGen.hpp
+21 -8
View File
@@ -39,8 +39,6 @@ struct PhysicalSlotInfo {
size_t size = 0;
};
using MemoryPlanArtifacts = std::string;
struct MemoryValueKey {
mlir::Value value;
std::optional<unsigned> lane;
@@ -48,15 +46,33 @@ struct MemoryValueKey {
bool operator==(const MemoryValueKey& other) const { return value == other.value && lane == other.lane; }
};
struct CompiledLocalMemoryEntry {
mlir::Value value;
MemEntry memory;
};
struct CompiledCoreMemoryPlan {
llvm::SmallVector<CompiledLocalMemoryEntry, 32> entries;
llvm::SmallVector<PhysicalSlotInfo, 32> slots;
uint64_t logicalBytes = 0;
uint64_t fallbackIntervals = 0;
uint64_t nestedSingleUseIntervals = 0;
};
struct MemoryReportRow {
uint64_t numAlloca = 0;
uint64_t sizeAlloca = 0;
uint64_t numGlobal = 0;
uint64_t sizeGlobal = 0;
uint64_t logicalAllocaBytes = 0;
uint64_t fallbackIntervals = 0;
uint64_t nestedSingleUseIntervals = 0;
bool operator==(const MemoryReportRow& other) const {
return numAlloca == other.numAlloca && sizeAlloca == other.sizeAlloca && numGlobal == other.numGlobal
&& sizeGlobal == other.sizeGlobal;
&& sizeGlobal == other.sizeGlobal && logicalAllocaBytes == other.logicalAllocaBytes
&& fallbackIntervals == other.fallbackIntervals
&& nestedSingleUseIntervals == other.nestedSingleUseIntervals;
}
};
@@ -93,26 +109,22 @@ class PimMemory {
llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap;
llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> ownedMemEntriesMap;
MemoryReportRow reportRow;
MemoryPlanArtifacts livenessArtifacts;
size_t minAlignment = 4;
size_t firstAvailableAddress = 0;
size_t nextPhysicalSlotId = 0;
MemEntry* gatherMemEntry(mlir::Value value, std::optional<unsigned> lane = std::nullopt);
size_t allocateAddress(size_t size, const MemoryValueKey& key);
void allocateGatheredMemory();
void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind);
PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key);
public:
PimMemory(llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap)
: globalMemEntriesMap(globalMemEntriesMap) {}
void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
void allocateCore(mlir::Operation* op, std::optional<unsigned> lane = std::nullopt);
void allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane = std::nullopt);
MemoryReportRow getReportRow() const;
const MemoryPlanArtifacts& getLivenessArtifacts() const { return livenessArtifacts; }
size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
MemEntry getMemEntry(const MemoryValueKey& key) const;
@@ -160,6 +172,7 @@ public:
struct CoreEmissionJob {
mlir::Operation* coreLikeOp = nullptr;
const CompiledCoreProgram* program = nullptr;
const CompiledCoreMemoryPlan* memoryPlan = nullptr;
size_t emittedCoreId = 0;
llvm::SmallVector<unsigned, 4> lanes;
std::optional<uint64_t> batchReportId;
src/PIM/Compiler/PimCompilerOptions.cpp
+3 -3
View File
@@ -74,7 +74,7 @@ llvm::cl::opt<bool>
llvm::cl::opt<bool>
pimDisableMemoryCoalescing("pim-disable-memory-coalescing",
llvm::cl::desc("Skip the PIM memory coalescing pass (developer diagnostic option)"),
llvm::cl::desc("Skip the early PIM IR memory coalescing pass (developer diagnostic)"),
llvm::cl::init(false),
llvm::cl::cat(OnnxMlirOptions));
@@ -124,10 +124,10 @@ llvm::cl::opt<bool> pimTraceCommunicationMaterialization(
llvm::cl::cat(OnnxMlirOptions));
llvm::cl::opt<size_t>
crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(2));
crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(128));
llvm::cl::opt<size_t>
crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(256));
crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(64));
llvm::cl::opt<long> coresCount("core-count",
llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."),
src/PIM/Compiler/PimCompilerOptions.hpp
+1 -1
View File
@@ -53,8 +53,8 @@ extern llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport;
extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;
extern llvm::cl::opt<bool> pimOnlyCodegen;
extern llvm::cl::opt<bool> pimDisableMemoryCoalescing;
extern llvm::cl::opt<bool> pimOnlyCodegen;
extern llvm::cl::opt<bool> useExperimentalConvImpl;
extern llvm::cl::opt<bool> pimEmitJson;
extern llvm::cl::opt<bool> pimReportConvLowering;
src/PIM/Compiler/PimCompilerUtils.cpp
+3
View File
@@ -21,6 +21,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
verifyExplicitPimCoreCount();
if (pimOnlyCodegen) {
pm.addPass(createPimLocalMemoryPlanningPass());
pm.addPass(createEmitPimCodePass());
return;
}
@@ -53,6 +54,8 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
pm.addPass(createMessagePass("Pim host constants folded"));
if (!pimDisableMemoryCoalescing)
pm.addPass(createPimMemoryCoalescingPass());
pm.addPass(createPimLocalMemoryPlanningPass());
pm.addPass(createMessagePass("Pim local memory planned"));
pm.addPass(createPimVerificationPass());
pm.addPass(createMessagePass("Pim verified"));
pm.addPass(createEmitPimCodePass());
src/PIM/Compiler/PimMemoryLiveness.cpp
-746
View File
@@ -1,746 +0,0 @@
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/raw_ostream.h"
#include <numeric>
#include <string>
#include <tuple>
#include <utility>
#include "Common/Support/CheckedArithmetic.hpp"
#include "Common/Support/ReportUtils.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
using namespace llvm;
using namespace mlir;
using namespace onnx_mlir;
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 MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional<unsigned> lane = std::nullopt) {
return {value, getLaneForMemoryValue(value, lane)};
}
struct MemoryTouchInterval {
uint64_t start = 0;
uint64_t end = 0;
Operation* firstTouchOp = nullptr;
Operation* lastTouchOp = nullptr;
uint64_t firstTouchPosition = 0;
uint64_t lastTouchPosition = 0;
bool hasRuntimeUse = false;
bool startUsedAllocFallback = false;
bool endUsedFallback = false;
bool escapesLoop = false;
std::string fallbackReason;
llvm::SmallVector<std::string, 8> aliasesFollowed;
};
struct OperationOrdering {
llvm::DenseMap<Operation*, uint64_t> position;
llvm::DenseMap<Operation*, uint64_t> subtreeEnd;
uint64_t nextPosition = 0;
};
static std::string printValueToString(mlir::Value value) {
std::string text;
llvm::raw_string_ostream os(text);
value.print(os);
os.flush();
return text;
}
static std::string printOperationToString(Operation* op) {
if (!op)
return "<none>";
std::string text;
llvm::raw_string_ostream os(text);
op->print(os);
os.flush();
return text;
}
static std::string printLocationToString(Location loc) {
std::string text;
llvm::raw_string_ostream os(text);
loc.print(os);
os.flush();
return text;
}
static std::string collapseWhitespace(StringRef text) {
std::string out;
out.reserve(text.size());
bool lastWasSpace = false;
for (char c : text) {
bool isSpace = c == ' ' || c == '\n' || c == '\t' || c == '\r';
if (isSpace) {
if (!lastWasSpace && !out.empty())
out.push_back(' ');
lastWasSpace = true;
continue;
}
out.push_back(c);
lastWasSpace = false;
}
return out;
}
static std::string abbreviate(StringRef text, size_t maxLen) {
if (text.size() <= maxLen)
return text.str();
return (text.take_front(maxLen - 3) + "...").str();
}
static std::string summarizeValue(mlir::Value value, size_t maxLen = 72) {
return abbreviate(collapseWhitespace(printValueToString(value)), maxLen);
}
static std::string summarizeOperation(Operation* op, size_t maxLen = 96) {
if (!op)
return "<none>";
std::string prefix = op->getName().getStringRef().str();
std::string full = collapseWhitespace(printOperationToString(op));
if (full == prefix)
return prefix;
return abbreviate(prefix + " :: " + full, maxLen);
}
static std::string summarizeLocation(Location loc, size_t maxLen = 88) {
return abbreviate(collapseWhitespace(printLocationToString(loc)), maxLen);
}
static void assignOperationOrdering(Operation* op, OperationOrdering& ordering) {
uint64_t position = ordering.nextPosition++;
ordering.position[op] = position;
uint64_t end = position;
for (Region& region : op->getRegions())
for (Block& block : region)
for (Operation& nestedOp : block) {
assignOperationOrdering(&nestedOp, ordering);
end = std::max(end, ordering.subtreeEnd.lookup(&nestedOp));
}
ordering.subtreeEnd[op] = end;
}
static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) {
OperationOrdering ordering;
if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
return ordering;
for (Operation& op : coreLikeOp->getRegion(0).front())
assignOperationOrdering(&op, ordering);
return ordering;
}
static bool isSupportedAliasOp(Operation* op) {
return isa<memref::SubViewOp,
memref::CastOp,
memref::CollapseShapeOp,
memref::ExpandShapeOp>(op);
}
static bool isRuntimeMemoryTouchOp(Operation* op) {
return isa<pim::PimMemCopyHostToDevOp,
pim::PimMemCopyDevToHostOp,
pim::PimMemCopyOp,
pim::PimReceiveOp,
pim::PimSendOp,
pim::PimConcatOp,
pim::PimVMMOp,
pim::PimTransposeOp,
pim::PimVVAddOp,
pim::PimVVSubOp,
pim::PimVVMulOp,
pim::PimVVMaxOp,
pim::PimVVDMulOp,
pim::PimVAvgOp,
pim::PimVReluOp,
pim::PimVTanhOp,
pim::PimVSigmOp,
pim::PimVSoftmaxOp>(op);
}
static bool isIgnoredLivenessUser(Operation* op) {
return isSupportedAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op) || isCoreStaticAddressOp(op);
}
static bool isWithin(mlir::Value value, Region* region) {
if (!region)
return false;
if (auto blockArg = dyn_cast<BlockArgument>(value))
return blockArg.getOwner()->getParent() == region;
if (Operation* definingOp = value.getDefiningOp())
return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion());
return false;
}
static bool isNestedAllocation(Operation* coreLikeOp, memref::AllocOp allocOp) {
if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
return false;
return allocOp->getBlock() != &coreLikeOp->getRegion(0).front();
}
static void addFallbackReason(std::string& reason, StringRef newReason) {
if (newReason.empty())
return;
if (!reason.empty())
reason += "; ";
reason += newReason.str();
}
static void appendAliasDescription(llvm::SmallVectorImpl<std::string>& aliases, mlir::Value value) {
std::string text = printValueToString(value);
if (!llvm::is_contained(aliases, text))
aliases.push_back(std::move(text));
}
struct OrderedTouchRange {
uint64_t start = 0;
uint64_t end = 0;
bool escapedLoop = false;
};
static OrderedTouchRange
getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const OperationOrdering& ordering) {
OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), false};
for (Operation* current = user; current; current = current->getParentOp()) {
auto forOp = dyn_cast<scf::ForOp>(current);
if (!forOp || isWithin(definingValue, &forOp.getRegion()))
continue;
range.start = std::min(range.start, ordering.position.lookup(forOp));
range.end = std::max(range.end, ordering.subtreeEnd.lookup(forOp));
range.escapedLoop = true;
}
return range;
}
static MemoryTouchInterval
computeMemoryTouchInterval(memref::AllocOp allocOp,
const OperationOrdering& ordering,
uint64_t fallbackEnd,
bool includeAliasDescriptions) {
MemoryTouchInterval interval;
interval.start = ordering.position.lookup(allocOp);
interval.end = interval.start;
auto recordAlias = [&](mlir::Value value) {
if (includeAliasDescriptions)
appendAliasDescription(interval.aliasesFollowed, value);
};
SmallPtrSet<mlir::Value, 16> visitedValues;
SmallPtrSet<Operation*, 32> visitedUsers;
SmallVector<mlir::Value> pendingValues;
pendingValues.push_back(allocOp.getResult());
auto parentLoop = allocOp->getParentOfType<scf::ForOp>();
while (!pendingValues.empty()) {
mlir::Value value = pendingValues.pop_back_val();
if (!visitedValues.insert(value).second)
continue;
for (Operation* user : value.getUsers()) {
if (!visitedUsers.insert(user).second)
continue;
if (isSupportedAliasOp(user)) {
for (mlir::Value result : user->getResults()) {
pendingValues.push_back(result);
recordAlias(result);
}
}
if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user)) {
for (OpResult result : user->getResults()) {
OpOperand* tiedOperand = dpsOp.getTiedOpOperand(result);
if (!tiedOperand || tiedOperand->get() != value)
continue;
pendingValues.push_back(result);
recordAlias(result);
}
}
if (auto forOp = dyn_cast<scf::ForOp>(user)) {
for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs())) {
if (initArg != value)
continue;
pendingValues.push_back(forOp.getRegionIterArgs()[index]);
pendingValues.push_back(forOp.getResult(index));
recordAlias(forOp.getRegionIterArgs()[index]);
recordAlias(forOp.getResult(index));
if (parentLoop && forOp != parentLoop)
interval.escapesLoop = true;
}
}
if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp());
auto indexSwitch = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp());
if (ifOp) {
for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
if (operand != value)
continue;
pendingValues.push_back(ifOp.getResult(index));
recordAlias(ifOp.getResult(index));
}
}
else if (indexSwitch) {
for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
if (operand != value)
continue;
pendingValues.push_back(indexSwitch.getResult(index));
recordAlias(indexSwitch.getResult(index));
}
}
else if (!forOp) {
addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
}
else {
for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
if (operand != value)
continue;
pendingValues.push_back(forOp.getResult(index));
recordAlias(forOp.getResult(index));
if (parentLoop && forOp == parentLoop)
interval.escapesLoop = true;
}
}
}
if (isRuntimeMemoryTouchOp(user)) {
uint64_t touchPosition = ordering.position.lookup(user);
if (!interval.hasRuntimeUse || touchPosition < interval.firstTouchPosition) {
interval.firstTouchPosition = touchPosition;
interval.firstTouchOp = user;
}
if (!interval.hasRuntimeUse || touchPosition > interval.lastTouchPosition) {
interval.lastTouchPosition = touchPosition;
interval.lastTouchOp = user;
}
OrderedTouchRange range = getEffectiveTouchRange(allocOp.getResult(), user, ordering);
interval.escapesLoop |= range.escapedLoop;
if (!interval.hasRuntimeUse) {
interval.start = range.start;
interval.end = range.end;
interval.hasRuntimeUse = true;
}
else {
if (range.start < interval.start) {
interval.start = range.start;
}
if (range.end > interval.end) {
interval.end = range.end;
}
}
continue;
}
if (isIgnoredLivenessUser(user))
continue;
addFallbackReason(interval.fallbackReason, "unhandled user op");
interval.endUsedFallback = true;
}
}
if (!interval.hasRuntimeUse) {
interval.startUsedAllocFallback = true;
interval.endUsedFallback = true;
interval.start = ordering.position.lookup(allocOp);
interval.end = fallbackEnd;
interval.firstTouchPosition = interval.start;
interval.lastTouchPosition = interval.end;
addFallbackReason(interval.fallbackReason, "no runtime memory touch");
return interval;
}
if (interval.endUsedFallback) {
interval.end = std::max(interval.end, fallbackEnd);
}
return interval;
}
static FailureOr<size_t> getAllocSizeBytes(memref::AllocOp allocOp) {
auto type = dyn_cast<ShapedType>(allocOp.getType());
if (!type)
return failure();
auto checkedBytes = pim::getCheckedShapedTypeSizeInBytes(type, allocOp, "memory allocation byte size");
if (failed(checkedBytes))
return failure();
return pim::checkedSize(*checkedBytes, allocOp, "memory allocation byte size");
}
static bool intervalsOverlap(const LocalAllocInterval& lhs, const LocalAllocInterval& rhs) {
return !(lhs.end < rhs.start || rhs.end < lhs.start);
}
static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, ArrayRef<LocalAllocInterval> intervals) {
uint64_t slotLogicalBytes = 0;
for (size_t intervalIndex : slot.intervalIndices)
slotLogicalBytes += intervals[intervalIndex].size;
return slotLogicalBytes;
}
} // namespace
SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp,
std::optional<unsigned> lane,
bool includeAliasDescriptions) {
SmallVector<LocalAllocInterval, 0> intervals;
OperationOrdering ordering = buildOperationOrdering(coreLikeOp);
if (ordering.position.empty())
return intervals;
uint64_t fallbackEnd = ordering.nextPosition == 0 ? 0 : ordering.nextPosition - 1;
size_t nextIntervalId = 0;
coreLikeOp->walk([&](memref::AllocOp allocOp) {
auto checkedSize = getAllocSizeBytes(allocOp);
if (failed(checkedSize)) {
llvm::errs() << "Failed to compute local allocation size for value: ";
allocOp.getResult().print(llvm::errs());
llvm::errs() << "\n";
llvm_unreachable("Failed to compute local allocation size");
}
MemoryTouchInterval touchInterval =
computeMemoryTouchInterval(allocOp, ordering, fallbackEnd, includeAliasDescriptions);
LocalAllocInterval interval;
interval.id = nextIntervalId++;
interval.alloc = allocOp;
interval.key = getMemoryValueKey(allocOp.getResult(), lane);
interval.start = touchInterval.start;
interval.end = touchInterval.end;
interval.size = *checkedSize;
interval.firstTouchOp = touchInterval.firstTouchOp;
interval.lastTouchOp = touchInterval.lastTouchOp;
interval.firstTouchPosition = touchInterval.firstTouchPosition;
interval.lastTouchPosition = touchInterval.lastTouchPosition;
interval.startUsedAllocFallback = touchInterval.startUsedAllocFallback;
interval.endUsedFallback = touchInterval.endUsedFallback;
interval.hasRuntimeUse = touchInterval.hasRuntimeUse;
interval.insideNestedRegion = isNestedAllocation(coreLikeOp, allocOp);
interval.escapesLoop = touchInterval.escapesLoop;
interval.fallbackReason = std::move(touchInterval.fallbackReason);
interval.aliasesFollowed = std::move(touchInterval.aliasesFollowed);
intervals.push_back(std::move(interval));
});
return intervals;
}
SmallVector<PlannedPhysicalSlot, 0> onnx_mlir::planPhysicalSlots(MutableArrayRef<LocalAllocInterval> intervals) {
SmallVector<PlannedPhysicalSlot, 0> slots;
SmallVector<size_t> intervalOrder(intervals.size());
std::iota(intervalOrder.begin(), intervalOrder.end(), 0);
llvm::stable_sort(intervalOrder, [&](size_t lhsIndex, size_t rhsIndex) {
const LocalAllocInterval& lhs = intervals[lhsIndex];
const LocalAllocInterval& rhs = intervals[rhsIndex];
if (lhs.size != rhs.size)
return lhs.size > rhs.size;
if (lhs.start != rhs.start)
return lhs.start < rhs.start;
if (lhs.end != rhs.end)
return lhs.end < rhs.end;
return lhs.id < rhs.id;
});
for (size_t intervalIndex : intervalOrder) {
LocalAllocInterval& interval = intervals[intervalIndex];
PlannedPhysicalSlot* bestSlot = nullptr;
auto bestKey = std::tuple<size_t, size_t, size_t, size_t>(std::numeric_limits<size_t>::max(),
std::numeric_limits<size_t>::max(),
std::numeric_limits<size_t>::max(),
std::numeric_limits<size_t>::max());
for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
PlannedPhysicalSlot& slot = slots[slotIndex];
bool compatible = true;
for (size_t otherIndex : slot.intervalIndices) {
if (intervalsOverlap(interval, intervals[otherIndex])) {
compatible = false;
break;
}
}
if (!compatible)
continue;
size_t resultingSize = std::max(slot.requiredSize, interval.size);
size_t growth = resultingSize - slot.requiredSize;
auto candidateKey =
std::tuple<size_t, size_t, size_t, size_t>(growth, resultingSize, slot.intervalIndices.size(), slot.id);
if (candidateKey < bestKey) {
bestKey = candidateKey;
bestSlot = &slot;
}
}
if (!bestSlot) {
slots.push_back({slots.size(), interval.size, interval.size, 0, {intervalIndex}});
interval.slotPlanIndex = slots.size() - 1;
interval.physicalSlotId = slots.back().id;
interval.physicalSlotSize = slots.back().requiredSize;
continue;
}
bestSlot->requiredSize = std::max(bestSlot->requiredSize, interval.size);
bestSlot->size = bestSlot->requiredSize;
bestSlot->intervalIndices.push_back(intervalIndex);
interval.slotPlanIndex = static_cast<size_t>(bestSlot - slots.data());
interval.physicalSlotId = bestSlot->id;
interval.physicalSlotSize = bestSlot->requiredSize;
}
return slots;
}
MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
std::optional<unsigned> lane,
ArrayRef<LocalAllocInterval> intervals,
ArrayRef<PlannedPhysicalSlot> slots,
size_t addressLimit,
PimMemoryReportLevel reportLevel) {
MemoryPlanArtifacts artifacts;
uint64_t totalLogicalBytes = 0;
uint64_t totalPhysicalBytes = 0;
uint64_t fallbackIntervals = 0;
uint64_t noRuntimeTouchIntervals = 0;
uint64_t reusedAllocations = 0;
uint64_t nestedIntervals = 0;
uint64_t loopEscapingIntervals = 0;
size_t largestLogicalAllocation = 0;
size_t largestPhysicalSlot = 0;
size_t maximumAssignedAddress = 0;
for (const LocalAllocInterval& interval : intervals) {
totalLogicalBytes += interval.size;
largestLogicalAllocation = std::max(largestLogicalAllocation, interval.size);
maximumAssignedAddress = std::max(maximumAssignedAddress, interval.assignedAddress + interval.physicalSlotSize);
if (interval.startUsedAllocFallback || interval.endUsedFallback)
++fallbackIntervals;
if (!interval.hasRuntimeUse)
++noRuntimeTouchIntervals;
if (interval.insideNestedRegion)
++nestedIntervals;
if (interval.escapesLoop)
++loopEscapingIntervals;
}
for (const PlannedPhysicalSlot& slot : slots) {
totalPhysicalBytes += slot.size;
largestPhysicalSlot = std::max(largestPhysicalSlot, slot.size);
if (slot.intervalIndices.size() > 1)
reusedAllocations += slot.intervalIndices.size() - 1;
}
uint64_t savedBytes = totalLogicalBytes >= totalPhysicalBytes ? totalLogicalBytes - totalPhysicalBytes : 0;
double savedPercent =
totalLogicalBytes == 0 ? 0.0 : 100.0 * static_cast<double>(savedBytes) / static_cast<double>(totalLogicalBytes);
raw_string_ostream os(artifacts);
os << "=== PIM Memory Liveness Report ===\n";
os << "Op: " << coreLikeOp->getName() << "\n";
if (lane)
os << "Lane: " << *lane << "\n";
os << "Summary:\n";
os << " logical allocation bytes: " << formatReportMemory(totalLogicalBytes) << " (" << totalLogicalBytes << ")\n";
os << " physical allocation bytes: " << formatReportMemory(totalPhysicalBytes) << " (" << totalPhysicalBytes
<< ")\n";
os << " saved bytes: " << formatReportMemory(savedBytes) << " (" << savedBytes << ")\n";
os << " saved percent: " << format("%.2f%%", savedPercent) << "\n";
os << " intervals: " << intervals.size() << "\n";
os << " physical slots: " << slots.size() << "\n";
os << " reused allocations: " << reusedAllocations << "\n";
os << " fallback intervals: " << fallbackIntervals << "\n";
os << " intervals with no runtime memory touch: " << noRuntimeTouchIntervals << "\n";
os << " nested allocations: " << nestedIntervals << "\n";
os << " loop-escaping allocations: " << loopEscapingIntervals << "\n";
os << " largest logical allocation: " << largestLogicalAllocation << "\n";
os << " largest physical slot: " << largestPhysicalSlot << "\n";
os << " address limit: " << addressLimit << "\n";
os << " peak physical memory: " << formatReportMemory(maximumAssignedAddress) << " (" << maximumAssignedAddress
<< ")\n";
os << " maximum assigned address: " << maximumAssignedAddress << "\n";
os << "\nHow To Read:\n";
os << " `summary` only shows the strongest reuse cases and the worst offenders.\n";
os << " Use `--pim-memory-report=full` when you need the complete slot-by-slot and interval-by-interval dump.\n";
os << " Large single-use slots, fallback intervals, and nested single-use allocations are the best places\n";
os << " to inspect if allocations should be moved, sunk, or made easier to coalesce earlier in the pipeline.\n";
SmallVector<const PlannedPhysicalSlot*> reusedSlots;
SmallVector<const PlannedPhysicalSlot*> singleUseSlots;
for (const PlannedPhysicalSlot& slot : slots)
if (slot.intervalIndices.size() > 1)
reusedSlots.push_back(&slot);
else
singleUseSlots.push_back(&slot);
llvm::stable_sort(reusedSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
uint64_t lhsLogicalBytes = getSlotLogicalBytes(*lhs, intervals);
uint64_t rhsLogicalBytes = getSlotLogicalBytes(*rhs, intervals);
if (lhs->intervalIndices.size() != rhs->intervalIndices.size())
return lhs->intervalIndices.size() > rhs->intervalIndices.size();
if (lhsLogicalBytes != rhsLogicalBytes)
return lhsLogicalBytes > rhsLogicalBytes;
if (lhs->size != rhs->size)
return lhs->size > rhs->size;
return lhs->id < rhs->id;
});
llvm::stable_sort(singleUseSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
if (lhs->size != rhs->size)
return lhs->size > rhs->size;
return lhs->id < rhs->id;
});
constexpr size_t kSummaryReuseLimit = 6;
constexpr size_t kSummaryOffenderLimit = 10;
os << "\nBest Reuse:\n";
if (reusedSlots.empty()) {
os << " no slots were shared by multiple intervals\n";
}
else {
for (const PlannedPhysicalSlot* slot : ArrayRef(reusedSlots).take_front(kSummaryReuseLimit)) {
uint64_t slotLogicalBytes = getSlotLogicalBytes(*slot, intervals);
os << " slot #" << slot->id << " addr=" << slot->address << " size=" << formatReportMemory(slot->size)
<< " intervals=" << slot->intervalIndices.size() << " logical_sum=" << formatReportMemory(slotLogicalBytes)
<< "\n";
for (size_t intervalIndex : slot->intervalIndices) {
const LocalAllocInterval& interval = intervals[intervalIndex];
os << " #" << interval.id << " [" << interval.start << "," << interval.end << "]"
<< " logical=" << formatReportMemory(interval.size)
<< " first=" << summarizeOperation(interval.firstTouchOp, 40)
<< " last=" << summarizeOperation(interval.lastTouchOp, 40) << "\n";
}
}
}
os << "\nTop Offenders:\n";
bool printedAttention = false;
for (const PlannedPhysicalSlot* slot : ArrayRef(singleUseSlots).take_front(kSummaryOffenderLimit)) {
const LocalAllocInterval& interval = intervals[slot->intervalIndices.front()];
printedAttention = true;
os << " slot #" << slot->id << " is single-use"
<< " size=" << formatReportMemory(slot->size) << " interval=#" << interval.id
<< " value=" << summarizeValue(interval.key.value, 56) << "\n";
os << " first=" << summarizeOperation(interval.firstTouchOp, 40)
<< " last=" << summarizeOperation(interval.lastTouchOp, 40)
<< " nested=" << (interval.insideNestedRegion ? "yes" : "no")
<< " escapes_loop=" << (interval.escapesLoop ? "yes" : "no") << "\n";
}
size_t fallbackPrinted = 0;
for (const LocalAllocInterval& interval : intervals) {
if (!(interval.startUsedAllocFallback || interval.endUsedFallback) || fallbackPrinted >= kSummaryOffenderLimit)
continue;
printedAttention = true;
++fallbackPrinted;
os << " fallback interval #" << interval.id << " size=" << formatReportMemory(interval.size)
<< " value=" << summarizeValue(interval.key.value, 56) << "\n";
os << " reason: " << (interval.fallbackReason.empty() ? "<none>" : interval.fallbackReason) << "\n";
}
size_t nestedPrinted = 0;
for (const LocalAllocInterval& interval : intervals) {
if (nestedPrinted >= kSummaryOffenderLimit)
break;
if (!(interval.insideNestedRegion && slots[interval.slotPlanIndex].intervalIndices.size() == 1))
continue;
printedAttention = true;
++nestedPrinted;
os << " nested single-use interval #" << interval.id << " slot #" << interval.physicalSlotId
<< " size=" << formatReportMemory(interval.size) << " value=" << summarizeValue(interval.key.value, 56)
<< "\n";
os << " hint: move or sink this alloc inside the nested region if the IR allows it.\n";
}
if (!printedAttention)
os << " no obvious blockers detected in this core\n";
if (reportLevel == PimMemoryReportFull) {
os << "\nSlot Reuse:\n";
for (const PlannedPhysicalSlot& slot : slots) {
uint64_t slotLogicalBytes = getSlotLogicalBytes(slot, intervals);
os << " slot #" << slot.id << " addr=" << slot.address << " size=" << formatReportMemory(slot.size) << " ("
<< slot.size << ")"
<< " intervals=" << slot.intervalIndices.size() << " logical_sum=" << formatReportMemory(slotLogicalBytes)
<< "\n";
for (size_t intervalIndex : slot.intervalIndices) {
const LocalAllocInterval& interval = intervals[intervalIndex];
mlir::Value allocValue = interval.key.value;
os << " [" << interval.start << "," << interval.end << "]"
<< " #" << interval.id << " logical=" << formatReportMemory(interval.size)
<< " nested=" << (interval.insideNestedRegion ? "yes" : "no")
<< " escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
<< " first=" << summarizeOperation(interval.firstTouchOp, 48)
<< " last=" << summarizeOperation(interval.lastTouchOp, 48) << "\n";
os << " value=" << summarizeValue(allocValue) << "\n";
}
}
}
if (reportLevel == PimMemoryReportFull) {
os << "\nInterval Details:\n";
for (const LocalAllocInterval& interval : intervals) {
const PlannedPhysicalSlot& slot = slots[interval.slotPlanIndex];
mlir::Value allocValue = interval.key.value;
Operation* definingOp = allocValue.getDefiningOp();
os << " #" << interval.id << " slot=" << slot.id << " live=[" << interval.start << "," << interval.end << "]"
<< " logical=" << formatReportMemory(interval.size)
<< " slot_size=" << formatReportMemory(interval.physicalSlotSize) << " addr=" << interval.assignedAddress
<< "\n";
os << " value=" << summarizeValue(allocValue, 88) << "\n";
os << " type=" << allocValue.getType() << "\n";
os << " loc="
<< summarizeLocation(definingOp ? definingOp->getLoc() : UnknownLoc::get(coreLikeOp->getContext())) << "\n";
os << " nested=" << (interval.insideNestedRegion ? "yes" : "no")
<< " escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
<< " start_fallback=" << (interval.startUsedAllocFallback ? "yes" : "no")
<< " end_fallback=" << (interval.endUsedFallback ? "yes" : "no") << "\n";
os << " first_use=" << summarizeOperation(interval.firstTouchOp) << " @" << interval.firstTouchPosition
<< "\n";
os << " last_use=" << summarizeOperation(interval.lastTouchOp) << " @" << interval.lastTouchPosition << "\n";
os << " slot_peers=";
bool first = true;
for (size_t otherIndex : slot.intervalIndices) {
if (intervals[otherIndex].id == interval.id)
continue;
if (!first)
os << ", ";
os << "#" << intervals[otherIndex].id;
first = false;
}
if (first)
os << "<none>";
os << "\n";
if (!interval.fallbackReason.empty())
os << " fallback_reason=" << interval.fallbackReason << "\n";
if (!interval.aliasesFollowed.empty()) {
os << " aliases_followed=" << interval.aliasesFollowed.size() << "\n";
for (const std::string& alias : interval.aliasesFollowed)
os << " - " << abbreviate(collapseWhitespace(alias), 108) << "\n";
}
}
}
os.flush();
return artifacts;
}
src/PIM/Compiler/PimMemoryLiveness.hpp
-62
View File
@@ -1,62 +0,0 @@
#pragma once
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include <limits>
#include <optional>
#include <string>
#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
namespace onnx_mlir {
struct LocalAllocInterval {
size_t id = 0;
mlir::memref::AllocOp alloc;
MemoryValueKey key;
uint64_t start = 0;
uint64_t end = 0;
size_t size = 0;
mlir::Operation* firstTouchOp = nullptr;
mlir::Operation* lastTouchOp = nullptr;
uint64_t firstTouchPosition = 0;
uint64_t lastTouchPosition = 0;
bool startUsedAllocFallback = false;
bool endUsedFallback = false;
bool hasRuntimeUse = false;
bool insideNestedRegion = false;
bool escapesLoop = false;
std::string fallbackReason;
llvm::SmallVector<std::string, 8> aliasesFollowed;
size_t slotPlanIndex = std::numeric_limits<size_t>::max();
size_t physicalSlotId = std::numeric_limits<size_t>::max();
size_t assignedAddress = 0;
size_t physicalSlotSize = 0;
};
struct PlannedPhysicalSlot {
size_t id = std::numeric_limits<size_t>::max();
size_t requiredSize = 0;
size_t size = 0;
size_t address = 0;
llvm::SmallVector<size_t, 8> intervalIndices;
};
llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp,
std::optional<unsigned> lane,
bool includeAliasDescriptions = true);
llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);
MemoryPlanArtifacts buildMemoryPlanArtifacts(mlir::Operation* coreLikeOp,
std::optional<unsigned> lane,
llvm::ArrayRef<LocalAllocInterval> intervals,
llvm::ArrayRef<PlannedPhysicalSlot> slots,
size_t addressLimit,
PimMemoryReportLevel reportLevel);
} // namespace onnx_mlir