diff --git a/.agents/invariants/PARALLELISM_AND_RUNTIME_INVARIANT.md b/.agents/invariants/PARALLELISM_AND_RUNTIME_INVARIANT.md deleted file mode 100644 index 5f76981..0000000 --- a/.agents/invariants/PARALLELISM_AND_RUNTIME_INVARIANT.md +++ /dev/null @@ -1,20 +0,0 @@ -# Parallelism and Runtime Invariant - -Compile-time optimization must not trade away execution parallelism or worsen -theoretical runtime behavior. The absence of a precise runtime measurement does -not make such a trade acceptable. - -In particular, a compile-time optimization must not: - -- serialize work that can execute independently; -- coarsen lane, core, batch, or scheduling granularity in a way that reduces - available execution parallelism; -- increase the theoretical runtime critical path, instruction count, memory - traffic, synchronization, or required copies merely to reduce compiler work; -- merge independently scheduled work when doing so may reduce concurrency. - -Compiler representations may share analysis, planning, or code-generation -work only when the emitted execution retains the same independent work, -schedule semantics, and available parallelism. If runtime impact cannot be -measured precisely, it must be established from the transformation's semantics; -uncertainty is not permission to accept a possible runtime regression. diff --git a/.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md b/.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md index 0d4c157..2b53e5b 100644 --- a/.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md +++ b/.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md @@ -11,7 +11,6 @@ This invariant applies to: - deferred communication; - Spatial-to-PIM lowering; - bufferization; -- memory coalescing; - liveness planning; - PIM code generation. @@ -19,12 +18,17 @@ This invariant applies to: A compiler compile-time or memory optimization must not reduce available hardware parallelism or worsen theoretical execution time. In the absence of a -precise runtime model, it is forbidden to serialize independent work, reduce +precise runtime model, uncertainty is not permission to accept a possible +runtime regression. It is forbidden to serialize independent work, reduce the number of simultaneously schedulable compute instances, replace parallel graph or core batches with sequential loops, introduce recomputation, add runtime copies, increase communication or instruction count, or lengthen the schedule critical path merely to reduce compiler time or memory usage. +Compiler representations may share analysis, planning, or code-generation +work only when the emitted execution retains the same independent work, +schedule semantics, scheduling granularity, and available parallelism. + An optimization is acceptable only when its static runtime proxies are equal or better. @@ -33,6 +37,8 @@ better. Do not: - use fewer parallel lanes or cores to simplify IR; +- coarsen lane, core, batch, or scheduling granularity when it may reduce + concurrency; - scalarize a valid graph or core batch; - replace independent operations with one sequential loop; - recompute values to reduce storage; diff --git a/AGENTS.md b/AGENTS.md index c9c6058..2572ece 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -14,6 +14,7 @@ Before modifying the relevant subsystem, read: * Use the debug build only when it is useful to obtain a clear stack trace with symbols, inspect names, place breakpoints, or test a small case interactively * The debug build is very slow, so use it only on small fast tests such as operation validations, not on network validations * Always prepend rtk to shell commands if missing and if rtk is available +* Always use the repository Python virtual environment (`.venv/bin/python` and its bundled tools) for Python commands # Core engineering philosophy diff --git a/README.md b/README.md index 595fecc..ff4b0f5 100644 --- a/README.md +++ b/README.md @@ -67,14 +67,11 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts Converts tensor-semantics PIM IR into memref-semantics PIM IR using MLIR's bufferization interfaces. -5. **Safe IR memory coalescing** - (`src/PIM/Dialect/Pim/Transforms/MemoryCoalescing`). - Normalizes compatible block-local allocations before final planning. -6. **PIM local-memory planning** +5. **PIM local-memory planning** (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`). - Computes whole-core lifetimes, coalesces non-overlapping allocations into - physical slots, assigns addresses, and records the explicit plan in PIM IR. -7. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and + Computes whole-core lifetimes, reuses addresses for non-overlapping + allocations, and records the explicit plan in PIM IR. +6. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and `src/PIM/Compiler`). Verifies the memory plan and other PIM invariants, then emits `.pim` core files, weights, and `memory.bin` / `config.json` without rerunning liveness. @@ -90,9 +87,10 @@ Supporting pieces: - `src/PIM/Pass` - pass registration and auxiliary passes. - `src/PIM/PimAccelerator.{cpp,hpp}` - ONNX-MLIR accelerator entry point. -## Key compiler options +## PIM compiler options -Pass these to `onnx-mlir` when compiling for PIM: +Pass these to `onnx-mlir` when compiling for PIM. These are all Raptor/PIM-specific +options; `onnx-mlir --help` lists the inherited ONNX-MLIR options. - `--maccel=PIM` - select the PIM accelerator. - `--EmitSpatial`, `--EmitPim`, `--EmitPimBufferized`, @@ -101,14 +99,30 @@ Pass these to `onnx-mlir` when compiling for PIM: - `--core-count=` - required positive core count for PIM compilation. - `--crossbar-size=` - crossbar width/height. Default in code is `128`. - `--crossbar-count=` - crossbars per core. Default in code is `64`. +- `--pim-memory-report=` - emit the concise combined memory report + under `reports/memory_report.txt`, or disable it. Default is `summary`. - `--pim-only-codegen` - assume input is already bufferized PIM IR and only run the codegen tail. - `--pim-emit-json` - also emit `core_*.json` instruction files alongside `core_*.pim`. -- `--pim-export-spatial-dataflow=` - control Spatial - dataflow CSV reports for the graph, trivially merged graph, scheduled, and - realized snapshots under `reports/`. +- `--pim-export-spatial-dataflow=` - + control Spatial dataflow CSV reports for the graph, trivially merged graph, + scheduled, and realized snapshots under `reports/`. Default is `none`. +- `--pim-conv-lowering=` - + select the convolution lowering strategy. Default is `auto`. +- `--pim-conv-im2col-max-elements=` - maximum globally materialized im2col + elements per convolution before streaming. Default is `1048576`. +- `--pim-conv-stream-chunk-positions=` - maximum output positions per + streamed convolution chunk. Default is `1024`. +- `--pim-report-conv-lowering=` - emit the bounded convolution + lowering report. Default is `true`. - `--use-experimental-conv-impl` - use the alternate convolution lowering. +- `--pim-detect-communication-deadlock` - statically simulate expanded + send/receive ordering and reject blocking deadlocks. Default is off. +- `--pim-materialize-scalar-fanout-global-order` - use the experimental, + expensive globally ordered scalar-fanout materializer. Default is off. +- `--pim-trace-communication-materialization` - emit verbose communication + materialization diagnostics and provenance attributes. Default is off. - `--ignore-concat-error` - soft-fail a ConcatOp corner case. ## Standard PIM hardware profile @@ -149,7 +163,7 @@ Python dependencies used by the validation scripts are `numpy`, `onnx`, and Per-operation validation from the repository root: ```bash -python validation/validate.py \ +.venv/bin/python validation/validate.py \ --raptor-path build_release/Release/bin/onnx-mlir \ --onnx-include-dir onnx-mlir/include \ --operations-dir validation/operations \ @@ -165,7 +179,7 @@ Validate one network or a subset by pointing `--operations-dir` at any directory containing `.onnx` files: ```bash -python validation/validate.py \ +.venv/bin/python validation/validate.py \ --raptor-path build_release/Release/bin/onnx-mlir \ --onnx-include-dir onnx-mlir/include \ --operations-dir validation/networks/yolo11n/depth_04 \ @@ -173,7 +187,6 @@ python validation/validate.py \ --crossbar-size 128 \ --core-count 144 \ --verbose \ - --raptor-extra-arg=--pim-memory-report=summary \ --raptor-extra-arg=--pim-detect-communication-deadlock \ --raptor-extra-arg=--pim-export-spatial-dataflow=none ``` @@ -200,7 +213,7 @@ Each validation run writes artifacts in the model workspace, for example under The compiler currently dumps dialect snapshots such as `spatial0.mlir`, `spatial1_graph.mlir`, `spatial2_trivial_merged.mlir`, `spatial3_scheduled_no_comm.mlir`, `spatial4_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`, -`pim2_folded.mlir`, `pim3_coalesced.mlir`, and `pim4_memory_planned.mlir` when an output directory is +`pim2_folded.mlir`, and `pim3_memory_planned.mlir` when an output directory is available. To rerun the simulator manually with tracing after validation has produced a @@ -229,7 +242,7 @@ Available operation suites under `validation/operations/`: `add`, `concat`, Generated operation tests can be regenerated with: ```bash -python3 validation/operations/gen_tests.py +.venv/bin/python validation/operations/gen_tests.py ``` ## Build diff --git a/src/PIM/CMakeLists.txt b/src/PIM/CMakeLists.txt index 2989ea0..e4af4ab 100644 --- a/src/PIM/CMakeLists.txt +++ b/src/PIM/CMakeLists.txt @@ -121,7 +121,6 @@ add_pim_library(OMPIMAccel OMPimCommon OMPimBufferization OMPimHostConstantFolding - OMPimMemoryCoalescing OMPimVerification MLIRTensorInferTypeOpInterfaceImpl ) diff --git a/src/PIM/Common/PimCommon.hpp b/src/PIM/Common/PimCommon.hpp index ba6b588..10b6bf6 100644 --- a/src/PIM/Common/PimCommon.hpp +++ b/src/PIM/Common/PimCommon.hpp @@ -23,6 +23,7 @@ #include "src/Compiler/CompilerOptions.hpp" #include +#include #include namespace onnx_mlir { @@ -30,11 +31,13 @@ namespace onnx_mlir { inline constexpr llvm::StringLiteral kCoreIdAttrName = "coreId"; inline constexpr llvm::StringLiteral kCoreIdsAttrName = "coreIds"; inline constexpr llvm::StringLiteral kLocalMemoryAddressAttrName = "pim.local_memory_address"; -inline constexpr llvm::StringLiteral kLocalMemorySlotAttrName = "pim.local_memory_slot"; -inline constexpr llvm::StringLiteral kLocalMemorySlotSizeAttrName = "pim.local_memory_slot_size"; -inline constexpr llvm::StringLiteral kLocalMemoryFallbackCountAttrName = "pim.local_memory_fallback_count"; -inline constexpr llvm::StringLiteral kLocalMemoryNestedSingleUseCountAttrName = - "pim.local_memory_nested_single_use_count"; +inline constexpr llvm::StringLiteral kLocalMemorySizeAttrName = "pim.local_memory_size"; +inline constexpr std::array kRemovedLocalMemoryPlanAttrNames = { + "pim.local_memory_slot", + "pim.local_memory_slot_size", + "pim.local_memory_fallback_count", + "pim.local_memory_nested_single_use_count", +}; inline constexpr size_t kPimLocalMemoryAddressLimit = static_cast(std::numeric_limits::max()); } // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index 6bd5fd4..4e2a0dc 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -141,6 +141,19 @@ static void printMemoryOverflowDiagnostic(const MemoryValueKey& key, } } +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) { @@ -193,12 +206,11 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE globalMemEntriesMap[key] = memEntry; switch (reportKind) { - case MemoryReportKind::Alloca: break; + case MemoryReportKind::Alloca: case MemoryReportKind::Global: - ++reportRow.numGlobal; - reportRow.sizeGlobal += memEntry.size; - break; case MemoryReportKind::Input: + ++reportRow.hostObjectCount; + break; case MemoryReportKind::None: break; } } @@ -242,12 +254,15 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) { } void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional lane) { - if (localPhysicalSlots.empty()) { - llvm::append_range(localPhysicalSlots, plan.slots); - reportRow.logicalAllocaBytes = plan.logicalBytes; - reportRow.fallbackIntervals = plan.fallbackIntervals; - reportRow.nestedSingleUseIntervals = plan.nestedSingleUseIntervals; + 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; @@ -255,49 +270,10 @@ void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional fields = { - {"Number of globals", std::to_string(row.numGlobal) }, - {"Global memory", formatReportMemory(row.sizeGlobal)} - }; - printReportFlatFields(os, fields); -} - -static void printCoreMemoryReportRow(raw_ostream& os, const MemoryReportEntry& entry) { - llvm::SmallVector fields = { - {"Number of allocas", std::to_string(entry.row.numAlloca) }, - {"Allocated memory", formatReportMemory(entry.row.sizeAlloca)} - }; - printReportFlatFields(os, fields); -} - -static void printBatchMemoryReportRow(raw_ostream& os, const MemoryReportEntry& entry) { - llvm::SmallVector perCoreFields = { - {"Number of allocas", std::to_string(entry.row.numAlloca) }, - {"Allocated memory", formatReportMemory(entry.row.sizeAlloca)} - }; - llvm::SmallVector totalFields = { - {"Number of allocas", std::to_string(entry.totalAllocaCount) }, - {"Batch memory", formatReportMemory(entry.totalAllocaBytes)} - }; - printReportPerCoreAndTotalFields(os, perCoreFields, totalFields); -} - -static MemoryReportRow addMemoryReportRows(const MemoryReportRow& lhs, const MemoryReportRow& rhs) { - MemoryReportRow result = lhs; - result.numAlloca += rhs.numAlloca; - result.sizeAlloca += rhs.sizeAlloca; - result.numGlobal += rhs.numGlobal; - result.sizeGlobal += rhs.sizeGlobal; - return result; -} - MemoryReportRow PimMemory::getReportRow() const { MemoryReportRow row = reportRow; - row.numAlloca = localPhysicalSlots.size(); - row.sizeAlloca = 0; - for (const PhysicalSlotInfo& slot : localPhysicalSlots) - row.sizeAlloca += slot.size; + row.hostBytes = firstAvailableAddress; + row.physicalLocalBytes = localArenaSize.value_or(0); return row; } @@ -379,29 +355,32 @@ llvm::FailureOr PimAcceleratorMemory::getIndexValue(mlir::Value value, return compiledIt->second.evaluate(knowledge); } +PimAcceleratorMemory::PimAcceleratorMemory() +: hostMem(memEntriesMap), fileReport(openMemoryReport(pimMemoryReport == PimMemoryReportSummary)) {} + +PimAcceleratorMemory::PimAcceleratorMemory( + const llvm::SmallDenseMap& 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, - row.numAlloca, - row.sizeAlloca}); + row}); } void PimAcceleratorMemory::recordBatchReport(uint64_t batchId, ArrayRef coreIds, - const MemoryReportRow& perCoreRow, - uint64_t totalAllocaCount, - uint64_t totalAllocaBytes) { + const MemoryReportRow& perCoreRow) { MemoryReportEntry entry; entry.kind = MemoryReportEntry::Kind::Batch; entry.id = batchId; llvm::append_range(entry.coreIds, coreIds); entry.row = perCoreRow; - entry.totalAllocaCount = totalAllocaCount; - entry.totalAllocaBytes = totalAllocaBytes; reportEntries.push_back(std::move(entry)); } @@ -410,93 +389,91 @@ void PimAcceleratorMemory::flushReport() { return; llvm::raw_os_ostream os(fileReport); - uint64_t totalGlobalMemory = hostReportRow.has_value() ? hostReportRow->sizeGlobal : 0; - uint64_t totalWeightsMemory = totalWeightBytes; - uint64_t totalCoresMemory = 0; - uint64_t totalLogicalCoreMemory = 0; - uint64_t totalFallbackIntervals = 0; - uint64_t totalNestedSingleUseIntervals = 0; - uint64_t maximumCorePeak = 0; - std::string worstEntry = ""; - 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(savedCoreMemory) / static_cast(totalLogicalCoreMemory); - - llvm::SmallVector 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) } + struct ReportGroup { + MemoryReportRow row; + SmallVector coreIds; + std::optional batchId; + }; + SmallVector 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(group.coreIds)); }; - printReportTotalsBlock(os, totalFields); - if (hostReportRow.has_value()) { - os << "\nHost:\n"; - printHostMemoryReportRow(os, *hostReportRow); + 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"; } - if (!reportEntries.empty()) { - if (hostReportRow.has_value()) - os << "\n"; - sortReportEntriesByFirstCore(reportEntries); - - for (size_t index = 0; index < reportEntries.size();) { - size_t runEnd = index + 1; - while (runEnd < reportEntries.size() && reportEntries[runEnd].kind == reportEntries[index].kind - && reportEntries[runEnd].row == reportEntries[index].row - && reportEntries[runEnd].totalAllocaCount == reportEntries[index].totalAllocaCount - && reportEntries[runEnd].totalAllocaBytes == reportEntries[index].totalAllocaBytes) { - ++runEnd; - } - - if (reportEntries[index].kind == MemoryReportEntry::Kind::Batch) { - os << "Batch "; - for (size_t batchIndex = index; batchIndex < runEnd; ++batchIndex) { - if (batchIndex != index) - os << ",\n "; - os << reportEntries[batchIndex].id << " (cores "; - printCompressedIntegerEntries(os, ArrayRef(reportEntries[batchIndex].coreIds)); - os << ")"; - } - } - else { - llvm::SmallVector coreIds; - for (size_t coreIndex = index; coreIndex < runEnd; ++coreIndex) - coreIds.push_back(reportEntries[coreIndex].coreIds.front()); - os << "Core "; - printCompressedIntegerEntries(os, ArrayRef(coreIds)); - } - os << ":\n"; - if (reportEntries[index].kind == MemoryReportEntry::Kind::Batch) - printBatchMemoryReportRow(os, reportEntries[index]); - else - printCoreMemoryReportRow(os, reportEntries[index]); - printReportEntrySeparator(os, runEnd < reportEntries.size()); - - index = runEnd; + 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(); @@ -834,25 +811,20 @@ static SmallVector collectTopLevelCoreLikeOps(func::FuncOp funcOp) { static FailureOr compileCoreMemoryPlan(Operation* coreLikeOp) { CompiledCoreMemoryPlan plan; - auto fallbackAttr = coreLikeOp->getAttrOfType(kLocalMemoryFallbackCountAttrName); - auto nestedSingleUseAttr = coreLikeOp->getAttrOfType(kLocalMemoryNestedSingleUseCountAttrName); - if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) { - coreLikeOp->emitError("requires complete PIM local-memory planning summary attributes before codegen"); + auto arenaAttr = coreLikeOp->getAttrOfType(kLocalMemorySizeAttrName); + if (!arenaAttr || arenaAttr.getInt() < 0 + || static_cast(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) { + coreLikeOp->emitError("requires a valid pim.local_memory_size attribute before codegen"); return failure(); } - plan.fallbackIntervals = static_cast(fallbackAttr.getInt()); - plan.nestedSingleUseIntervals = static_cast(nestedSingleUseAttr.getInt()); - SmallDenseMap slotIndices; + plan.arenaSize = static_cast(arenaAttr.getInt()); bool hasFailure = false; coreLikeOp->walk([&](memref::AllocOp allocOp) { if (hasFailure) return; auto addressAttr = allocOp->getAttrOfType(kLocalMemoryAddressAttrName); - auto slotAttr = allocOp->getAttrOfType(kLocalMemorySlotAttrName); - auto slotSizeAttr = allocOp->getAttrOfType(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"); + if (!addressAttr || addressAttr.getInt() < 0) { + allocOp.emitOpError("requires a non-negative pim.local_memory_address attribute before codegen"); hasFailure = true; return; } @@ -864,25 +836,23 @@ static FailureOr compileCoreMemoryPlan(Operation* coreLi return; } size_t address = static_cast(addressAttr.getInt()); - size_t slotId = static_cast(slotAttr.getInt()); - size_t slotSize = static_cast(slotSizeAttr.getInt()); - plan.entries.push_back({allocOp.getResult(), {address, static_cast(*checkedSize)}}); - plan.logicalBytes += *checkedSize; - - auto [slotIt, inserted] = slotIndices.try_emplace(slotId, plan.slots.size()); - if (inserted) { - plan.slots.push_back({slotId, 0, slotSize}); + if (address > plan.arenaSize || *checkedSize > plan.arenaSize - address) { + allocOp.emitOpError("has a planned local-memory range outside its core arena"); + hasFailure = true; return; } - const PhysicalSlotInfo& existing = plan.slots[slotIt->second]; - if (existing.size != slotSize) { - allocOp.emitOpError("has a PIM local-memory arena inconsistent with another allocation"); + plan.entries.push_back({allocOp.getResult(), {address, static_cast(*checkedSize)}}); + auto logicalBytes = pim::checkedAdd( + static_cast(plan.logicalBytes), static_cast(*checkedSize), allocOp, "logical local bytes"); + if (failed(logicalBytes)) { hasFailure = true; + return; } + plan.logicalBytes = *logicalBytes; + ++plan.logicalAllocationCount; }); if (hasFailure) return failure(); - llvm::sort(plan.slots, [](const PhysicalSlotInfo& lhs, const PhysicalSlotInfo& rhs) { return lhs.id < rhs.id; }); return plan; } @@ -1462,7 +1432,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: for (const SmallVector& group : batchJobIndices) { SmallVector reportedCoreIds; - MemoryReportRow batchRow; std::optional batchPerCoreRow; for (size_t jobIndex : group) { @@ -1475,15 +1444,13 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: reportedCoreIds.push_back(pim::checkedI32OrCrash(job.emittedCoreId, "batch report core id")); if (!batchPerCoreRow) batchPerCoreRow = result.reportRow; - batchRow = addMemoryReportRows(batchRow, 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 {}), - batchRow.numAlloca, - batchRow.sizeAlloca); + memory.recordBatchReport( + batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {})); } maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1; diff --git a/src/PIM/Compiler/PimCodeGen.hpp b/src/PIM/Compiler/PimCodeGen.hpp index 385d6b7..41bc893 100644 --- a/src/PIM/Compiler/PimCodeGen.hpp +++ b/src/PIM/Compiler/PimCodeGen.hpp @@ -33,12 +33,6 @@ struct MemEntry { size_t size; }; -struct PhysicalSlotInfo { - size_t id = 0; - size_t address = 0; - size_t size = 0; -}; - struct MemoryValueKey { mlir::Value value; std::optional lane; @@ -53,26 +47,22 @@ struct CompiledLocalMemoryEntry { struct CompiledCoreMemoryPlan { llvm::SmallVector entries; - llvm::SmallVector slots; + uint64_t logicalAllocationCount = 0; uint64_t logicalBytes = 0; - uint64_t fallbackIntervals = 0; - uint64_t nestedSingleUseIntervals = 0; + size_t arenaSize = 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; + uint64_t hostObjectCount = 0; + uint64_t hostBytes = 0; + uint64_t logicalLocalAllocationCount = 0; + uint64_t logicalLocalBytes = 0; + uint64_t physicalLocalBytes = 0; bool operator==(const MemoryReportRow& other) const { - return numAlloca == other.numAlloca && sizeAlloca == other.sizeAlloca && numGlobal == other.numGlobal - && sizeGlobal == other.sizeGlobal && logicalAllocaBytes == other.logicalAllocaBytes - && fallbackIntervals == other.fallbackIntervals - && nestedSingleUseIntervals == other.nestedSingleUseIntervals; + return hostObjectCount == other.hostObjectCount && hostBytes == other.hostBytes + && logicalLocalAllocationCount == other.logicalLocalAllocationCount + && logicalLocalBytes == other.logicalLocalBytes && physicalLocalBytes == other.physicalLocalBytes; } }; @@ -99,16 +89,14 @@ struct MemoryReportEntry { uint64_t id = 0; llvm::SmallVector coreIds; MemoryReportRow row; - uint64_t totalAllocaCount = 0; - uint64_t totalAllocaBytes = 0; }; class PimMemory { llvm::SmallVector memEntries; - llvm::SmallVector localPhysicalSlots; llvm::SmallDenseMap& globalMemEntriesMap; llvm::SmallDenseMap ownedMemEntriesMap; MemoryReportRow reportRow; + std::optional localArenaSize; size_t minAlignment = 4; size_t firstAvailableAddress = 0; @@ -145,12 +133,9 @@ private: mutable llvm::DenseMap compiledAddressExprs; public: - PimAcceleratorMemory() - : hostMem(memEntriesMap), fileReport(openReportFile("memory_report")) {} - PimAcceleratorMemory(const llvm::SmallDenseMap& initialMemEntries, bool enableReport) - : memEntriesMap(initialMemEntries), - hostMem(memEntriesMap), - fileReport(enableReport ? openReportFile("memory_report") : std::fstream()) {} + PimAcceleratorMemory(); + PimAcceleratorMemory( + const llvm::SmallDenseMap& initialMemEntries, bool enableReport); PimMemory& getOrCreateDeviceMem(size_t id); @@ -160,11 +145,8 @@ public: llvm::FailureOr getIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge = {}) const; void reportHost(); void recordCoreReport(size_t coreId, const MemoryReportRow& row); - void recordBatchReport(uint64_t batchId, - llvm::ArrayRef coreIds, - const MemoryReportRow& perCoreRow, - uint64_t totalAllocaCount, - uint64_t totalAllocaBytes); + void recordBatchReport( + uint64_t batchId, llvm::ArrayRef coreIds, const MemoryReportRow& perCoreRow); void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; } void flushReport(); }; diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index 5cf8a33..5e50aa7 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -19,10 +19,8 @@ llvm::cl::opt pimMemoryReport( "pim-memory-report", llvm::cl::desc("Emit a human-readable PIM memory planning report"), llvm::cl::values(clEnumValN(PimMemoryReportNone, "none", "Do not emit any PIM memory planning report")), - llvm::cl::values( - clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise slot reuse report with key offenders")), - llvm::cl::values(clEnumValN(PimMemoryReportFull, "full", "Emit the full detailed PIM memory planning report")), - llvm::cl::init(PimMemoryReportNone), + llvm::cl::values(clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise PIM memory summary")), + llvm::cl::init(PimMemoryReportSummary), llvm::cl::cat(OnnxMlirOptions)); llvm::cl::opt pimConvLowering( @@ -72,12 +70,6 @@ llvm::cl::opt llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions)); -llvm::cl::opt - pimDisableMemoryCoalescing("pim-disable-memory-coalescing", - llvm::cl::desc("Skip the early PIM IR memory coalescing pass (developer diagnostic)"), - llvm::cl::init(false), - llvm::cl::cat(OnnxMlirOptions)); - llvm::cl::opt useExperimentalConvImpl("use-experimental-conv-impl", llvm::cl::desc("Use experimental implementation for convolution"), llvm::cl::init(false), diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index d109ccc..5e58c23 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -23,7 +23,6 @@ typedef enum { typedef enum { PimMemoryReportNone = 0, PimMemoryReportSummary = 1, - PimMemoryReportFull = 2, } PimMemoryReportLevel; typedef enum { @@ -53,7 +52,6 @@ extern llvm::cl::opt pimMemoryReport; extern llvm::cl::opt pimConvLowering; extern llvm::cl::opt pimExportSpatialDataflow; -extern llvm::cl::opt pimDisableMemoryCoalescing; extern llvm::cl::opt pimOnlyCodegen; extern llvm::cl::opt useExperimentalConvImpl; extern llvm::cl::opt pimEmitJson; diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index 54a4517..ad48d03 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -22,6 +22,7 @@ void addPassesPim(OwningOpRef& module, if (pimOnlyCodegen) { pm.addPass(createPimLocalMemoryPlanningPass()); + pm.addPass(createPimVerificationPass()); pm.addPass(createEmitPimCodePass()); return; } @@ -52,8 +53,6 @@ void addPassesPim(OwningOpRef& module, pm.addPass(mlir::createLowerAffinePass()); pm.addPass(createPimHostConstantFoldingPass()); 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()); diff --git a/src/PIM/Dialect/Pim/Analysis/CMakeLists.txt b/src/PIM/Dialect/Pim/Analysis/CMakeLists.txt index 43336bf..d7452e3 100644 --- a/src/PIM/Dialect/Pim/Analysis/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/Analysis/CMakeLists.txt @@ -5,4 +5,8 @@ add_pim_library(OMPimLocalMemoryLifetimeAnalysis INCLUDE_DIRS PUBLIC ${PIM_PUBLIC_INCLUDE_DIRS} + + LINK_LIBS PUBLIC + OMPimCommon + PimOps ) diff --git a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp b/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp index 0bf3da7..ea8013c 100644 --- a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp +++ b/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp @@ -2,70 +2,204 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Interfaces/DestinationStyleOpInterface.h" +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" #include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" using namespace mlir; -namespace onnx_mlir { -namespace pim { +namespace onnx_mlir::pim { +namespace { + +struct OperationOrdering { + DenseMap position; + DenseMap subtreeEnd; + uint64_t next = 0; +}; + +static void orderOperation(Operation* op, OperationOrdering& ordering) { + uint64_t end = ordering.next; + ordering.position[op] = ordering.next++; + for (Region& region : op->getRegions()) + for (Block& block : region) + for (Operation& nested : block) { + orderOperation(&nested, ordering); + end = std::max(end, ordering.subtreeEnd.lookup(&nested)); + } + ordering.subtreeEnd[op] = end; +} + +static OperationOrdering buildOrdering(Operation* coreLikeOp) { + OperationOrdering ordering; + if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty()) + return ordering; + for (Operation& op : coreLikeOp->getRegion(0).front()) + orderOperation(&op, ordering); + return ordering; +} + +static bool isRuntimeMemoryTouch(Operation* op) { + return isa(op); +} + +static bool isIgnoredUser(Operation* op) { + return isLocalMemoryAliasOp(op) || isa(op) + || isCoreStaticAddressOp(op); +} + +static bool isWithin(Value value, Region* region) { + if (auto argument = dyn_cast(value)) + return argument.getOwner()->getParent() == region; + if (Operation* definingOp = value.getDefiningOp()) + return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion()); + return false; +} + +static std::pair getTouchRange(Value definingValue, Operation* user, const OperationOrdering& ordering) { + uint64_t start = ordering.position.lookup(user); + uint64_t end = start; + for (Operation* current = user; current; current = current->getParentOp()) + if (auto forOp = dyn_cast(current); forOp && !isWithin(definingValue, &forOp.getRegion())) { + start = std::min(start, ordering.position.lookup(forOp)); + end = std::max(end, ordering.subtreeEnd.lookup(forOp)); + } + return {start, end}; +} + +static FailureOr getAllocationSize(memref::AllocOp allocation) { + auto type = dyn_cast(allocation.getType()); + if (!type) + return failure(); + auto bytes = getCheckedShapedTypeSizeInBytes(type, allocation, "memory allocation byte size"); + if (failed(bytes)) + return failure(); + return checkedSize(*bytes, allocation, "memory allocation byte size"); +} + +} // namespace bool isLocalMemoryAliasOp(Operation* op) { return isa(op); } -LogicalResult walkLocalMemoryUses(Value root, - llvm::function_ref visitUser) { +LogicalResult walkLocalMemoryUses(Value root, LocalMemoryUseCallback callback) { llvm::SmallPtrSet visitedValues; - llvm::SmallPtrSet visitedUsers; - llvm::SmallVector pendingValues {root}; - auto addAlias = [&](Value value) { pendingValues.push_back(value); }; - - while (!pendingValues.empty()) { - Value value = pendingValues.pop_back_val(); + llvm::SmallPtrSet visitedUses; + SmallVector pending = {root}; + while (!pending.empty()) { + Value value = pending.pop_back_val(); if (!visitedValues.insert(value).second) continue; - for (Operation* user : value.getUsers()) { - if (!visitedUsers.insert(user).second) - continue; - if (failed(visitUser(value, user))) + for (OpOperand& use : value.getUses()) { + if (!visitedUses.insert(&use).second || failed(callback(use))) return failure(); + Operation* user = use.getOwner(); - if (isLocalMemoryAliasOp(user)) - for (Value result : user->getResults()) - addAlias(result); + if (isLocalMemoryAliasOp(user) && use.getOperandNumber() == 0) { + if (user->getNumResults() != 1) + return failure(); + pending.push_back(user->getResult(0)); + } - if (auto dpsOp = dyn_cast(user)) - for (OpResult result : user->getResults()) - if (OpOperand* tied = dpsOp.getTiedOpOperand(result); tied && tied->get() == value) - addAlias(result); + if (auto dpsOp = dyn_cast(user); dpsOp && dpsOp.isDpsInit(&use)) + pending.push_back(dpsOp.getTiedOpResult(&use)); if (auto forOp = dyn_cast(user)) - for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs())) - if (initArg == value) { - addAlias(forOp.getRegionIterArgs()[index]); - addAlias(forOp.getResult(index)); + for (auto [index, init] : llvm::enumerate(forOp.getInitArgsMutable())) + if (&init == &use) { + pending.push_back(forOp.getRegionIterArgs()[index]); + pending.push_back(forOp.getResult(index)); } - auto yieldOp = dyn_cast(user); - if (!yieldOp) - continue; - for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) { - if (operand != value) - continue; - if (auto forOp = dyn_cast(yieldOp->getParentOp())) - addAlias(forOp.getResult(index)); - else if (auto ifOp = dyn_cast(yieldOp->getParentOp())) - addAlias(ifOp.getResult(index)); - else if (auto switchOp = dyn_cast(yieldOp->getParentOp())) - addAlias(switchOp.getResult(index)); + if (auto yieldOp = dyn_cast(user)) { + unsigned index = use.getOperandNumber(); + Operation* parent = yieldOp->getParentOp(); + if (auto forOp = dyn_cast(parent)) + pending.push_back(forOp.getResult(index)); + else if (auto ifOp = dyn_cast(parent)) + pending.push_back(ifOp.getResult(index)); + else if (auto switchOp = dyn_cast(parent)) + pending.push_back(switchOp.getResult(index)); + else + return failure(); } } } return success(); } -} // namespace pim -} // namespace onnx_mlir +FailureOr> analyzeLocalMemoryLifetimes(Operation* coreLikeOp) { + SmallVector intervals; + OperationOrdering ordering = buildOrdering(coreLikeOp); + if (ordering.position.empty()) + return intervals; + uint64_t fallbackEnd = ordering.next - 1; + bool failedSize = false; + coreLikeOp->walk([&](memref::AllocOp allocation) { + auto size = getAllocationSize(allocation); + if (failed(size)) { + failedSize = true; + return; + } + + LocalMemoryInterval interval {allocation, ordering.position.lookup(allocation), ordering.position.lookup(allocation), *size}; + bool hasRuntimeTouch = false; + bool needsFallback = false; + if (failed(walkLocalMemoryUses(allocation.getResult(), [&](OpOperand& use) { + Operation* user = use.getOwner(); + if (isRuntimeMemoryTouch(user)) { + auto [start, end] = getTouchRange(allocation.getResult(), user, ordering); + if (!hasRuntimeTouch) { + interval.start = start; + interval.end = end; + hasRuntimeTouch = true; + } + else { + interval.start = std::min(interval.start, start); + interval.end = std::max(interval.end, end); + } + } + else if (!isIgnoredUser(user)) { + needsFallback = true; + } + return success(); + }))) + needsFallback = true; + if (!hasRuntimeTouch || needsFallback) { + interval.start = std::min(interval.start, ordering.position.lookup(allocation)); + interval.end = fallbackEnd; + } + intervals.push_back(interval); + }); + if (failedSize) + return failure(); + return intervals; +} + +bool localMemoryLifetimesOverlap(const LocalMemoryInterval& lhs, const LocalMemoryInterval& rhs) { + return !(lhs.end < rhs.start || rhs.end < lhs.start); +} + +} // namespace onnx_mlir::pim diff --git a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp b/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp index e9b1147..f40e39a 100644 --- a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp +++ b/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp @@ -1,17 +1,25 @@ #pragma once -#include "mlir/IR/Operation.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "llvm/ADT/STLFunctionalExtras.h" +#include "llvm/ADT/SmallVector.h" -namespace onnx_mlir { -namespace pim { +namespace onnx_mlir::pim { + +struct LocalMemoryInterval { + mlir::memref::AllocOp allocation; + uint64_t start = 0; + uint64_t end = 0; + size_t size = 0; +}; + +using LocalMemoryUseCallback = llvm::function_ref; bool isLocalMemoryAliasOp(mlir::Operation* op); +mlir::LogicalResult walkLocalMemoryUses(mlir::Value root, LocalMemoryUseCallback callback); +mlir::FailureOr> +analyzeLocalMemoryLifetimes(mlir::Operation* coreLikeOp); +bool localMemoryLifetimesOverlap(const LocalMemoryInterval& lhs, const LocalMemoryInterval& rhs); -mlir::LogicalResult walkLocalMemoryUses( - mlir::Value root, - llvm::function_ref visitUser); - -} // namespace pim -} // namespace onnx_mlir +} // namespace onnx_mlir::pim diff --git a/src/PIM/Dialect/Pim/CMakeLists.txt b/src/PIM/Dialect/Pim/CMakeLists.txt index 9badbab..b25921d 100644 --- a/src/PIM/Dialect/Pim/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/CMakeLists.txt @@ -4,7 +4,6 @@ add_onnx_mlir_dialect_doc(pim Pim.td) add_subdirectory(Analysis) add_subdirectory(Transforms/Bufferization) add_subdirectory(Transforms/HostConstantFolding) -add_subdirectory(Transforms/MemoryCoalescing) add_subdirectory(Transforms/LocalMemoryPlanning) add_subdirectory(Transforms/Verification) diff --git a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt index d6e09bb..43a4530 100644 --- a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt @@ -8,7 +8,6 @@ add_pim_library(OMPimLocalMemoryPlanning LINK_LIBS PUBLIC OMPimCommon - OMPimCompilerOptions OMPimLocalMemoryLifetimeAnalysis PimOps ) diff --git a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp index d2a9993..4a85565 100644 --- a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp +++ b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp @@ -1,668 +1,169 @@ -#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Pass/Pass.h" -#include "mlir/IR/Value.h" -#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" -#include "llvm/Support/MathExtras.h" -#include "llvm/Support/FileSystem.h" -#include "llvm/Support/raw_os_ostream.h" -#include "llvm/Support/raw_ostream.h" -#include -#include #include -#include #include -#include -#include "Common/Support/CheckedArithmetic.hpp" -#include "Common/Support/ReportUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" #include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp" #include "src/Accelerators/PIM/Pass/PIMPasses.h" using namespace llvm; using namespace mlir; -using namespace onnx_mlir; - -namespace { - -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; -}; - -struct OperationOrdering { - llvm::DenseMap position; - llvm::DenseMap subtreeEnd; - uint64_t nextPosition = 0; -}; - -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) { - std::string text; - llvm::raw_string_ostream os(text); - if (auto result = dyn_cast(value)) - os << result.getOwner()->getName() << '#' << result.getResultNumber(); - else if (auto blockArg = dyn_cast(value)) - os << "block_arg#" << blockArg.getArgNumber(); - else - os << ""; - os << " : " << value.getType(); - os.flush(); - return abbreviate(text, maxLen); -} - -static std::string summarizeOperation(Operation* op, size_t maxLen = 96) { - if (!op) - return ""; - return abbreviate(op->getName().getStringRef(), 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 isRuntimeMemoryTouchOp(Operation* op) { - return isa(op); -} - -static bool isIgnoredLivenessUser(Operation* op) { - return pim::isLocalMemoryAliasOp(op) || isa(op) - || isCoreStaticAddressOp(op); -} - -static bool isWithin(mlir::Value value, Region* region) { - if (!region) - return false; - if (auto blockArg = dyn_cast(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(); -} - -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(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) { - MemoryTouchInterval interval; - interval.start = ordering.position.lookup(allocOp); - interval.end = interval.start; - - auto parentLoop = allocOp->getParentOfType(); - (void) pim::walkLocalMemoryUses( - allocOp.getResult(), - [&](mlir::Value value, Operation* user) { - if (auto forOp = dyn_cast(user); - forOp && parentLoop && forOp != parentLoop && llvm::is_contained(forOp.getInitArgs(), value)) - interval.escapesLoop = true; - if (auto yieldOp = dyn_cast(user)) { - auto forOp = dyn_cast(yieldOp->getParentOp()); - auto ifOp = dyn_cast(yieldOp->getParentOp()); - auto indexSwitch = dyn_cast(yieldOp->getParentOp()); - if (!forOp && !ifOp && !indexSwitch) - addFallbackReason(interval.fallbackReason, "yield without scf.for parent"); - else if (forOp && parentLoop && forOp == parentLoop && llvm::is_contained(yieldOp.getOperands(), value)) - 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; - } - return success(); - } - - if (isIgnoredLivenessUser(user)) - return success(); - - addFallbackReason(interval.fallbackReason, "unhandled user op"); - interval.endUsedFallback = true; - return success(); - }); - - 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 getAllocSizeBytes(memref::AllocOp allocOp) { - auto type = dyn_cast(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 intervals) { - uint64_t slotLogicalBytes = 0; - for (size_t intervalIndex : slot.intervalIndices) - slotLogicalBytes += intervals[intervalIndex].size; - return slotLogicalBytes; -} - -static bool placementsOverlap(const PlannedPhysicalSlot& lhs, const PlannedPhysicalSlot& rhs) { - return lhs.address < rhs.address + rhs.requiredSize && rhs.address < lhs.address + lhs.requiredSize; -} - -static bool hasAddressReuse(size_t slotIndex, ArrayRef slots) { - if (slots[slotIndex].intervalIndices.size() > 1) - return true; - return llvm::any_of(llvm::enumerate(slots), [&](auto indexedSlot) { - return indexedSlot.index() != slotIndex && placementsOverlap(slots[slotIndex], indexedSlot.value()); - }); -} - -} // namespace - -SmallVector onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp) { - SmallVector 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); - LocalAllocInterval interval; - interval.id = nextIntervalId++; - interval.alloc = allocOp; - 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.valueSummary = summarizeValue(allocOp.getResult(), 88); - interval.firstTouchSummary = summarizeOperation(touchInterval.firstTouchOp); - interval.lastTouchSummary = summarizeOperation(touchInterval.lastTouchOp); - intervals.push_back(std::move(interval)); - }); - - return intervals; -} - -SmallVector onnx_mlir::planPhysicalSlots(MutableArrayRef intervals) { - SmallVector slots; - SmallVector 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]; - SmallVector conflictingSlots; - SmallVector candidateAddresses = {0}; - size_t currentPeak = 0; - for (const PlannedPhysicalSlot& slot : slots) { - currentPeak = std::max(currentPeak, slot.address + slot.requiredSize); - if (llvm::any_of(slot.intervalIndices, [&](size_t otherIndex) { - return intervalsOverlap(interval, intervals[otherIndex]); - })) { - conflictingSlots.push_back(&slot); - size_t candidate = llvm::alignTo(slot.address + slot.requiredSize, 4); - if (!llvm::is_contained(candidateAddresses, candidate)) - candidateAddresses.push_back(candidate); - } - } - - size_t bestAddress = std::numeric_limits::max(); - auto bestKey = std::tuple(std::numeric_limits::max(), - std::numeric_limits::max()); - for (size_t candidate : candidateAddresses) { - bool overlaps = llvm::any_of(conflictingSlots, [&](const PlannedPhysicalSlot* slot) { - return candidate < slot->address + slot->requiredSize && slot->address < candidate + interval.size; - }); - if (overlaps) - continue; - auto candidateKey = std::tuple(std::max(currentPeak, candidate + interval.size), candidate); - if (candidateKey < bestKey) { - bestKey = candidateKey; - bestAddress = candidate; - } - } - assert(bestAddress != std::numeric_limits::max() && "address after all conflicts must be available"); - - auto reusable = llvm::find_if(slots, [&](const PlannedPhysicalSlot& slot) { - return slot.address == bestAddress && slot.requiredSize == interval.size; - }); - if (reusable != slots.end()) { - reusable->intervalIndices.push_back(intervalIndex); - interval.slotPlanIndex = static_cast(reusable - slots.begin()); - } - else { - slots.push_back({slots.size(), interval.size, bestAddress, {intervalIndex}}); - interval.slotPlanIndex = slots.size() - 1; - } - } - - return slots; -} - -CoreMemoryPlan onnx_mlir::buildCoreMemoryPlan(Operation* coreLikeOp) { - CoreMemoryPlan plan; - plan.coreLikeOp = coreLikeOp; - plan.intervals = buildLocalAllocIntervals(coreLikeOp); - plan.slots = planPhysicalSlots(plan.intervals); - return plan; -} - -LogicalResult onnx_mlir::assignPhysicalSlotAddresses(CoreMemoryPlan& plan, size_t addressLimit) { - SmallVector slotOrder(plan.slots.size()); - std::iota(slotOrder.begin(), slotOrder.end(), 0); - llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) { - const PlannedPhysicalSlot& lhs = plan.slots[lhsIndex]; - const PlannedPhysicalSlot& rhs = plan.slots[rhsIndex]; - if (lhs.requiredSize != rhs.requiredSize) - return lhs.requiredSize > rhs.requiredSize; - return lhs.id < rhs.id; - }); - - size_t nextSlotId = 0; - for (size_t slotIndex : slotOrder) { - PlannedPhysicalSlot& slot = plan.slots[slotIndex]; - if (slot.address > addressLimit || slot.requiredSize > addressLimit - slot.address) { - plan.coreLikeOp->emitError() << "PIM local memory plan exceeds the signed int32 address range while assigning " - "a " - << slot.requiredSize << " byte physical placement at address " << slot.address; - return failure(); - } - - slot.id = nextSlotId++; - } - return success(); -} - -std::string onnx_mlir::buildMemoryPlanReport(Operation* coreLikeOp, - ArrayRef intervals, - ArrayRef slots, - size_t addressLimit, - PimMemoryReportLevel reportLevel) { - std::string 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); - if (interval.startUsedAllocFallback || interval.endUsedFallback) - ++fallbackIntervals; - if (!interval.hasRuntimeUse) - ++noRuntimeTouchIntervals; - if (interval.insideNestedRegion) - ++nestedIntervals; - if (interval.escapesLoop) - ++loopEscapingIntervals; - } - for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) { - const PlannedPhysicalSlot& slot = slots[slotIndex]; - largestPhysicalSlot = std::max(largestPhysicalSlot, slot.requiredSize); - maximumAssignedAddress = std::max(maximumAssignedAddress, slot.address + slot.requiredSize); - if (hasAddressReuse(slotIndex, slots)) - reusedAllocations += slot.intervalIndices.size(); - } - totalPhysicalBytes = maximumAssignedAddress; - - uint64_t savedBytes = totalLogicalBytes >= totalPhysicalBytes ? totalLogicalBytes - totalPhysicalBytes : 0; - double savedPercent = - totalLogicalBytes == 0 ? 0.0 : 100.0 * static_cast(savedBytes) / static_cast(totalLogicalBytes); - - raw_string_ostream os(artifacts); - os << "=== PIM Memory Liveness Report ===\n"; - os << "Op: " << coreLikeOp->getName() << "\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 << " address placements: " << slots.size() << "\n"; - os << " allocations sharing address space: " << 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 address placement: " << largestPhysicalSlot << "\n"; - os << " address limit: " << addressLimit << "\n"; - os << " peak physical memory: " << formatReportMemory(maximumAssignedAddress) << " (" << maximumAssignedAddress - << ")\n"; - os << " maximum assigned address: " << maximumAssignedAddress << "\n"; - - SmallVector reusedSlots; - SmallVector singleUseSlots; - for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) { - if (hasAddressReuse(slotIndex, slots)) - reusedSlots.push_back(&slots[slotIndex]); - else - singleUseSlots.push_back(&slots[slotIndex]); - } - - 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->requiredSize != rhs->requiredSize) - return lhs->requiredSize > rhs->requiredSize; - return lhs->id < rhs->id; - }); - llvm::stable_sort(singleUseSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) { - if (lhs->requiredSize != rhs->requiredSize) - return lhs->requiredSize > rhs->requiredSize; - return lhs->id < rhs->id; - }); - - constexpr size_t kSummaryReuseLimit = 6; - constexpr size_t kSummaryOffenderLimit = 10; - - os << "\nBest Address Reuse:\n"; - if (reusedSlots.empty()) { - os << " no address ranges were reused\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->requiredSize) - << " intervals=" << slot->intervalIndices.size() << " logical_sum=" << formatReportMemory(slotLogicalBytes) - << "\n"; - } - } - - os << "\nTop Offenders:\n"; - for (const PlannedPhysicalSlot* slot : ArrayRef(singleUseSlots).take_front(kSummaryOffenderLimit)) { - const LocalAllocInterval& interval = intervals[slot->intervalIndices.front()]; - os << " slot #" << slot->id << " is single-use" - << " size=" << formatReportMemory(slot->requiredSize) << " interval=#" << interval.id - << " value=" << abbreviate(interval.valueSummary, 56) << "\n"; - os << " first=" << abbreviate(interval.firstTouchSummary, 40) - << " last=" << abbreviate(interval.lastTouchSummary, 40) - << " nested=" << (interval.insideNestedRegion ? "yes" : "no") - << " escapes_loop=" << (interval.escapesLoop ? "yes" : "no") << "\n"; - } - if (singleUseSlots.empty()) - 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.requiredSize) - << " (" << slot.requiredSize << ")" - << " intervals=" << slot.intervalIndices.size() << " logical_sum=" << formatReportMemory(slotLogicalBytes) - << "\n"; - for (size_t intervalIndex : slot.intervalIndices) { - const LocalAllocInterval& interval = intervals[intervalIndex]; - os << " [" << interval.start << "," << interval.end << "]" - << " #" << interval.id << " logical=" << formatReportMemory(interval.size) << " (" << interval.size - << ")" - << " nested=" << (interval.insideNestedRegion ? "yes" : "no") - << " escapes_loop=" << (interval.escapesLoop ? "yes" : "no") - << " first=" << abbreviate(interval.firstTouchSummary, 48) - << " last=" << abbreviate(interval.lastTouchSummary, 48) << "\n"; - os << " value=" << abbreviate(interval.valueSummary, 72) << "\n"; - if (!interval.fallbackReason.empty()) - os << " fallback_reason=" << interval.fallbackReason << "\n"; - } - } - } - os.flush(); - - return artifacts; -} namespace onnx_mlir { namespace { +static bool rangesOverlap(size_t lhsAddress, size_t lhsSize, size_t rhsAddress, size_t rhsSize) { + return lhsAddress < rhsAddress + rhsSize && rhsAddress < lhsAddress + lhsSize; +} + +static FailureOr alignedEnd(size_t address, size_t size, size_t limit) { + if (address > limit || size > limit - address) + return failure(); + size_t end = address + size; + size_t padding = (4 - end % 4) % 4; + if (padding > limit - end) + return failure(); + return end + padding; +} + struct PimLocalMemoryPlanningPass : PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimLocalMemoryPlanningPass) StringRef getArgument() const override { return "pim-local-memory-planning"; } StringRef getDescription() const override { - return "Plan liveness-based physical slots and addresses for PIM core-local memory"; + return "Plan liveness-based addresses for PIM core-local memory"; } void runOnOperation() override { - std::fstream reportFile; - std::unique_ptr report; - if (pimMemoryReport != PimMemoryReportNone) { - reportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt"); - if (reportFile.is_open()) - report = std::make_unique(reportFile); - } - else if (std::string outputRoot = getOutputDir(); !outputRoot.empty()) { - sys::fs::remove(outputRoot + "/reports/pim_memory_liveness_report.txt"); - } - Builder builder(&getContext()); - uint64_t nextBatchId = 0; - bool failedPlanning = false; + bool hasFailure = false; getOperation().walk([&](Operation* op) { - if (failedPlanning || !isa(op)) + if (hasFailure || !isa(op)) return; - CoreMemoryPlan plan = buildCoreMemoryPlan(op); - if (failed(assignPhysicalSlotAddresses(plan, kPimLocalMemoryAddressLimit))) { - failedPlanning = true; - return; - } - - size_t arenaSize = 0; - for (const PlannedPhysicalSlot& placement : plan.slots) - arenaSize = std::max(arenaSize, placement.address + placement.requiredSize); - for (const LocalAllocInterval& interval : plan.intervals) { - const PlannedPhysicalSlot& slot = plan.slots[interval.slotPlanIndex]; - interval.alloc->setAttr(kLocalMemoryAddressAttrName, builder.getI64IntegerAttr(slot.address)); - interval.alloc->setAttr(kLocalMemorySlotAttrName, builder.getI64IntegerAttr(0)); - interval.alloc->setAttr(kLocalMemorySlotSizeAttrName, builder.getI64IntegerAttr(arenaSize)); - } - uint64_t fallbackIntervals = llvm::count_if(plan.intervals, [](const LocalAllocInterval& interval) { - return interval.startUsedAllocFallback || interval.endUsedFallback; + op->removeAttr(kLocalMemorySizeAttrName); + for (StringRef name : kRemovedLocalMemoryPlanAttrNames) + op->removeAttr(name); + op->walk([&](memref::AllocOp allocation) { + allocation->removeAttr(kLocalMemoryAddressAttrName); + for (StringRef name : kRemovedLocalMemoryPlanAttrNames) + allocation->removeAttr(name); }); - uint64_t nestedSingleUseIntervals = llvm::count_if(plan.intervals, [&](const LocalAllocInterval& interval) { - return interval.insideNestedRegion && !hasAddressReuse(interval.slotPlanIndex, plan.slots); - }); - op->setAttr(kLocalMemoryFallbackCountAttrName, builder.getI64IntegerAttr(fallbackIntervals)); - op->setAttr(kLocalMemoryNestedSingleUseCountAttrName, builder.getI64IntegerAttr(nestedSingleUseIntervals)); - if (!report) - return; - std::string planReport = buildMemoryPlanReport( - op, plan.intervals, plan.slots, kPimLocalMemoryAddressLimit, pimMemoryReport); - if (auto coreOp = dyn_cast(op)) { - *report << "Core " << coreOp.getCoreId() << ":\n"; - *report << planReport; + auto plan = buildCoreMemoryPlan(op); + if (failed(plan)) { + hasFailure = true; return; } - - SmallVector coreIds = getBatchCoreIds(cast(op)); - llvm::sort(coreIds); - coreIds.erase(std::unique(coreIds.begin(), coreIds.end()), coreIds.end()); - uint64_t batchId = nextBatchId++; - for (int32_t coreId : coreIds) - *report << "Batch " << batchId << " core " << coreId << ":\n" << planReport; + op->setAttr(kLocalMemorySizeAttrName, builder.getI64IntegerAttr(plan->arenaSize)); + for (const LocalMemoryPlacement& placement : plan->placements) + for (size_t intervalIndex : placement.intervalIndices) + plan->intervals[intervalIndex].allocation->setAttr( + kLocalMemoryAddressAttrName, builder.getI64IntegerAttr(placement.address)); }); - - if (report) { - report->flush(); - reportFile.close(); - } - if (failedPlanning) { + if (hasFailure) { signalPassFailure(); return; } - dumpModule(getOperation(), "pim4_memory_planned"); + dumpModule(getOperation(), "pim3_memory_planned"); } }; } // namespace +FailureOr> +planLocalMemoryPlacements(ArrayRef intervals, size_t addressLimit) { + SmallVector placements; + SmallVector order(intervals.size()); + std::iota(order.begin(), order.end(), 0); + llvm::stable_sort(order, [&](size_t lhsIndex, size_t rhsIndex) { + const auto& lhs = intervals[lhsIndex]; + const auto& 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 lhsIndex < rhsIndex; + }); + + for (size_t intervalIndex : order) { + const auto& interval = intervals[intervalIndex]; + SmallVector conflicts; + SmallVector candidates = {0}; + size_t currentPeak = 0; + for (const LocalMemoryPlacement& placement : placements) { + auto end = alignedEnd(placement.address, placement.size, addressLimit); + if (failed(end)) + return failure(); + currentPeak = std::max(currentPeak, *end); + if (!llvm::any_of(placement.intervalIndices, [&](size_t otherIndex) { + return pim::localMemoryLifetimesOverlap(interval, intervals[otherIndex]); + })) + continue; + conflicts.push_back(&placement); + if (!llvm::is_contained(candidates, *end)) + candidates.push_back(*end); + } + + size_t bestAddress = std::numeric_limits::max(); + auto bestKey = std::tuple(std::numeric_limits::max(), bestAddress); + for (size_t candidate : candidates) { + auto end = alignedEnd(candidate, interval.size, addressLimit); + if (failed(end) || llvm::any_of(conflicts, [&](const LocalMemoryPlacement* placement) { + return rangesOverlap(candidate, interval.size, placement->address, placement->size); + })) + continue; + auto key = std::tuple(std::max(currentPeak, *end), candidate); + if (key < bestKey) { + bestKey = key; + bestAddress = candidate; + } + } + if (bestAddress == std::numeric_limits::max()) + return failure(); + + auto reusable = llvm::find_if(placements, [&](const LocalMemoryPlacement& placement) { + return placement.address == bestAddress && placement.size == interval.size; + }); + if (reusable != placements.end()) + reusable->intervalIndices.push_back(intervalIndex); + else + placements.push_back({bestAddress, interval.size, {intervalIndex}}); + } + return placements; +} + +FailureOr buildCoreMemoryPlan(Operation* coreLikeOp) { + CoreMemoryPlan plan; + plan.coreLikeOp = coreLikeOp; + auto intervals = pim::analyzeLocalMemoryLifetimes(coreLikeOp); + if (failed(intervals)) + return failure(); + plan.intervals = std::move(*intervals); + auto placements = planLocalMemoryPlacements(plan.intervals, kPimLocalMemoryAddressLimit); + if (failed(placements)) { + coreLikeOp->emitError("PIM local-memory plan exceeds the signed int32 address range"); + return failure(); + } + plan.placements = std::move(*placements); + for (const LocalMemoryPlacement& placement : plan.placements) { + auto end = alignedEnd(placement.address, placement.size, kPimLocalMemoryAddressLimit); + if (failed(end)) { + coreLikeOp->emitError("PIM local-memory plan has invalid address arithmetic"); + return failure(); + } + plan.arenaSize = std::max(plan.arenaSize, *end); + } + return plan; +} + std::unique_ptr createPimLocalMemoryPlanningPass() { return std::make_unique(); } diff --git a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp index 80ae74e..76d3e13 100644 --- a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp +++ b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp @@ -1,64 +1,24 @@ #pragma once -#include "mlir/Dialect/MemRef/IR/MemRef.h" - -#include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/SmallVector.h" - -#include -#include - -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" namespace onnx_mlir { -struct LocalAllocInterval { - size_t id = 0; - mlir::memref::AllocOp alloc; - 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; - std::string valueSummary; - std::string firstTouchSummary; - std::string lastTouchSummary; - size_t slotPlanIndex = std::numeric_limits::max(); -}; - -struct PlannedPhysicalSlot { - size_t id = std::numeric_limits::max(); - size_t requiredSize = 0; +struct LocalMemoryPlacement { size_t address = 0; + size_t size = 0; llvm::SmallVector intervalIndices; }; struct CoreMemoryPlan { mlir::Operation* coreLikeOp = nullptr; - llvm::SmallVector intervals; - llvm::SmallVector slots; + llvm::SmallVector intervals; + llvm::SmallVector placements; + size_t arenaSize = 0; }; -llvm::SmallVector buildLocalAllocIntervals(mlir::Operation* coreLikeOp); - -llvm::SmallVector planPhysicalSlots(llvm::MutableArrayRef intervals); - -CoreMemoryPlan buildCoreMemoryPlan(mlir::Operation* coreLikeOp); - -mlir::LogicalResult assignPhysicalSlotAddresses(CoreMemoryPlan& plan, size_t addressLimit); - -std::string buildMemoryPlanReport(mlir::Operation* coreLikeOp, - llvm::ArrayRef intervals, - llvm::ArrayRef slots, - size_t addressLimit, - PimMemoryReportLevel reportLevel); +mlir::FailureOr> +planLocalMemoryPlacements(llvm::ArrayRef intervals, size_t addressLimit); +mlir::FailureOr buildCoreMemoryPlan(mlir::Operation* coreLikeOp); } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/CMakeLists.txt b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/CMakeLists.txt deleted file mode 100644 index 3bf4699..0000000 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/CMakeLists.txt +++ /dev/null @@ -1,15 +0,0 @@ -add_pim_library(OMPimMemoryCoalescing - MemoryCoalescing.cpp - MemoryCoalescing.hpp - MemoryCoalescingPass.cpp - - EXCLUDE_FROM_OM_LIBS - - INCLUDE_DIRS PUBLIC - ${PIM_PUBLIC_INCLUDE_DIRS} - - LINK_LIBS PUBLIC - OMPimCommon - OMPimLocalMemoryLifetimeAnalysis - PimOps -) diff --git a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.cpp b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.cpp deleted file mode 100644 index ffa4ae8..0000000 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.cpp +++ /dev/null @@ -1,194 +0,0 @@ -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" - -#include - -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp" - -using namespace mlir; - -namespace onnx_mlir { -namespace pim { - -namespace { - -static bool isCandidateAllocType(MemRefType type) { - return type && type.hasStaticShape() && type.getLayout().isIdentity() - && hasByteSizedElementType(type.getElementType()); -} - -static uint64_t getTypeSizeBytes(MemRefType type) { - return static_cast(type.getNumElements() * getElementTypeSizeInBytes(type.getElementType())); -} - -static Operation* getTopLevelAncestorInBlock(Operation* op, Block& block) { - Operation* current = op; - while (current && current->getBlock() != &block) - current = current->getParentOp(); - return current; -} - -static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis); - -static FailureOr -getLastUseInstruction(memref::AllocOp allocOp, Block& scopeBlock, const DenseMap& opOrder) { - uint64_t endInstruction = opOrder.lookup(allocOp); - if (failed(walkLocalMemoryUses(allocOp.getResult(), [&](Value, Operation* user) { - if (auto yieldOp = dyn_cast(user)) { - if (!isa(yieldOp->getParentOp())) - return failure(); - } - Operation* orderedUser = getTopLevelAncestorInBlock(user, scopeBlock); - if (!orderedUser) - return failure(); - auto order = opOrder.find(orderedUser); - if (order == opOrder.end()) - return failure(); - endInstruction = std::max(endInstruction, order->second); - return success(); - }))) - return failure(); - - return endInstruction; -} - -static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis) { - for (Operation& op : block) - for (Region& region : op.getRegions()) - for (Block& nestedBlock : region) - analyzeBlock(nestedBlock, analysis); - - DenseMap opOrder; - uint64_t nextInstruction = 0; - for (Operation& op : block) - opOrder.try_emplace(&op, nextInstruction++); - - MemoryCoalescingBlockAnalysis blockAnalysis; - blockAnalysis.block = █ - - for (Operation& op : block) { - auto allocOp = dyn_cast(&op); - if (!allocOp) - continue; - - auto allocType = dyn_cast(allocOp.getType()); - if (!isCandidateAllocType(allocType)) { - ++blockAnalysis.skippedAllocations; - continue; - } - - auto endInstruction = getLastUseInstruction(allocOp, block, opOrder); - if (failed(endInstruction)) { - ++blockAnalysis.skippedAllocations; - continue; - } - - blockAnalysis.candidates.push_back( - AllocationCandidate {allocOp, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)}); - } - - analysis.skippedAllocations += blockAnalysis.skippedAllocations; - if (!blockAnalysis.candidates.empty() || blockAnalysis.skippedAllocations != 0) - analysis.blocks.push_back(std::move(blockAnalysis)); -} - -} // namespace - -uint64_t MemoryCoalescingAnalysis::getCandidateCount() const { - uint64_t total = 0; - for (const MemoryCoalescingBlockAnalysis& block : blocks) - total += block.candidates.size(); - return total; -} - -uint64_t MemoryCoalescingAnalysis::getCandidateBytes() const { - uint64_t total = 0; - for (const MemoryCoalescingBlockAnalysis& block : blocks) - for (const AllocationCandidate& candidate : block.candidates) - total += candidate.sizeBytes; - return total; -} - -MemoryCoalescingAnalysis analyzeMemoryCoalescingCandidates(Operation* coreLikeOp) { - MemoryCoalescingAnalysis analysis; - if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty()) - return analysis; - - analyzeBlock(coreLikeOp->getRegion(0).front(), analysis); - return analysis; -} - -MemoryCoalescingStats -coalesceMemory(Operation* coreLikeOp, const MemoryCoalescingAnalysis& analysis, RewriterBase& rewriter) { - (void) coreLikeOp; - - MemoryCoalescingStats stats; - stats.skippedAllocations = analysis.skippedAllocations; - - for (const MemoryCoalescingBlockAnalysis& blockAnalysis : analysis.blocks) { - auto candidates = blockAnalysis.candidates; - llvm::sort(candidates, [](const AllocationCandidate& lhs, const AllocationCandidate& rhs) { - if (lhs.startInstruction != rhs.startInstruction) - return lhs.startInstruction < rhs.startInstruction; - return lhs.endInstruction < rhs.endInstruction; - }); - - struct ActiveStorage { - memref::AllocOp root; - uint64_t endInstruction = 0; - }; - - SmallVector active; - SmallVector freeList; - - for (AllocationCandidate& candidate : candidates) { - for (auto it = active.begin(); it != active.end();) { - if (it->endInstruction < candidate.startInstruction) { - freeList.push_back(it->root); - it = active.erase(it); - continue; - } - ++it; - } - - auto bestFit = freeList.end(); - uint64_t bestFitBytes = std::numeric_limits::max(); - auto candidateType = cast(candidate.alloc.getType()); - for (auto it = freeList.begin(); it != freeList.end(); ++it) { - auto freeType = cast((*it).getType()); - if (freeType != candidateType) - continue; - - uint64_t freeBytes = getTypeSizeBytes(freeType); - if (freeBytes < candidate.sizeBytes || freeBytes >= bestFitBytes) - continue; - - bestFit = it; - bestFitBytes = freeBytes; - } - - if (bestFit == freeList.end()) { - active.push_back(ActiveStorage {candidate.alloc, candidate.endInstruction}); - continue; - } - - memref::AllocOp root = *bestFit; - freeList.erase(bestFit); - candidate.alloc.getResult().replaceAllUsesWith(root.getResult()); - rewriter.eraseOp(candidate.alloc); - active.push_back(ActiveStorage {root, candidate.endInstruction}); - ++stats.removedAllocs; - stats.savedBytes += candidate.sizeBytes; - } - } - - return stats; -} - -} // namespace pim -} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp deleted file mode 100644 index 197cda2..0000000 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp +++ /dev/null @@ -1,44 +0,0 @@ -#pragma once - -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/IR/PatternMatch.h" - -#include "llvm/ADT/SmallVector.h" - -namespace onnx_mlir { -namespace pim { - -struct AllocationCandidate { - mlir::memref::AllocOp alloc; - uint64_t startInstruction = 0; - uint64_t endInstruction = 0; - uint64_t sizeBytes = 0; -}; - -struct MemoryCoalescingBlockAnalysis { - mlir::Block* block = nullptr; - llvm::SmallVector candidates; - uint64_t skippedAllocations = 0; -}; - -struct MemoryCoalescingAnalysis { - llvm::SmallVector blocks; - uint64_t skippedAllocations = 0; - - uint64_t getCandidateCount() const; - uint64_t getCandidateBytes() const; -}; - -struct MemoryCoalescingStats { - uint64_t removedAllocs = 0; - uint64_t savedBytes = 0; - uint64_t skippedAllocations = 0; -}; - -MemoryCoalescingAnalysis analyzeMemoryCoalescingCandidates(mlir::Operation* coreLikeOp); - -MemoryCoalescingStats -coalesceMemory(mlir::Operation* coreLikeOp, const MemoryCoalescingAnalysis& analysis, mlir::RewriterBase& rewriter); - -} // namespace pim -} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp deleted file mode 100644 index 6042af9..0000000 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp +++ /dev/null @@ -1,101 +0,0 @@ -#include "mlir/Pass/Pass.h" - -#include "llvm/Support/raw_os_ostream.h" - -#include - -#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" - -using namespace mlir; - -namespace onnx_mlir { -namespace { - -struct CoalescingReportRow { - uint64_t candidates = 0; - uint64_t logicalBytes = 0; - uint64_t skipped = 0; - uint64_t removed = 0; - uint64_t savedBytes = 0; -}; - -struct CoalescingReportEntry { - std::string label; - uint64_t coreCount = 1; - CoalescingReportRow row; -}; - -static void emitReport(ArrayRef entries) { - std::fstream file = openReportFile("memory_coalescing_report"); - if (!file.is_open()) - return; - llvm::raw_os_ostream os(file); - CoalescingReportRow total; - for (const CoalescingReportEntry& entry : entries) { - total.candidates += entry.row.candidates * entry.coreCount; - total.logicalBytes += entry.row.logicalBytes * entry.coreCount; - total.skipped += entry.row.skipped * entry.coreCount; - total.removed += entry.row.removed * entry.coreCount; - total.savedBytes += entry.row.savedBytes * entry.coreCount; - } - printReportTotalsBlock(os, - {{"Number of candidates", std::to_string(total.candidates)}, - {"Logical candidate bytes", formatReportMemory(total.logicalBytes)}, - {"Skipped allocations", std::to_string(total.skipped)}, - {"Removed allocations", std::to_string(total.removed)}, - {"Saved memory", formatReportMemory(total.savedBytes)}}); - for (const CoalescingReportEntry& entry : entries) { - os << "\n" << entry.label << ":\n"; - printReportFlatFields(os, - {{"Number of candidates", std::to_string(entry.row.candidates)}, - {"Logical candidate bytes", formatReportMemory(entry.row.logicalBytes)}, - {"Skipped allocations", std::to_string(entry.row.skipped)}, - {"Removed allocations", std::to_string(entry.row.removed)}, - {"Saved memory", formatReportMemory(entry.row.savedBytes)}}); - } - os.flush(); -} - -struct PimMemoryCoalescingPass : PassWrapper> { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimMemoryCoalescingPass) - - StringRef getArgument() const override { return "pim-memory-coalescing"; } - StringRef getDescription() const override { return "Safely coalesce compatible block-local PIM allocations"; } - - void runOnOperation() override { - IRRewriter rewriter(&getContext()); - SmallVector reportEntries; - uint64_t nextBatchId = 0; - getOperation().walk([&](Operation* op) { - if (!isa(op)) - return; - auto analysis = pim::analyzeMemoryCoalescingCandidates(op); - auto stats = pim::coalesceMemory(op, analysis, rewriter); - CoalescingReportRow row {analysis.getCandidateCount(), - analysis.getCandidateBytes(), - stats.skippedAllocations, - stats.removedAllocs, - stats.savedBytes}; - if (auto coreOp = dyn_cast(op)) { - reportEntries.push_back({"Core " + std::to_string(coreOp.getCoreId()), 1, row}); - return; - } - SmallVector coreIds = getBatchCoreIds(cast(op)); - reportEntries.push_back( - {"Batch " + std::to_string(nextBatchId++), std::max(1, coreIds.size()), row}); - }); - emitReport(reportEntries); - dumpModule(getOperation(), "pim3_coalesced"); - } -}; - -} // namespace - -std::unique_ptr createPimMemoryCoalescingPass() { return std::make_unique(); } - -} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt b/src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt index 3afe2bf..de55435 100644 --- a/src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt @@ -7,6 +7,7 @@ add_pim_library(OMPimVerification OMPimCommon OMPimCompilerOptions OMPimBufferization + OMPimLocalMemoryLifetimeAnalysis PimOps SpatialOps ) diff --git a/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp b/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp index 34bc33d..7a3cae7 100644 --- a/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp @@ -8,6 +8,7 @@ #include "llvm/Support/FormatVariadic.h" #include "llvm/Support/raw_ostream.h" +#include #include #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" @@ -18,6 +19,7 @@ #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" #include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" @@ -109,61 +111,127 @@ static bool isCoreWeightBlockArgument(Value value) { return false; } -struct VerifiedLocalMemorySlot { - uint64_t size = 0; -}; - static LogicalResult verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diagnostics) { - DenseMap slots; bool hasFailure = false; - auto fallbackAttr = coreLikeOp->getAttrOfType(kLocalMemoryFallbackCountAttrName); - auto nestedSingleUseAttr = coreLikeOp->getAttrOfType(kLocalMemoryNestedSingleUseCountAttrName); - if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) { - diagnostics.report(coreLikeOp, [&](Operation* op) { - op->emitError("requires complete non-negative PIM local-memory planning summary attributes"); + for (StringRef name : kRemovedLocalMemoryPlanAttrNames) + if (coreLikeOp->hasAttr(name)) { + diagnostics.report(coreLikeOp, [name](Operation* op) { + op->emitError() << "contains removed PIM local-memory planning attribute '" << name << "'"; + }); + hasFailure = true; + } + + auto arenaAttr = coreLikeOp->getAttrOfType(kLocalMemorySizeAttrName); + if (!arenaAttr || arenaAttr.getInt() < 0 + || static_cast(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) { + diagnostics.report(coreLikeOp, [](Operation* op) { + op->emitError("requires one valid non-negative pim.local_memory_size attribute"); }); hasFailure = true; } - coreLikeOp->walk([&](memref::AllocOp allocOp) { - auto addressAttr = allocOp->getAttrOfType(kLocalMemoryAddressAttrName); - auto slotAttr = allocOp->getAttrOfType(kLocalMemorySlotAttrName); - auto slotSizeAttr = allocOp->getAttrOfType(kLocalMemorySlotSizeAttrName); - if (!addressAttr || !slotAttr || !slotSizeAttr || addressAttr.getInt() < 0 || slotAttr.getInt() < 0 - || slotSizeAttr.getInt() < 0) { - diagnostics.report(allocOp, [&](Operation*) { - allocOp.emitOpError("requires a complete non-negative PIM local-memory plan"); + if (hasFailure) + return failure(); + uint64_t arenaSize = static_cast(arenaAttr.getInt()); + + auto analyzed = pim::analyzeLocalMemoryLifetimes(coreLikeOp); + if (failed(analyzed)) { + diagnostics.report(coreLikeOp, [](Operation* op) { + op->emitError("cannot analyze PIM local-memory lifetimes for plan verification"); + }); + return failure(); + } + + struct Event { + uint64_t position; + bool start; + size_t intervalIndex; + }; + SmallVector events; + SmallVector addresses(analyzed->size()); + for (auto indexed : llvm::enumerate(*analyzed)) { + size_t index = indexed.index(); + auto& interval = indexed.value(); + memref::AllocOp allocation = interval.allocation; + for (StringRef name : kRemovedLocalMemoryPlanAttrNames) + if (allocation->hasAttr(name)) { + diagnostics.report(allocation, [name](Operation* op) { + op->emitOpError() << "contains removed PIM local-memory planning attribute '" << name << "'"; + }); + hasFailure = true; + } + auto addressAttr = allocation->getAttrOfType(kLocalMemoryAddressAttrName); + if (!addressAttr || addressAttr.getInt() < 0) { + diagnostics.report(allocation, [](Operation* op) { + op->emitOpError("requires one non-negative pim.local_memory_address attribute"); }); hasFailure = true; - return; + continue; } - uint64_t address = static_cast(addressAttr.getInt()); - uint64_t slotId = static_cast(slotAttr.getInt()); - uint64_t slotSize = static_cast(slotSizeAttr.getInt()); - auto logicalSize = pim::getCheckedShapedTypeSizeInBytes( - cast(allocOp.getType()), allocOp, "planned local allocation byte size"); - bool invalidRange = failed(logicalSize) || address > slotSize || *logicalSize > slotSize - address - || slotSize > kPimLocalMemoryAddressLimit || address % 4 != 0 || slotId != 0; - if (invalidRange) { - diagnostics.report(allocOp, [&](Operation*) { - allocOp.emitOpError() << "has an invalid PIM local-memory range: address=" << address - << ", logical size=" << (succeeded(logicalSize) ? *logicalSize : 0) - << ", slot size=" << slotSize; + if (address % 4 != 0 || address > arenaSize || interval.size > arenaSize - address) { + diagnostics.report(allocation, [&](Operation* op) { + op->emitOpError() << "has invalid PIM local-memory range [" << address << ", " + << (address <= arenaSize && interval.size <= arenaSize - address + ? address + interval.size + : arenaSize) + << ") for arena size " << arenaSize; }); hasFailure = true; - return; + continue; + } + addresses[index] = static_cast(address); + events.push_back({interval.start, true, index}); + events.push_back({interval.end, false, index}); + } + if (hasFailure) + return failure(); + + llvm::sort(events, [](const Event& lhs, const Event& rhs) { + return lhs.position != rhs.position ? lhs.position < rhs.position : lhs.start > rhs.start; + }); + std::map active; + for (const Event& event : events) { + size_t index = event.intervalIndex; + size_t address = addresses[index]; + if (!event.start) { + auto it = active.find(address); + if (it != active.end() && it->second == index) + active.erase(it); + continue; } - auto [slotIt, inserted] = slots.try_emplace(slotId, VerifiedLocalMemorySlot {slotSize}); - if (!inserted && slotIt->second.size != slotSize) { - diagnostics.report(allocOp, [&](Operation*) { - allocOp.emitOpError() << "has inconsistent size for PIM local-memory arena " << slotId; + const auto& interval = (*analyzed)[index]; + auto successor = active.lower_bound(address); + auto conflict = [&](std::map::iterator it) { + if (it == active.end()) + return false; + const auto& other = (*analyzed)[it->second]; + return address < it->first + other.size && it->first < address + interval.size; + }; + auto predecessor = successor; + bool hasPredecessor = predecessor != active.begin(); + if (hasPredecessor) + --predecessor; + auto conflicting = conflict(successor) ? successor + : hasPredecessor && conflict(predecessor) ? predecessor : active.end(); + if (conflicting != active.end()) { + const auto& other = (*analyzed)[conflicting->second]; + memref::AllocOp allocation = interval.allocation; + memref::AllocOp otherAllocation = other.allocation; + diagnostics.report(allocation, [&](Operation*) { + auto diagnostic = allocation.emitOpError() + << "PIM local-memory plan assigns simultaneously live allocations to overlapping ranges; first range [" + << conflicting->first << ", " << conflicting->first + other.size << "), second range [" << address + << ", " << address + interval.size << "), live positions overlap at [" + << std::max(interval.start, other.start) << ", " << std::min(interval.end, other.end) << "]"; + diagnostic.attachNote(otherAllocation.getLoc()) << "first allocation"; }); - hasFailure = true; + return failure(); } - }); - return success(!hasFailure); + active.emplace(address, index); + } + return success(); } static bool isSupportedCoreInstructionOp(Operation* op) { diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp index cee4178..bd1735f 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp @@ -16,6 +16,11 @@ namespace spatial { namespace { +// Pressure means distinct weights exceed half the fleet's one-copy capacity. +// The reserved headroom conservatively avoids greedy capacity fragmentation; +// makespan, communication, instruction-path, and local-memory proxies remain stable. +constexpr size_t kHighCrossbarPressureCapacityDivisor = 2; + struct ScheduledTask { size_t processor = std::numeric_limits::max(); Time startTime = 0; @@ -136,7 +141,7 @@ void verifyOctTableSize(size_t nodeCount, size_t processorCount) { bool hasHighCrossbarPressure(const ComputeGraph& graph, size_t processorCount, size_t crossbarCapacity) { if (crossbarCapacity > std::numeric_limits::max() / processorCount) return false; - const size_t threshold = processorCount * crossbarCapacity / 2; + const size_t threshold = processorCount * crossbarCapacity / kHighCrossbarPressureCapacityDivisor; CrossbarUsage distinctWeights; for (const ComputeGraphNode& node : graph.nodes) { for (const CrossbarWeight& weight : node.crossbarUsage) { diff --git a/src/PIM/Pass/PIMPasses.h b/src/PIM/Pass/PIMPasses.h index 62cae80..85b399b 100644 --- a/src/PIM/Pass/PIMPasses.h +++ b/src/PIM/Pass/PIMPasses.h @@ -15,8 +15,6 @@ std::unique_ptr createSpatialToPimPass(); std::unique_ptr createPimBufferizationPass(); -std::unique_ptr createPimMemoryCoalescingPass(); - std::unique_ptr createMergeComputeNodesPass(); std::unique_ptr createTrivialGraphComputeMergePass(); diff --git a/src/PIM/PimAccelerator.cpp b/src/PIM/PimAccelerator.cpp index f9be891..b6c8198 100644 --- a/src/PIM/PimAccelerator.cpp +++ b/src/PIM/PimAccelerator.cpp @@ -76,7 +76,6 @@ void PimAccelerator::registerPasses(int optLevel) const { registerPass(createLowerSpatialPlansPass); registerPass(createSpatialToPimPass); registerPass(createPimBufferizationPass); - registerPass(createPimMemoryCoalescingPass); registerPass(createTrivialGraphComputeMergePass); registerPass(createMergeComputeNodesPass); registerPass(createPimHostConstantFoldingPass); diff --git a/test/PIM/PimMemoryLivenessPlannerTest.cpp b/test/PIM/PimMemoryLivenessPlannerTest.cpp index 2f6eb4c..2400526 100644 --- a/test/PIM/PimMemoryLivenessPlannerTest.cpp +++ b/test/PIM/PimMemoryLivenessPlannerTest.cpp @@ -4,103 +4,97 @@ #include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp" -using onnx_mlir::LocalAllocInterval; -using onnx_mlir::CoreMemoryPlan; -using onnx_mlir::assignPhysicalSlotAddresses; -using onnx_mlir::planPhysicalSlots; +using onnx_mlir::LocalMemoryPlacement; +using onnx_mlir::pim::LocalMemoryInterval; +using onnx_mlir::planLocalMemoryPlacements; namespace { -LocalAllocInterval makeInterval(size_t id, size_t size, uint64_t start, uint64_t end) { - LocalAllocInterval interval; - interval.id = id; +LocalMemoryInterval makeInterval(size_t size, uint64_t start, uint64_t end) { + LocalMemoryInterval interval; interval.size = size; interval.start = start; interval.end = end; return interval; } -size_t getPeak(llvm::ArrayRef slots) { +size_t getPeak(llvm::ArrayRef placements) { size_t peak = 0; - for (const auto& slot : slots) - peak = std::max(peak, slot.address + slot.requiredSize); + for (const auto& placement : placements) + peak = std::max(peak, placement.address + placement.size); return peak; } -void assertSingleSlotCase(LocalAllocInterval a, LocalAllocInterval b, size_t expectedSlotSize) { - llvm::SmallVector intervals = {a, b}; - auto slots = planPhysicalSlots(intervals); - assert(slots.size() == 1); - assert(slots.front().requiredSize == expectedSlotSize); - assert(intervals[0].slotPlanIndex == intervals[1].slotPlanIndex); +void assertSinglePlacementCase(LocalMemoryInterval a, LocalMemoryInterval b, size_t expectedSize) { + llvm::SmallVector intervals = {a, b}; + auto placements = planLocalMemoryPlacements(intervals, 1024); + assert(mlir::succeeded(placements)); + assert(placements->size() == 1); + assert(placements->front().size == expectedSize); + assert(placements->front().intervalIndices.size() == 2); } int testSameSizeNonOverlap() { std::cout << "testSameSizeNonOverlap:" << std::endl; - assertSingleSlotCase(makeInterval(0, 64, 0, 10), makeInterval(1, 64, 11, 20), 64); + assertSinglePlacementCase(makeInterval(64, 0, 10), makeInterval(64, 11, 20), 64); return 0; } int testLargerFirst() { std::cout << "testLargerFirst:" << std::endl; - llvm::SmallVector intervals = { - makeInterval(0, 100, 0, 10), makeInterval(1, 40, 11, 20)}; - auto slots = planPhysicalSlots(intervals); - assert(slots.size() == 2); - assert(getPeak(slots) == 100); + llvm::SmallVector intervals = { + makeInterval(100, 0, 10), makeInterval(40, 11, 20)}; + auto placements = planLocalMemoryPlacements(intervals, 1024); + assert(mlir::succeeded(placements) && placements->size() == 2); + assert(getPeak(*placements) == 100); return 0; } int testSmallerFirst() { std::cout << "testSmallerFirst:" << std::endl; - llvm::SmallVector intervals = { - makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20)}; - auto slots = planPhysicalSlots(intervals); - assert(slots.size() == 2); - assert(getPeak(slots) == 100); + llvm::SmallVector intervals = { + makeInterval(40, 0, 10), makeInterval(100, 11, 20)}; + auto placements = planLocalMemoryPlacements(intervals, 1024); + assert(mlir::succeeded(placements) && placements->size() == 2); + assert(getPeak(*placements) == 100); return 0; } int testOverlapNeedsTwoSlots() { std::cout << "testOverlapNeedsTwoSlots:" << std::endl; - llvm::SmallVector intervals = { - makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)}; - auto slots = planPhysicalSlots(intervals); - assert(slots.size() == 2); - assert(intervals[0].slotPlanIndex != intervals[1].slotPlanIndex); + llvm::SmallVector intervals = { + makeInterval(100, 0, 20), makeInterval(40, 10, 30)}; + auto placements = planLocalMemoryPlacements(intervals, 1024); + assert(mlir::succeeded(placements) && placements->size() == 2); + assert((*placements)[0].address != (*placements)[1].address); return 0; } int testReuseChain() { std::cout << "testReuseChain:" << std::endl; - llvm::SmallVector intervals = { - makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), makeInterval(2, 20, 21, 30)}; - auto slots = planPhysicalSlots(intervals); - assert(slots.size() == 3); - assert(getPeak(slots) == 100); + llvm::SmallVector intervals = { + makeInterval(40, 0, 10), makeInterval(100, 11, 20), makeInterval(20, 21, 30)}; + auto placements = planLocalMemoryPlacements(intervals, 1024); + assert(mlir::succeeded(placements) && placements->size() == 3); + assert(getPeak(*placements) == 100); return 0; } int testPartialAddressReuse() { std::cout << "testPartialAddressReuse:" << std::endl; - llvm::SmallVector intervals = { - makeInterval(0, 100, 0, 10), makeInterval(1, 60, 11, 20), makeInterval(2, 40, 11, 20)}; - auto slots = planPhysicalSlots(intervals); - assert(getPeak(slots) == 100); + llvm::SmallVector intervals = { + makeInterval(100, 0, 10), makeInterval(60, 11, 20), makeInterval(40, 11, 20)}; + auto placements = planLocalMemoryPlacements(intervals, 1024); + assert(mlir::succeeded(placements)); + assert(getPeak(*placements) == 100); return 0; } -int testAddressAssignment() { - std::cout << "testAddressAssignment:" << std::endl; - CoreMemoryPlan plan; - plan.intervals = {makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)}; - plan.slots = planPhysicalSlots(plan.intervals); - assert(mlir::succeeded(assignPhysicalSlotAddresses(plan, 1024))); - assert(plan.slots.size() == 2); - const auto& firstSlot = plan.slots[plan.intervals[0].slotPlanIndex]; - const auto& secondSlot = plan.slots[plan.intervals[1].slotPlanIndex]; - assert(firstSlot.address == 0 && firstSlot.requiredSize == 100); - assert(secondSlot.address == 100 && secondSlot.requiredSize == 40); +int testAddressLimit() { + std::cout << "testAddressLimit:" << std::endl; + llvm::SmallVector intervals = { + makeInterval(100, 0, 20), makeInterval(40, 10, 30)}; + assert(mlir::failed(planLocalMemoryPlacements(intervals, 128))); return 0; } @@ -117,7 +111,7 @@ int main(int argc, char *argv[]) { failures += testOverlapNeedsTwoSlots(); failures += testReuseChain(); failures += testPartialAddressReuse(); - failures += testAddressAssignment(); + failures += testAddressLimit(); if (failures != 0) { std::cerr << failures << " test failures\n"; return EXIT_FAILURE; diff --git a/validation/raptor.py b/validation/raptor.py index ac64911..a97f2bb 100644 --- a/validation/raptor.py +++ b/validation/raptor.py @@ -11,7 +11,6 @@ PIM_PASS_LABELS = ( ("SpatialToPimPass", "Spatial to PIM"), ("PimBufferizationPass", "Bufferize PIM"), ("HostConstantFoldingPass", "Fold Host Constants"), - ("PimMemoryCoalescingPass", "Coalesce Local Memory"), ("PimLocalMemoryPlanningPass", "Plan Local Memory"), ("VerificationPass", "Verify PIM"), ("EmitPimCodePass", "Emit PIM Code"), @@ -43,7 +42,7 @@ def _format_command(cmd): def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path, crossbar_size, crossbar_count, core_count=None, - pim_memory_report="none", raptor_extra_args=None, cwd=None, verbose=False, + raptor_extra_args=None, cwd=None, verbose=False, reporter=None, timeout_sec=None): # Define the arguments, with the possibility to set crossbar size and count args = [ @@ -57,8 +56,6 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path, ] if core_count is not None: args.append(f"--core-count={core_count}") - if pim_memory_report != "none": - args.append(f"--pim-memory-report={pim_memory_report}") if raptor_extra_args: args.extend(str(arg) for arg in raptor_extra_args) if verbose: diff --git a/validation/validate.py b/validation/validate.py index b58f3f3..29d5be3 100755 --- a/validation/validate.py +++ b/validation/validate.py @@ -75,8 +75,6 @@ def main(): ap.add_argument("--crossbar-size", type=int, default=128) ap.add_argument("--crossbar-count", type=int, default=64) ap.add_argument("--core-count", type=int, default=144) - ap.add_argument("--pim-memory-report", choices=("none", "summary", "full"), default="none", - help="Emit a human-readable PIM memory planning report during codegen.") ap.add_argument("--raptor-extra-arg", action="append", default=[], help="Additional argument to pass through to the Raptor compiler. Repeat as needed.") ap.add_argument("--command-timeout-seconds", type=float, default=1000000.0, @@ -144,7 +142,6 @@ def main(): result = validate_network( onnx_path, a.raptor_path, a.onnx_include_dir, simulator_dir, crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count, core_count=a.core_count, - pim_memory_report=a.pim_memory_report, raptor_extra_args=a.raptor_extra_arg, command_timeout_seconds=a.command_timeout_seconds, threshold=a.threshold, diff --git a/validation/validate_one.py b/validation/validate_one.py index a1316c7..9f663c6 100644 --- a/validation/validate_one.py +++ b/validation/validate_one.py @@ -312,7 +312,7 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1 def validate_network(network_onnx_path, raptor_path, onnx_include_dir, simulator_dir, crossbar_size=128, crossbar_count=64, core_count=144, - pim_memory_report="none", raptor_extra_args=None, + raptor_extra_args=None, threshold=1e-3, rtol=1e-5, seed=0, reporter=None, model_index=1, model_total=1, verbose=False, command_timeout_seconds=60.0, mode=MODE_FULL): @@ -367,7 +367,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir, pim_pass_timings = compile_with_raptor( network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count, core_count=core_count, - pim_memory_report=pim_memory_report, raptor_extra_args=raptor_extra_args, + raptor_extra_args=raptor_extra_args, cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds) print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}") reporter.advance() @@ -407,7 +407,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir, pim_pass_timings = compile_with_raptor( network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count, core_count=core_count, - pim_memory_report=pim_memory_report, raptor_extra_args=raptor_extra_args, + raptor_extra_args=raptor_extra_args, cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds) print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}") reporter.advance()