diff --git a/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md b/.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md similarity index 98% rename from invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md rename to .agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md index 95f0f87..c969d1b 100644 --- a/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md +++ b/.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md @@ -8,10 +8,10 @@ Every developer or coding agent modifying Spatial graph construction, graph verification, Blueprint handling, or `MergeComputeNodes` must read this file after `README.md` and `AGENTS.md`. -`AGENTS.md` must contain this instruction: +`AGENTS.md` must reference this invariant: ```text -* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes. +* `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md` ``` ## Scope diff --git a/invariants/PARALLELISM_AND_RUNTIME_INVARIANT.md b/.agents/invariants/PARALLELISM_AND_RUNTIME_INVARIANT.md similarity index 100% rename from invariants/PARALLELISM_AND_RUNTIME_INVARIANT.md rename to .agents/invariants/PARALLELISM_AND_RUNTIME_INVARIANT.md diff --git a/.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md b/.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md new file mode 100644 index 0000000..0d4c157 --- /dev/null +++ b/.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md @@ -0,0 +1,67 @@ +# Compiler Performance Optimization Invariant + +## Scope + +This invariant applies to: + +- ONNX-to-Spatial lowering; +- graph and trivial merge; +- PEFT scheduling; +- scheduled materialization; +- deferred communication; +- Spatial-to-PIM lowering; +- bufferization; +- memory coalescing; +- liveness planning; +- PIM code generation. + +## Invariant + +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 +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. + +An optimization is acceptable only when its static runtime proxies are equal or +better. + +## Forbidden optimization trades + +Do not: + +- use fewer parallel lanes or cores to simplify IR; +- scalarize a valid graph or core batch; +- replace independent operations with one sequential loop; +- recompute values to reduce storage; +- add device/device or host/device copies to reduce peak memory; +- change hardware parameters to make a benchmark easier; +- reduce communication concurrency by imposing a global serial order; +- increase emitted instruction count to reduce compiler analysis; +- move PIM work to the host. + +## Required proof + +Every performance change must report before and after: + +- compiler wall time; +- compiler peak RSS; +- IR operation and value counts at the changed stage; +- logical compute-instance count; +- graph compute batch and scheduled batch lane counts; +- PEFT class and core assignment; +- maximum scheduled critical-path or step count; +- MVM/VMM and vector instruction counts; +- send and receive counts and bytes; +- total emitted instruction count; +- maximum instructions assigned to one core; +- number and size of runtime memory copies; +- final host, weights, logical-core, physical-core, and maximum-core memory; +- numerical validation. + +## Stop rule + +If compile time or memory improves but a runtime proxy worsens, stop and reject +the change. diff --git a/AGENTS.md b/AGENTS.md index 99de7f5..c9c6058 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,5 +1,11 @@ * Always read the full README.md before doing anything -* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes. + +# Required project invariants + +Before modifying the relevant subsystem, read: + +* `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md` +* `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md` * Build commands: * `cmake --build ./build_release` * `cmake --build ./build_debug` @@ -165,6 +171,10 @@ Do not refactor when: * If a change adds a walk, cache, analysis, or structural traversal, justify why it is needed * For hot paths, prefer preserving existing asymptotic behavior unless a better structure is part of the requested change * If performance may change, mention the expected impact and suggest a targeted timing check +* Compile-time and compiler-memory improvements must not trade away hardware + parallelism or theoretical runtime. Preserve or improve the static runtime + proxies defined in + `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md`. # Goal-driven execution diff --git a/README.md b/README.md index ecb7258..595fecc 100644 --- a/README.md +++ b/README.md @@ -67,20 +67,22 @@ 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. **Static memory coalescing** - (`src/PIM/Dialect/Pim/Transforms/StaticMemoryCoalescing`). - Reuses compatible local memref allocations inside PIM cores before codegen. - -6. **PIM code generation** (`src/PIM/Pass/PimCodegen` and +5. **Safe IR memory coalescing** + (`src/PIM/Dialect/Pim/Transforms/MemoryCoalescing`). + Normalizes compatible block-local allocations before final planning. +6. **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 `src/PIM/Compiler`). - Folds host constants, materializes remaining host constants, verifies PIM IR, - emits `.pim` core files, writes weights, and writes `memory.bin` / - `config.json`. + Verifies the memory plan and other PIM invariants, then emits `.pim` core + files, weights, and `memory.bin` / `config.json` without rerunning liveness. Supporting pieces: - `src/PIM/Common` - shared IR, filesystem, diagnostics, reports, and utility helpers. -- `src/PIM/Compiler` - PIM compiler options, memory/address planning, binary +- `src/PIM/Compiler` - PIM compiler options, planned-address materialization, binary instruction format, artifact writing, weight emission, and codegen entry points. - `src/PIM/Conversion/SpatialToGraphviz` - optional Spatial graphviz conversion @@ -97,8 +99,8 @@ Pass these to `onnx-mlir` when compiling for PIM: `--EmitPimCodegen` - stop the PIM pipeline at the requested stage. The PIM default is `--EmitPimCodegen`. - `--core-count=` - required positive core count for PIM compilation. -- `--crossbar-size=` - crossbar width/height. Default in code is `2`. -- `--crossbar-count=` - crossbars per core. Default in code is `256`. +- `--crossbar-size=` - crossbar width/height. Default in code is `128`. +- `--crossbar-count=` - crossbars per core. Default in code is `64`. - `--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 @@ -109,12 +111,28 @@ Pass these to `onnx-mlir` when compiling for PIM: - `--use-experimental-conv-impl` - use the alternate convolution lowering. - `--ignore-concat-error` - soft-fail a ConcatOp corner case. +## Standard PIM hardware profile + +Raptor's standard development and YOLO validation profile is: + +| Parameter | Value | +| --- | ---: | +| Cores | 144 | +| Crossbars per core | 64 | +| Crossbar size | 128 × 128 | + +Canonical compiler flags: + +`--crossbar-count=64 --crossbar-size=128 --core-count=144` + +`--core-count` remains mandatory and must be passed explicitly to the compiler. + Example: ```bash ./build_release/Release/bin/onnx-mlir model.onnx -o /tmp/raptor/model \ --maccel=PIM --EmitPimCodegen \ - --crossbar-size=2048 --crossbar-count=256 --core-count=1000 + --crossbar-count=64 --crossbar-size=128 --core-count=144 ``` This writes PIM artifacts under `/tmp/raptor/pim/`. @@ -131,21 +149,33 @@ Python dependencies used by the validation scripts are `numpy`, `onnx`, and Per-operation validation from the repository root: ```bash -python3 validation/validate.py \ +python validation/validate.py \ --raptor-path build_release/Release/bin/onnx-mlir \ --onnx-include-dir onnx-mlir/include \ - --core-count 1000 + --operations-dir validation/operations \ + --crossbar-count 64 \ + --crossbar-size 128 \ + --core-count 144 \ + --verbose \ + --raptor-extra-arg=--pim-detect-communication-deadlock \ + --raptor-extra-arg=--pim-export-spatial-dataflow=none ``` Validate one network or a subset by pointing `--operations-dir` at any directory containing `.onnx` files: ```bash -python3 validation/validate.py \ +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 \ - --crossbar-size 2048 --crossbar-count 256 --core-count 1000 + --crossbar-count 64 \ + --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 ``` Useful validation options: @@ -170,7 +200,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`, and `pim3_coalesced.mlir` when an output directory is +`pim2_folded.mlir`, `pim3_coalesced.mlir`, and `pim4_memory_planned.mlir` when an output directory is available. To rerun the simulator manually with tracing after validation has produced a diff --git a/src/PIM/CMakeLists.txt b/src/PIM/CMakeLists.txt index 86c13c3..2989ea0 100644 --- a/src/PIM/CMakeLists.txt +++ b/src/PIM/CMakeLists.txt @@ -120,8 +120,8 @@ add_pim_library(OMPIMAccel OMSpatialToPim OMPimCommon OMPimBufferization - OMPimMemoryCoalescing OMPimHostConstantFolding + OMPimMemoryCoalescing OMPimVerification MLIRTensorInferTypeOpInterfaceImpl ) diff --git a/src/PIM/Common/IR/CoreBlockUtils.cpp b/src/PIM/Common/IR/CoreBlockUtils.cpp index c369411..814b197 100644 --- a/src/PIM/Common/IR/CoreBlockUtils.cpp +++ b/src/PIM/Common/IR/CoreBlockUtils.cpp @@ -3,6 +3,7 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp" @@ -12,8 +13,10 @@ namespace onnx_mlir { bool isCoreStaticAddressOp(mlir::Operation* op) { if (mlir::isa; +enum class CoreWalkMode { + ExecuteCommunication, + StructuralExtremes +}; +using CoreWalkCallback = llvm::function_ref; -static void propagateRegionResults(mlir::ValueRange results, - mlir::Region& region, - StaticValueKnowledge& knowledge) { +static void propagateRegionResults(mlir::ValueRange results, mlir::Region& region, StaticValueKnowledge& knowledge) { if (region.empty()) return; auto yield = mlir::cast(region.front().getTerminator()); @@ -50,23 +53,38 @@ static void propagateRegionResults(mlir::ValueRange results, static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block, const StaticValueKnowledge& initialKnowledge, CoreWalkMode mode, + const PimCoreCommunicationPlan* communicationPlan, CoreWalkCallback callback) { bool hasFailure = false; StaticValueKnowledge knowledge = initialKnowledge; - llvm::StringRef purpose = mode == CoreWalkMode::ExecuteAllIterations ? "codegen" : "verification"; - for (mlir::Operation& op : block) { + llvm::StringRef purpose = mode == CoreWalkMode::ExecuteCommunication ? "communication verification" : "verification"; + llvm::SmallVector structuralOperations; + llvm::ArrayRef operations; + if (communicationPlan) { + auto it = communicationPlan->find(&block); + if (it != communicationPlan->end()) + operations = it->second; + } + else { + structuralOperations.reserve(block.getOperations().size()); + for (mlir::Operation& op : block) + structuralOperations.push_back(&op); + operations = structuralOperations; + } + + for (mlir::Operation* operation : operations) { + mlir::Operation& op = *operation; if (mlir::isa(op) || isCoreStaticAddressOp(&op)) continue; if (auto loadOp = mlir::dyn_cast(op); loadOp && succeeded(resolveIndexValue(loadOp.getResult(), knowledge))) continue; - if (auto forOp = mlir::dyn_cast(op)) { auto lower = resolveIndexValue(forOp.getLowerBound(), knowledge); auto upper = resolveIndexValue(forOp.getUpperBound(), knowledge); auto step = resolveIndexValue(forOp.getStep(), knowledge); if (failed(lower) || failed(upper) || failed(step) - || (mode == CoreWalkMode::ExecuteAllIterations && *step <= 0)) { + || (mode == CoreWalkMode::ExecuteCommunication && *step <= 0)) { forOp.emitOpError() << "requires statically evaluable scf.for bounds for PIM " << purpose; hasFailure = true; continue; @@ -84,13 +102,13 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block, loopKnowledge.indexValues[forOp.getInductionVar()] = induction; for (auto [index, iterArg] : llvm::enumerate(forOp.getRegionIterArgs())) loopKnowledge.aliases[iterArg] = carryValues ? iterValues[index] : forOp.getInitArgs()[index]; - hasFailure |= failed(walkPimCoreBlockImpl(body, loopKnowledge, mode, callback)); + hasFailure |= failed(walkPimCoreBlockImpl(body, loopKnowledge, mode, communicationPlan, callback)); auto yield = mlir::cast(body.getTerminator()); for (auto [index, yielded] : llvm::enumerate(yield.getOperands())) iterValues[index] = resolveLoopCarriedAlias(yielded, loopKnowledge); }; - if (mode == CoreWalkMode::ExecuteAllIterations) { + if (mode == CoreWalkMode::ExecuteCommunication) { for (int64_t induction = *lower; induction < *upper; induction += *step) visitIteration(induction, true); } @@ -113,12 +131,14 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block, continue; } mlir::Region& selected = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion(); - if (mode == CoreWalkMode::ExecuteAllIterations) { - hasFailure |= !selected.empty() && failed(walkPimCoreBlockImpl(selected.front(), knowledge, mode, callback)); + if (mode == CoreWalkMode::ExecuteCommunication) { + hasFailure |= !selected.empty() + && failed(walkPimCoreBlockImpl(selected.front(), knowledge, mode, communicationPlan, callback)); } else { for (mlir::Region* region : {&ifOp.getThenRegion(), &ifOp.getElseRegion()}) - hasFailure |= !region->empty() && failed(walkPimCoreBlockImpl(region->front(), knowledge, mode, callback)); + hasFailure |= !region->empty() + && failed(walkPimCoreBlockImpl(region->front(), knowledge, mode, communicationPlan, callback)); } propagateRegionResults(ifOp.getResults(), selected, knowledge); continue; @@ -138,33 +158,54 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block, break; } for (mlir::Region& region : switchOp->getRegions()) { - if (mode == CoreWalkMode::ExecuteAllIterations && ®ion != selected) + if (mode == CoreWalkMode::ExecuteCommunication && ®ion != selected) continue; - hasFailure |= failed(walkPimCoreBlockImpl(region.front(), knowledge, mode, callback)); + hasFailure |= failed(walkPimCoreBlockImpl(region.front(), knowledge, mode, communicationPlan, callback)); } propagateRegionResults(switchOp.getResults(), *selected, knowledge); continue; } - hasFailure |= failed(callback(op, knowledge)); + if (mode != CoreWalkMode::ExecuteCommunication || mlir::isa(op)) + hasFailure |= failed(callback(op, knowledge)); } return mlir::success(!hasFailure); } } // namespace -mlir::LogicalResult -walkPimCoreBlock(mlir::Block& block, - const StaticValueKnowledge& knowledge, - llvm::function_ref callback) { - return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::ExecuteAllIterations, callback); +PimCoreCommunicationPlan buildPimCoreCommunicationPlan(mlir::Block& block) { + llvm::DenseSet communicationAncestors; + block.walk([&](mlir::Operation* op) { + if (!mlir::isa(op)) + return; + for (mlir::Operation* parent = op->getParentOp(); parent; parent = parent->getParentOp()) + communicationAncestors.insert(parent); + }); + + PimCoreCommunicationPlan plan; + block.walk([&](mlir::Operation* op) { + bool isCommunication = mlir::isa(op); + bool isControlFlow = mlir::isa(op); + if (isCommunication || (isControlFlow && (op->getNumResults() != 0 || communicationAncestors.contains(op)))) + plan[op->getBlock()].push_back(op); + }); + return plan; +} + +mlir::LogicalResult walkPimCoreCommunicationBlock( + mlir::Block& block, + const PimCoreCommunicationPlan& plan, + const StaticValueKnowledge& knowledge, + llvm::function_ref callback) { + return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::ExecuteCommunication, &plan, callback); } mlir::LogicalResult walkPimCoreBlockStructurally( mlir::Block& block, const StaticValueKnowledge& knowledge, llvm::function_ref callback) { - return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::StructuralExtremes, callback); + return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::StructuralExtremes, nullptr, callback); } } // namespace onnx_mlir diff --git a/src/PIM/Common/IR/CoreBlockUtils.hpp b/src/PIM/Common/IR/CoreBlockUtils.hpp index 3a9c56e..d41828c 100644 --- a/src/PIM/Common/IR/CoreBlockUtils.hpp +++ b/src/PIM/Common/IR/CoreBlockUtils.hpp @@ -3,23 +3,30 @@ #include "mlir/IR/Block.h" #include "mlir/Support/LogicalResult.h" +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLFunctionalExtras.h" +#include "llvm/ADT/SmallVector.h" #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp" namespace onnx_mlir { +using PimCoreCommunicationPlan = llvm::DenseMap>; + /// Returns true for ops in a `pim.core` body that only participate in static /// address or index computation and therefore do not emit PIM instructions. bool isCoreStaticAddressOp(mlir::Operation* op); -/// Walks a `pim.core` body, statically unrolling nested `scf.for` loops when -/// their bounds are known and invoking `callback` only on instruction-emitting -/// operations. -mlir::LogicalResult -walkPimCoreBlock(mlir::Block& block, - const StaticValueKnowledge& knowledge, - llvm::function_ref callback); +/// Walks a `pim.core` body's communication stream, statically unrolling +/// control flow that contains send/receive operations and invoking `callback` +/// on those operations in execution order. +PimCoreCommunicationPlan buildPimCoreCommunicationPlan(mlir::Block& block); + +mlir::LogicalResult walkPimCoreCommunicationBlock( + mlir::Block& block, + const PimCoreCommunicationPlan& plan, + const StaticValueKnowledge& knowledge, + llvm::function_ref callback); /// Walks a `pim.core`-like body structurally for verification without /// enumerating full loop trip counts. Loop bounds must still be statically diff --git a/src/PIM/Common/PimCommon.hpp b/src/PIM/Common/PimCommon.hpp index e65019a..ba6b588 100644 --- a/src/PIM/Common/PimCommon.hpp +++ b/src/PIM/Common/PimCommon.hpp @@ -22,9 +22,19 @@ #include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp" #include "src/Compiler/CompilerOptions.hpp" +#include +#include + 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 size_t kPimLocalMemoryAddressLimit = static_cast(std::numeric_limits::max()); } // namespace onnx_mlir diff --git a/src/PIM/Compiler/CMakeLists.txt b/src/PIM/Compiler/CMakeLists.txt index dd37a9d..95e2015 100644 --- a/src/PIM/Compiler/CMakeLists.txt +++ b/src/PIM/Compiler/CMakeLists.txt @@ -18,7 +18,6 @@ add_pim_library(OMPimCompilerUtils PimArtifactWriter.cpp PimCodeGen.cpp PimCoreProgram.cpp - PimMemoryLiveness.cpp PimWeightEmitter.cpp EXCLUDE_FROM_OM_LIBS @@ -31,8 +30,8 @@ add_pim_library(OMPimCompilerUtils OMPimCompilerOptions OMPimCommon OMPimBufferization - OMPimMemoryCoalescing OMPimHostConstantFolding + OMPimLocalMemoryPlanning OMPimVerification OMPimPasses OMONNXToSpatial diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index 25bfcef..6bd5fd4 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -15,6 +15,7 @@ #include "llvm/ADT/StringExtras.h" #include "llvm/Support/Debug.h" #include "llvm/Support/FileSystem.h" +#include "llvm/Support/FormatVariadic.h" #include "llvm/Support/JSON.h" #include "llvm/Support/raw_ostream.h" @@ -31,8 +32,8 @@ #include "Common/IR/CompactAsmUtils.hpp" #include "Common/PimCommon.hpp" -#include "Common/Support/Diagnostics.hpp" #include "Common/Support/CheckedArithmetic.hpp" +#include "Common/Support/Diagnostics.hpp" #include "Common/Support/ReportUtils.hpp" #include "Conversion/ONNXToSpatial/Common/Common.hpp" #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" @@ -43,7 +44,6 @@ #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp" -#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp" #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" @@ -110,8 +110,6 @@ static Operation* getDiagnosticAnchor(mlir::Value value) { // PIM instruction immediates are serialized as signed int32_t fields today // (`sldi` goes through checkedI32OrCrash), so local addresses must stay within // the non-negative int32_t range. -static constexpr size_t kPimAddressLimit = static_cast(std::numeric_limits::max()); - static FailureOr checkedAlignTo(size_t value, size_t alignment, Operation* anchor, StringRef fieldName) { if (alignment == 0) return value; @@ -129,7 +127,7 @@ static void printMemoryOverflowDiagnostic(const MemoryValueKey& key, llvm::errs() << "Requested allocation size: " << requestedSize << " bytes\n"; llvm::errs() << "Current firstAvailableAddress: " << currentFirstAvailableAddress << "\n"; llvm::errs() << "Aligned end address: " << alignedEndAddress << "\n"; - llvm::errs() << "Address limit: " << kPimAddressLimit << " (signed int32_t immediate range)\n"; + llvm::errs() << "Address limit: " << kPimLocalMemoryAddressLimit << " (signed int32_t immediate range)\n"; if (key.lane) llvm::errs() << "Lane: " << *key.lane << "\n"; llvm::errs() << "Value: "; @@ -152,10 +150,13 @@ size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) { FailureOr checkedAlignedEnd = failure(); if (succeeded(checkedEnd)) checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment"); - if (address > kPimAddressLimit || failed(checkedEnd) || *checkedEnd > kPimAddressLimit - || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimAddressLimit) { + if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit + || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimLocalMemoryAddressLimit) { printMemoryOverflowDiagnostic( - key, size, firstAvailableAddress, succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit); + key, + size, + firstAvailableAddress, + succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit); llvm_unreachable("PIM local memory allocation overflow"); } firstAvailableAddress = *checkedAlignedEnd; @@ -202,15 +203,6 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE } } -PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key) { - PhysicalSlotInfo slot; - slot.id = nextPhysicalSlotId++; - slot.address = allocateAddress(slotSize, key); - slot.size = slotSize; - localPhysicalSlots.push_back(slot); - return slot; -} - void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) { SmallDenseMap globalConstants; SmallVector, 16> globalAliases; @@ -249,71 +241,17 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) { globalMemEntriesMap[getMemoryValueKey(alias)] = getMemEntry(getMemoryValueKey(original)); } -void PimMemory::allocateCore(Operation* op, std::optional lane) { - auto intervals = buildLocalAllocIntervals(op, lane, pimMemoryReport == PimMemoryReportFull); - SmallVector plannedSlots = planPhysicalSlots(intervals); - - SmallVector slotOrder(plannedSlots.size()); - std::iota(slotOrder.begin(), slotOrder.end(), 0); - llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) { - const PlannedPhysicalSlot& lhs = plannedSlots[lhsIndex]; - const PlannedPhysicalSlot& rhs = plannedSlots[rhsIndex]; - if (lhs.requiredSize != rhs.requiredSize) - return lhs.requiredSize > rhs.requiredSize; - return lhs.id < rhs.id; - }); - - SmallVector usedExistingSlots(localPhysicalSlots.size(), false); - for (size_t slotIndex : slotOrder) { - PlannedPhysicalSlot& slot = plannedSlots[slotIndex]; - size_t bestExistingIndex = std::numeric_limits::max(); - auto bestKey = std::tuple( - std::numeric_limits::max(), std::numeric_limits::max(), std::numeric_limits::max()); - - for (size_t existingIndex = 0; existingIndex < localPhysicalSlots.size(); ++existingIndex) { - if (usedExistingSlots[existingIndex]) - continue; - const PhysicalSlotInfo& existingSlot = localPhysicalSlots[existingIndex]; - if (existingSlot.size < slot.requiredSize) - continue; - auto candidateKey = - std::tuple(existingSlot.size - slot.requiredSize, existingSlot.size, existingSlot.id); - if (candidateKey < bestKey) { - bestKey = candidateKey; - bestExistingIndex = existingIndex; - } - } - - if (bestExistingIndex != std::numeric_limits::max()) { - const PhysicalSlotInfo& existingSlot = localPhysicalSlots[bestExistingIndex]; - slot.id = existingSlot.id; - slot.address = existingSlot.address; - slot.size = existingSlot.size; - usedExistingSlots[bestExistingIndex] = true; - } - else { - PhysicalSlotInfo newSlot = allocatePhysicalSlot(slot.requiredSize, intervals[slot.intervalIndices.front()].key); - slot.id = newSlot.id; - slot.address = newSlot.address; - slot.size = newSlot.size; - usedExistingSlots.push_back(true); - } - - for (size_t intervalIndex : slot.intervalIndices) { - LocalAllocInterval& interval = intervals[intervalIndex]; - interval.physicalSlotId = slot.id; - interval.assignedAddress = slot.address; - interval.physicalSlotSize = slot.size; - MemEntry memEntry {slot.address, interval.size}; - ownedMemEntriesMap[interval.key] = memEntry; - globalMemEntriesMap[interval.key] = memEntry; - } +void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional lane) { + if (localPhysicalSlots.empty()) { + llvm::append_range(localPhysicalSlots, plan.slots); + reportRow.logicalAllocaBytes = plan.logicalBytes; + reportRow.fallbackIntervals = plan.fallbackIntervals; + reportRow.nestedSingleUseIntervals = plan.nestedSingleUseIntervals; } - - if (pimMemoryReport != PimMemoryReportNone) { - MemoryPlanArtifacts artifacts = - buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport); - livenessArtifacts += artifacts; + for (const CompiledLocalMemoryEntry& entry : plan.entries) { + MemoryValueKey key = getMemoryValueKey(entry.value, lane); + ownedMemEntriesMap[key] = entry.memory; + globalMemEntriesMap[key] = entry.memory; } } @@ -475,13 +413,41 @@ void PimAcceleratorMemory::flushReport() { uint64_t totalGlobalMemory = hostReportRow.has_value() ? hostReportRow->sizeGlobal : 0; uint64_t totalWeightsMemory = totalWeightBytes; uint64_t totalCoresMemory = 0; - for (const MemoryReportEntry& entry : reportEntries) + uint64_t totalLogicalCoreMemory = 0; + uint64_t totalFallbackIntervals = 0; + uint64_t totalNestedSingleUseIntervals = 0; + uint64_t maximumCorePeak = 0; + std::string worstEntry = ""; + 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)}, - {"Cores memory", formatReportMemory(totalCoresMemory) } + 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) } }; printReportTotalsBlock(os, totalFields); @@ -552,19 +518,14 @@ void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instructio void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const { switch (instruction.opcode) { - case pim_binary::Opcode::sldi: - scalarRegisterValues[instruction.rd] = instruction.r2OrImm; - break; + case pim_binary::Opcode::sldi: scalarRegisterValues[instruction.rd] = instruction.r2OrImm; break; case pim_binary::Opcode::sld: case pim_binary::Opcode::sadd: case pim_binary::Opcode::ssub: case pim_binary::Opcode::smul: case pim_binary::Opcode::saddi: - case pim_binary::Opcode::smuli: - scalarRegisterValues[instruction.rd].reset(); - break; - default: - break; + case pim_binary::Opcode::smuli: scalarRegisterValues[instruction.rd].reset(); break; + default: break; } } @@ -580,8 +541,8 @@ void PimCodeGen::genSetRegisterImmediate(uint8_t registerNumber, int32_t immedia } void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const { - genSetRegisterImmediate( - pim::checkedU8OrCrash(registerNumber, "register number"), pim::checkedI32OrCrash(immediate, "register immediate")); + genSetRegisterImmediate(pim::checkedU8OrCrash(registerNumber, "register number"), + pim::checkedI32OrCrash(immediate, "register immediate")); } void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const { @@ -729,8 +690,7 @@ void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKno for (size_t outerIndex = 0; outerIndex < outerCount; ++outerIndex) { size_t dstOffset = (outerIndex * outputConcatDim + concatOffset) * innerCount * elementSize; size_t srcOffset = outerIndex * inputConcatDim * innerCount * elementSize; - emitMemCopyOp( - pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len"); + emitMemCopyOp(pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len"); } concatOffset += inputConcatDim; @@ -788,10 +748,11 @@ void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const Sta size_t maxElementOffset = plan.totalBytes == 0 ? 0 : plan.totalBytes - plan.elementBytes; int32_t maxSourceAddress = pim::checkedI32OrCrash( pim::checkedAddOrCrash(srcAddr, maxElementOffset, "transpose source address"), "transpose source address"); - int32_t maxDestinationAddress = pim::checkedI32OrCrash( - pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"), "transpose destination address"); - (void)maxSourceAddress; - (void)maxDestinationAddress; + int32_t maxDestinationAddress = + pim::checkedI32OrCrash(pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"), + "transpose destination address"); + (void) maxSourceAddress; + (void) maxDestinationAddress; pim_binary::InstructionRecord copyInstruction; copyInstruction.opcode = pim_binary::Opcode::lmv; @@ -871,6 +832,60 @@ static SmallVector collectTopLevelCoreLikeOps(func::FuncOp funcOp) { return coreLikeOps; } +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"); + return failure(); + } + plan.fallbackIntervals = static_cast(fallbackAttr.getInt()); + plan.nestedSingleUseIntervals = static_cast(nestedSingleUseAttr.getInt()); + SmallDenseMap slotIndices; + 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"); + hasFailure = true; + return; + } + + auto checkedSize = pim::getCheckedShapedTypeSizeInBytes( + cast(allocOp.getType()), allocOp, "planned local allocation byte size"); + if (failed(checkedSize)) { + hasFailure = true; + 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}); + return; + } + const PhysicalSlotInfo& existing = plan.slots[slotIt->second]; + if (existing.size != slotSize) { + allocOp.emitOpError("has a PIM local-memory arena inconsistent with another allocation"); + hasFailure = true; + } + }); + if (hasFailure) + return failure(); + llvm::sort(plan.slots, [](const PhysicalSlotInfo& lhs, const PhysicalSlotInfo& rhs) { return lhs.id < rhs.id; }); + return plan; +} + struct CoreEmissionResult { static constexpr size_t kMaxStoredCodegenDiagnostics = 8; @@ -882,10 +897,8 @@ struct CoreEmissionResult { OnnxMlirCompilerErrorCodes status = CompilerSuccess; MemoryReportRow reportRow; llvm::SmallVector usedWeights; - MemoryPlanArtifacts livenessArtifacts; llvm::SmallVector diagnostics; size_t diagnosticCount = 0; - void recordDiagnostic(Operation* op, StringRef message) { ++diagnosticCount; if (diagnostics.size() < kMaxStoredCodegenDiagnostics) @@ -1012,13 +1025,21 @@ static LogicalResult executeCompiledCorePlan( } auto emitBinary = [&](auto op, pim_binary::Opcode opcode) { - coreCodeGen.emitBinaryVectorOp(opcode, op.getOutputBuffer(), op.getLhs(), op.getRhs(), - getVectorByteSizeOrCrash(cast(op.getLhs().getType())), knowledge); + coreCodeGen.emitBinaryVectorOp(opcode, + op.getOutputBuffer(), + op.getLhs(), + op.getRhs(), + getVectorByteSizeOrCrash(cast(op.getLhs().getType())), + knowledge); }; auto emitUnary = [&](auto op, pim_binary::Opcode opcode, int32_t r2OrImm, int32_t generic1) { - coreCodeGen.emitUnaryVectorOp(opcode, op.getOutputBuffer(), op.getInput(), + coreCodeGen.emitUnaryVectorOp(opcode, + op.getOutputBuffer(), + op.getInput(), getVectorByteSizeOrCrash(cast(op.getInput().getType())), - knowledge, r2OrImm, generic1); + knowledge, + r2OrImm, + generic1); }; switch (node.opKind) { @@ -1038,20 +1059,19 @@ static LogicalResult executeCompiledCorePlan( else return failure(); break; - case CompiledCoreOpKind::Transpose: - coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); - break; - case CompiledCoreOpKind::VVAdd: emitBinary(cast(node.op), pim_binary::Opcode::vvadd); break; - case CompiledCoreOpKind::VVSub: emitBinary(cast(node.op), pim_binary::Opcode::vvsub); break; - case CompiledCoreOpKind::VVMul: emitBinary(cast(node.op), pim_binary::Opcode::vvmul); break; - case CompiledCoreOpKind::VVMax: emitBinary(cast(node.op), pim_binary::Opcode::vvmax); break; - case CompiledCoreOpKind::VVDMul: emitBinary(cast(node.op), pim_binary::Opcode::vvdmul); break; - case CompiledCoreOpKind::VAvg: emitUnary(cast(node.op), pim_binary::Opcode::vavg, 1, 1); break; - case CompiledCoreOpKind::VRelu: emitUnary(cast(node.op), pim_binary::Opcode::vrelu, 0, 0); break; - case CompiledCoreOpKind::VTanh: emitUnary(cast(node.op), pim_binary::Opcode::vtanh, 0, 0); break; - case CompiledCoreOpKind::VSigm: emitUnary(cast(node.op), pim_binary::Opcode::vsigm, 0, 0); break; + case CompiledCoreOpKind::Transpose: coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break; + case CompiledCoreOpKind::VVAdd: emitBinary(cast(node.op), pim_binary::Opcode::vvadd); break; + case CompiledCoreOpKind::VVSub: emitBinary(cast(node.op), pim_binary::Opcode::vvsub); break; + case CompiledCoreOpKind::VVMul: emitBinary(cast(node.op), pim_binary::Opcode::vvmul); break; + case CompiledCoreOpKind::VVMax: emitBinary(cast(node.op), pim_binary::Opcode::vvmax); break; + case CompiledCoreOpKind::VVDMul: emitBinary(cast(node.op), pim_binary::Opcode::vvdmul); break; + case CompiledCoreOpKind::VAvg: emitUnary(cast(node.op), pim_binary::Opcode::vavg, 1, 1); break; + case CompiledCoreOpKind::VRelu: emitUnary(cast(node.op), pim_binary::Opcode::vrelu, 0, 0); break; + case CompiledCoreOpKind::VTanh: emitUnary(cast(node.op), pim_binary::Opcode::vtanh, 0, 0); break; + case CompiledCoreOpKind::VSigm: emitUnary(cast(node.op), pim_binary::Opcode::vsigm, 0, 0); break; case CompiledCoreOpKind::VSoftmax: - emitUnary(cast(node.op), pim_binary::Opcode::vsoftmax, 0, 0); break; + emitUnary(cast(node.op), pim_binary::Opcode::vsoftmax, 0, 0); + break; } } return success(); @@ -1103,8 +1123,7 @@ static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp, } static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) { - std::string outputCorePath = - (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str(); + std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str(); std::error_code errorCode; raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None); if (errorCode) { @@ -1128,8 +1147,7 @@ static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath if (!pimEmitJson.getValue()) return CompilerSuccess; - std::string outputCoreJsonPath = - (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str(); + std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str(); errorCode = std::error_code(); raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode); if (errorCode) { @@ -1169,18 +1187,27 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: SmallDenseMap materializedHostGlobals = collectMaterializedHostGlobals(moduleOp, funcOp, memory); llvm::DenseMap> compiledPrograms; + llvm::DenseMap> memoryPlans; for (Operation* op : coreLikeOps) { auto program = std::make_unique(); if (failed(compileCoreProgram(op, *program))) return CompilerFailure; compiledPrograms.try_emplace(op, std::move(program)); + auto memoryPlan = compileCoreMemoryPlan(op); + if (failed(memoryPlan)) + return CompilerFailure; + memoryPlans.try_emplace(op, std::make_unique(std::move(*memoryPlan))); } - auto getCompiledProgram = [&](Operation* op) { auto it = compiledPrograms.find(op); assert(it != compiledPrograms.end() && "missing compiled PIM core program"); return it->second.get(); }; + auto getMemoryPlan = [&](Operation* op) { + auto it = memoryPlans.find(op); + assert(it != memoryPlans.end() && "missing PIM core memory plan"); + return it->second.get(); + }; llvm::DenseMap emittedCoreIds; size_t nextEmittedCoreId = 0; @@ -1210,6 +1237,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: CoreEmissionJob job; job.coreLikeOp = coreOp; job.program = getCompiledProgram(op); + job.memoryPlan = getMemoryPlan(op); job.emittedCoreId = emittedCoreIds.lookup(originalCoreId); jobs.push_back(std::move(job)); continue; @@ -1229,6 +1257,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: CoreEmissionJob job; job.coreLikeOp = coreBatchOp; job.program = getCompiledProgram(op); + job.memoryPlan = getMemoryPlan(op); job.emittedCoreId = emittedCoreIds.lookup(originalCoreId); job.lanes = lanesByCoreId.lookup(originalCoreId); job.batchReportId = nextBatchReportId; @@ -1332,17 +1361,16 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: if (auto coreOp = dyn_cast(job.coreLikeOp)) { aliasMaterializedHostGlobals(coreOp, moduleOp, materializedHostGlobals, jobMemory); auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId); - deviceMemory.allocateCore(coreOp); + deviceMemory.allocateCore(*job.memoryPlan); StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp); if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) { - (void)finalizeInstructions(); + (void) finalizeInstructions(); result.status = CompilerFailure; return result; } result.reportRow = deviceMemory.getReportRow(); result.usedWeights = std::move(usedWeights); - result.livenessArtifacts = deviceMemory.getLivenessArtifacts(); } else { auto coreBatchOp = cast(job.coreLikeOp); @@ -1352,10 +1380,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: for (unsigned lane : job.lanes) { StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane); - deviceMemory.allocateCore(coreBatchOp, lane); + deviceMemory.allocateCore(*job.memoryPlan, lane); coreCodeGen.setBatchLane(lane); if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) { - (void)finalizeInstructions(); + (void) finalizeInstructions(); result.status = CompilerFailure; return result; } @@ -1363,7 +1391,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: result.reportRow = deviceMemory.getReportRow(); result.usedWeights = std::move(usedWeights); - result.livenessArtifacts = deviceMemory.getLivenessArtifacts(); } if (!finalizeInstructions()) { @@ -1388,9 +1415,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: for (const CoreEmissionResult& result : jobResults) { if (!summaryAnchor && !result.diagnostics.empty()) summaryAnchor = result.diagnostics.front().op; - for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics) { + for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics) diagnostics.report(diagnostic.op, [&](Operation* op) { op->emitError() << diagnostic.message; }); - } size_t unreportedCount = result.diagnosticCount - result.diagnostics.size(); diagnostics.noteFailures(static_cast(unreportedCount)); } @@ -1420,19 +1446,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath); memory.setTotalWeightBytes(weightEmission.totalWeightBytes); auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName; - if (std::string reportsRoot = getOutputDir(); !reportsRoot.empty()) { - std::string reportsDir = reportsRoot + "/reports"; - sys::fs::remove(reportsDir + "/pim_memory_liveness_report.txt"); - sys::fs::remove(reportsDir + "/pim_memory_liveness_report.json"); - sys::fs::remove(reportsDir + "/pim_memory_liveness_timeline.dot"); - } - std::fstream livenessReportFile; - std::unique_ptr livenessReportOs; - if (pimMemoryReport != PimMemoryReportNone) { - livenessReportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt"); - livenessReportOs = std::make_unique(livenessReportFile); - } - for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) { const CoreEmissionJob& job = jobs[jobIndex]; const CoreEmissionResult& result = jobResults[jobIndex]; @@ -1443,8 +1456,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: return err; xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup); memory.recordCoreReport(job.emittedCoreId, result.reportRow); - if (livenessReportFile.is_open()) - *livenessReportOs << "Core " << job.emittedCoreId << ":\n" << result.livenessArtifacts; continue; } } @@ -1473,18 +1484,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: batchPerCoreRow.value_or(MemoryReportRow {}), batchRow.numAlloca, batchRow.sizeAlloca); - if (livenessReportFile.is_open()) - for (size_t jobIndex : group) - *livenessReportOs << "Batch " << batchReportId << " core " << jobs[jobIndex].emittedCoreId << ":\n" - << jobResults[jobIndex].livenessArtifacts; } maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1; - if (livenessReportFile.is_open()) { - livenessReportOs->flush(); - livenessReportFile.close(); - } memory.flushReport(); return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath); } diff --git a/src/PIM/Compiler/PimCodeGen.hpp b/src/PIM/Compiler/PimCodeGen.hpp index 1a31ce8..385d6b7 100644 --- a/src/PIM/Compiler/PimCodeGen.hpp +++ b/src/PIM/Compiler/PimCodeGen.hpp @@ -39,8 +39,6 @@ struct PhysicalSlotInfo { size_t size = 0; }; -using MemoryPlanArtifacts = std::string; - struct MemoryValueKey { mlir::Value value; std::optional lane; @@ -48,15 +46,33 @@ struct MemoryValueKey { bool operator==(const MemoryValueKey& other) const { return value == other.value && lane == other.lane; } }; +struct CompiledLocalMemoryEntry { + mlir::Value value; + MemEntry memory; +}; + +struct CompiledCoreMemoryPlan { + llvm::SmallVector entries; + llvm::SmallVector slots; + uint64_t logicalBytes = 0; + uint64_t fallbackIntervals = 0; + uint64_t nestedSingleUseIntervals = 0; +}; + struct MemoryReportRow { uint64_t numAlloca = 0; uint64_t sizeAlloca = 0; uint64_t numGlobal = 0; uint64_t sizeGlobal = 0; + uint64_t logicalAllocaBytes = 0; + uint64_t fallbackIntervals = 0; + uint64_t nestedSingleUseIntervals = 0; bool operator==(const MemoryReportRow& other) const { return numAlloca == other.numAlloca && sizeAlloca == other.sizeAlloca && numGlobal == other.numGlobal - && sizeGlobal == other.sizeGlobal; + && sizeGlobal == other.sizeGlobal && logicalAllocaBytes == other.logicalAllocaBytes + && fallbackIntervals == other.fallbackIntervals + && nestedSingleUseIntervals == other.nestedSingleUseIntervals; } }; @@ -93,26 +109,22 @@ class PimMemory { llvm::SmallDenseMap& globalMemEntriesMap; llvm::SmallDenseMap ownedMemEntriesMap; MemoryReportRow reportRow; - MemoryPlanArtifacts livenessArtifacts; size_t minAlignment = 4; size_t firstAvailableAddress = 0; - size_t nextPhysicalSlotId = 0; MemEntry* gatherMemEntry(mlir::Value value, std::optional lane = std::nullopt); size_t allocateAddress(size_t size, const MemoryValueKey& key); void allocateGatheredMemory(); void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind); - PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key); public: PimMemory(llvm::SmallDenseMap& globalMemEntriesMap) : globalMemEntriesMap(globalMemEntriesMap) {} void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp); - void allocateCore(mlir::Operation* op, std::optional lane = std::nullopt); + void allocateCore(const CompiledCoreMemoryPlan& plan, std::optional lane = std::nullopt); MemoryReportRow getReportRow() const; - const MemoryPlanArtifacts& getLivenessArtifacts() const { return livenessArtifacts; } size_t getFirstAvailableAddress() const { return firstAvailableAddress; } MemEntry getMemEntry(const MemoryValueKey& key) const; @@ -160,6 +172,7 @@ public: struct CoreEmissionJob { mlir::Operation* coreLikeOp = nullptr; const CompiledCoreProgram* program = nullptr; + const CompiledCoreMemoryPlan* memoryPlan = nullptr; size_t emittedCoreId = 0; llvm::SmallVector lanes; std::optional batchReportId; diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index b50dc9e..5cf8a33 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -74,7 +74,7 @@ llvm::cl::opt llvm::cl::opt pimDisableMemoryCoalescing("pim-disable-memory-coalescing", - llvm::cl::desc("Skip the PIM memory coalescing pass (developer diagnostic option)"), + llvm::cl::desc("Skip the early PIM IR memory coalescing pass (developer diagnostic)"), llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions)); @@ -124,10 +124,10 @@ llvm::cl::opt pimTraceCommunicationMaterialization( llvm::cl::cat(OnnxMlirOptions)); llvm::cl::opt - crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(2)); + crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(128)); llvm::cl::opt - crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(256)); + crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(64)); llvm::cl::opt coresCount("core-count", llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."), diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index 0f782a2..d109ccc 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -53,8 +53,8 @@ extern llvm::cl::opt pimMemoryReport; extern llvm::cl::opt pimConvLowering; extern llvm::cl::opt pimExportSpatialDataflow; -extern llvm::cl::opt pimOnlyCodegen; extern llvm::cl::opt pimDisableMemoryCoalescing; +extern llvm::cl::opt pimOnlyCodegen; extern llvm::cl::opt useExperimentalConvImpl; extern llvm::cl::opt pimEmitJson; extern llvm::cl::opt pimReportConvLowering; diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index 6c9e946..54a4517 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -21,6 +21,7 @@ void addPassesPim(OwningOpRef& module, verifyExplicitPimCoreCount(); if (pimOnlyCodegen) { + pm.addPass(createPimLocalMemoryPlanningPass()); pm.addPass(createEmitPimCodePass()); return; } @@ -53,6 +54,8 @@ void addPassesPim(OwningOpRef& module, pm.addPass(createMessagePass("Pim host constants folded")); if (!pimDisableMemoryCoalescing) pm.addPass(createPimMemoryCoalescingPass()); + pm.addPass(createPimLocalMemoryPlanningPass()); + pm.addPass(createMessagePass("Pim local memory planned")); pm.addPass(createPimVerificationPass()); pm.addPass(createMessagePass("Pim verified")); pm.addPass(createEmitPimCodePass()); diff --git a/src/PIM/Dialect/Pim/Analysis/CMakeLists.txt b/src/PIM/Dialect/Pim/Analysis/CMakeLists.txt new file mode 100644 index 0000000..43336bf --- /dev/null +++ b/src/PIM/Dialect/Pim/Analysis/CMakeLists.txt @@ -0,0 +1,8 @@ +add_pim_library(OMPimLocalMemoryLifetimeAnalysis + LocalMemoryLifetimeAnalysis.cpp + + EXCLUDE_FROM_OM_LIBS + + INCLUDE_DIRS PUBLIC + ${PIM_PUBLIC_INCLUDE_DIRS} +) diff --git a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp b/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp new file mode 100644 index 0000000..0bf3da7 --- /dev/null +++ b/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp @@ -0,0 +1,71 @@ +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Interfaces/DestinationStyleOpInterface.h" + +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" + +#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" + +using namespace mlir; + +namespace onnx_mlir { +namespace pim { + +bool isLocalMemoryAliasOp(Operation* op) { + return isa(op); +} + +LogicalResult walkLocalMemoryUses(Value root, + llvm::function_ref visitUser) { + 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(); + if (!visitedValues.insert(value).second) + continue; + for (Operation* user : value.getUsers()) { + if (!visitedUsers.insert(user).second) + continue; + if (failed(visitUser(value, user))) + return failure(); + + if (isLocalMemoryAliasOp(user)) + for (Value result : user->getResults()) + addAlias(result); + + 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 forOp = dyn_cast(user)) + for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs())) + if (initArg == value) { + addAlias(forOp.getRegionIterArgs()[index]); + addAlias(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)); + } + } + } + return success(); +} + +} // namespace pim +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp b/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp new file mode 100644 index 0000000..e9b1147 --- /dev/null +++ b/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "mlir/IR/Operation.h" + +#include "llvm/ADT/STLFunctionalExtras.h" + +namespace onnx_mlir { +namespace pim { + +bool isLocalMemoryAliasOp(mlir::Operation* op); + +mlir::LogicalResult walkLocalMemoryUses( + mlir::Value root, + llvm::function_ref visitUser); + +} // namespace pim +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/CMakeLists.txt b/src/PIM/Dialect/Pim/CMakeLists.txt index 6eebe9a..9badbab 100644 --- a/src/PIM/Dialect/Pim/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/CMakeLists.txt @@ -1,9 +1,11 @@ add_onnx_mlir_dialect(Pim pim) add_onnx_mlir_dialect_doc(pim Pim.td) +add_subdirectory(Analysis) add_subdirectory(Transforms/Bufferization) -add_subdirectory(Transforms/MemoryCoalescing) add_subdirectory(Transforms/HostConstantFolding) +add_subdirectory(Transforms/MemoryCoalescing) +add_subdirectory(Transforms/LocalMemoryPlanning) add_subdirectory(Transforms/Verification) add_pim_library(PimOps diff --git a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt new file mode 100644 index 0000000..d6e09bb --- /dev/null +++ b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt @@ -0,0 +1,14 @@ +add_pim_library(OMPimLocalMemoryPlanning + LocalMemoryPlanning.cpp + + EXCLUDE_FROM_OM_LIBS + + INCLUDE_DIRS PUBLIC + ${PIM_PUBLIC_INCLUDE_DIRS} + + LINK_LIBS PUBLIC + OMPimCommon + OMPimCompilerOptions + OMPimLocalMemoryLifetimeAnalysis + PimOps +) diff --git a/src/PIM/Compiler/PimMemoryLiveness.cpp b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp similarity index 51% rename from src/PIM/Compiler/PimMemoryLiveness.cpp rename to src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp index a53a998..d2a9993 100644 --- a/src/PIM/Compiler/PimMemoryLiveness.cpp +++ b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp @@ -1,14 +1,18 @@ #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 "mlir/Interfaces/DestinationStyleOpInterface.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallPtrSet.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 @@ -16,9 +20,13 @@ #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/PimMemoryLiveness.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; @@ -26,19 +34,6 @@ using namespace onnx_mlir; namespace { -static std::optional getLaneForMemoryValue(mlir::Value value, std::optional lane) { - if (!lane) - return std::nullopt; - auto allocOp = value.getDefiningOp(); - if (!allocOp || !allocOp->getParentOfType()) - return std::nullopt; - return lane; -} - -static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional lane = std::nullopt) { - return {value, getLaneForMemoryValue(value, lane)}; -} - struct MemoryTouchInterval { uint64_t start = 0; uint64_t end = 0; @@ -51,7 +46,6 @@ struct MemoryTouchInterval { bool endUsedFallback = false; bool escapesLoop = false; std::string fallbackReason; - llvm::SmallVector aliasesFollowed; }; struct OperationOrdering { @@ -60,50 +54,6 @@ struct OperationOrdering { uint64_t nextPosition = 0; }; -static std::string printValueToString(mlir::Value value) { - std::string text; - llvm::raw_string_ostream os(text); - value.print(os); - os.flush(); - return text; -} - -static std::string printOperationToString(Operation* op) { - if (!op) - return ""; - std::string text; - llvm::raw_string_ostream os(text); - op->print(os); - os.flush(); - return text; -} - -static std::string printLocationToString(Location loc) { - std::string text; - llvm::raw_string_ostream os(text); - loc.print(os); - os.flush(); - return text; -} - -static std::string collapseWhitespace(StringRef text) { - std::string out; - out.reserve(text.size()); - bool lastWasSpace = false; - for (char c : text) { - bool isSpace = c == ' ' || c == '\n' || c == '\t' || c == '\r'; - if (isSpace) { - if (!lastWasSpace && !out.empty()) - out.push_back(' '); - lastWasSpace = true; - continue; - } - out.push_back(c); - lastWasSpace = false; - } - return out; -} - static std::string abbreviate(StringRef text, size_t maxLen) { if (text.size() <= maxLen) return text.str(); @@ -111,21 +61,23 @@ static std::string abbreviate(StringRef text, size_t maxLen) { } static std::string summarizeValue(mlir::Value value, size_t maxLen = 72) { - return abbreviate(collapseWhitespace(printValueToString(value)), maxLen); + 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 ""; - std::string prefix = op->getName().getStringRef().str(); - std::string full = collapseWhitespace(printOperationToString(op)); - if (full == prefix) - return prefix; - return abbreviate(prefix + " :: " + full, maxLen); -} - -static std::string summarizeLocation(Location loc, size_t maxLen = 88) { - return abbreviate(collapseWhitespace(printLocationToString(loc)), maxLen); + return abbreviate(op->getName().getStringRef(), maxLen); } static void assignOperationOrdering(Operation* op, OperationOrdering& ordering) { @@ -151,13 +103,6 @@ static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) { return ordering; } -static bool isSupportedAliasOp(Operation* op) { - return isa(op); -} - static bool isRuntimeMemoryTouchOp(Operation* op) { return isa(op) || isCoreStaticAddressOp(op); + return pim::isLocalMemoryAliasOp(op) || isa(op) + || isCoreStaticAddressOp(op); } static bool isWithin(mlir::Value value, Region* region) { @@ -207,12 +153,6 @@ static void addFallbackReason(std::string& reason, StringRef newReason) { reason += newReason.str(); } -static void appendAliasDescription(llvm::SmallVectorImpl& aliases, mlir::Value value) { - std::string text = printValueToString(value); - if (!llvm::is_contained(aliases, text)) - aliases.push_back(std::move(text)); -} - struct OrderedTouchRange { uint64_t start = 0; uint64_t end = 0; @@ -233,97 +173,28 @@ getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const Operati return range; } -static MemoryTouchInterval -computeMemoryTouchInterval(memref::AllocOp allocOp, - const OperationOrdering& ordering, - uint64_t fallbackEnd, - bool includeAliasDescriptions) { +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 recordAlias = [&](mlir::Value value) { - if (includeAliasDescriptions) - appendAliasDescription(interval.aliasesFollowed, value); - }; - SmallPtrSet visitedValues; - SmallPtrSet visitedUsers; - SmallVector pendingValues; - pendingValues.push_back(allocOp.getResult()); auto parentLoop = allocOp->getParentOfType(); - - while (!pendingValues.empty()) { - mlir::Value value = pendingValues.pop_back_val(); - if (!visitedValues.insert(value).second) - continue; - - for (Operation* user : value.getUsers()) { - if (!visitedUsers.insert(user).second) - continue; - - if (isSupportedAliasOp(user)) { - for (mlir::Value result : user->getResults()) { - pendingValues.push_back(result); - recordAlias(result); - } - } - - if (auto dpsOp = dyn_cast(user)) { - for (OpResult result : user->getResults()) { - OpOperand* tiedOperand = dpsOp.getTiedOpOperand(result); - if (!tiedOperand || tiedOperand->get() != value) - continue; - pendingValues.push_back(result); - recordAlias(result); - } - } - - if (auto forOp = dyn_cast(user)) { - for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs())) { - if (initArg != value) - continue; - pendingValues.push_back(forOp.getRegionIterArgs()[index]); - pendingValues.push_back(forOp.getResult(index)); - recordAlias(forOp.getRegionIterArgs()[index]); - recordAlias(forOp.getResult(index)); - if (parentLoop && forOp != parentLoop) - interval.escapesLoop = true; - } - } - + (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 (ifOp) { - for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) { - if (operand != value) - continue; - pendingValues.push_back(ifOp.getResult(index)); - recordAlias(ifOp.getResult(index)); - } - } - else if (indexSwitch) { - for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) { - if (operand != value) - continue; - pendingValues.push_back(indexSwitch.getResult(index)); - recordAlias(indexSwitch.getResult(index)); - } - } - else if (!forOp) { + if (!forOp && !ifOp && !indexSwitch) addFallbackReason(interval.fallbackReason, "yield without scf.for parent"); - } - else { - for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) { - if (operand != value) - continue; - pendingValues.push_back(forOp.getResult(index)); - recordAlias(forOp.getResult(index)); - if (parentLoop && forOp == parentLoop) - interval.escapesLoop = true; - } - } + else if (forOp && parentLoop && forOp == parentLoop && llvm::is_contained(yieldOp.getOperands(), value)) + interval.escapesLoop = true; } if (isRuntimeMemoryTouchOp(user)) { @@ -345,23 +216,21 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, interval.hasRuntimeUse = true; } else { - if (range.start < interval.start) { + if (range.start < interval.start) interval.start = range.start; - } - if (range.end > interval.end) { + if (range.end > interval.end) interval.end = range.end; - } } - continue; + return success(); } if (isIgnoredLivenessUser(user)) - continue; + return success(); addFallbackReason(interval.fallbackReason, "unhandled user op"); interval.endUsedFallback = true; - } - } + return success(); + }); if (!interval.hasRuntimeUse) { interval.startUsedAllocFallback = true; @@ -374,9 +243,8 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, return interval; } - if (interval.endUsedFallback) { + if (interval.endUsedFallback) interval.end = std::max(interval.end, fallbackEnd); - } return interval; } @@ -402,11 +270,21 @@ static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, 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, - std::optional lane, - bool includeAliasDescriptions) { +SmallVector onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp) { SmallVector intervals; OperationOrdering ordering = buildOperationOrdering(coreLikeOp); if (ordering.position.empty()) @@ -423,12 +301,10 @@ SmallVector onnx_mlir::buildLocalAllocIntervals(Operation llvm_unreachable("Failed to compute local allocation size"); } - MemoryTouchInterval touchInterval = - computeMemoryTouchInterval(allocOp, ordering, fallbackEnd, includeAliasDescriptions); + MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd); LocalAllocInterval interval; interval.id = nextIntervalId++; interval.alloc = allocOp; - interval.key = getMemoryValueKey(allocOp.getResult(), lane); interval.start = touchInterval.start; interval.end = touchInterval.end; interval.size = *checkedSize; @@ -442,7 +318,9 @@ SmallVector onnx_mlir::buildLocalAllocIntervals(Operation interval.insideNestedRegion = isNestedAllocation(coreLikeOp, allocOp); interval.escapesLoop = touchInterval.escapesLoop; interval.fallbackReason = std::move(touchInterval.fallbackReason); - interval.aliasesFollowed = std::move(touchInterval.aliasesFollowed); + interval.valueSummary = summarizeValue(allocOp.getResult(), 88); + interval.firstTouchSummary = summarizeOperation(touchInterval.firstTouchOp); + interval.lastTouchSummary = summarizeOperation(touchInterval.lastTouchOp); intervals.push_back(std::move(interval)); }); @@ -467,60 +345,94 @@ SmallVector onnx_mlir::planPhysicalSlots(MutableArrayRef for (size_t intervalIndex : intervalOrder) { LocalAllocInterval& interval = intervals[intervalIndex]; - PlannedPhysicalSlot* bestSlot = nullptr; - auto bestKey = std::tuple(std::numeric_limits::max(), - std::numeric_limits::max(), - std::numeric_limits::max(), - std::numeric_limits::max()); - - for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) { - PlannedPhysicalSlot& slot = slots[slotIndex]; - bool compatible = true; - for (size_t otherIndex : slot.intervalIndices) { - if (intervalsOverlap(interval, intervals[otherIndex])) { - compatible = false; - break; - } + 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); } - if (!compatible) - continue; + } - size_t resultingSize = std::max(slot.requiredSize, interval.size); - size_t growth = resultingSize - slot.requiredSize; - auto candidateKey = - std::tuple(growth, resultingSize, slot.intervalIndices.size(), slot.id); + 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; - bestSlot = &slot; + bestAddress = candidate; } } + assert(bestAddress != std::numeric_limits::max() && "address after all conflicts must be available"); - if (!bestSlot) { - slots.push_back({slots.size(), interval.size, interval.size, 0, {intervalIndex}}); - interval.slotPlanIndex = slots.size() - 1; - interval.physicalSlotId = slots.back().id; - interval.physicalSlotSize = slots.back().requiredSize; - continue; + 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; } - - bestSlot->requiredSize = std::max(bestSlot->requiredSize, interval.size); - bestSlot->size = bestSlot->requiredSize; - bestSlot->intervalIndices.push_back(intervalIndex); - interval.slotPlanIndex = static_cast(bestSlot - slots.data()); - interval.physicalSlotId = bestSlot->id; - interval.physicalSlotSize = bestSlot->requiredSize; } return slots; } -MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp, - std::optional lane, - ArrayRef intervals, - ArrayRef slots, - size_t addressLimit, - PimMemoryReportLevel reportLevel) { - MemoryPlanArtifacts artifacts; +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; @@ -536,7 +448,6 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp, for (const LocalAllocInterval& interval : intervals) { totalLogicalBytes += interval.size; largestLogicalAllocation = std::max(largestLogicalAllocation, interval.size); - maximumAssignedAddress = std::max(maximumAssignedAddress, interval.assignedAddress + interval.physicalSlotSize); if (interval.startUsedAllocFallback || interval.endUsedFallback) ++fallbackIntervals; if (!interval.hasRuntimeUse) @@ -546,12 +457,14 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp, if (interval.escapesLoop) ++loopEscapingIntervals; } - for (const PlannedPhysicalSlot& slot : slots) { - totalPhysicalBytes += slot.size; - largestPhysicalSlot = std::max(largestPhysicalSlot, slot.size); - if (slot.intervalIndices.size() > 1) - reusedAllocations += slot.intervalIndices.size() - 1; + 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 = @@ -560,8 +473,6 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp, raw_string_ostream os(artifacts); os << "=== PIM Memory Liveness Report ===\n"; os << "Op: " << coreLikeOp->getName() << "\n"; - if (lane) - os << "Lane: " << *lane << "\n"; os << "Summary:\n"; os << " logical allocation bytes: " << formatReportMemory(totalLogicalBytes) << " (" << totalLogicalBytes << ")\n"; os << " physical allocation bytes: " << formatReportMemory(totalPhysicalBytes) << " (" << totalPhysicalBytes @@ -569,32 +480,27 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp, os << " saved bytes: " << formatReportMemory(savedBytes) << " (" << savedBytes << ")\n"; os << " saved percent: " << format("%.2f%%", savedPercent) << "\n"; os << " intervals: " << intervals.size() << "\n"; - os << " physical slots: " << slots.size() << "\n"; - os << " reused allocations: " << reusedAllocations << "\n"; + os << " 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 physical slot: " << largestPhysicalSlot << "\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"; - os << "\nHow To Read:\n"; - os << " `summary` only shows the strongest reuse cases and the worst offenders.\n"; - os << " Use `--pim-memory-report=full` when you need the complete slot-by-slot and interval-by-interval dump.\n"; - os << " Large single-use slots, fallback intervals, and nested single-use allocations are the best places\n"; - os << " to inspect if allocations should be moved, sunk, or made easier to coalesce earlier in the pipeline.\n"; - SmallVector reusedSlots; SmallVector singleUseSlots; - for (const PlannedPhysicalSlot& slot : slots) - if (slot.intervalIndices.size() > 1) - reusedSlots.push_back(&slot); + for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) { + if (hasAddressReuse(slotIndex, slots)) + reusedSlots.push_back(&slots[slotIndex]); else - singleUseSlots.push_back(&slot); + singleUseSlots.push_back(&slots[slotIndex]); + } llvm::stable_sort(reusedSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) { uint64_t lhsLogicalBytes = getSlotLogicalBytes(*lhs, intervals); @@ -603,140 +509,66 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp, return lhs->intervalIndices.size() > rhs->intervalIndices.size(); if (lhsLogicalBytes != rhsLogicalBytes) return lhsLogicalBytes > rhsLogicalBytes; - if (lhs->size != rhs->size) - return lhs->size > rhs->size; + 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->size != rhs->size) - return lhs->size > rhs->size; + 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 Reuse:\n"; + os << "\nBest Address Reuse:\n"; if (reusedSlots.empty()) { - os << " no slots were shared by multiple intervals\n"; + 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->size) + os << " slot #" << slot->id << " addr=" << slot->address << " size=" << formatReportMemory(slot->requiredSize) << " intervals=" << slot->intervalIndices.size() << " logical_sum=" << formatReportMemory(slotLogicalBytes) << "\n"; - for (size_t intervalIndex : slot->intervalIndices) { - const LocalAllocInterval& interval = intervals[intervalIndex]; - os << " #" << interval.id << " [" << interval.start << "," << interval.end << "]" - << " logical=" << formatReportMemory(interval.size) - << " first=" << summarizeOperation(interval.firstTouchOp, 40) - << " last=" << summarizeOperation(interval.lastTouchOp, 40) << "\n"; - } } } os << "\nTop Offenders:\n"; - bool printedAttention = false; for (const PlannedPhysicalSlot* slot : ArrayRef(singleUseSlots).take_front(kSummaryOffenderLimit)) { const LocalAllocInterval& interval = intervals[slot->intervalIndices.front()]; - printedAttention = true; os << " slot #" << slot->id << " is single-use" - << " size=" << formatReportMemory(slot->size) << " interval=#" << interval.id - << " value=" << summarizeValue(interval.key.value, 56) << "\n"; - os << " first=" << summarizeOperation(interval.firstTouchOp, 40) - << " last=" << summarizeOperation(interval.lastTouchOp, 40) + << " 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"; } - size_t fallbackPrinted = 0; - for (const LocalAllocInterval& interval : intervals) { - if (!(interval.startUsedAllocFallback || interval.endUsedFallback) || fallbackPrinted >= kSummaryOffenderLimit) - continue; - printedAttention = true; - ++fallbackPrinted; - os << " fallback interval #" << interval.id << " size=" << formatReportMemory(interval.size) - << " value=" << summarizeValue(interval.key.value, 56) << "\n"; - os << " reason: " << (interval.fallbackReason.empty() ? "" : interval.fallbackReason) << "\n"; - } - size_t nestedPrinted = 0; - for (const LocalAllocInterval& interval : intervals) { - if (nestedPrinted >= kSummaryOffenderLimit) - break; - if (!(interval.insideNestedRegion && slots[interval.slotPlanIndex].intervalIndices.size() == 1)) - continue; - printedAttention = true; - ++nestedPrinted; - os << " nested single-use interval #" << interval.id << " slot #" << interval.physicalSlotId - << " size=" << formatReportMemory(interval.size) << " value=" << summarizeValue(interval.key.value, 56) - << "\n"; - os << " hint: move or sink this alloc inside the nested region if the IR allows it.\n"; - } - if (!printedAttention) + 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.size) << " (" - << slot.size << ")" + 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]; - mlir::Value allocValue = interval.key.value; os << " [" << interval.start << "," << interval.end << "]" - << " #" << interval.id << " logical=" << formatReportMemory(interval.size) + << " #" << interval.id << " logical=" << formatReportMemory(interval.size) << " (" << interval.size + << ")" << " nested=" << (interval.insideNestedRegion ? "yes" : "no") << " escapes_loop=" << (interval.escapesLoop ? "yes" : "no") - << " first=" << summarizeOperation(interval.firstTouchOp, 48) - << " last=" << summarizeOperation(interval.lastTouchOp, 48) << "\n"; - os << " value=" << summarizeValue(allocValue) << "\n"; - } - } - } - - if (reportLevel == PimMemoryReportFull) { - os << "\nInterval Details:\n"; - for (const LocalAllocInterval& interval : intervals) { - const PlannedPhysicalSlot& slot = slots[interval.slotPlanIndex]; - mlir::Value allocValue = interval.key.value; - Operation* definingOp = allocValue.getDefiningOp(); - os << " #" << interval.id << " slot=" << slot.id << " live=[" << interval.start << "," << interval.end << "]" - << " logical=" << formatReportMemory(interval.size) - << " slot_size=" << formatReportMemory(interval.physicalSlotSize) << " addr=" << interval.assignedAddress - << "\n"; - os << " value=" << summarizeValue(allocValue, 88) << "\n"; - os << " type=" << allocValue.getType() << "\n"; - os << " loc=" - << summarizeLocation(definingOp ? definingOp->getLoc() : UnknownLoc::get(coreLikeOp->getContext())) << "\n"; - os << " nested=" << (interval.insideNestedRegion ? "yes" : "no") - << " escapes_loop=" << (interval.escapesLoop ? "yes" : "no") - << " start_fallback=" << (interval.startUsedAllocFallback ? "yes" : "no") - << " end_fallback=" << (interval.endUsedFallback ? "yes" : "no") << "\n"; - os << " first_use=" << summarizeOperation(interval.firstTouchOp) << " @" << interval.firstTouchPosition - << "\n"; - os << " last_use=" << summarizeOperation(interval.lastTouchOp) << " @" << interval.lastTouchPosition << "\n"; - os << " slot_peers="; - bool first = true; - for (size_t otherIndex : slot.intervalIndices) { - if (intervals[otherIndex].id == interval.id) - continue; - if (!first) - os << ", "; - os << "#" << intervals[otherIndex].id; - first = false; - } - if (first) - os << ""; - os << "\n"; - if (!interval.fallbackReason.empty()) - os << " fallback_reason=" << interval.fallbackReason << "\n"; - if (!interval.aliasesFollowed.empty()) { - os << " aliases_followed=" << interval.aliasesFollowed.size() << "\n"; - for (const std::string& alias : interval.aliasesFollowed) - os << " - " << abbreviate(collapseWhitespace(alias), 108) << "\n"; + << " 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"; } } } @@ -744,3 +576,95 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp, return artifacts; } + +namespace onnx_mlir { +namespace { + +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"; + } + + 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; + getOperation().walk([&](Operation* op) { + if (failedPlanning || !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; + }); + 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; + 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; + }); + + if (report) { + report->flush(); + reportFile.close(); + } + if (failedPlanning) { + signalPassFailure(); + return; + } + dumpModule(getOperation(), "pim4_memory_planned"); + } +}; + +} // namespace + +std::unique_ptr createPimLocalMemoryPlanningPass() { + return std::make_unique(); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimMemoryLiveness.hpp b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp similarity index 54% rename from src/PIM/Compiler/PimMemoryLiveness.hpp rename to src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp index 5a25174..80ae74e 100644 --- a/src/PIM/Compiler/PimMemoryLiveness.hpp +++ b/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp @@ -6,10 +6,8 @@ #include "llvm/ADT/SmallVector.h" #include -#include #include -#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" namespace onnx_mlir { @@ -17,7 +15,6 @@ namespace onnx_mlir { struct LocalAllocInterval { size_t id = 0; mlir::memref::AllocOp alloc; - MemoryValueKey key; uint64_t start = 0; uint64_t end = 0; size_t size = 0; @@ -31,32 +28,37 @@ struct LocalAllocInterval { bool insideNestedRegion = false; bool escapesLoop = false; std::string fallbackReason; - llvm::SmallVector aliasesFollowed; + std::string valueSummary; + std::string firstTouchSummary; + std::string lastTouchSummary; size_t slotPlanIndex = std::numeric_limits::max(); - size_t physicalSlotId = std::numeric_limits::max(); - size_t assignedAddress = 0; - size_t physicalSlotSize = 0; }; struct PlannedPhysicalSlot { size_t id = std::numeric_limits::max(); size_t requiredSize = 0; - size_t size = 0; size_t address = 0; llvm::SmallVector intervalIndices; }; -llvm::SmallVector buildLocalAllocIntervals(mlir::Operation* coreLikeOp, - std::optional lane, - bool includeAliasDescriptions = true); +struct CoreMemoryPlan { + mlir::Operation* coreLikeOp = nullptr; + llvm::SmallVector intervals; + llvm::SmallVector slots; +}; + +llvm::SmallVector buildLocalAllocIntervals(mlir::Operation* coreLikeOp); llvm::SmallVector planPhysicalSlots(llvm::MutableArrayRef intervals); -MemoryPlanArtifacts buildMemoryPlanArtifacts(mlir::Operation* coreLikeOp, - std::optional lane, - llvm::ArrayRef intervals, - llvm::ArrayRef slots, - size_t addressLimit, - PimMemoryReportLevel reportLevel); +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); } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/CMakeLists.txt b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/CMakeLists.txt index d746087..3bf4699 100644 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/CMakeLists.txt @@ -10,5 +10,6 @@ add_pim_library(OMPimMemoryCoalescing 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 index c697040..ffa4ae8 100644 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.cpp +++ b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.cpp @@ -1,15 +1,13 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Interfaces/DestinationStyleOpInterface.h" - #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallPtrSet.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; @@ -19,10 +17,6 @@ namespace pim { namespace { -static bool isSupportedAliasOp(Operation* op) { - return isa(op); -} - static bool isCandidateAllocType(MemRefType type) { return type && type.hasStaticShape() && type.getLayout().isIdentity() && hasByteSizedElementType(type.getElementType()); @@ -44,59 +38,21 @@ static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis); static FailureOr getLastUseInstruction(memref::AllocOp allocOp, Block& scopeBlock, const DenseMap& opOrder) { uint64_t endInstruction = opOrder.lookup(allocOp); - SmallPtrSet visitedValues; - SmallPtrSet visitedUsers; - SmallVector pendingValues; - pendingValues.push_back(allocOp.getResult()); - - while (!pendingValues.empty()) { - Value value = pendingValues.pop_back_val(); - if (!visitedValues.insert(value).second) - continue; - - for (Operation* user : value.getUsers()) { - if (!visitedUsers.insert(user).second) - continue; - - if (isSupportedAliasOp(user)) - llvm::append_range(pendingValues, user->getResults()); - - if (auto dpsOp = dyn_cast(user)) { - for (OpResult result : user->getResults()) { - OpOperand* tiedOperand = dpsOp.getTiedOpOperand(result); - if (tiedOperand && tiedOperand->get() == value) - pendingValues.push_back(result); - } - } - - if (auto forOp = dyn_cast(user)) { - for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs())) { - if (initArg != value) - continue; - pendingValues.push_back(forOp.getRegionIterArgs()[index]); - pendingValues.push_back(forOp.getResult(index)); - } - } - + if (failed(walkLocalMemoryUses(allocOp.getResult(), [&](Value, Operation* user) { if (auto yieldOp = dyn_cast(user)) { - auto forOp = dyn_cast(yieldOp->getParentOp()); - if (!forOp) + if (!isa(yieldOp->getParentOp())) return failure(); - for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) - if (operand == value) - pendingValues.push_back(forOp.getResult(index)); } - 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; } @@ -133,7 +89,7 @@ static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis) { } blockAnalysis.candidates.push_back( - AllocationCandidate {allocOp, &block, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)}); + AllocationCandidate {allocOp, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)}); } analysis.skippedAllocations += blockAnalysis.skippedAllocations; @@ -150,6 +106,14 @@ uint64_t MemoryCoalescingAnalysis::getCandidateCount() const { 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()) diff --git a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp index 67e1b1d..197cda2 100644 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp +++ b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp @@ -10,7 +10,6 @@ namespace pim { struct AllocationCandidate { mlir::memref::AllocOp alloc; - mlir::Block* scopeBlock = nullptr; uint64_t startInstruction = 0; uint64_t endInstruction = 0; uint64_t sizeBytes = 0; @@ -27,6 +26,7 @@ struct MemoryCoalescingAnalysis { uint64_t skippedAllocations = 0; uint64_t getCandidateCount() const; + uint64_t getCandidateBytes() const; }; struct MemoryCoalescingStats { diff --git a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp index 2025cfc..6042af9 100644 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp @@ -1,207 +1,94 @@ -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" -#include "llvm/ADT/SmallVector.h" #include "llvm/Support/raw_os_ostream.h" #include -#include "Common/IR/CompactAsmUtils.hpp" #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" -#include "src/Accelerators/PIM/Common/Support/DebugDump.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; -using namespace onnx_mlir::compact_asm; namespace onnx_mlir { namespace { -// This pass is an IR cleanup step after bufferization. It only rewrites -// obviously compatible local allocations with non-overlapping lifetimes inside -// the same block and leaves the final physical memory planning to codegen. - struct CoalescingReportRow { - uint64_t numCandidates = 0; - uint64_t numSkipped = 0; - uint64_t numRemoved = 0; + uint64_t candidates = 0; + uint64_t logicalBytes = 0; + uint64_t skipped = 0; + uint64_t removed = 0; uint64_t savedBytes = 0; - - bool operator==(const CoalescingReportRow& other) const { - return numCandidates == other.numCandidates && numSkipped == other.numSkipped && numRemoved == other.numRemoved - && savedBytes == other.savedBytes; - } }; struct CoalescingReportEntry { - enum class Kind { - Core, - Batch - }; - - Kind kind = Kind::Core; - uint64_t id = 0; - llvm::SmallVector coreIds; + std::string label; + uint64_t coreCount = 1; CoalescingReportRow row; }; -static std::string formatMemory(uint64_t bytes) { return formatReportMemory(bytes); } - -static void printReportRow(raw_ostream& os, const CoalescingReportRow& row) { - llvm::SmallVector fields = { - {"Number of candidates", std::to_string(row.numCandidates)}, - {"Skipped allocations", std::to_string(row.numSkipped) }, - {"Removed allocations", std::to_string(row.numRemoved) }, - {"Saved memory", formatMemory(row.savedBytes) } - }; - printReportFlatFields(os, fields); -} - -static CoalescingReportRow getTotalRow(const CoalescingReportEntry& entry) { - uint64_t factor = std::max(1, entry.coreIds.size()); - return {entry.row.numCandidates * factor, - entry.row.numSkipped * factor, - entry.row.numRemoved * factor, - entry.row.savedBytes * factor}; -} - static void emitReport(ArrayRef entries) { std::fstream file = openReportFile("memory_coalescing_report"); if (!file.is_open()) return; - llvm::raw_os_ostream os(file); - CoalescingReportRow totalRow; + CoalescingReportRow total; for (const CoalescingReportEntry& entry : entries) { - CoalescingReportRow entryTotal = getTotalRow(entry); - totalRow.numCandidates += entryTotal.numCandidates; - totalRow.numSkipped += entryTotal.numSkipped; - totalRow.numRemoved += entryTotal.numRemoved; - totalRow.savedBytes += entryTotal.savedBytes; + 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; } - - llvm::SmallVector totalFields = { - {"Number of candidates", std::to_string(totalRow.numCandidates)}, - {"Skipped allocations", std::to_string(totalRow.numSkipped) }, - {"Removed allocations", std::to_string(totalRow.numRemoved) }, - {"Saved memory", formatMemory(totalRow.savedBytes) } - }; - printReportTotalsBlock(os, totalFields); - if (!entries.empty()) - os << "\n"; - - llvm::SmallVector sortedEntries(entries.begin(), entries.end()); - sortReportEntriesByFirstCore(sortedEntries); - - for (size_t index = 0; index < sortedEntries.size();) { - size_t runEnd = index + 1; - while (runEnd < sortedEntries.size() && sortedEntries[runEnd].kind == sortedEntries[index].kind - && sortedEntries[runEnd].row == sortedEntries[index].row) { - ++runEnd; - } - - if (sortedEntries[index].kind == CoalescingReportEntry::Kind::Batch) { - os << "Batch "; - for (size_t batchIndex = index; batchIndex < runEnd; ++batchIndex) { - if (batchIndex != index) - os << ",\n "; - os << sortedEntries[batchIndex].id << " (cores "; - printCompressedIntegerEntries(os, ArrayRef(sortedEntries[batchIndex].coreIds)); - os << ")"; - } - } - else { - llvm::SmallVector coreIds; - for (size_t coreIndex = index; coreIndex < runEnd; ++coreIndex) - coreIds.push_back(sortedEntries[coreIndex].coreIds.front()); - os << "Core "; - printCompressedIntegerEntries(os, ArrayRef(coreIds)); - } - - os << ":\n"; - if (sortedEntries[index].kind == CoalescingReportEntry::Kind::Batch) { - llvm::SmallVector perCoreFields = { - {"Number of candidates", std::to_string(sortedEntries[index].row.numCandidates)}, - {"Skipped allocations", std::to_string(sortedEntries[index].row.numSkipped) }, - {"Removed allocations", std::to_string(sortedEntries[index].row.numRemoved) }, - {"Saved memory", formatMemory(sortedEntries[index].row.savedBytes) } - }; - CoalescingReportRow totalRow = getTotalRow(sortedEntries[index]); - llvm::SmallVector totalFields = { - {"Number of candidates", std::to_string(totalRow.numCandidates)}, - {"Skipped allocations", std::to_string(totalRow.numSkipped) }, - {"Removed allocations", std::to_string(totalRow.numRemoved) }, - {"Saved memory", formatMemory(totalRow.savedBytes) } - }; - printReportPerCoreAndTotalFields(os, perCoreFields, totalFields); - } - else { - printReportRow(os, sortedEntries[index].row); - } - printReportEntrySeparator(os, runEnd < sortedEntries.size()); - index = runEnd; + 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(); - file.close(); } struct PimMemoryCoalescingPass : PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimMemoryCoalescingPass) StringRef getArgument() const override { return "pim-memory-coalescing"; } - StringRef getDescription() const override { return "Analyze local PIM memory reuse opportunities"; } - - PimMemoryCoalescingPass() = default; - PimMemoryCoalescingPass(const PimMemoryCoalescingPass& pass) {} + StringRef getDescription() const override { return "Safely coalesce compatible block-local PIM allocations"; } void runOnOperation() override { IRRewriter rewriter(&getContext()); SmallVector reportEntries; uint64_t nextBatchId = 0; - bool hasFailure = false; - getOperation().walk([&](Operation* op) { - if (hasFailure || !isa(op)) + if (!isa(op)) return; - auto analysis = pim::analyzeMemoryCoalescingCandidates(op); auto stats = pim::coalesceMemory(op, analysis, rewriter); - CoalescingReportRow row { - analysis.getCandidateCount(), stats.skippedAllocations, stats.removedAllocs, stats.savedBytes}; - + CoalescingReportRow row {analysis.getCandidateCount(), + analysis.getCandidateBytes(), + stats.skippedAllocations, + stats.removedAllocs, + stats.savedBytes}; if (auto coreOp = dyn_cast(op)) { - auto checkedCoreId = - pim::checkedI32(static_cast(coreOp.getCoreId()), coreOp, "memory coalescing core id"); - if (failed(checkedCoreId)) { - hasFailure = true; - return; - } - reportEntries.push_back( - {CoalescingReportEntry::Kind::Core, static_cast(coreOp.getCoreId()), {*checkedCoreId}, row}); + reportEntries.push_back({"Core " + std::to_string(coreOp.getCoreId()), 1, row}); return; } - - auto coreIds = getBatchCoreIds(cast(op)); - CoalescingReportEntry entry; - entry.kind = CoalescingReportEntry::Kind::Batch; - entry.id = nextBatchId++; - llvm::append_range(entry.coreIds, coreIds); - entry.row = row; - reportEntries.push_back(std::move(entry)); + SmallVector coreIds = getBatchCoreIds(cast(op)); + reportEntries.push_back( + {"Batch " + std::to_string(nextBatchId++), std::max(1, coreIds.size()), row}); }); - - if (hasFailure) { - signalPassFailure(); - return; - } - emitReport(reportEntries); dumpModule(getOperation(), "pim3_coalesced"); } diff --git a/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp b/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp index 9612d25..34bc33d 100644 --- a/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp @@ -15,8 +15,9 @@ #include "src/Accelerators/PIM/Common/IR/SubviewUtils.hpp" #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" +#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/PimOps.hpp" #include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" @@ -108,6 +109,63 @@ 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"); + }); + 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"); + }); + hasFailure = true; + return; + } + + 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; + }); + hasFailure = true; + return; + } + + 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; + }); + hasFailure = true; + } + }); + return success(!hasFailure); +} + static bool isSupportedCoreInstructionOp(Operation* op) { return isa(base.getDefiningOp()); } - -enum class CommunicationEventKind { Send, Receive }; +enum class CommunicationEventKind { + Send, + Receive +}; struct CommunicationEvent { CommunicationEventKind kind = CommunicationEventKind::Send; @@ -157,19 +217,6 @@ struct CommunicationEvent { int64_t peerCoreId = 0; int64_t size = 0; uint64_t ordinal = 0; - std::optional minChannelId; - std::string materializer; - std::optional traceId; - std::optional commOrder; - std::optional traceClassId; - std::optional traceBlockOrdinal; - std::string traceKind; - std::string tracePhase; - std::string traceClassKind; - std::string tracePayload; - std::string traceMessages; - std::string tracePrevOp; - std::string traceNextOp; Operation* op = nullptr; }; @@ -183,7 +230,6 @@ constexpr StringLiteral kRaptorMinChannelIdAttr = "raptor.min_channel_id"; constexpr StringLiteral kRaptorMaterializerAttr = "raptor.materializer"; constexpr StringLiteral kRaptorCommOrderAttr = "raptor.comm_order"; constexpr StringLiteral kRaptorCommTraceIdAttr = "raptor.comm_trace_id"; -constexpr StringLiteral kRaptorCommTraceKindAttr = "raptor.comm_trace_kind"; constexpr StringLiteral kRaptorCommTracePhaseAttr = "raptor.comm_trace_phase"; constexpr StringLiteral kRaptorCommTraceClassIdAttr = "raptor.comm_trace_class_id"; constexpr StringLiteral kRaptorCommTraceClassKindAttr = "raptor.comm_trace_class_kind"; @@ -225,31 +271,44 @@ static std::string formatOperationSummary(Operation* op) { } static std::string formatCommunicationEvent(const CommunicationEvent& event) { + std::optional minChannelId = getNearestIntegerAttr(event.op, kRaptorMinChannelIdAttr); + std::optional commOrder = getNearestIntegerAttr(event.op, kRaptorCommOrderAttr); + std::optional traceId = getNearestIntegerAttr(event.op, kRaptorCommTraceIdAttr); + std::optional traceClassId = getNearestIntegerAttr(event.op, kRaptorCommTraceClassIdAttr); + std::optional traceBlockOrdinal = getNearestIntegerAttr(event.op, kRaptorCommTraceBlockOrdinalAttr); + std::string materializer = getNearestStringAttr(event.op, kRaptorMaterializerAttr); + std::string tracePhase = getNearestStringAttr(event.op, kRaptorCommTracePhaseAttr); + std::string traceClassKind = getNearestStringAttr(event.op, kRaptorCommTraceClassKindAttr); + std::string tracePayload = getNearestStringAttr(event.op, kRaptorCommTracePayloadAttr); + std::string traceMessages = getNearestStringAttr(event.op, kRaptorCommTraceMessagesAttr); + std::string tracePrevOp = getNearestStringAttr(event.op, kRaptorCommTracePrevOpAttr); + std::string traceNextOp = getNearestStringAttr(event.op, kRaptorCommTraceNextOpAttr); + std::string text; llvm::raw_string_ostream os(text); os << "core " << event.coreId << " " << getCommunicationEventKindName(event.kind) << " " - << (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId - << " size " << event.size << "B ordinal " << event.ordinal; - if (event.minChannelId) - os << " min_channel " << *event.minChannelId; - if (event.commOrder) - os << " comm_order " << *event.commOrder; - if (!event.materializer.empty()) - os << " materializer " << event.materializer; - if (event.traceId) - os << " trace#" << *event.traceId; - if (!event.tracePhase.empty()) - os << " phase " << event.tracePhase; - if (event.traceClassId) - os << " class " << event.traceClassKind << "#" << *event.traceClassId; - if (event.traceBlockOrdinal) - os << " block_ordinal " << *event.traceBlockOrdinal; - if (!event.tracePayload.empty()) - os << " payload " << event.tracePayload; - if (!event.traceMessages.empty()) - os << " messages {" << event.traceMessages << "}"; - if (!event.tracePrevOp.empty() || !event.traceNextOp.empty()) - os << " inserted_between [" << event.tracePrevOp << " | " << event.traceNextOp << "]"; + << (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId << " size " + << event.size << "B ordinal " << event.ordinal; + if (minChannelId) + os << " min_channel " << *minChannelId; + if (commOrder) + os << " comm_order " << *commOrder; + if (!materializer.empty()) + os << " materializer " << materializer; + if (traceId) + os << " trace#" << *traceId; + if (!tracePhase.empty()) + os << " phase " << tracePhase; + if (traceClassId) + os << " class " << traceClassKind << "#" << *traceClassId; + if (traceBlockOrdinal) + os << " block_ordinal " << *traceBlockOrdinal; + if (!tracePayload.empty()) + os << " payload " << tracePayload; + if (!traceMessages.empty()) + os << " messages {" << traceMessages << "}"; + if (!tracePrevOp.empty() || !traceNextOp.empty()) + os << " inserted_between [" << tracePrevOp << " | " << traceNextOp << "]"; return os.str(); } @@ -261,10 +320,8 @@ static bool areMatchedCommunicationEvents(const CommunicationEvent& lhs, const C || (lhs.kind == CommunicationEventKind::Receive && rhs.kind == CommunicationEventKind::Send); } - -static std::optional findMatchingCounterpartIndex(const CommunicationEventVector& events, - const CommunicationEvent& event, - size_t begin) { +static std::optional +findMatchingCounterpartIndex(const CommunicationEventVector& events, const CommunicationEvent& event, size_t begin) { for (size_t index = begin; index < events.size(); ++index) if (areMatchedCommunicationEvents(event, events[index])) return index; @@ -288,8 +345,7 @@ static void printCounterpartProbe(llvm::raw_ostream& os, peerPc = peerPcIt->second; os << " counterpart probe for " << formatCommunicationEvent(blockedEvent) << "\n"; - os << " peer core " << blockedEvent.peerCoreId << " current pc " << peerPc << " of " << peerEvents.size() - << "\n"; + os << " peer core " << blockedEvent.peerCoreId << " current pc " << peerPc << " of " << peerEvents.size() << "\n"; std::optional nextMatch = findMatchingCounterpartIndex(peerEvents, blockedEvent, peerPc); std::optional anyMatch = findMatchingCounterpartIndex(peerEvents, blockedEvent, 0); @@ -317,7 +373,8 @@ static void printCounterpartProbe(llvm::raw_ostream& os, return; os << " peer operations blocking before that counterpart:\n"; - size_t end = std::min(peerEvents.size(), std::min(*nextMatch + static_cast(1), peerPc + static_cast(12))); + size_t end = + std::min(peerEvents.size(), std::min(*nextMatch + static_cast(1), peerPc + static_cast(12))); for (size_t index = peerPc; index < end; ++index) { os << (index == peerPc ? " pc => " : " ") << "#" << index << " " << formatCommunicationEvent(peerEvents[index]) << "\n"; @@ -327,77 +384,58 @@ static void printCounterpartProbe(llvm::raw_ostream& os, os << " ... " << (*nextMatch - end + 1) << " more peer communication event(s) before the counterpart\n"; } -static CommunicationEvent makeCommunicationEvent(CommunicationEventKind kind, - int64_t coreId, - int64_t peerCoreId, - int64_t size, - uint64_t ordinal, - Operation* op) { - return CommunicationEvent {kind, - coreId, - peerCoreId, - size, - ordinal, - getNearestIntegerAttr(op, kRaptorMinChannelIdAttr), - getNearestStringAttr(op, kRaptorMaterializerAttr), - getNearestIntegerAttr(op, kRaptorCommTraceIdAttr), - getNearestIntegerAttr(op, kRaptorCommOrderAttr), - getNearestIntegerAttr(op, kRaptorCommTraceClassIdAttr), - getNearestIntegerAttr(op, kRaptorCommTraceBlockOrdinalAttr), - getNearestStringAttr(op, kRaptorCommTraceKindAttr), - getNearestStringAttr(op, kRaptorCommTracePhaseAttr), - getNearestStringAttr(op, kRaptorCommTraceClassKindAttr), - getNearestStringAttr(op, kRaptorCommTracePayloadAttr), - getNearestStringAttr(op, kRaptorCommTraceMessagesAttr), - getNearestStringAttr(op, kRaptorCommTracePrevOpAttr), - getNearestStringAttr(op, kRaptorCommTraceNextOpAttr), - op}; +static CommunicationEvent makeCommunicationEvent( + CommunicationEventKind kind, int64_t coreId, int64_t peerCoreId, int64_t size, uint64_t ordinal, Operation* op) { + return CommunicationEvent {kind, coreId, peerCoreId, size, ordinal, op}; } static LogicalResult appendCoreCommunicationEvents(Block& block, + const PimCoreCommunicationPlan& plan, int64_t coreId, const StaticValueKnowledge& initialKnowledge, SmallVectorImpl& events, pim::CappedDiagnosticReporter& diagnostics) { - return walkPimCoreBlock(block, initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) { - if (auto sendOp = dyn_cast(&op)) { - auto targetCoreId = resolveIndexValue(sendOp.getTargetCoreId(), knowledge); - if (failed(targetCoreId)) { - diagnostics.report(&op, [](Operation* illegalOp) { - illegalOp->emitOpError("cannot statically resolve send target core for PIM communication deadlock check"); - }); - return failure(); + return walkPimCoreCommunicationBlock( + block, plan, initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) { + if (auto sendOp = dyn_cast(&op)) { + auto targetCoreId = resolveIndexValue(sendOp.getTargetCoreId(), knowledge); + if (failed(targetCoreId)) { + diagnostics.report(&op, [](Operation* illegalOp) { + illegalOp->emitOpError("cannot statically resolve send target core for PIM communication deadlock check"); + }); + return failure(); + } + + events.push_back(makeCommunicationEvent(CommunicationEventKind::Send, + coreId, + *targetCoreId, + sendOp.getSize(), + static_cast(events.size()), + &op)); + return success(); } - events.push_back(makeCommunicationEvent(CommunicationEventKind::Send, - coreId, - *targetCoreId, - sendOp.getSize(), - static_cast(events.size()), - &op)); - return success(); - } + if (auto receiveOp = dyn_cast(&op)) { + auto sourceCoreId = resolveIndexValue(receiveOp.getSourceCoreId(), knowledge); + if (failed(sourceCoreId)) { + diagnostics.report(&op, [](Operation* illegalOp) { + illegalOp->emitOpError( + "cannot statically resolve receive source core for PIM communication deadlock check"); + }); + return failure(); + } - if (auto receiveOp = dyn_cast(&op)) { - auto sourceCoreId = resolveIndexValue(receiveOp.getSourceCoreId(), knowledge); - if (failed(sourceCoreId)) { - diagnostics.report(&op, [](Operation* illegalOp) { - illegalOp->emitOpError("cannot statically resolve receive source core for PIM communication deadlock check"); - }); - return failure(); + events.push_back(makeCommunicationEvent(CommunicationEventKind::Receive, + coreId, + *sourceCoreId, + receiveOp.getSize(), + static_cast(events.size()), + &op)); + return success(); } - events.push_back(makeCommunicationEvent(CommunicationEventKind::Receive, - coreId, - *sourceCoreId, - receiveOp.getSize(), - static_cast(events.size()), - &op)); return success(); - } - - return success(); - }); + }); } static void printCommunicationWindow(llvm::raw_ostream& os, @@ -466,9 +504,10 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp, ArrayRef cycle) { printCommunicationDeadlockReport(coreEvents, programCounters, cycle); - auto diagnostic = moduleOp.emitError() - << "PIM communication deadlock check found a blocking send/receive cycle while statically simulating the " - "expanded per-core communication streams; see the PIM static communication deadlock report above"; + auto diagnostic = + moduleOp.emitError() + << "PIM communication deadlock check found a blocking send/receive cycle while statically simulating the " + "expanded per-core communication streams; see the PIM static communication deadlock report above"; for (int64_t coreId : cycle) { auto eventsIt = coreEvents.find(coreId); @@ -479,14 +518,15 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp, const CommunicationEvent& event = eventsIt->second[pcIt->second]; Diagnostic& note = diagnostic.attachNote(event.op->getLoc()); note << formatCommunicationEvent(event); - if (!event.materializer.empty()) - note << " emitted by " << event.materializer; + std::string materializer = getNearestStringAttr(event.op, kRaptorMaterializerAttr); + if (!materializer.empty()) + note << " emitted by " << materializer; } } -static FailureOr> findCommunicationWaitCycle( - const DenseMap& coreEvents, - const DenseMap& programCounters) { +static FailureOr> +findCommunicationWaitCycle(const DenseMap& coreEvents, + const DenseMap& programCounters) { for (const auto& [startCoreId, events] : coreEvents) { auto startPcIt = programCounters.find(startCoreId); if (startPcIt == programCounters.end() || startPcIt->second >= events.size()) @@ -519,7 +559,7 @@ static FailureOr> findCommunicationWaitCycle( } static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp, - pim::CappedDiagnosticReporter& diagnostics) { + pim::CappedDiagnosticReporter& diagnostics) { DenseMap coreEvents; bool hasFailure = false; @@ -530,14 +570,20 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp, for (Operation& op : funcOp.getBody().front().getOperations()) { if (auto coreOp = dyn_cast(&op)) { int64_t coreId = coreOp.getCoreId(); - if (failed(appendCoreCommunicationEvents( - coreOp.getBody().front(), coreId, StaticValueKnowledge {}, coreEvents[coreId], diagnostics))) + PimCoreCommunicationPlan plan = buildPimCoreCommunicationPlan(coreOp.getBody().front()); + if (failed(appendCoreCommunicationEvents(coreOp.getBody().front(), + plan, + coreId, + StaticValueKnowledge {}, + coreEvents[coreId], + diagnostics))) hasFailure = true; continue; } if (auto coreBatchOp = dyn_cast(&op)) { SmallVector coreIds = getBatchCoreIds(coreBatchOp); + PimCoreCommunicationPlan plan = buildPimCoreCommunicationPlan(coreBatchOp.getBody().front()); size_t laneCount = static_cast(coreBatchOp.getLaneCount()); for (size_t lane = 0; lane < laneCount; ++lane) { StaticValueKnowledge laneKnowledge; @@ -546,11 +592,14 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp, laneKnowledge.aliases[coreBatchOp.getInputArgument(inputIndex)] = coreBatchOp.getInputs()[inputIndex]; SmallVector laneCoreIds = getLaneChunkCoreIds(coreIds, laneCount, static_cast(lane)); - for (int32_t coreId : laneCoreIds) { - if (failed(appendCoreCommunicationEvents( - coreBatchOp.getBody().front(), coreId, laneKnowledge, coreEvents[coreId], diagnostics))) + for (int32_t coreId : laneCoreIds) + if (failed(appendCoreCommunicationEvents(coreBatchOp.getBody().front(), + plan, + coreId, + laneKnowledge, + coreEvents[coreId], + diagnostics))) hasFailure = true; - } } } } @@ -607,9 +656,10 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp, return failure(); } - auto diagnostic = moduleOp.emitError() - << "PIM communication deadlock check stalled without finding a closed wait cycle; this usually means a " - "send/receive peer is missing or ordered after a finished core"; + auto diagnostic = + moduleOp.emitError() + << "PIM communication deadlock check stalled without finding a closed wait cycle; this usually means a " + "send/receive peer is missing or ordered after a finished core"; for (const auto& [coreId, events] : coreEvents) { size_t pc = programCounters[coreId]; if (pc >= events.size()) @@ -652,6 +702,7 @@ struct VerificationPass : PassWrapper> for (Operation& op : funcOp.getBody().front().getOperations()) { if (auto coreOp = dyn_cast(&op)) { (void) verifyCoreWeights(moduleOp, coreOp, diagnostics); + (void) verifyLocalMemoryPlan(coreOp, diagnostics); StaticValueKnowledge knowledge; (void) verifyCoreLikeOperands(coreOp, knowledge, diagnostics); continue; @@ -659,6 +710,7 @@ struct VerificationPass : PassWrapper> if (auto coreBatchOp = dyn_cast(&op)) { (void) verifyCoreWeights(moduleOp, coreBatchOp, diagnostics); + (void) verifyLocalMemoryPlan(coreBatchOp, diagnostics); llvm::SmallVector lanes; lanes.push_back(0); if (coreBatchOp.getLaneCount() > 1) diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp index 6011cec..cee4178 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp @@ -133,6 +133,22 @@ 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; + CrossbarUsage distinctWeights; + for (const ComputeGraphNode& node : graph.nodes) { + for (const CrossbarWeight& weight : node.crossbarUsage) { + if (!containsCrossbarWeight(distinctWeights, weight)) + distinctWeights.push_back(weight); + if (distinctWeights.size() > threshold) + return true; + } + } + return false; +} + } // namespace Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount) { @@ -145,6 +161,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu if (processorCount == 0) llvm::report_fatal_error("PEFT scheduler: processor count must be positive"); MeshModel mesh = MeshModel::infer(processorCount); + const bool preferCrossbarReuse = hasHighCrossbarPressure(graph, processorCount, options.crossbarCapacity); verifyOctTableSize(nodeCount, processorCount); std::vector> reverseLevels = buildReverseLevels(graph); @@ -282,6 +299,8 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu Time eft = addOrMax(est, computeCost); Time oeft = addOrMax(eft, oct[task * processorCount + processor]); size_t centerDistance = mesh.getCenterDistance(processor); + bool betterCrossbarChoice = + preferCrossbarReuse ? overlapCount > bestOverlapCount : overlapCount < bestOverlapCount; if (oeft < bestOeft || (oeft == bestOeft && eft < bestEft) || (oeft == bestOeft && eft == bestEft && est < bestEst)) { @@ -302,7 +321,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu bestCenterDistance = centerDistance; } else if (oeft == bestOeft && eft == bestEft && est == bestEst - && centerDistance == bestCenterDistance && overlapCount < bestOverlapCount) { + && centerDistance == bestCenterDistance && betterCrossbarChoice) { bestProcessor = processor; bestEst = est; bestEft = eft; diff --git a/src/PIM/Pass/PIMPasses.h b/src/PIM/Pass/PIMPasses.h index b2a8d8a..62cae80 100644 --- a/src/PIM/Pass/PIMPasses.h +++ b/src/PIM/Pass/PIMPasses.h @@ -23,6 +23,8 @@ std::unique_ptr createTrivialGraphComputeMergePass(); std::unique_ptr createPimHostConstantFoldingPass(); +std::unique_ptr createPimLocalMemoryPlanningPass(); + std::unique_ptr createPimVerificationPass(); std::unique_ptr createEmitPimCodePass(); diff --git a/src/PIM/PimAccelerator.cpp b/src/PIM/PimAccelerator.cpp index a623d96..f9be891 100644 --- a/src/PIM/PimAccelerator.cpp +++ b/src/PIM/PimAccelerator.cpp @@ -80,6 +80,7 @@ void PimAccelerator::registerPasses(int optLevel) const { registerPass(createTrivialGraphComputeMergePass); registerPass(createMergeComputeNodesPass); registerPass(createPimHostConstantFoldingPass); + registerPass(createPimLocalMemoryPlanningPass); registerPass(createPimVerificationPass); registerPass(createEmitPimCodePass); } diff --git a/test/PIM/PimMemoryLivenessPlannerTest.cpp b/test/PIM/PimMemoryLivenessPlannerTest.cpp index 246d722..2f6eb4c 100644 --- a/test/PIM/PimMemoryLivenessPlannerTest.cpp +++ b/test/PIM/PimMemoryLivenessPlannerTest.cpp @@ -2,9 +2,11 @@ #include #include -#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp" +#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; namespace { @@ -18,12 +20,19 @@ LocalAllocInterval makeInterval(size_t id, size_t size, uint64_t start, uint64_t return interval; } +size_t getPeak(llvm::ArrayRef slots) { + size_t peak = 0; + for (const auto& slot : slots) + peak = std::max(peak, slot.address + slot.requiredSize); + 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].physicalSlotId == intervals[1].physicalSlotId); + assert(intervals[0].slotPlanIndex == intervals[1].slotPlanIndex); } int testSameSizeNonOverlap() { @@ -34,13 +43,21 @@ int testSameSizeNonOverlap() { int testLargerFirst() { std::cout << "testLargerFirst:" << std::endl; - assertSingleSlotCase(makeInterval(0, 100, 0, 10), makeInterval(1, 40, 11, 20), 100); + 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); return 0; } int testSmallerFirst() { std::cout << "testSmallerFirst:" << std::endl; - assertSingleSlotCase(makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), 100); + 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); return 0; } @@ -50,7 +67,7 @@ int testOverlapNeedsTwoSlots() { makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)}; auto slots = planPhysicalSlots(intervals); assert(slots.size() == 2); - assert(intervals[0].physicalSlotId != intervals[1].physicalSlotId); + assert(intervals[0].slotPlanIndex != intervals[1].slotPlanIndex); return 0; } @@ -59,10 +76,31 @@ int testReuseChain() { 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() == 1); - assert(slots.front().requiredSize == 100); - assert(intervals[0].physicalSlotId == intervals[1].physicalSlotId); - assert(intervals[1].physicalSlotId == intervals[2].physicalSlotId); + assert(slots.size() == 3); + assert(getPeak(slots) == 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); + 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); return 0; } @@ -78,6 +116,8 @@ int main(int argc, char *argv[]) { failures += testSmallerFirst(); failures += testOverlapNeedsTwoSlots(); failures += testReuseChain(); + failures += testPartialAddressReuse(); + failures += testAddressAssignment(); if (failures != 0) { std::cerr << failures << " test failures\n"; return EXIT_FAILURE; diff --git a/validation/raptor.py b/validation/raptor.py index cbb0e6d..ac64911 100644 --- a/validation/raptor.py +++ b/validation/raptor.py @@ -11,6 +11,8 @@ 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"), ) diff --git a/validation/tools/classification_local_image_validation.py b/validation/tools/classification_local_image_validation.py index a440913..7223540 100644 --- a/validation/tools/classification_local_image_validation.py +++ b/validation/tools/classification_local_image_validation.py @@ -250,9 +250,9 @@ def main(): parser.add_argument("--raptor-path", type=Path, default=defaults["raptor_path"]) parser.add_argument("--onnx-include-dir", type=Path, default=defaults["onnx_include_dir"]) parser.add_argument("--simulator-dir", type=Path, default=defaults["simulator_dir"]) - parser.add_argument("--crossbar-size", type=int, default=2048) - parser.add_argument("--crossbar-count", type=int, default=256) - parser.add_argument("--core-count", type=int, default=1000) + parser.add_argument("--crossbar-size", type=int, default=128) + parser.add_argument("--crossbar-count", type=int, default=64) + parser.add_argument("--core-count", type=int, default=144) parser.add_argument("--top-k", type=int, default=5) parser.add_argument("--command-timeout-seconds", type=float, default=7200.0) parser.add_argument("--skip-compile", action="store_true") diff --git a/validation/tools/yolo_local_image_validation.py b/validation/tools/yolo_local_image_validation.py index 6aac69b..f73c825 100644 --- a/validation/tools/yolo_local_image_validation.py +++ b/validation/tools/yolo_local_image_validation.py @@ -198,9 +198,9 @@ def main(): parser.add_argument("--raptor-path", type=Path, default=defaults["raptor_path"]) parser.add_argument("--onnx-include-dir", type=Path, default=defaults["onnx_include_dir"]) parser.add_argument("--simulator-dir", type=Path, default=defaults["simulator_dir"]) - parser.add_argument("--crossbar-size", type=int, default=2048) - parser.add_argument("--crossbar-count", type=int, default=256) - parser.add_argument("--core-count", type=int, default=1000) + parser.add_argument("--crossbar-size", type=int, default=128) + parser.add_argument("--crossbar-count", type=int, default=64) + parser.add_argument("--core-count", type=int, default=144) parser.add_argument("--command-timeout-seconds", type=float, default=7200.0) parser.add_argument("--skip-compile", action="store_true") parser.add_argument("--verbose", action="store_true") diff --git a/validation/tools/yolo_real_image_validation.py b/validation/tools/yolo_real_image_validation.py index 173989c..52b3fce 100644 --- a/validation/tools/yolo_real_image_validation.py +++ b/validation/tools/yolo_real_image_validation.py @@ -399,9 +399,9 @@ def main(): parser.add_argument("--remote-project", default="/home/gmagnani/Project/Raptor") parser.add_argument("--remote-python", default="/home/gmagnani/venv/bin/python") parser.add_argument("--network-dir", default="validation/networks/yolo11n/depth_51") - parser.add_argument("--crossbar-size", type=int, default=2048) - parser.add_argument("--crossbar-count", type=int, default=256) - parser.add_argument("--core-count", type=int, default=1000) + parser.add_argument("--crossbar-size", type=int, default=128) + parser.add_argument("--crossbar-count", type=int, default=64) + parser.add_argument("--core-count", type=int, default=144) parser.add_argument("--command-timeout-seconds", type=int, default=7200) parser.add_argument("--skip-compile", action="store_true") parser.add_argument("--annotated-dir", default="validation/networks/yolo11n/depth_51/real_image_validation/annotated") diff --git a/validation/validate.py b/validation/validate.py index cb4f47e..b58f3f3 100755 --- a/validation/validate.py +++ b/validation/validate.py @@ -72,10 +72,9 @@ def main(): ap.add_argument("--relative-threshold", type=float, default=1e-5, help="Relative tolerance for per-element output comparison.") ap.add_argument("--seed", type=int, default=0, help="RNG seed for generated validation inputs.") - ap.add_argument("--crossbar-size", type=int, default=64) - ap.add_argument("--crossbar-count", type=int, default=8) - ap.add_argument("--core-count", type=int, default=None, - help="Core count to pass to Raptor. Required for PIM validation.") + 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=[], @@ -120,8 +119,6 @@ def main(): missing_args.append("--raptor-path") if not a.onnx_include_dir: missing_args.append("--onnx-include-dir") - if a.core_count is None: - missing_args.append("--core-count") if missing_args: ap.error("the following arguments are required unless --clean is used: " + ", ".join(missing_args)) diff --git a/validation/validate_one.py b/validation/validate_one.py index 9c81752..a1316c7 100644 --- a/validation/validate_one.py +++ b/validation/validate_one.py @@ -311,7 +311,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=64, crossbar_count=8, core_count=None, + simulator_dir, crossbar_size=128, crossbar_count=64, core_count=144, pim_memory_report="none", raptor_extra_args=None, threshold=1e-3, rtol=1e-5, seed=0, reporter=None, model_index=1, model_total=1, verbose=False,