diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index 95ed52e005..63a449f286 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -67,9 +67,7 @@ std::unique_ptr createConvertToPTOOpPass(); std::unique_ptr createInferPTOMemScopePass(); std::unique_ptr -createPlanMemoryPass(const PlanMemoryOptions &options = {}); -std::unique_ptr -createPlanMemoryModernPass(const PlanMemoryOptions &options); +createPlanMemoryModernPass(const PlanMemoryOptions &options = {}); std::unique_ptr createPTORemoveRedundantBarrierPass(); std::unique_ptr createPTOValidateIntToPtrUsesPass(); std::unique_ptr createPTORematerializeFixpipeVectorQuantPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 5edf8ca953..c1e70d3f0b 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -179,7 +179,7 @@ def PTOMaterializeImplicitTmp def PlanMemory : Pass<"pto-plan-memory", "ModuleOp"> { let summary = "Plan memory for PTO Ops"; - let constructor = "mlir::pto::createPlanMemoryPass()"; + let constructor = "mlir::pto::createPlanMemoryModernPass()"; let options = [ Option<"memMode", "mem-mode", "std::string", "\"local\"", diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index efd02b5d2c..882c9160c4 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -119,7 +119,6 @@ add_mlir_dialect_library(PTOTransforms Utils.cpp OptMemPlanForPipeline.cpp InferPTOMemScope.cpp - PTOPlanMemory.cpp PTOPlanMemoryModern.cpp PTORemoveRedundantBarrier.cpp InferPTOLayout.cpp diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp deleted file mode 100644 index a0b9c18594..0000000000 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ /dev/null @@ -1,2876 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -//===- PlanMemory.cpp ----Plan Buffer Memory Address ----------------------===// -//===----------------------------------------------------------------------===// - -#include "PTO/Support/CodeConstants.h" -#include "PTOPlanMemory.h" - -#include "PTO/IR/PTOMultiBuffer.h" -#include "PTO/IR/PTOTypeUtils.h" -#include "Utils.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/GPU/IR/GPUDialect.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "mlir/IR/AsmState.h" -#include "mlir/Transforms/DialectConversion.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -#include "llvm/Support/Debug.h" -#include "llvm/Support/ErrorHandling.h" -#include "llvm/Support/raw_ostream.h" - -#include -#include -#include -#include -#include - -#define DEBUG_TYPE "pto-plan-memory" -#define LDBG(X) LLVM_DEBUG(llvm::dbgs() << X) - -namespace mlir { -namespace pto { -#define GEN_PASS_DEF_PLANMEMORY -#include "PTO/Transforms/Passes.h.inc" -} // namespace pto -} // namespace mlir - -using namespace mlir; -using namespace pto; - -namespace { - -constexpr int64_t kBitsPerByte = 8; -constexpr unsigned kI32BitWidth = 32; -constexpr unsigned kMemoryEffectReserveSize = 8; -constexpr int kSingleBufferCount = 1; -constexpr int kDoubleBufferCount = 2; -constexpr int64_t kA5VecLocalMemBits = 2031616; -constexpr int64_t kA3VecLocalMemBits = 1572864; -constexpr int64_t kMatLocalMemBits = 4194304; -constexpr int64_t kLocalMemAlignmentBytes = 256; - -struct LocalMemSpec { - int64_t capacityBits = 0; - int64_t alignBytes = 1; -}; - -static std::optional getTileBufferFootprintBytes(TileBufType type) { - ArrayRef shape = type.getShape(); - unsigned elemBytes = getPTOStorageElemByteSize(type.getElementType()); - if (elemBytes == 0) { - return std::nullopt; - } - - if (type.getCompactModeI32() != - static_cast(pto::CompactMode::RowPlusOne)) { - std::optional totalStaticSize = getStaticTotalSize(shape); - if (!totalStaticSize.has_value()) { - return std::nullopt; - } - return totalStaticSize.value() * static_cast(elemBytes); - } - - if (shape.size() != mlir::pto::kValue2 || llvm::is_contained(shape, ShapedType::kDynamic)) { - return std::nullopt; - } - - bool rowMajor = - type.getBLayoutValueI32() == static_cast(pto::BLayout::RowMajor); - int64_t major = rowMajor ? shape[0] : shape[1]; - int64_t minor = rowMajor ? shape[1] : shape[0]; - if (major == 0 || minor == 0) { - return 0; - } - return ((major - 1) * (minor + 1) + minor) * - static_cast(elemBytes); -} - -static int64_t ceilDivBitsToBytes(int64_t bits) { - return (bits + kBitsPerByte - 1) / kBitsPerByte; -} - -static int64_t alignUpBytes(int64_t value, int64_t align) { - int64_t safeAlign = std::max(align, 1); - if (safeAlign == 1) { - return value; - } - int64_t rem = value % safeAlign; - if (rem == 0) { - return value; - } - return value + (safeAlign - rem); -} - -static size_t plannerAlignBitsFromBytes(size_t alignBytes) { - return std::max(alignBytes, 1) * kBitsPerByte; -} - -static LocalMemSpec getLocalMemSpec(Operation *op, AddressSpace as) { - switch (as) { - case AddressSpace::VEC: - return isTargetArchA5(op) - ? LocalMemSpec{kA5VecLocalMemBits, kLocalMemAlignmentBytes} - : LocalMemSpec{kA3VecLocalMemBits, kLocalMemAlignmentBytes}; - case AddressSpace::MAT: - return LocalMemSpec{kMatLocalMemBits, kLocalMemAlignmentBytes}; - default: - return LocalMemSpec{}; - } -} - -static bool isNameIn(StringRef name, ArrayRef names) { - return llvm::is_contained(names, name); -} - -static bool isIgnoredA5TmpOperandUse(OpOperand &use) { - Operation *owner = use.getOwner(); - unsigned operandNo = use.getOperandNumber(); - StringRef name = owner->getName().getStringRef(); - - if (auto dpsOp = dyn_cast(owner)) { - if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) { - return false; - } - } else if (auto dpsOp = dyn_cast(owner)) { - if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) { - return false; - } - } - - if (isNameIn(name, {"pto.trowargmax", "pto.trowargmin", "pto.trowmax", - "pto.trowmin", "pto.trowsum", "pto.trowprod"})) { - return operandNo == 1; - } - - if (name == "pto.txors") { - return operandNo == mlir::pto::kValue2; - } - - if (isNameIn(name, {"pto.tprelu", "pto.txor", "pto.tsels", - "pto.trowexpand", "pto.tcolexpand", - "pto.trowexpandadd", "pto.trowexpanddiv", - "pto.trowexpandexpdif", "pto.trowexpandmax", - "pto.trowexpandmin", "pto.trowexpandmul", - "pto.trowexpandsub", "pto.tcolexpandadd", - "pto.tcolexpanddiv", "pto.tcolexpandexpdif", - "pto.tcolexpandmax", "pto.tcolexpandmin", - "pto.tcolexpandmul", "pto.tcolexpandsub"})) { - return operandNo == mlir::pto::kValue2; - } - - if (name == "pto.tsel") { - return operandNo == mlir::pto::kValue3; - } - - return false; -} - -static bool isA5IgnoredTmpAlloc(pto::AllocTileOp allocTile) { - if (getTargetArch(allocTile.getOperation()) != PTOArch::A5) { - return false; - } - - Value value = allocTile.getResult(); - if (value.use_empty()) { - return false; - } - - for (OpOperand &use : value.getUses()) { - if (!isIgnoredA5TmpOperandUse(use)) { - return false; - } - } - return true; -} - -static void collectStableValueOrder(Region ®ion, - AsmState &asmState, - DenseMap &stableValueKeys, - SmallVectorImpl &seenValues) { - auto recordValue = [&asmState, &seenValues, &stableValueKeys](Value value) { - if (stableValueKeys.find(value) != stableValueKeys.end()) { - return; - } - std::string key; - llvm::raw_string_ostream os(key); - value.printAsOperand(os, asmState); - stableValueKeys[value] = os.str(); - seenValues.push_back(value); - }; - - for (Block &block : region) { - for (BlockArgument blockArg : block.getArguments()) { - recordValue(blockArg); - } - for (Operation &op : block) { - for (Value result : op.getResults()) { - recordValue(result); - } - for (Region &nestedRegion : op.getRegions()) { - collectStableValueOrder(nestedRegion, asmState, stableValueKeys, - seenValues); - } - } - } -} - -static StableValueOrderMap buildStableValueOrder(func::FuncOp func) { - DenseMap stableValueKeys; - SmallVector seenValues; - AsmState asmState(func); - collectStableValueOrder(func.getBody(), asmState, stableValueKeys, seenValues); - - llvm::sort(seenValues, [&](Value lhs, Value rhs) { - const std::string &lhsKey = stableValueKeys.find(lhs)->second; - const std::string &rhsKey = stableValueKeys.find(rhs)->second; - if (lhsKey != rhsKey) { - return lhsKey < rhsKey; - } - return isLessValue(lhs, rhs); - }); - - StableValueOrderMap stableValueOrder; - for (auto [index, value] : llvm::enumerate(seenValues)) { - stableValueOrder[value] = index; - } - return stableValueOrder; -} - -static uint32_t lookupStableValueOrder( - Value value, const StableValueOrderMap &stableValueOrder) { - auto it = stableValueOrder.find(value); - if (it != stableValueOrder.end()) { - return it->second; - } - return std::numeric_limits::max(); -} - -static void sortValuesByStableOrder( - SmallVectorImpl &values, - const StableValueOrderMap &stableValueOrder) { - llvm::sort(values, [&](Value lhs, Value rhs) { - uint32_t lhsOrder = lookupStableValueOrder(lhs, stableValueOrder); - uint32_t rhsOrder = lookupStableValueOrder(rhs, stableValueOrder); - if (lhsOrder != rhsOrder) { - return lhsOrder < rhsOrder; - } - return isLessValue(lhs, rhs); - }); -} - -static SmallVector getScratchBuffersFromEffects(Operation *op, - ValueRange dpsInits, - const StableValueOrderMap &stableValueOrder) { - SmallVector scratchBuffers; - auto memEffect = dyn_cast(op); - if (!memEffect) { - return scratchBuffers; - } - - SmallVector, - kMemoryEffectReserveSize> - effects; - memEffect.getEffects(effects); - for (const auto &effect : effects) { - if (!isa(effect.getEffect())) { - continue; - } - Value value = effect.getValue(); - if (!value) { - continue; - } - if (!llvm::is_contained(op->getOperands(), value)) { - continue; - } - if (llvm::is_contained(dpsInits, value)) { - continue; - } - if (!llvm::is_contained(scratchBuffers, value)) { - scratchBuffers.push_back(value); - } - } - sortValuesByStableOrder(scratchBuffers, stableValueOrder); - return scratchBuffers; -} - -static SmallVector -getMemoryEffectBufferOperands(Operation *op, - const StableValueOrderMap &stableValueOrder) { - SmallVector buffers; - auto memEffect = dyn_cast(op); - if (!memEffect) { - return buffers; - } - - SmallVector, - kMemoryEffectReserveSize> - effects; - memEffect.getEffects(effects); - for (const auto &effect : effects) { - if (!isa(effect.getEffect())) { - continue; - } - Value value = effect.getValue(); - if (!value || !GetBufferSpaceAttr(value)) { - continue; - } - if (!llvm::is_contained(buffers, value)) { - buffers.push_back(value); - } - } - sortValuesByStableOrder(buffers, stableValueOrder); - return buffers; -} - -static SmallVector -getScratchConflictPairsFromEffects(Operation *op, ValueRange dpsInits, - const StableValueOrderMap &stableValueOrder) { - SmallVector conflictPairs; - SmallVector scratchBuffers = - getScratchBuffersFromEffects(op, dpsInits, stableValueOrder); - for (Value scratch : scratchBuffers) { - for (Value dst : dpsInits) { - if (!scratch || !dst || scratch == dst) { - continue; - } - conflictPairs.emplace_back(scratch, dst); - } - } - return conflictPairs; -} - -enum class ReserveBufferMode { - None, - Auto, - Manual, -}; - -struct ReserveBufferPlan { - ReserveBufferMode mode = ReserveBufferMode::None; - ReserveBufferOp reserveOp; - AddressSpace addressSpace = AddressSpace::Zero; - int64_t sizeBytes = 0; - int64_t capacityBytes = 0; - int64_t alignBytes = 1; -}; - -using ReserveBufferPlans = SmallVector; - -static LogicalResult analyzeReserveBufferPlans(func::FuncOp funcOp, - ReserveBufferPlans &plans) { - SmallVector reserveOps; - funcOp.walk( - [&](ReserveBufferOp reserveOp) { reserveOps.push_back(reserveOp); }); - - if (reserveOps.empty()) { - return success(); - } - - for (ReserveBufferOp reserveOp : reserveOps) { - AddressSpace as = reserveOp.getLocation().getAddressSpace(); - auto spec = getLocalMemSpec(reserveOp.getOperation(), as); - if (spec.capacityBits <= 0 || spec.alignBytes <= 0) { - return reserveOp.emitOpError("unsupported reserve_buffer location"); - } - - int64_t capacityBytes = spec.capacityBits / kBitsPerByte; - int64_t sizeBytes = reserveOp.getSize(); - bool autoAlloc = reserveOp.getAutoAlloc(); - - ReserveBufferPlan &plan = plans.emplace_back(); - plan.mode = autoAlloc ? ReserveBufferMode::Auto : ReserveBufferMode::Manual; - plan.reserveOp = reserveOp; - plan.addressSpace = as; - plan.sizeBytes = sizeBytes; - plan.capacityBytes = capacityBytes; - plan.alignBytes = spec.alignBytes; - - // Auto mode only declares that one contiguous region must be reserved. - // The concrete base is filled later from a hole in the target local space. - if (autoAlloc) { - if (reserveOp.getBaseAttr()) { - return reserveOp.emitOpError( - "expects 'base' to be absent when 'auto' is true"); - } - continue; - } - - // In manual mode, reserve_buffer.base is already fixed by the frontend or - // an earlier stage. Only basic validation is needed here. - auto baseAttr = reserveOp.getBaseAttr(); - if (!baseAttr) { - return reserveOp.emitOpError("expects 'base' when 'auto' is false"); - } - - int64_t baseBytes = baseAttr.getInt(); - if (baseBytes % spec.alignBytes != 0) { - return reserveOp.emitOpError( - "expects 'base' to satisfy the address-space alignment"); - } - if (baseBytes + sizeBytes > capacityBytes) { - return reserveOp.emitOpError("exceeds available local memory capacity"); - } - } - - return success(); -} - -struct OccupiedByteRange { - int64_t begin = 0; - int64_t end = 0; -}; - -static LogicalResult assignAutoReserveBufferBases( - ReserveBufferPlans &plans, - const std::map &bufferInfos, - const DenseMap> &buffer2Offsets) { - std::map> occupiedByAddressSpace; - for (const auto &it : bufferInfos) { - Value buffer = it.first; - const BufferInfo &bufferInfo = it.second; - - auto offsetsIt = buffer2Offsets.find(buffer); - if (offsetsIt == buffer2Offsets.end()) { - continue; - } - - // Reserve-buffer allocation intentionally happens after normal MemPlan. - // Reconstruct the already occupied byte ranges from the planned local - // buffers, then place reserve_buffer into the first aligned hole. - int64_t occupiedSizeBytes = ceilDivBitsToBytes(bufferInfo.constBits); - if (bufferInfo.operation) { - auto spec = getLocalMemSpec(bufferInfo.operation, bufferInfo.bufferScope); - occupiedSizeBytes = - alignUpBytes(occupiedSizeBytes, std::max(spec.alignBytes, 1)); - } - for (uint64_t offsetBytes : offsetsIt->second) { - occupiedByAddressSpace[bufferInfo.bufferScope].push_back( - OccupiedByteRange{static_cast(offsetBytes), - static_cast(offsetBytes) + occupiedSizeBytes}); - } - } - - auto normalizeRanges = [](SmallVector &ranges) { - llvm::sort(ranges, - [](const OccupiedByteRange &lhs, const OccupiedByteRange &rhs) { - return lhs.begin < rhs.begin; - }); - - SmallVector merged; - for (const OccupiedByteRange &range : ranges) { - if (merged.empty() || range.begin > merged.back().end) { - merged.push_back(range); - continue; - } - merged.back().end = std::max(merged.back().end, range.end); - } - ranges.swap(merged); - }; - - for (auto &it : occupiedByAddressSpace) { - normalizeRanges(it.second); - } - - for (ReserveBufferPlan &plan : plans) { - if (plan.mode != ReserveBufferMode::Auto || !plan.reserveOp) { - continue; - } - - SmallVector &occupied = - occupiedByAddressSpace[plan.addressSpace]; - - // First-fit search: try address 0 first, then keep moving the candidate to - // the end of the current occupied interval until a large-enough aligned - // hole is found. - int64_t candidateBase = 0; - for (const OccupiedByteRange &range : occupied) { - candidateBase = alignUpBytes(candidateBase, plan.alignBytes); - if (candidateBase + plan.sizeBytes <= range.begin) { - break; - } - candidateBase = std::max(candidateBase, range.end); - } - candidateBase = alignUpBytes(candidateBase, plan.alignBytes); - if (candidateBase + plan.sizeBytes > plan.capacityBytes) { - return plan.reserveOp.emitOpError( - "failed to allocate local memory hole for reserve_buffer"); - } - - plan.reserveOp->setAttr( - "base", - IntegerAttr::get( - IntegerType::get(plan.reserveOp.getContext(), kI32BitWidth), - candidateBase)); - occupied.push_back( - OccupiedByteRange{candidateBase, candidateBase + plan.sizeBytes}); - normalizeRanges(occupied); - } - return success(); -} - -} // namespace - -void MemLivenessAnalysis::build() { - Region &funcRegion = func_.getBody(); - stableValueOrder = buildStableValueOrder(func_); - Liveness live(func_); - // Recursively obtaining IR information. - RecursionIR(&funcRegion, live); - // the lifetime of the buffer. - GenerateBufferLife(); -} - -bool MemLivenessAnalysis::isLocalMemPlan() const { - return planMode == MemPlanMode::LOCAL_MEM_PLAN; -} - -bool MemLivenessAnalysis::isGlobalWorkSpaceMemPlan() const { - return planMode == MemPlanMode::GLOBAL_WORKSPACE_PLAN; -} - -void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { - auto result = region->walk([&](Operation *op) { - // recursive control flow - if (auto ifOp = dyn_cast(op)) { - RecursiveIfOp(ifOp, live); - return WalkResult::skip(); - } else if (auto forOp = dyn_cast(op)) { - RecursiveForOp(forOp, live); - return WalkResult::skip(); - } else if (auto whileOp = dyn_cast(op)) { - RecursiveWhileOp(whileOp, live); - return WalkResult::skip(); - } else if (auto fusionRegion = dyn_cast(op)) { - RecursiveFusionRegionOp(fusionRegion, live); - return WalkResult::skip(); - } else if (auto section = dyn_cast(op)) { - // Sections are transparent containers. Analyze tileop-local cube work - // in place so it contributes to liveness. - UpdateLinearOperation(section.getOperation()); - RecursionIR(§ion.getBody(), live); - auto sectionEnd = UpdateLinearOperation(section.getOperation()); - OpKillHandle(sectionEnd, live, section->getBlock()); - return WalkResult::skip(); - } else if (auto section = dyn_cast(op)) { - // Keep vector sections transparent for the same reason as cube sections. - UpdateLinearOperation(section.getOperation()); - RecursionIR(§ion.getBody(), live); - auto sectionEnd = UpdateLinearOperation(section.getOperation()); - OpKillHandle(sectionEnd, live, section->getBlock()); - return WalkResult::skip(); - } - - // process operation - auto curOpInfo = UpdateLinearOperation(op); - auto mayAliasOp = getOperationAliasInfo(op); - if (mayAliasOp.has_value()) { - auto aliasPair = mayAliasOp.value(); - UpdateBufferAlias(aliasPair.first, aliasPair.second); - } else if (isa(op)) { - // Runtime-bound tile handles do not allocate static local storage. - return WalkResult::advance(); - } else if (isLocalMemPlan() && dyn_cast(op)) { - auto allocTileOp = cast(op); - if (allocTileOp.getAddr()) { - return WalkResult::advance(); - } - if (isA5IgnoredTmpAlloc(allocTileOp)) { - return WalkResult::advance(); - } - auto memorySpaceAttr = GetBufferSpaceAttr(allocTileOp.getResult()); - if (!isLocalBuffer(memorySpaceAttr)) { - allocTileOp.emitError("Alloc tile buffer not at local space"); - return WalkResult::interrupt(); - } - if (auto attr = allocTileOp->getAttrOfType( - mlir::pto::kPtoMultiBufferAttrName)) { - uint64_t n = attr.getValue().getZExtValue(); - if (n < mlir::pto::kPtoMultiBufferMinNum || - n > mlir::pto::kPtoMultiBufferMaxNum) { - allocTileOp.emitError() - << "pto.multi_buffer must be in [" - << mlir::pto::kPtoMultiBufferMinNum << ", " - << mlir::pto::kPtoMultiBufferMaxNum << "] (got " << n << ")"; - return WalkResult::interrupt(); - } - buffer2MultiNum[allocTileOp.getResult()] = static_cast(n); - } - UpdateOpBufferInfo(op, op->getResults()); - return WalkResult::advance(); - } else if (isLocalMemPlan() && dyn_cast(op)) { - auto allocMultiOp = cast(op); - if (allocMultiOp.getAddr()) { - return WalkResult::advance(); - } - auto memorySpaceAttr = GetBufferSpaceAttr(allocMultiOp.getResult()); - if (!isLocalBuffer(memorySpaceAttr)) { - allocMultiOp.emitError("Alloc multi tile buffer not at local space"); - return WalkResult::interrupt(); - } - uint32_t count = allocMultiOp.getResult().getType().getCount(); - buffer2MultiNum[allocMultiOp.getResult()] = count; - UpdateOpBufferInfo(op, op->getResults()); - return WalkResult::advance(); - } else if (auto tprintOp = dyn_cast(op)) { - // TPrintOp only reads from buffer, similar to LoadOp - UpdateOpGenInfo(curOpInfo, llvm::to_vector(op->getOperands())); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto tgetvalOp = dyn_cast(op)) { - (void)tgetvalOp; - UpdateOpGenInfo(curOpInfo, llvm::to_vector(op->getOperands())); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto setValidShapeOp = dyn_cast(op)) { - (void)setValidShapeOp; - // Metadata-only update on an existing tile handle. Keep the source buffer - // alive through this operation, but do not model it as producing a new - // alias/result buffer. - UpdateOpGenInfo(curOpInfo, ValueRange{op->getOperand(0)}); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto getValidShapeOp = dyn_cast(op)) { - (void)getValidShapeOp; - // Metadata-only read from an existing tile handle. It touches the source - // buffer for liveness, but the scalar row/col results are not buffers. - UpdateOpGenInfo(curOpInfo, ValueRange{op->getOperand(0)}); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto ptoDpsOp = dyn_cast(op)) { - // PTO ops with destination (tile_buf, partition_view, etc.). - SmallVector genBuffers = llvm::to_vector(op->getOperands()); - auto scratchBuffers = getScratchBuffersFromEffects( - op, ptoDpsOp.getDpsInits(), stableValueOrder); - genBuffers.append(scratchBuffers.begin(), scratchBuffers.end()); - UpdateOpGenInfo(curOpInfo, genBuffers); - for (const auto &conflictPair : - getScratchConflictPairsFromEffects(op, ptoDpsOp.getDpsInits(), - stableValueOrder)) { - RecordSemanticConflict(conflictPair.first, conflictPair.second); - } - for (const auto &[lhs, rhs] : getSemanticNoAliasPairs(op)) - RecordSemanticConflict(lhs, rhs); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto dstStyleOp = dyn_cast(op)) { - // Preserve generic destination-style aliasing for non-tile tensors. - UpdateInitAndResAlias(dstStyleOp); - UpdateOpGenInfo(curOpInfo, llvm::to_vector(dstStyleOp.getDpsInits())); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto selectOp = dyn_cast(op)) { - UpdateBufferAlias(selectOp.getResult(), selectOp.getTrueValue(), true); - UpdateBufferAlias(selectOp.getResult(), selectOp.getFalseValue(), true); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto tileBufAddrOp = dyn_cast(op)) { - UpdateOpGenInfo(curOpInfo, ValueRange{tileBufAddrOp.getSrc()}); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto callOp = dyn_cast(op)) { - UpdateOpGenInfo(curOpInfo, llvm::to_vector(callOp->getOperands())); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto simtLaunch = dyn_cast(op)) { - // A launched SIMT helper reads its explicit argument list. Treat it as - // a call for liveness so tile-derived local memrefs remain live through - // the launch. - UpdateOpGenInfo(curOpInfo, llvm::to_vector(simtLaunch->getOperands())); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (isa(op)) { - UpdateOpGenInfo(curOpInfo, llvm::to_vector(op->getOperands())); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto gpuLaunchOp = dyn_cast(op)) { - UpdateOpGenInfo(curOpInfo, llvm::to_vector(gpuLaunchOp->getOperands())); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (isa(op)) { - // Vector/Cube micro-ops such as vlds/vsts are not destination-style - // ops, but their explicit memory effects identify the local buffers - // that must remain live through the instruction. - SmallVector buffers = - getMemoryEffectBufferOperands(op, stableValueOrder); - UpdateOpGenInfo(curOpInfo, buffers); - OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (failed(CheckIfUnknownOpTouchBuffer(op))) { - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - if (result == WalkResult::interrupt()) { - llvm_unreachable("PlanMemory Traverse IR Failed! "); - } -} - -void MemLivenessAnalysis::UpdateInitAndResAlias( - DestinationStyleOpInterface dstStyleOp) { - auto results = dstStyleOp->getResults(); - if (results.empty()) { - return; - } - for (auto [res, init] : - llvm::zip(dstStyleOp->getResults(), dstStyleOp.getDpsInits())) { - auto iter = buffer2AliasVec.find(init); - if (iter == buffer2AliasVec.end()) { - continue; - } - auto tensorType = dyn_cast_or_null(res.getType()); - if (tensorType) { - UpdateBufferAlias(res, init); - } - } -} - -OpInfo *MemLivenessAnalysis::UpdateLinearOperation(Operation *op) { - auto opInfo = std::make_unique(op, seqIndex++); - auto curOpInfo = opInfo.get(); - linearOperation.push_back(std::move(opInfo)); - return curOpInfo; -} - -void MemLivenessAnalysis::UpdateForOpBufferAlias(scf::ForOp forOp) { - if (forOp.getResults().empty()) { - return; - } - if (!forOp.getRegionIterArgs().empty()) { - if (forOp.getYieldedValues().size() != forOp.getRegionIterArgs().size() || - forOp.getInitArgs().size() != forOp.getRegionIterArgs().size()) { - llvm::report_fatal_error("scf.for alias sizes are inconsistent"); - } - for (auto [i, arg] : llvm::enumerate(forOp.getRegionIterArgs())) { - // yielded values alias region iter args. - UpdateBufferAlias(forOp.getYieldedValues()[i], arg); - } - } - if (forOp->getResults().size() != forOp.getYieldedValues().size()) { - llvm::report_fatal_error("scf.for result/yield sizes are inconsistent"); - } - for (auto [i, arg] : llvm::enumerate(forOp.getYieldedValues())) { - // forOp result values alias region iter yielded values. - UpdateBufferAlias(forOp->getResult(i), arg); - } -} - -void MemLivenessAnalysis::UpdateWhileOpBufferAlias(scf::WhileOp whileOp) { - if (whileOp.getResults().size() != whileOp.getYieldedValues().size()) - llvm::report_fatal_error("scf.while result/yield sizes are inconsistent"); - - // The before region receives the initial values and the after region - // receives the values forwarded by scf.condition. The after-region yield - // is the back-edge which becomes the next before-region argument and the - // scf.while result on exit. - for (auto [i, init] : llvm::enumerate(whileOp.getInits())) { - if (i < whileOp.getBeforeArguments().size()) - UpdateBufferAlias(whileOp.getBeforeArguments()[i], init); - } - - auto conditionArgs = whileOp.getConditionOp().getArgs(); - for (auto [i, arg] : llvm::enumerate(whileOp.getAfterArguments())) { - if (i < conditionArgs.size()) - UpdateBufferAlias(arg, conditionArgs[i]); - } - - for (auto [i, yielded] : llvm::enumerate(whileOp.getYieldedValues())) { - if (i < whileOp.getBeforeArguments().size()) - UpdateBufferAlias(whileOp.getBeforeArguments()[i], yielded); - if (i < whileOp.getConditionOp().getArgs().size()) - UpdateBufferAlias(whileOp.getResult(i), - whileOp.getConditionOp().getArgs()[i]); - if (i < whileOp.getResults().size()) - UpdateBufferAlias(whileOp.getResult(i), yielded); - } -} - -void MemLivenessAnalysis::RecursiveForOp(scf::ForOp forOp, Liveness live) { - // Model loop-carried tile handles as aliases of the yielded roots. - auto forBeginSeq = UpdateLinearOperation(forOp.getOperation()); - UpdateOpGenInfo(forBeginSeq, GetLiveBuffersInLoop(forOp, live)); - UpdateForOpInitArgsAlias(forOp); - RecursionIR(&forOp.getRegion(), live); - UpdateForOpBufferAlias(forOp); - auto forEndSeq = UpdateLinearOperation(forOp.getOperation()); - OpKillHandle(forEndSeq, live, forOp->getBlock()); -} - -void MemLivenessAnalysis::RecursiveWhileOp(scf::WhileOp whileOp, - Liveness live) { - auto whileBeginSeq = UpdateLinearOperation(whileOp.getOperation()); - UpdateOpGenInfo(whileBeginSeq, GetLiveBuffersInLoop(whileOp.getOperation(), live)); - - UpdateWhileOpBufferAlias(whileOp); - RecursionIR(&whileOp.getBefore(), live); - RecursionIR(&whileOp.getAfter(), live); - UpdateWhileOpBufferAlias(whileOp); - - auto whileEndSeq = UpdateLinearOperation(whileOp.getOperation()); - OpKillHandle(whileEndSeq, live, whileOp->getBlock()); -} - -void MemLivenessAnalysis::UpdateForOpInitArgsAlias(scf::ForOp forOp) { - if (forOp.getInitArgs().empty()) { - return; - } - if (forOp.getInitArgs().size() != forOp.getRegionIterArgs().size()) { - llvm::report_fatal_error("scf.for init/iter-arg sizes are inconsistent"); - } - for (auto [i, arg] : llvm::enumerate(forOp.getInitArgs())) { - // init args alias region iter args. - UpdateBufferAlias(forOp.getRegionIterArgs()[i], arg); - } -} - -void MemLivenessAnalysis::UpdateIfOpBufferAlias(scf::IfOp ifOp, - scf::YieldOp yieldOp) { - if (ifOp.getResults().empty()) { - return; - } - if (ifOp->getResults().size() != yieldOp->getOperands().size()) { - llvm::report_fatal_error("scf.if result/yield sizes are inconsistent"); - } - for (auto [i, arg] : llvm::enumerate(yieldOp->getOperands())) { - // Multiple buffers involved, requiring one-to-one correspondence. - UpdateBufferAlias(ifOp->getResult(i), arg); - } -} - -void MemLivenessAnalysis::RecursiveIfOp(scf::IfOp ifOp, Liveness live) { - // Join the tile roots yielded by each branch into the if result. - UpdateLinearOperation(ifOp.getOperation()); - RecursionIR(&ifOp.getThenRegion(), live); - auto curIfElse = UpdateLinearOperation(ifOp.getOperation()); - UpdateIfOpBufferAlias(ifOp, ifOp.thenYield()); - - auto curIfEnd = curIfElse; - if (ifOp.elseBlock()) { - RecursionIR(&ifOp.getElseRegion(), live); - curIfEnd = UpdateLinearOperation(ifOp.getOperation()); - UpdateIfOpBufferAlias(ifOp, ifOp.elseYield()); - } - OpKillHandle(curIfEnd, live, ifOp->getBlock()); -} - -void MemLivenessAnalysis::UpdateFusionRegionBufferAlias( - pto::FusionRegionOp fusionRegion, pto::YieldOp yieldOp) { - if (fusionRegion.getResults().empty()) { - return; - } - if (fusionRegion->getResults().size() != yieldOp->getOperands().size()) { - llvm::report_fatal_error( - "pto.fusion_region result/yield sizes are inconsistent"); - } - for (auto [i, yielded] : llvm::enumerate(yieldOp->getOperands())) { - UpdateBufferAlias(fusionRegion->getResult(i), yielded); - } -} - -void MemLivenessAnalysis::RecursiveFusionRegionOp(pto::FusionRegionOp fusionRegion, - Liveness live) { - UpdateLinearOperation(fusionRegion.getOperation()); - RecursionIR(&fusionRegion.getBody(), live); - - auto yieldOp = - dyn_cast(fusionRegion.getBody().front().getTerminator()); - if (!yieldOp) { - llvm::report_fatal_error("pto.fusion_region must terminate with pto.yield"); - } - UpdateFusionRegionBufferAlias(fusionRegion, yieldOp); - - auto regionEnd = UpdateLinearOperation(fusionRegion.getOperation()); - OpKillHandle(regionEnd, live, fusionRegion->getBlock()); -} - -SmallVector MemLivenessAnalysis::GetLiveBuffersInLoop(Operation *loopOp, - Liveness live) { - SmallVector allocBeforeLoopBuffers; - const auto *liveBlockInfo = live.getLiveness(loopOp->getBlock()); - auto currentLiveValues = - liveBlockInfo->currentlyLiveValues(loopOp); - if (currentLiveValues.empty()) { - return allocBeforeLoopBuffers; - } - // The gen buffer of the same operation must ensure the order of priority. - SetVector currentLiveValuesOrder; - for (auto buffer : currentLiveValues) { - currentLiveValuesOrder.insert(buffer); - } - for (const Value &operand : currentLiveValuesOrder) { - auto aliasBuffers = GetAliasBuffers(operand); - aliasBuffers.insert(operand); - for (auto Buffer : aliasBuffers) { - auto iter = buffer2status.find(Buffer); - if (iter != buffer2status.end()) { - allocBeforeLoopBuffers.push_back(Buffer); - } - } - } - sortValuesByStableOrder(allocBeforeLoopBuffers, stableValueOrder); - return allocBeforeLoopBuffers; -} - -bool MemLivenessAnalysis::isSkippableOp(Operation *op) const { - // Call-like ops are still modeled explicitly. Only pure terminators and - // dim queries are skipped here. - return isa(op); -} - -LogicalResult -MemLivenessAnalysis::CheckIfUnknownOpTouchBuffer(Operation *op) const { - if (isSkippableOp(op) || isGlobalWorkSpaceMemPlan()) { - // This scene can be ignored. - return success(); - } - if (isOpTouchLocalBuffer(op)) { - op->emitError("PlanMemory Fail : Unrecognized type of Operation touches " - "local buffer!"); - return failure(); - } - return success(); -} - -void MemLivenessAnalysis::UpdateBufferAlias(Value buffer, Value aliasBuffer, - bool isIgnoreInplace) { - // union all alias buffers about `aliasBuffer` and `buffer` - auto unionAliasSet = - Union(GetAliasBuffers(aliasBuffer), GetAliasBuffers(buffer)); - unionAliasSet.insert(buffer); - unionAliasSet.insert(aliasBuffer); - - // update alias map info for each buffer - // e.g. if A alias B, C alias D, now update: - // A alias B,C,D; B alias A,C,D; C alias A,B,D; D alias A,B,C - for (auto buf : unionAliasSet) { - // remove buf self from union alias set - auto clonedAliasSet = unionAliasSet; - clonedAliasSet.remove(buf); - - buffer2AliasVec[buf] = clonedAliasSet; - } - - // Mark aliases that must not participate in inplace merging. - auto it = bufferInfos.find(aliasBuffer); - if (isIgnoreInplace && it != bufferInfos.end()) { - it->second.ignoreInplace = true; - } -} - -SetVector MemLivenessAnalysis::Union(SetVector set1, - SetVector set2) { - SetVector unionSet; - unionSet.insert(set1.begin(), set1.end()); - unionSet.insert(set2.begin(), set2.end()); - return unionSet; -} - -SetVector MemLivenessAnalysis::GetAliasBuffers(Value aliasBuffer) { - if (!aliasBuffer) { - return {}; - } - - auto trueVar = buffer2AliasVec.find(aliasBuffer); - if (trueVar != buffer2AliasVec.end()) { - return trueVar->second; - } - return {}; -} - -void MemLivenessAnalysis::UpdateOpBufferInfo(Operation *op, - const ValueRange &results) { - for (const Value &operand : results) { - auto it = buffer2status.find(operand); - if (it != buffer2status.end()) { - continue; - } - bufferInfos[operand] = GenerateBufferInfo(op, operand); - buffer2status[operand] = BufferStatus::DEFFINED; - } -} - -void MemLivenessAnalysis::UpdateOpGenInfo(OpInfo *opInfo, - const ValueRange &results) { - if (results.empty()) { - return; - } - for (Value operand : results) { - auto aliasBuffers = GetAliasBuffers(operand); - aliasBuffers.insert(operand); - for (auto buffer : aliasBuffers) { - UpdateOperandGenInfo(opInfo, buffer); - } - } -} - -void MemLivenessAnalysis::UpdateOperandGenInfo(OpInfo *opInfo, Value operand) { - auto iter_buffer = buffer2status.find(operand); - if (iter_buffer == buffer2status.end()) { - return; - } - if (iter_buffer->second == BufferStatus::DEFFINED) { - genKillMap[opInfo].gen.push_back(operand); - buffer2status[iter_buffer->first] = BufferStatus::GENED; - } else if (iter_buffer->second == BufferStatus::KILLED) { - llvm_unreachable("The buffer memory has been released and cannot be used " - "again! "); - } -} - -void MemLivenessAnalysis::OpKillHandle(OpInfo *opInfo, Liveness live, - Block *block) { - const auto *liveBlockInfo = live.getLiveness(block); - auto currentLiveValues = - liveBlockInfo->currentlyLiveValues(opInfo->operation); - if (currentLiveValues.empty()) { - return; - } - SmallVector liveValues(currentLiveValues.begin(), - currentLiveValues.end()); - sortValuesByStableOrder(liveValues, stableValueOrder); - for (const Value &operand : liveValues) { - UpdateOpKillInfo(opInfo, operand, live); - } -} - -void MemLivenessAnalysis::UpdateOpKillInfo(OpInfo *opInfo, Value operand, - Liveness live) { - auto aliasBuffers = GetAliasBuffers(operand); - aliasBuffers.insert(operand); - for (Value aliasBuffer : aliasBuffers) { - auto iterBuffer = buffer2status.find(aliasBuffer); - if (iterBuffer == buffer2status.end()) { - return; - } - if (iterBuffer->second == BufferStatus::GENED && - IsInSameBlock(iterBuffer->first.getDefiningOp(), opInfo->operation) && - AllDeadAfter(opInfo->operation, aliasBuffers, live)) { - genKillMap[opInfo].kill.push_back(aliasBuffer); - buffer2status[iterBuffer->first] = BufferStatus::KILLED; - } - } -} - -bool MemLivenessAnalysis::IsInSameBlock(Operation *op1, Operation *op2) const { - return op1->getBlock() == op2->getBlock(); -} - -bool MemLivenessAnalysis::AllDeadAfter(Operation *op, SetVector aliasVec, - Liveness live) const { - for (auto aliasBuffer : aliasVec) { - if (!live.isDeadAfter(aliasBuffer, op)) { - return false; - } - } - return true; -} - -void MemLivenessAnalysis::RecordSemanticConflict(Value lhs, Value rhs) { - SetVector lhsAliases = GetAliasBuffers(lhs); - lhsAliases.insert(lhs); - SetVector rhsAliases = GetAliasBuffers(rhs); - rhsAliases.insert(rhs); - - auto appendUniquePair = [this](Value a, Value b) { - if (!a || !b || a == b) { - return; - } - ValuePair pair = isLessValue(a, b) ? ValuePair(a, b) : ValuePair(b, a); - if (!llvm::is_contained(semanticConflictPairs, pair)) { - semanticConflictPairs.push_back(pair); - } - }; - - for (Value a : lhsAliases) { - for (Value b : rhsAliases) { - appendUniquePair(a, b); - } - } -} - -BufferInfo MemLivenessAnalysis::GenerateBufferInfo(Operation *op, - Value operand) { - auto memorySpaceAttr = GetBufferSpaceAttr(operand); - if (isLocalMemPlan() && isLocalBuffer(memorySpaceAttr)) { - if (!memorySpaceAttr.has_value()) { - llvm::report_fatal_error("local buffer must have memory space"); - } - return GetBufferInfo(op, operand, - memorySpaceAttr.value().getAddressSpace()); - } - llvm_unreachable("buffer must has BufferInfo !"); -} - -BufferInfo MemLivenessAnalysis::GetBufferInfo(Operation *op, Value operand, - pto::AddressSpace bufferScope) { - BufferInfo bufferInfo; - bufferInfo.operation = op; - bufferInfo.bufferScope = bufferScope; - // get buffer size, now for static shape - Type elementType; - std::optional footprintBytes; - if (auto tileType = dyn_cast(operand.getType())) { - elementType = tileType.getElementType(); - footprintBytes = getTileBufferFootprintBytes(tileType); - } else if (auto multiType = dyn_cast(operand.getType())) { - TileBufType slotType = multiType.getSlotType(); - elementType = slotType.getElementType(); - footprintBytes = getTileBufferFootprintBytes(slotType); - } else { - llvm_unreachable("local memory planner expects tile buffer roots"); - } - bufferInfo.bufferType = elementType; - if (!footprintBytes.has_value()) { - llvm::report_fatal_error( - "failed to obtain buffer static physical footprint"); - } - bufferInfo.constBits = footprintBytes.value() * kBitsPerByte; - return bufferInfo; -} - -void MemLivenessAnalysis::GenerateBufferLife() { - int scopeTime = 0; - for (size_t i = 0; i < linearOperation.size(); ++i) { - auto it = genKillMap.find(linearOperation[i].get()); - if (it == genKillMap.end()) { - scopeTime++; - continue; - } - // Time given to buffer start. - for (const Value &genBuffer : it->second.gen) { - std::unique_ptr bufferLife = - std::make_unique(genBuffer); - bufferLife->allocTime = scopeTime; - buffer2Life[genBuffer] = std::move(bufferLife); - } - // Time given to buffer end. - for (const Value &killBuffer : it->second.kill) { - auto iter = buffer2Life.find(killBuffer); - if (iter == buffer2Life.end()) { - llvm::report_fatal_error("buffer lifetime killed before generation"); - } - iter->second->freeTime = scopeTime; - } - scopeTime++; - } -} - -std::shared_ptr -StorageEntry::GetBufferLifeByValue(const Value v) const { - auto find = std::find_if( - bufferLifeVec.begin(), bufferLifeVec.end(), - [v](std::shared_ptr life) { return life->buffer == v; }); - if (find != bufferLifeVec.end()) { - return *find; - } - return nullptr; -} - -bool MemPlan::IsReusePTOOp(Operation *op) const { - if (restrictInplaceAsISA) { - return false; - } - - // not in ISA but confirmed with hardware developers: - // elementwise ops with the same shape and the same bitwidth operands can also - // do memory inplace for src and dst - return false; -} - -SmallVector MemPlan::GenerateInplaceList() { - SmallVector inplaceList; - DenseMap hasTouchOp; - inplaceList.insert(inplaceList.end(), inplacePairList.begin(), - inplacePairList.end()); - for (auto &operationSeq : linearOperation) { - auto it = genKillMap.find(operationSeq.get()); - if (it == genKillMap.end()) { - continue; - } - if (hasTouchOp[operationSeq->operation]) { - continue; - } - - SmallVector genBuffers(it->second.gen.begin(), it->second.gen.end()); - SmallVector killBuffers(it->second.kill.begin(), it->second.kill.end()); - sortValuesByStableOrder(genBuffers, stableValueOrder); - sortValuesByStableOrder(killBuffers, stableValueOrder); - - for (const Value &genBuffer : genBuffers) { - auto genBufferIter = bufferInfos.find(genBuffer); - if (genBufferIter == bufferInfos.end()) { - llvm::report_fatal_error("gen buffer missing from buffer info map"); - } - if (genBufferIter->second.ignoreInplace) { - continue; - } - - for (const Value &killBuffer : killBuffers) { - auto killBufferIter = bufferInfos.find(killBuffer); - if (killBufferIter == bufferInfos.end()) { - llvm::report_fatal_error("kill buffer missing from buffer info map"); - } - if (killBufferIter->second.ignoreInplace) { - continue; - } - - bool bufferSizeMatch = - killBufferIter->second.constBits >= genBufferIter->second.constBits; - bool isResuableOp = IsReusePTOOp(it->first->operation); - bool canInplace = bufferSizeMatch && isResuableOp; - if (canInplace) { - inplaceList.emplace_back(std::make_pair(genBuffer, killBuffer)); - break; - } - } - } - // Nodes in inplace are only processed once. - hasTouchOp[operationSeq->operation] = true; - } - return inplaceList; -} - -void MemPlan::EmitPlanMemoryFailureInfo() { - if (failApplyBufferInfo.empty()) { - return; - } - for (auto &iter : failApplyBufferInfo) { - AddressSpace space = iter.first; - func_.emitError() << stringifyEnum(space) << " overflow, requires " - << iter.second << " bits while " - << GetBufferSpaceInfo(space).second << " bits available!"; - } -} - -bool MemPlan::RecordOverflowIfAny() { - if (!failApplyBufferInfo.empty()) { - return true; - } - if (planMode != MemPlanMode::LOCAL_MEM_PLAN || - memscope2rootStorageEntry.empty()) { - return false; - } - - for (auto &it : memscope2rootStorageEntry) { - auto *rootStorageEntry = it.second; - if (!rootStorageEntry) { - continue; - } - auto bufferSpaceInfo = - GetBufferSpaceInfo(rootStorageEntry->bufInfo->bufferScope); - size_t maxBits = bufferSpaceInfo.second; - uint64_t maxAllocBits = rootStorageEntry->alignedConstBits; - for (auto *child : rootStorageEntry->mergedChildren) { - maxAllocBits = - std::max(maxAllocBits, child->bitsOffset + child->alignedConstBits); - } - if (maxAllocBits > maxBits) { - failApplyBufferInfo[rootStorageEntry->bufInfo->bufferScope] = - maxAllocBits; - } - } - - return !failApplyBufferInfo.empty(); -} - -bool MemPlan::HasSemanticConflict(const StorageEntry *entry, - const BufferLifeVec &bufferLives) const { - if (!entry || semanticConflictPairs.empty() || bufferLives.empty()) { - return false; - } - - auto containsPair = [this](Value lhs, Value rhs) { - ValuePair pair = isLessValue(lhs, rhs) ? ValuePair(lhs, rhs) - : ValuePair(rhs, lhs); - return llvm::is_contained(semanticConflictPairs, pair); - }; - - for (Value entryBuffer : entry->inplaceBuffers) { - for (const auto &life : bufferLives) { - if (!life) { - continue; - } - Value otherBuffer = life->buffer; - if (!otherBuffer || entryBuffer == otherBuffer) { - continue; - } - if (containsPair(entryBuffer, otherBuffer)) { - return true; - } - } - } - return false; -} - -// Plan Memory algorithm. -LogicalResult MemPlan::plan() { - // Construct StorageEntry structure. - GenerateStorageEntry(); - // Plan memory address. - PlanStatus as = planMode == MemPlanMode::LOCAL_MEM_PLAN - ? PlanLocalMemAddress() - : PlanWorkSpaceMemAddress(); - if (as == PlanStatus::PLAN_FAILED) { - EmitPlanMemoryFailureInfo(); - return failure(); - } - if (RecordOverflowIfAny()) { - EmitPlanMemoryFailureInfo(); - return failure(); - } - auto hasAddressOverlap = [](const StorageEntry *lhs, const StorageEntry *rhs) { - uint64_t lhsBegin = lhs->bitsOffset; - uint64_t lhsEnd = lhs->bitsOffset + lhs->alignedConstBits; - uint64_t rhsBegin = rhs->bitsOffset; - uint64_t rhsEnd = rhs->bitsOffset + rhs->alignedConstBits; - return lhsBegin < rhsEnd && rhsBegin < lhsEnd; - }; - SmallVector plannedEntries; - plannedEntries.reserve(StorageEntryVec.size() + pingEntry2RelationPongEntry.size()); - for (const auto &entry : StorageEntryVec) { - plannedEntries.push_back(entry.get()); - } - for (const auto &entry : pingEntry2RelationPongEntry) { - plannedEntries.push_back(entry.second.get()); - } - for (size_t i = 0; i < plannedEntries.size(); ++i) { - for (size_t j = i + 1; j < plannedEntries.size(); ++j) { - const StorageEntry *lhs = plannedEntries[i]; - const StorageEntry *rhs = plannedEntries[j]; - if (!lhs || !rhs) { - continue; - } - if (lhs->bufInfo->bufferScope != rhs->bufInfo->bufferScope) { - continue; - } - if (!hasAddressOverlap(lhs, rhs)) { - continue; - } - bool lifeOverlap = - !GetOverlapBufferLife(lhs->bufferLifeVec, rhs->bufferLifeVec).empty(); - bool semanticConflict = HasSemanticConflict(lhs, rhs->bufferLifeVec); - if (!lifeOverlap && !semanticConflict) { - continue; - } - func_.emitError() - << "PlanMemory produced overlapping local buffers in " - << stringifyEnum(lhs->bufInfo->bufferScope) - << " at offsets " << lhs->bitsOffset << " and " << rhs->bitsOffset; - return failure(); - } - } - // Update the address information of each buffer after memory buffer. - UpdateBuffer2Offsets(); - if (enablePrintMemoryAllocatedSize) { - PrintSuccessfulAllocatedMaxBits(); - } - return success(); -} - -void MemPlan::GenerateStorageEntry() { - // create new storage entry. - for (auto &operation : linearOperation) { - auto it = genKillMap.find(operation.get()); - if (it == genKillMap.end()) { - continue; - } - SmallVector genBuffers(it->second.gen.begin(), it->second.gen.end()); - sortValuesByStableOrder(genBuffers, stableValueOrder); - for (const Value &genBuffer : genBuffers) { - auto iter = bufferInfos.find(genBuffer); - if (iter == bufferInfos.end()) { - continue; - } - const std::shared_ptr &bufLife = buffer2Life.at(genBuffer); - std::unique_ptr entry = std::make_unique(); - entry->bufInfo = &iter->second; - entry->bufferLifeVec.emplace_back(bufLife); - entry->inplaceBuffers.emplace_back(iter->first); - auto multiBuffer = buffer2MultiNum.find(genBuffer); - if (multiBuffer != buffer2MultiNum.end()) { - entry->multiBufferNum = multiBuffer->second; - } - buffer2storageEntry[genBuffer] = entry.get(); - // Verify the validity of parameters after initialization. - ValidateParameters(entry); - StorageEntryVec.emplace_back(std::move(entry)); - } - } -} - -void MemPlan::PrintSuccessfulAllocatedMaxBits() { - auto it = memscope2rootStorageEntry.find(pto::AddressSpace::VEC); - if (it != memscope2rootStorageEntry.end()) { - if (!it->second) { - llvm::report_fatal_error("missing root storage entry for VEC scope"); - } - uint64_t ubAllocBits = it->second->alignedConstBits + it->second->bitsOffset; - for (auto& child : it->second->mergedChildren) { - ubAllocBits = std::max(ubAllocBits, child->bitsOffset + child->alignedConstBits); - } - llvm::outs() << "[PTOPlanMemory] Allocated UB size = " << ubAllocBits - << " bits\n"; - } -} - -void MemPlan::ValidateParameters(std::unique_ptr &e) const { - if (!e->bufInfo->operation) { - llvm::report_fatal_error("storage entry missing defining operation"); - } - if (e->bufInfo->constBits < 0U) { - llvm::report_fatal_error("storage entry has invalid memory size"); - } - if (e->bufferLifeVec.empty()) { - llvm::report_fatal_error("storage entry missing lifetime information"); - } -} - -void MemPlan::UpdateBuffer2Offsets() { - for (auto &e : StorageEntryVec) { - // Skip sibling (slot >= 1) entries -- their offsets are written via the - // primary entry's `relationOtherBuffers` traversal below. Without this - // skip the sibling offsets would be appended in StorageEntryVec order - // rather than slot order, breaking the runtime contract that - // `buffer2Offsets[buffer][k]` is slot k's physical offset. - if (e->isMultiBufferSlot) { - continue; - } - for (Value &buffer : e->inplaceBuffers) { - buffer2Offsets[buffer].push_back( - (e->bitsOffset + kBitsToByte - 1) / kBitsToByte); - // Multi-buffer primary: append sibling offsets in slot order so the - // final offsets list is [slot0, slot1, ..., slotN-1]. - for (auto *sibling : e->relationOtherBuffers) { - if (!sibling) { - continue; - } - buffer2Offsets[buffer].push_back( - (sibling->bitsOffset + kBitsToByte - 1) / kBitsToByte); - } - } - } - // In the MultiBuffer scenario, single reuse db will result in additional - // storageEntry. - UpdateMultiBufferReuseExtraOffset(); -} - -void MemPlan::UpdateMultiBufferReuseExtraOffset() { - if (pingEntry2RelationPongEntry.empty()) { - return; - } - - for (auto &relationEntry : pingEntry2RelationPongEntry) { - for (Value &buffer : relationEntry.second->inplaceBuffers) { - // MultiBuffer can cause multiple addrs. - buffer2Offsets[buffer].push_back( - (relationEntry.second->bitsOffset + kBitsToByte - 1) / - kBitsToByte); - } - } -} - -void MemPlan::MergeInplaceSE() { - // get the list of inplace value pair. - SmallVector inplaceList = GenerateInplaceList(); - // try to merge storage entries. genSE is replaced by KillSE. - for (const auto &pairIter : inplaceList) { - const StorageEntry *genSE = buffer2storageEntry[pairIter.first]; - StorageEntry *killSE = buffer2storageEntry[pairIter.second]; - if (genSE == killSE) { - // already same storageEntry, no need to inplace. - continue; - } - if (genSE == nullptr || killSE == nullptr) { - llvm::report_fatal_error("invalid storage entry during inplace merge"); - } - BufferLifeVec mergedBufferLifeVec; - mergedBufferLifeVec.insert(mergedBufferLifeVec.end(), - genSE->bufferLifeVec.begin(), - genSE->bufferLifeVec.end()); - mergedBufferLifeVec.insert(mergedBufferLifeVec.end(), - killSE->bufferLifeVec.begin(), - killSE->bufferLifeVec.end()); - MergeBufferVec(mergedBufferLifeVec); - killSE->bufferLifeVec.swap(mergedBufferLifeVec); - // merge allocs of two storage entry and update inplaceBuffers. - killSE->inplaceBuffers.insert(killSE->inplaceBuffers.begin(), - genSE->inplaceBuffers.begin(), - genSE->inplaceBuffers.end()); - - // Take the maximum value for the inplace scene. - killSE->multiBufferNum = - std::max(genSE->multiBufferNum, killSE->multiBufferNum); - - // all buffers have same storage entry after merging - for (auto &buffer : genSE->inplaceBuffers) { - buffer2storageEntry[buffer] = killSE; - } - // remove the alloc info of dst after successful merging - auto e = std::find_if(StorageEntryVec.begin(), StorageEntryVec.end(), - [genSE](std::unique_ptr &se) { - return se.get() == genSE; - }); - StorageEntryVec.erase(e); - } -} - -PlanStatus MemPlan::PlanLocalMemAddress() { - // merge from the first storage entry - MergeInplaceSE(); - dmaFirstPipelineOpt.build(func_); - ExpandMultiBufferStorageEntry(); - MergeSameScopeSE(); - return PlanMemAddressOfWholeLocalBuffer(); -} - -PlanStatus MemPlan::PlanWorkSpaceMemAddress() { - // merge from the first storage entry - MergeInplaceSE(); - ExpandMultiBufferStorageEntry(); - return PlanMemOffsetOfWholeWorkSpace(); -} - -PlanStatus MemPlan::PlanMemOffsetOfWholeWorkSpace() { - for (auto &it : workSpaceArg2rootStorageEntry) { - StorageEntry *rootStorageEntry = it.second; - if (!enableGlobalReuse) { - GlobalWorkspaceNoReuse(rootStorageEntry); - continue; - } - MemBoundList outline; - PlanRecHis history; - SpecInfo si; - // Can be reuse without conflicting life intervals. - si.specLevel = si.minLevel; - int childrenNum = static_cast(rootStorageEntry->mergedChildren.size()); - outline.push_back(std::make_shared( - BufferLifeVec(), 0, std::numeric_limits::max(), nullptr)); - - // The initial value is rootStorageEntry. - StorageEntry *curEntry = rootStorageEntry; - while (si.childIdx < childrenNum) { - curEntry->alignedConstBits = - static_cast(curEntry->bufInfo->constBits); - curEntry->childIdx = si.childIdx; - LogicalResult planResult = MultiSpecPlan(si, outline, history, curEntry); - if (failed(planResult)) { - return PlanStatus::PLAN_FAILED; - } - if (si.childIdx >= childrenNum) { - break; - } - curEntry = rootStorageEntry->mergedChildren[si.childIdx]; - } - } - planStatus = PlanStatus::PLAN_SUCCESS; - return planStatus; -} - -void MemPlan::GlobalWorkspaceNoReuse(StorageEntry *rootStorageEntry) { - rootStorageEntry->bitsOffset = 0; - uint64_t offset = static_cast(rootStorageEntry->bufInfo->constBits); - for (StorageEntry *child : rootStorageEntry->mergedChildren) { - child->bitsOffset = offset; - offset += static_cast(child->bufInfo->constBits); - } -} - -void MemPlan::ExpandMultiBufferStorageEntry() { - // For each multi-buffer primary entry, create (N - 1) sibling entries so - // the planner can lay out one physical slot per sibling. Siblings are - // pushed into `StorageEntryVec` and participate in normal Stage0/Stage2 - // address allocation. The primary keeps `relationOtherBuffers` pointing - // at the siblings in slot order (slot 1..N-1), and `relationPongEntry` - // aliases the first sibling so existing N == 2 codepaths keep working. - size_t size = StorageEntryVec.size(); - for (size_t i = 0; i < size; i++) { - auto *primary = StorageEntryVec[i].get(); - if (primary->multiBufferNum <= 1) { - continue; - } - uint32_t n = primary->multiBufferNum; - for (uint32_t slot = 1; slot < n; ++slot) { - auto entry = std::make_unique(); - entry->bufInfo = primary->bufInfo; - entry->bufferLifeVec = primary->bufferLifeVec; - entry->alignedConstBits = primary->alignedConstBits; - entry->inplaceBuffers = primary->inplaceBuffers; - entry->multiBufferNum = n; - entry->isMultiBufferSlot = true; - primary->relationOtherBuffers.push_back(entry.get()); - StorageEntryVec.push_back(std::move(entry)); - } - if (!primary->relationOtherBuffers.empty()) { - primary->relationPongEntry = primary->relationOtherBuffers.front(); - } - } -} - -bool MemPlan::IsEnoughForBuffersNoReuse(StorageEntry *rootStorageEntry, - size_t restBufferSize, - size_t alignUnit) { - auto iter = - bufferScope2RequiredSize.find(rootStorageEntry->bufInfo->bufferScope); - if (iter == bufferScope2RequiredSize.end()) { - llvm::report_fatal_error("missing required-size entry for buffer scope"); - } - if (iter->second < restBufferSize) { - // Even when the scope fits without reuse (no peak to save), honor - // largest-first placement so the option means the same thing on both paths: - // a deterministic decreasing-size layout regardless of whether reuse kicks - // in. Stable sort keeps uniform-size scopes byte-identical to the default. - if (orderBySize) { - rootStorageEntry = GetSizeOrderedRootStorageEntry(rootStorageEntry); - } - PlanBuffersWithoutReuse(rootStorageEntry, alignUnit); - return true; - } - return false; -} - -void MemPlan::PlanBuffersWithoutReuse(StorageEntry *rootStorageEntry, - size_t alignUnit) { - uint offset = 0; - rootStorageEntry->bitsOffset = offset; - offset = AlignUp(rootStorageEntry->bufInfo->constBits, alignUnit); - rootStorageEntry->alignedConstBits = offset; - for (StorageEntry *child : rootStorageEntry->mergedChildren) { - child->bitsOffset = offset; - uint64_t alignedBits = AlignUp(child->bufInfo->constBits, alignUnit); - offset += alignedBits; - child->alignedConstBits = alignedBits; - } -} - -void MemPlan::MergeSameScopeSE() { - // Construct root StorageEntry and collect the same scope StorageEntry - for (auto &iter : StorageEntryVec) { - auto iter_scope = - memscope2rootStorageEntry.find(iter->bufInfo->bufferScope); - if (iter_scope == memscope2rootStorageEntry.end()) { - memscope2rootStorageEntry[iter->bufInfo->bufferScope] = iter.get(); - } else { - iter_scope->second->mergedChildren.push_back(iter.get()); - } - } - - // set bufferScope2RequiredSize for all StorageEntry - for (auto &rootStorageEntry : memscope2rootStorageEntry) { - auto bufferSpaceInfo = GetBufferSpaceInfo(rootStorageEntry.first); - size_t accumulateSize = AlignUp(rootStorageEntry.second->bufInfo->constBits, - bufferSpaceInfo.first); - for (auto &childrenStorageEntry : rootStorageEntry.second->mergedChildren) { - size_t curStorageSize = AlignUp(childrenStorageEntry->bufInfo->constBits, - bufferSpaceInfo.first); - accumulateSize = accumulateSize + curStorageSize; - } - bufferScope2RequiredSize[rootStorageEntry.first] = accumulateSize; - } -} - -void MemPlan::PlanMemAddressForLevel0( - StorageEntry *rootStorageEntry) { - // get the buffer info for a given scope. - auto bufferSpaceInfo = - GetBufferSpaceInfo(rootStorageEntry->bufInfo->bufferScope); - size_t align = bufferSpaceInfo.first; - size_t maxBits = UINT64_MAX; - rootStorageEntry = GetReorderRootStorageEntry(rootStorageEntry); - // memory outline in a given buffer scope. - MemBoundList outline; - PlanRecHis history; - SpecInfo si; - si.specLevel = SPEC_LEVEL_0; - si.maxLevel = SPEC_LEVEL_0; - int childrenNum = static_cast(rootStorageEntry->mergedChildren.size()); - outline.push_back( - std::make_shared(BufferLifeVec(), 0, maxBits, nullptr)); - - // The initial value is rootStorageEntry. - StorageEntry *curEntry = rootStorageEntry; - while (si.childIdx < childrenNum) { - uint64_t needBits = static_cast(curEntry->bufInfo->constBits); - curEntry->alignedConstBits = AlignUp(needBits, align); - curEntry->childIdx = si.childIdx; - (void)MultiSpecPlan(si, outline, history, curEntry); - if (si.childIdx >= childrenNum) { - break; - } - curEntry = rootStorageEntry->mergedChildren[si.childIdx]; - } - // Find the max appled bits from all children and root, which is the max - // memory applied in this buffer space. - uint64_t maxAllocBits = rootStorageEntry->alignedConstBits; - auto children = rootStorageEntry->mergedChildren; - for (auto *child : children) { - maxAllocBits = - std::max(maxAllocBits, child->bitsOffset + child->alignedConstBits); - } - failApplyBufferInfo[rootStorageEntry->bufInfo->bufferScope] = maxAllocBits; -} - -PlanStatus MemPlan::PlanMemAddressOfWholeLocalBuffer() { - // Start plan - for (auto &it : memscope2rootStorageEntry) { - StorageEntry *rootStorageEntry = it.second; - // get the buffer info for a given scope. - auto bufferSpaceInfo = - GetBufferSpaceInfo(rootStorageEntry->bufInfo->bufferScope); - size_t align = bufferSpaceInfo.first; - size_t maxBits = bufferSpaceInfo.second; - if (rootStorageEntry->mergedChildren.empty()) { - PlanStatus status = PlanSingleLocalBuffer(rootStorageEntry, align, maxBits); - if (status != PlanStatus::PLAN_SUCCESS) { - return status; - } - continue; - } - if (IsEnoughForBuffersNoReuse(rootStorageEntry, maxBits, align)) { - continue; - } - PlanStatus status = PlanReusableLocalBuffer(rootStorageEntry, align, maxBits); - if (status != PlanStatus::PLAN_SUCCESS) { - return status; - } - } - planStatus = PlanStatus::PLAN_SUCCESS; - return planStatus; -} - -PlanStatus MemPlan::PlanSingleLocalBuffer(StorageEntry *rootStorageEntry, - size_t align, size_t maxBits) { - uint64_t needAlignedBits = AlignUp(rootStorageEntry->bufInfo->constBits, align); - if (needAlignedBits > maxBits) { - failApplyBufferInfo[rootStorageEntry->bufInfo->bufferScope] = - needAlignedBits; - return PlanStatus::PLAN_FAILED; - } - rootStorageEntry->bitsOffset = 0; - rootStorageEntry->alignedConstBits = needAlignedBits; - return PlanStatus::PLAN_SUCCESS; -} - -PlanStatus MemPlan::PlanReusableLocalBuffer(StorageEntry *rootStorageEntry, - size_t align, size_t maxBits) { - rootStorageEntry = GetReorderRootStorageEntry(rootStorageEntry); - ReportMemLifeDebugInfo(rootStorageEntry); - - MemBoundList outline; - PlanRecHis history; - SpecInfo si; - si.specLevel = si.maxLevel; - int childrenNum = static_cast(rootStorageEntry->mergedChildren.size()); - outline.push_back( - std::make_shared(BufferLifeVec(), 0, maxBits, nullptr)); - - StorageEntry *curEntry = rootStorageEntry; - while (si.childIdx < childrenNum) { - uint64_t needBits = static_cast(curEntry->bufInfo->constBits); - curEntry->alignedConstBits = AlignUp(needBits, align); - curEntry->childIdx = si.childIdx; - LDBG("\n"); - LDBG("----------Need-Plan-CurEntry---------\n"); - ReportCurEntryDebugInfo(curEntry); - LDBG("\n"); - LogicalResult planResult = MultiSpecPlan(si, outline, history, curEntry); - if (failed(planResult)) { - StatusWrapper statusWrapper = {false, curEntry->alignedConstBits, - &si, outline, - history, rootStorageEntry}; - LDBG("\n"); - LDBG("----------ApplyFailStrategy---------\n"); - ReportCurEntryDebugInfo(curEntry); - LDBG("\n"); - PlanStatus status = ApplyFailStrategy(statusWrapper, maxBits); - if (status == PlanStatus::RESTART_NEW_PLAN) { - si = SpecInfo(); - curEntry = rootStorageEntry; - continue; - } - if (status == PlanStatus::PLAN_FAILED) { - ReportAllocatedEntryDebugInfo(rootStorageEntry); - PlanMemAddressForLevel0(rootStorageEntry); - return status; - } - } - if (si.childIdx >= childrenNum) { - break; - } - curEntry = rootStorageEntry->mergedChildren[si.childIdx]; - } - return PlanStatus::PLAN_SUCCESS; -} - -void MemPlan::ReportMemLifeDebugInfo(StorageEntry *rootStorageEntry) { - LDBG("-------------------------- Buffer2Life --------------------------\n"); - MemLifeDebugInfo(rootStorageEntry); - for (auto &StorageEntry : rootStorageEntry->mergedChildren) { - MemLifeDebugInfo(StorageEntry); - } -} - -void MemPlan::MemLifeDebugInfo(StorageEntry *storageEntry) { - for (auto &buffer : storageEntry->inplaceBuffers) { - (void)buffer; - LDBG("Buffer : " << buffer << "\n"); - } - for (auto &bufferLife : storageEntry->bufferLifeVec) { - (void)bufferLife; - LDBG("bufferLife : " - << "allocTime : " << bufferLife->allocTime - << " , freeTime : " << bufferLife->freeTime << "\n"); - } - LDBG("\n"); -} - -void MemPlan::ReportCurEntryDebugInfo(const StorageEntry *curEntry) { - for (auto &buffer : curEntry->inplaceBuffers) { - (void)buffer; - LDBG("buffer : " << buffer); - } -} - -StorageEntry * -MemPlan::GetReorderRootStorageEntry(StorageEntry *rootStorageEntry) { - if (orderBySize) { - return GetSizeOrderedRootStorageEntry(rootStorageEntry); - } - if (rootStorageEntry->bufInfo->bufferScope != pto::AddressSpace::VEC) { - return rootStorageEntry; - } - SmallVector origStorageEntryVec = {rootStorageEntry}; - origStorageEntryVec.insert(origStorageEntryVec.end(), - rootStorageEntry->mergedChildren.begin(), - rootStorageEntry->mergedChildren.end()); - - // reorder storage entrys: dma touched buffers + other buffers + scalar - // touched buffers - SmallVector reorderedStorageEntryVec; - SmallVector touchPipeScalarStorageEntryVec; - for (auto &storageEntry : origStorageEntryVec) { - for (auto &buffer : storageEntry->inplaceBuffers) { - if (dmaFirstPipelineOpt.IsDmaBuffer(buffer)) { - reorderedStorageEntryVec.push_back(storageEntry); - break; - } - if (dmaFirstPipelineOpt.IsScalarBuffer(buffer)) { - touchPipeScalarStorageEntryVec.push_back(storageEntry); - break; - } - } - } - for (auto &storageEntry : origStorageEntryVec) { - auto it1 = std::find(reorderedStorageEntryVec.begin(), - reorderedStorageEntryVec.end(), storageEntry); - auto it2 = std::find(touchPipeScalarStorageEntryVec.begin(), - touchPipeScalarStorageEntryVec.end(), storageEntry); - if (it1 == reorderedStorageEntryVec.end() && - it2 == touchPipeScalarStorageEntryVec.end()) { - reorderedStorageEntryVec.push_back(storageEntry); - } - } - - reorderedStorageEntryVec.insert(reorderedStorageEntryVec.end(), - touchPipeScalarStorageEntryVec.begin(), - touchPipeScalarStorageEntryVec.end()); - - // Ensure that ping pong is continuously plan mem in the multi buffer. - ReorderContinuousPingPongEntry(reorderedStorageEntryVec); - StorageEntry *reorderedRootStorageEntry = reorderedStorageEntryVec[0]; - reorderedRootStorageEntry->mergedChildren.clear(); - for (size_t j = 1; j < reorderedStorageEntryVec.size(); ++j) { - reorderedRootStorageEntry->mergedChildren.push_back( - reorderedStorageEntryVec[j]); - } - return reorderedRootStorageEntry; -} - -void MemPlan::ReorderContinuousPingPongEntry( - SmallVector &storageEntryVec) { - SmallVector reorderedStorageEntryVec; - for (auto &storageEntry : storageEntryVec) { - auto it = std::find(reorderedStorageEntryVec.begin(), - reorderedStorageEntryVec.end(), storageEntry); - if (it == reorderedStorageEntryVec.end()) { - reorderedStorageEntryVec.push_back(storageEntry); - if (storageEntry->multiBufferNum == kDoubleBufferCount && - storageEntry->relationPongEntry) { - // Ping Pong continuous save. - reorderedStorageEntryVec.push_back(storageEntry->relationPongEntry); - } - } - } - reorderedStorageEntryVec.swap(storageEntryVec); -} - -StorageEntry * -MemPlan::GetSizeOrderedRootStorageEntry(StorageEntry *rootStorageEntry) { - // First-fit-decreasing: place the largest buffers first. For the heterogeneous - // buffer sizes that real kernels produce, decreasing-size order packs tighter - // than an arbitrary/DMA-first order (this is the ordering XLA, TVM and SOMAS - // all use). Applies to every memory space, unlike the DMA-first reorder which - // is VEC-only. - SmallVector entries = {rootStorageEntry}; - entries.insert(entries.end(), rootStorageEntry->mergedChildren.begin(), - rootStorageEntry->mergedChildren.end()); - - // Stable sort by decreasing buffer size. Stable keeps the original order among - // equal-size buffers, so uniform-size instances (e.g. the plan_memory_* tests) - // are left untouched. - std::stable_sort(entries.begin(), entries.end(), - [](const StorageEntry *a, const StorageEntry *b) { - return a->bufInfo->constBits > b->bufInfo->constBits; - }); - - // Keep ping-pong (double-buffer) pairs contiguous so double-buffering is - // preserved (same post-processing the DMA-first path applies). - ReorderContinuousPingPongEntry(entries); - - // Rebuild the flat root -> children structure around the new (largest) root. - // Clear every entry's child list first: when the root changes, the previous - // root would otherwise keep its stale child list (forming a cycle), and only - // the new root should carry the flat list of the others. - StorageEntry *reorderedRootStorageEntry = entries[0]; - for (StorageEntry *entry : entries) { - entry->mergedChildren.clear(); - } - for (size_t j = 1; j < entries.size(); ++j) { - reorderedRootStorageEntry->mergedChildren.push_back(entries[j]); - } - // Keep the scope -> root map consistent so later consumers (RecordOverflowIfAny, - // PrintSuccessfulAllocatedMaxBits) read the new root and its full child list. - // This must accompany the clear above: clearing children without updating the - // map would leave the stale root pointing at an empty child list. - memscope2rootStorageEntry[reorderedRootStorageEntry->bufInfo->bufferScope] = - reorderedRootStorageEntry; - return reorderedRootStorageEntry; -} - -std::pair -MemPlan::GetBufferSpaceInfo(pto::AddressSpace &space) const { - switch (space) { - case pto::AddressSpace::VEC: - return std::make_pair(plannerAlignBitsFromBytes(ubAlignSize), ubSpaceSize); - case pto::AddressSpace::MAT: - return std::make_pair(plannerAlignBitsFromBytes(l1AlignSize), l1SpaceSize); - case pto::AddressSpace::ACC: - return std::make_pair(plannerAlignBitsFromBytes(l0cAlignSize), - l0cSpaceSize); - case pto::AddressSpace::LEFT: - return std::make_pair(plannerAlignBitsFromBytes(l0aAlignSize), - l0aSpaceSize); - case pto::AddressSpace::RIGHT: - return std::make_pair(plannerAlignBitsFromBytes(l0bAlignSize), - l0bSpaceSize); - case pto::AddressSpace::BIAS: - return std::make_pair(plannerAlignBitsFromBytes(biasAlignSize), - biasSpaceSize); - case pto::AddressSpace::SCALING: - return std::make_pair(plannerAlignBitsFromBytes(scalingAlignSize), - scalingSpaceSize); - case pto::AddressSpace::Zero: - case pto::AddressSpace::GM: - return std::make_pair(size_t{0}, size_t{0}); - } - - llvm_unreachable("Temporarily unsupported memory buffer space !"); -} - -LogicalResult MemPlan::MultiSpecPlan(SpecInfo &si, MemBoundList &outline, - PlanRecHis &history, StorageEntry *entry) { - LogicalResult planResult = failure(); - for (int i = si.specLevel; i >= si.minLevel; i--) { - planResult = SpecAlloc(outline, history, entry, si, i); - if (succeeded(planResult)) { - if (si.childIdx == si.specStartIdx) { - // In roll back plan, when the specified specStartIdx is reached, - // the subsequent plan still adopts the maxLevel strategy. - si.specLevel = si.maxLevel; - } - si.childIdx++; - break; - } - } - return planResult; -} - -LogicalResult MemPlan::SpecAlloc(MemBoundList &outline, PlanRecHis &his, - StorageEntry *e, const SpecInfo &si, - int localLevel) { - if (std::any_of(his.begin(), his.end(), - [e](PlanRecord &r) { return r.entry && r.entry == e; })) { - // If the plan has already been completed, return success directly. - return success(); - } - for (MemBoundListConstIter start = outline.begin(); start != outline.end(); - ++start) { - uint64_t size = 0; - uint64_t allocOffset = (*start)->offset; - for (MemBoundListConstIter end = start; end != outline.end(); ++end) { - std::shared_ptr last = *end; - size += last->extent; - // if index & addr are as same as last rollback result, - // continue to find next result - if (IsSamePlanAsLastRollBack(allocOffset, e->childIdx, si) || - VerifyConflictStage0(e, last)) { - start = end; - break; - } - if (size < e->alignedConstBits) { - continue; - } - // If SPEC_LEVEL_1, then the address of pong Offset address needs to be - // allocated. - uint64_t pongOffset{0}; - if (localLevel == SPEC_LEVEL_1 && - VerifyConflictStage1(outline, his, e, - OutlineSectionInfo(start, end, size, false), - pongOffset)) { - break; - } - - if (VerifyConflictStage2(his, e, localLevel, start, outline)) { - break; - } - e->bitsOffset = allocOffset; - UpdateOutline(outline, his, e, - OutlineSectionInfo(start, end, size, false), localLevel); - - if (localLevel == SPEC_LEVEL_1) { - // There is no conflict with the historical plan of buffer life, and - // the address of the Pong can be assigned. - PlanRelationPongEntryAddress(pongOffset, e); - SpecAllocRelationPongEntry(outline, his, e, pongOffset); - } - LDBG("APPLY_SPEC_LEVEL: " << localLevel << "\n"); - bool needRecord = - allocatedEntry.end() == - std::find(allocatedEntry.begin(), allocatedEntry.end(), e); - if (needRecord) { - allocatedEntry.push_back(e); - } - return success(); - } - } - return failure(); -} - -LoopLikeOpInterface -MemPlan::GetBufferParentLoop(const SmallVector &buffers) { - llvm::SmallSet parentLoopVec; - for (auto buffer : buffers) { - if (!buffer.getDefiningOp()) { - if (!isa( - buffer.getParentBlock()->getParentOp())) { - llvm::report_fatal_error("expected loop-carried block argument"); - } - // Init args and region iter arg are inplace, ignore Region Iter Arg - // without DefineOp. - continue; - } - LoopLikeOpInterface bufferParentLoop = getParentLoop(buffer); - if (bufferParentLoop) { - parentLoopVec.insert(bufferParentLoop); - } else { - return nullptr; - } - } - if (parentLoopVec.size() == 1) { - return *parentLoopVec.begin(); - } - return nullptr; -} - -bool MemPlan::VerifyConflictStage1(MemBoundList &outline, PlanRecHis &his, - StorageEntry *e, - const OutlineSectionInfo &outlineInfo, - uint64_t &pongOffset) { - if (outlineInfo.mem_start != outlineInfo.mem_end) { - return true; - } - auto reuseBoundStorageEntry = (*outlineInfo.mem_start)->lastStorageEntry; - if (!reuseBoundStorageEntry) { - // This area has not been planed, so there is no need to consider it. - return true; - } - - StorageEntry *multiRelationPongEntry = - GetMultiRelationPongEntry(reuseBoundStorageEntry); - if (multiRelationPongEntry) { - if (e->multiBufferNum == kSingleBufferCount || - (e->multiBufferNum == kDoubleBufferCount && e->relationPongEntry && - (e->relationPongEntry->bitsOffset != 0))) { - auto parentLoop1 = GetBufferParentLoop(e->inplaceBuffers); - auto parentLoop2 = - GetBufferParentLoop(reuseBoundStorageEntry->inplaceBuffers); - if (!(parentLoop1 != nullptr && parentLoop2 != nullptr && - parentLoop1 == parentLoop2)) { - // Cannot be reused under the same for. - return true; - } - // There are two situations: - // 1. Single reuse DB. - // 2. DB reuse DB. - pongOffset = multiRelationPongEntry->bitsOffset; - bool conflict = std::any_of( - his.begin(), his.end(), [pongOffset, e, this](PlanRecord &r) { - return this->IsBufferLifeVecConflict(r, pongOffset, e); - }); - if (!conflict) { - return false; - } - } - } - return true; -} - -StorageEntry * -MemPlan::GetMultiRelationPongEntry(const StorageEntry *reuseBoundStorageEntry) { - if (reuseBoundStorageEntry->multiBufferNum == kDoubleBufferCount && - reuseBoundStorageEntry->relationPongEntry && - (reuseBoundStorageEntry->relationPongEntry->bitsOffset != 0)) { - // If the reuseBoundStorageEntry itself requires db, directly match and - // return relationPongEntry. - return reuseBoundStorageEntry->relationPongEntry; - } - auto iter = pingEntry2RelationPongEntry.find(reuseBoundStorageEntry); - if (iter != pingEntry2RelationPongEntry.end()) { - // If the reuseBoundStorageEntry itself is single, but has already been - // reused with db and has an extra pong StorageEntry is added. - return iter->second.get(); - } - return nullptr; -} - -void MemPlan::SpecAllocRelationPongEntry(MemBoundList &outline, PlanRecHis &his, - StorageEntry *e, uint64_t offset) { - for (MemBoundListConstIter start = outline.begin(); start != outline.end(); - ++start) { - uint64_t size = 0; - // Find the MemBound corresponding to the Pong offset. - if ((*start)->offset != offset) { - continue; - } - for (MemBoundListConstIter end = start; end != outline.end(); ++end) { - std::shared_ptr last = *end; - size += last->extent; - if (size < e->alignedConstBits) { - continue; - } - StorageEntry *pongStorageEntry = nullptr; - auto iter = pingEntry2RelationPongEntry.find(e); - if (iter != pingEntry2RelationPongEntry.end()) { - pongStorageEntry = iter->second.get(); - } - if (e->multiBufferNum == kDoubleBufferCount && e->relationPongEntry) { - pongStorageEntry = e->relationPongEntry; - } - if (!pongStorageEntry) { - llvm::report_fatal_error("pong storage entry not found"); - } - UpdateOutline(outline, his, pongStorageEntry, - OutlineSectionInfo(start, end, size, true), SPEC_LEVEL_1); - return; - } - } -} - -bool MemPlan::IsBufferLifeVecConflict(PlanRecord &r, uint64_t offset, - const StorageEntry *e) const { - if ((r.firstMemBound->offset + r.allExtent > offset) && - (r.firstMemBound->offset < offset + e->alignedConstBits)) { - if (HasSemanticConflict(e, r.firstMemBound->bufferLifeVec)) { - return true; - } - DenseMap intersection = - GetOverlapBufferLife(r.entry->bufferLifeVec, e->bufferLifeVec); - return !intersection.empty(); - } - return false; -} - -void MemPlan::PlanRelationPongEntryAddress(uint64_t offset, StorageEntry *e) { - if (e->multiBufferNum == kSingleBufferCount) { - std::unique_ptr entry = std::make_unique(); - entry->bufInfo = e->bufInfo; - entry->bufferLifeVec = e->bufferLifeVec; - entry->alignedConstBits = e->alignedConstBits; - entry->inplaceBuffers = e->inplaceBuffers; - entry->multiBufferNum = e->multiBufferNum; - entry->bitsOffset = offset; - pingEntry2RelationPongEntry[e] = std::move(entry); - } else if (e->multiBufferNum == kDoubleBufferCount) { - e->relationPongEntry->bitsOffset = offset; - } - // N > 2: the Stage1 "place ping next to a free pong slot" optimization is - // not modeled for the general N-way case in this release. Sibling entries - // get their own addresses via the normal Stage0/Stage2 paths in - // `PlanReusableLocalBuffer` / `PlanSingleLocalBuffer`. This branch is a - // no-op rather than an unreachable so the planner can keep making forward - // progress on N > 2 inputs. -} - -bool MemPlan::VerifyConflictStage2(PlanRecHis &his, const StorageEntry *e, - int specLevel, MemBoundListConstIter &start, - const MemBoundList &outline) { - if (specLevel != SPEC_LEVEL_2) { - return false; - } - bool touchMemCanUse = false; - MemBoundListConstIter foundMem; - - for (auto iter = start; iter != outline.end(); ++iter) { - uint64_t offset = (*iter)->offset; - bool conflict = - std::any_of(his.begin(), his.end(), [offset, e, this](PlanRecord &r) { - return (r.firstMemBound->offset + r.allExtent > offset) && - (r.firstMemBound->offset < offset + e->alignedConstBits) && - this->PipeConflict(r.entry, e, this->pipeDmaConflictMap); - }); - // if conflict, continue finding the first bound that has no conflict - // if last bound do not meet the size, continue - if (conflict || - ((*iter == outline.back()) && (*iter)->extent < e->alignedConstBits)) { - continue; - } - touchMemCanUse = true; - foundMem = iter; - break; - } - - if (touchMemCanUse) { - bool conflict = (foundMem != start); - start = conflict ? --foundMem : start; - return conflict; - } - // if cannot find a bound that has no conflict with current entry, - return true; -} - -bool MemPlan::PipeConflict(const StorageEntry *e1, const StorageEntry *e2, - DenseMap &conflictMap) { - if (e1 == nullptr || e2 == nullptr) { - return false; - } - auto sePair = std::make_pair(e1, e2); - auto [iter, isInserted] = conflictMap.try_emplace(sePair, false); - if (!isInserted) { - return iter->second; - } - - for (const Value var1 : e1->inplaceBuffers) { - for (const Value var2 : e2->inplaceBuffers) { - bool conflict = dmaFirstPipelineOpt.BufferPipeConflict(var1, var2); - if (conflict) { - iter->second = true; - return true; - } - } - } - return false; -} - -void MemPlan::UpdateOutline(MemBoundList &outline, PlanRecHis &his, - StorageEntry *e, - const OutlineSectionInfo &outlineInfo, - int localLevel) const { - auto start = outlineInfo.mem_start; - MemBoundListConstIter end = outlineInfo.mem_end; - // outline: - // |-------start+end-------------| - // |--head--|--split e--|--tail--| - uint64_t res = outlineInfo.size - e->alignedConstBits; - std::shared_ptr last = *end; - ++end; - std::shared_ptr bound; - SmallVector> splitBound; - // split e, to get Boundbound - if (splitOutline) { - // add splitBound by splitting e to section - AddMemBoundInSectionalWay(e, start, end, splitBound); - } else { - // origin outline - BufferLifeVec life(e->bufferLifeVec.begin(), e->bufferLifeVec.end()); - MergeBufferLife(start, end, life); - splitBound.emplace_back(std::make_shared( - life, e->bitsOffset, e->alignedConstBits, e)); - } - - // insert tail memory bound - if (res > 0) { - bound = std::make_shared(last->bufferLifeVec, - last->offset + last->extent - res, - res, last->lastStorageEntry); - end = outline.insert(end, bound); - } - // insert split e memory bound - for (int i = static_cast(splitBound.size()) - 1; i >= 0; --i) { - end = outline.insert(end, splitBound[i]); - } - // record the current plan of first split entry in his - his.emplace_back(PlanRecord{localLevel, - e->childIdx, - res > 0, - false, - splitBound.size(), - e, - e->alignedConstBits, - splitBound[0], - {}, - outlineInfo.isDirectlyRollback}); - PlanRecord &r = his.back(); - r.replaced.splice(r.replaced.begin(), outline, start, end); -} - -void MemPlan::AddMemBoundInSectionalWay( - StorageEntry *e, MemBoundListConstIter start, MemBoundListConstIter end, - SmallVector> &splitBound) const { - // |--outline1--|--outline2--|--outline3--| - // |---------e------------ | - // |--split e1 -|-split e2-| - for (auto iter = start; iter != end; ++iter) { - BufferLifeVec life(e->bufferLifeVec.begin(), e->bufferLifeVec.end()); - life.insert(life.end(), (*iter)->bufferLifeVec.begin(), - (*iter)->bufferLifeVec.end()); - // merge the buffer life - MergeBufferVec(life); - // get the extent - uint64_t size = 0; - if (std::distance(start, iter) == std::distance(start, end) - 1) { - // deal with the last split e2 - size = e->bitsOffset + e->alignedConstBits - (*iter)->offset; - } else { - size = (*iter)->extent; - } - splitBound.emplace_back( - std::make_shared(life, (*iter)->offset, size, e)); - } -} - -inline void MemPlan::MergeBufferLife(MemBoundList::const_iterator start, - MemBoundList::const_iterator end, - BufferLifeVec &newLife) const { - size_t size = 0; - for (auto it = start; it != end; ++it) { - size += (*it)->bufferLifeVec.size(); - } - newLife.reserve(size); - for (auto it = start; it != end; ++it) { - newLife.insert(newLife.end(), (*it)->bufferLifeVec.begin(), - (*it)->bufferLifeVec.end()); - } - MergeBufferVec(newLife); -} - -void MemPlan::MergeBufferVec(BufferLifeVec &bufferLife) const { - if (bufferLife.empty()) { - return; - } - BufferLifeVec mergedLife; - mergedLife.reserve(bufferLife.size()); - // sort life by alloc and free time - std::sort(bufferLife.begin(), bufferLife.end(), CompareBufferLife()); - int start = bufferLife[0]->allocTime; - int end = bufferLife[0]->freeTime; - auto buffer = bufferLife[0]->buffer; - // merge life - for (size_t i = 1; i < bufferLife.size(); i++) { - auto &life = bufferLife[i]; - if (life->allocTime <= end + 1) { - end = end < life->freeTime ? life->freeTime : end; - } else { - mergedLife.emplace_back(std::make_unique(buffer, start, end)); - buffer = life->buffer; - start = life->allocTime; - end = life->freeTime; - } - } - mergedLife.emplace_back(std::make_unique(buffer, start, end)); - bufferLife.swap(mergedLife); -} - -bool MemPlan::IsSamePlanAsLastRollBack(uint64_t allocOffset, int curChildIdx, - const SpecInfo &si) const { - return curChildIdx == si.rollbackIdx && allocOffset == si.rollbackAddr; -} - -// spec_level == SPEC_LEVEL_0 -inline bool -MemPlan::VerifyConflictStage0(StorageEntry *e, - const std::shared_ptr &last) { - if (HasSemanticConflict(e, last->bufferLifeVec)) { - return true; - } - // level_0: offset = 0, offset means life distance - DenseMap intersection = - GetOverlapBufferLife(e->bufferLifeVec, last->bufferLifeVec); - return !intersection.empty(); -} - -// verify two buffer life vectors is conflict or not -// The key pair looks like the following diagram -// indicate that var1 is generated later than var2. -// buffer2 -// --- PLAN_time -// buffer1 intersected | | -// buffer_life | | -// PLAN_time --- lo --- --- -// | | --- | | -// | | --- | | -// --- hi --- --- free_time -// | | -// | | -// free_time --- -// Meantime, the overlap is the intersected buffer_life. -DenseMap -MemPlan::GetOverlapBufferLife(const BufferLifeVec &b1, - const BufferLifeVec &b2) const { - DenseMap intersection; - size_t i = 0; - size_t j = 0; - size_t b1Len = b1.size(); - size_t b2Len = b2.size(); - if (b1Len == 0 || b2Len == 0) { - return intersection; - } - while (i < b1Len && j < b2Len) { - auto lo = std::max(b1[i]->allocTime, b2[j]->allocTime); - auto hi = std::min(b1[i]->freeTime, b2[j]->freeTime); - if (lo <= hi) { - BufferLife bufferLife(nullptr, lo, hi); - ValuePair key = - lo == b1[i]->allocTime && hi == b2[j]->freeTime - ? std::make_pair(b1[i]->buffer, - b2[j]->buffer) // case in the diagram - : std::make_pair(b2[j]->buffer, b1[i]->buffer); // opposing case - intersection.try_emplace(key, bufferLife); - } - if (b1[i]->freeTime < b2[j]->freeTime) { - i += 1; - } else { - j += 1; - } - } - return intersection; -} - -PlanStatus MemPlan::ApplyFailStrategy(StatusWrapper &statusWrapper, - const size_t maxBits) { - RollBackForAllocFail(statusWrapper, maxBits); - // second class rollback, level 1 --> 0 - if (statusWrapper.si->specLevel > SPEC_LEVEL_0 && - statusWrapper.si->childIdx >= 0) { - statusWrapper.si->specLevel--; - return PlanStatus::CONTINUE_PLAN; - } - if (!splitOutline) { - // roll back to origin again, enable split outline. - splitOutline = true; - return PlanStatus::RESTART_NEW_PLAN; - } - return PlanStatus::PLAN_FAILED; -} - -void MemPlan::ReportAllocatedEntryDebugInfo(StorageEntry *rootStorageEntry) { - auto printRecord = [this](const StorageEntry *entry) { - uint64_t needByte = - (entry->alignedConstBits + kBitsToByte - 1) / kBitsToByte; - uint64_t offsetByte = - (entry->bitsOffset + kBitsToByte - 1) / kBitsToByte; - (void)needByte; - (void)offsetByte; - ReportCurEntryDebugInfo(entry); - LDBG(", offset: " << offsetByte); - LDBG(", extent: " << needByte); - LDBG(", buffer life: "); - for (auto &bufferLife : entry->bufferLifeVec) { - (void)bufferLife; - LDBG("[" << bufferLife->allocTime << "-" << bufferLife->freeTime - << "], "); - } - }; - LDBG("--------------------------BUFFER ALLOCATE " - "START-------------------------- " - << "\n" - << "\n"); - LDBG(" BUFFER ALLOCATE START: UB" - << "\n"); - if (!allocatedEntry.empty()) { - for (auto &entry : allocatedEntry) { - printRecord(entry); - LDBG("\n"); - } - size_t num = allocatedEntry.size() - 1; - if (rootStorageEntry->mergedChildren.size() <= num) { - llvm::report_fatal_error("missing failed storage entry"); - } - const StorageEntry *failedSe = rootStorageEntry->mergedChildren[num]; - printRecord(failedSe); - LDBG("alloc fail,because exceed bound of memory \n" - << " BUFFER ALLOCATE END \n"); - LDBG("\n" - << "--------------------------BUFFER ALLOCATE " - "END-------------------------- " - << "\n"); - } -} - -LogicalResult MemPlan::InitMemSpecsFromModule(func::FuncOp funcOp) { - struct MemSpec { - int ubSpaceSize; - int l1SpaceSize; - int l0aSpaceSize; - int l0bSpaceSize; - int l0cSpaceSize; - int ubAlignSize; - int l1AlignSize; - int l0cAlignSize; - int l0aAlignSize; - int l0bAlignSize; - int biasAlignSize; - int biasSpaceSize; - int scalingAlignSize; - int scalingSpaceSize; - }; - - const MemSpec kA3 = { - 1572864, 4194304, 524288, 524288, 1048576, 256, 256, - 4096, 4096, 4096, 256, 524288, 256, 1572864}; - const MemSpec kA5 = { - 2031616, 4194304, 524288, 524288, 2097152, 256, 256, - 4096, 4096, 4096, 256, 524288, 256, 2031616}; - - auto applySpec = [this](const MemSpec &spec) { - ubSpaceSize = spec.ubSpaceSize; - l1SpaceSize = spec.l1SpaceSize; - l0aSpaceSize = spec.l0aSpaceSize; - l0bSpaceSize = spec.l0bSpaceSize; - l0cSpaceSize = spec.l0cSpaceSize; - ubAlignSize = spec.ubAlignSize; - l1AlignSize = spec.l1AlignSize; - l0cAlignSize = spec.l0cAlignSize; - l0aAlignSize = spec.l0aAlignSize; - l0bAlignSize = spec.l0bAlignSize; - biasAlignSize = spec.biasAlignSize; - biasSpaceSize = spec.biasSpaceSize; - scalingAlignSize = spec.scalingAlignSize; - scalingSpaceSize = spec.scalingSpaceSize; - }; - - // Default to a3. - applySpec(kA3); - - // --pto-arch options: - // a3 -> default memory spec - // a5 -> override memory spec - if (isTargetArchA5(getTopLevelModuleOp(funcOp))) { - applySpec(kA5); - } - return success(); -} - -void MemPlan::RollBackForAllocFail(StatusWrapper &statusWrapper, - const size_t maxBits) { - while (ContinueRollBack(statusWrapper)) { - RollBackForAllocFailInner(statusWrapper, maxBits); - } -} - -bool MemPlan::ContinueRollBack(const StatusWrapper &statusWrapper) const { - return (!statusWrapper.hasEnoughRollBackSize) && - (!statusWrapper.history.empty() && (!statusWrapper.outline.empty())); -} - -void MemPlan::RollBackForAllocFailInner(StatusWrapper &statusWrapper, - const size_t maxBits) { - auto &si = statusWrapper.si; - if (si->childIdx > si->specStartIdx) { - si->specStartIdx = si->childIdx; - } - // Check whether the container is empty before accessing "history" - while (!statusWrapper.history.empty()) { - PlanRecord r = - RollbackOutline(statusWrapper.history, statusWrapper.outline); - auto iter = pingEntry2RelationPongEntry.find(r.entry); - if (iter != pingEntry2RelationPongEntry.end()) { - pingEntry2RelationPongEntry.erase(iter); - } - if (r.isDirectlyRollback || - (r.entry->multiBufferNum == kDoubleBufferCount && - !r.entry->relationPongEntry)) { - continue; - } - si->childIdx = r.childIdx; - si->specLevel = r.specLevel; - if (si->specLevel > si->minLevel) { - // record rollback info: index and address - si->rollbackAddr = - si->childIdx == -1 - ? UINT64_MAX - : statusWrapper.RootE->mergedChildren[si->childIdx]->bitsOffset; - si->rollbackIdx = si->childIdx; - if (statusWrapper.si->rollbackAddr + statusWrapper.alignedConstBits > - maxBits) { - continue; - } - statusWrapper.hasEnoughRollBackSize = true; - break; - } - } -} - -PlanRecord MemPlan::RollbackOutline(PlanRecHis &history, - MemBoundList &outline) const { - auto r = history.back(); - auto it = std::find(outline.begin(), outline.end(), r.firstMemBound); - // |--head--|--split entry--|--tail--| - // erase head - if (r.headed) { - it--; - it = outline.erase(it); - } - // erase split entry - for (size_t i = 0; i < r.splitNums; i++) { - it = outline.erase(it); - } - // erase tail - if (r.tailed) { - it = outline.erase(it); - } - // restore outline and replaced - outline.splice(it, r.replaced); - history.pop_back(); - return r; -} - -namespace { - -class LegacyAllocTileOpAddPlannedAddressPattern - : public OpRewritePattern { -public: - explicit LegacyAllocTileOpAddPlannedAddressPattern( - MLIRContext *context, - DenseMap> buffer2Offsets) - : OpRewritePattern(context), - buffer2Offsets(std::move(buffer2Offsets)) {} - - LogicalResult matchAndRewrite(pto::AllocTileOp op, - PatternRewriter &rewriter) const override { - if (op.getAddr()) { - return failure(); - } - - auto tileType = dyn_cast(op.getResult().getType()); - if (!tileType) { - return failure(); - } - - auto it = buffer2Offsets.find(op.getResult()); - if (it == buffer2Offsets.end() || it->second.empty()) { - return failure(); - } - - if (it->second.size() != 1) { - return rewriter.notifyMatchFailure( - op, "single alloc_tile root expects exactly one planned address"); - } - - Value addr = rewriter.create(op.getLoc(), - it->second.front(), 64); - auto planned = rewriter.create( - op.getLoc(), tileType, addr, - op.getValidRow() ? op.getValidRow() : Value(), - op.getValidCol() ? op.getValidCol() : Value()); - for (NamedAttribute attr : op->getAttrs()) { - if (attr.getName().getValue() == "operandSegmentSizes") { - continue; - } - planned->setAttr(attr.getName(), attr.getValue()); - } - - rewriter.replaceOp(op, planned.getResult()); - return success(); - } - -private: - DenseMap> buffer2Offsets; -}; - -class LegacyAllocMultiTileOpAddPlannedAddressesPattern - : public OpRewritePattern { -public: - explicit LegacyAllocMultiTileOpAddPlannedAddressesPattern( - MLIRContext *context, - DenseMap> buffer2Offsets) - : OpRewritePattern(context), - buffer2Offsets(std::move(buffer2Offsets)) {} - - LogicalResult matchAndRewrite(pto::AllocMultiTileOp op, - PatternRewriter &rewriter) const override { - if (op.getAddr() || op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) { - return failure(); - } - auto it = buffer2Offsets.find(op.getResult()); - if (it == buffer2Offsets.end() || it->second.empty()) { - return failure(); - } - if (it->second.size() != op.getResult().getType().getCount()) { - return rewriter.notifyMatchFailure( - op, "planned address count does not match multi_tile_buf count"); - } - - SmallVector addrs; - addrs.reserve(it->second.size()); - for (uint64_t offset : it->second) { - addrs.push_back(static_cast(offset)); - } - rewriter.modifyOpInPlace(op, [&] { - op->setAttr(pto::kPtoMultiBufferAddrsAttrName, - rewriter.getDenseI64ArrayAttr(addrs)); - }); - return success(); - } - -private: - DenseMap> buffer2Offsets; -}; - -static FailureOr parseLegacyMemPlanMode(func::FuncOp func, - llvm::StringRef memMode) { - if (memMode.equals_insensitive("local") || - memMode.equals_insensitive("local-mem-plan")) { - return MemPlanMode::LOCAL_MEM_PLAN; - } - if (memMode.equals_insensitive("global-work-space-plan")) { - return MemPlanMode::GLOBAL_WORKSPACE_PLAN; - } - func.emitError("unsupported mem-mode '") - << memMode << "'; only 'local' is supported by the PTOAS pipeline"; - return failure(); -} - -struct PlanMemoryPass : public mlir::pto::impl::PlanMemoryBase { -public: - PlanMemoryPass() = default; - explicit PlanMemoryPass(const mlir::pto::PlanMemoryOptions &planMemoryOption) - : PlanMemoryBase(planMemoryOption) {} - - void runOnOperation() override; - -private: - void populateBufferAddressToAllocOp( - RewritePatternSet &patterns, MemPlanMode mode, - DenseMap> buffer2Offsets) { - if (mode == MemPlanMode::LOCAL_MEM_PLAN) { - patterns.add( - patterns.getContext(), buffer2Offsets); - patterns.add( - patterns.getContext(), buffer2Offsets); - } - } -}; -} // namespace - -void PlanMemoryPass::runOnOperation() { - ModuleOp moduleOp = getOperation(); - SmallVector funcs; - moduleOp.walk([&](func::FuncOp funcOp) { - // TileOp helpers only contain compute code and deliberately do not own - // alloc_tile/reserve_buffer lifetimes. All other functions, including - // ordinary functions in backend child modules, must be planned. - if (!funcOp->hasAttr("pto.tileop.helper")) { - funcs.push_back(funcOp); - } - }); - - for (func::FuncOp funcOp : funcs) { - auto parsedMode = parseLegacyMemPlanMode(funcOp, this->memMode); - if (failed(parsedMode)) { - return signalPassFailure(); - } - MemPlanMode mode = *parsedMode; - ReserveBufferPlans reservePlans; - if (mode == MemPlanMode::LOCAL_MEM_PLAN && - failed(analyzeReserveBufferPlans(funcOp, reservePlans))) { - return signalPassFailure(); - } - if (mode == MemPlanMode::LOCAL_MEM_PLAN) { - for (ReserveBufferPlan &reservePlan : reservePlans) { - if (reservePlan.mode != ReserveBufferMode::Manual) { - continue; - } - reservePlan.reserveOp.emitOpError( - "pto.reserve_buffer with explicit 'base' (auto = false) is not " - "supported in PlanMemory; use --pto-level=level3 or set auto = true"); - return signalPassFailure(); - } - } - - MemLivenessAnalysis memLiveness(funcOp, mode); - memLiveness.build(); - - constexpr bool enableGlobalReuse = false; - constexpr bool enablePrintMemoryAllocatedSize = false; - constexpr bool restrictInplaceAsISA = false; - MemPlan memPlan(mode, enableGlobalReuse, enablePrintMemoryAllocatedSize, - restrictInplaceAsISA, this->orderBySize); - if (failed(memPlan.InitMemSpecsFromModule(funcOp))) { - return signalPassFailure(); - } - memPlan.func_ = funcOp; - memPlan.SetLinearOperation(memLiveness.linearOperation); - memPlan.SetBufferInfos(memLiveness.bufferInfos); - memPlan.SetBuffer2Life(memLiveness.buffer2Life); - memPlan.SetGenKillMap(memLiveness.genKillMap); - memPlan.SetBuffer2MultiNum(memLiveness.buffer2MultiNum); - memPlan.SetInplacePairList(memLiveness.inplacePairList); - memPlan.SetSemanticConflictPairs(memLiveness.semanticConflictPairs); - memPlan.SetStableValueOrder(std::move(memLiveness.stableValueOrder)); - if (failed(memPlan.plan())) { - return signalPassFailure(); - } - // Keep reserve_buffer allocation outside the core MemPlan algorithm: - // normal local buffers are planned first, then reserve_buffer claims one - // aligned hole in its target address space. - if (mode == MemPlanMode::LOCAL_MEM_PLAN && - failed(assignAutoReserveBufferBases( - reservePlans, memLiveness.bufferInfos, memPlan.GetBuffer2Offsets()))) { - return signalPassFailure(); - } - - RewritePatternSet patterns(&getContext()); - populateBufferAddressToAllocOp(patterns, mode, memPlan.GetBuffer2Offsets()); - if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns)))) { - return signalPassFailure(); - } - if (failed(verifySemanticNoAliasRanges(funcOp))) - return signalPassFailure(); - - bool hasUnplannedAllocTile = false; - funcOp.walk([&](pto::AllocTileOp op) { - if (op.getAddr()) { - return; - } - if (op->use_empty()) { - return; - } - if (isA5IgnoredTmpAlloc(op)) { - return; - } - op.emitError( - "PTOPlanMemory failed to assign an address to pto.alloc_tile"); - hasUnplannedAllocTile = true; - }); - if (hasUnplannedAllocTile) { - return signalPassFailure(); - } - } -} - -std::unique_ptr -mlir::pto::createPlanMemoryPass(const PlanMemoryOptions &options) { - return std::make_unique(options); -} diff --git a/lib/PTO/Transforms/PTOPlanMemory.h b/lib/PTO/Transforms/PTOPlanMemory.h deleted file mode 100644 index 175fe25ed7..0000000000 --- a/lib/PTO/Transforms/PTOPlanMemory.h +++ /dev/null @@ -1,819 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -//===- PlanMemory.h ----Plan Buffer Memory Address ------------------------===// -//===----------------------------------------------------------------------===// -#ifndef PTO_PLAN_MEMORY_H -#define PTO_PLAN_MEMORY_H - -#include -#include "OptMemPlanForPipeline.h" -#include "PTO/IR/PTO.h" -#include "PTO/Transforms/Passes.h" -#include "mlir/Analysis/Liveness.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "llvm/ADT/SmallSet.h" - - -namespace mlir { -namespace pto { - -// Value comparator for std::map -inline bool isLessValue(const Value &a, const Value &b) { - return a.getImpl() < b.getImpl(); -} - -struct ValueComparator { - bool operator()(const Value &a, const Value &b) const { - return isLessValue(a, b); - } -}; - -using StableValueOrderMap = DenseMap; - -/// Various states when collecting gen-kill. -enum class BufferStatus { UNDEFFINED = 0, DEFFINED, GENED, KILLED }; - -/// Pair of inplace Value. -using ValuePair = std::pair; - -enum class MemPlanMode { - LOCAL_MEM_PLAN, - GLOBAL_WORKSPACE_PLAN, -}; - -/// Result status after plan memory. -enum class PlanStatus { - PLAN_SUCCESS = 0, - RESTART_NEW_PLAN, - CONTINUE_PLAN, - PLAN_FAILED -}; - -/// Memory reuse plan mode can be achieved without conflicting life -/// intervals, offset = 0. -constexpr const int SPEC_LEVEL_0 = 0; - -/// By increasing the lifespan by 1 without conflict, -/// memory reuse plan mode can be implemented to avoid dependency on -/// continuous instructions caused by plan, offset = 1. -constexpr const int SPEC_LEVEL_1 = 1; - -/// pipe conflict opt. -constexpr const int SPEC_LEVEL_2 = 2; - -/// plan information of alloc buffer. -struct BufferInfo { - /// Alloc operation of buffer. - Operation *operation{nullptr}; - /// Space corresponding to buffer. - pto::AddressSpace bufferScope; - /// The size required for the buffer. - int64_t constBits{0}; - /// The type of element in the buffer. - Type bufferType; - /// Alias buffer does not participate in inplace. - /// e.g : - /// alloc A - /// for(arg = A) : - /// alloc B - /// ... - /// alloc C - /// vadd ins(B, D), outs(C) - /// scf.yield C - /// Put (A, C) in inplacePairList and inplace them together in next plan - /// memory. Because here do not union the lifetime of A and C, just set the - /// ignoreInplace of A and C to be true so that A will not be inplaced with - /// other buffer due to wrong lifetime. - /// Extending the lifetime union of A and C would allow further inplace reuse. - bool ignoreInplace{false}; -}; - -/// linear operation info. -struct OpInfo { - OpInfo(Operation *operation, int index) - : operation(operation), index(index) {} - Operation *operation{nullptr}; - int index{0}; -}; - -struct GenKillEntry { - /// record the gen operands, namely the operand buffer that is firstly written - /// by operation. - SmallVector gen; - - /// record the kill operands, namely the operand buffer that is last read by - /// operation. - SmallVector kill; -}; - -/// Record buffer life interval information. -struct BufferLife { - BufferLife(Value buffer, int64_t start, int64_t end) - : buffer(buffer), allocTime(start), freeTime(end) {} - explicit BufferLife(Value buffer) : buffer(buffer) {} - /// buffer value. - Value buffer; - /// the buffer allocate time. - int64_t allocTime{-1}; - /// the buffer free time. - int64_t freeTime{-1}; -}; - -/// a list of buffer life for a given storage entry -using BufferLifeVec = SmallVector>; - -struct StorageEntry { - /// The the buffer plan info. - BufferInfo *bufInfo{nullptr}; - - /// The lifespan of a buffer. - BufferLifeVec bufferLifeVec; - - /// The children of this entry, not including itself. - SmallVector mergedChildren; - - /// The current entry needs to be planed an aligned size. - uint64_t alignedConstBits{0}; - - /// The current entry's child index. - int childIdx; - - /// The starting address after the current entry allocation. - uint64_t bitsOffset{0}; - - /// Allocs that inplace buffer this entry. - SmallVector inplaceBuffers; - - /// multiBuffer relation StorageEntry. - /// For N >= 2 this aliases `relationOtherBuffers.front()` -- kept around so - /// existing N == 2 code paths can keep using the single-sibling field. - StorageEntry *relationPongEntry{nullptr}; - - /// Sibling slot entries for multi-buffer (N - 1 entries for slot 1..N-1). - /// The primary entry occupies slot 0; siblings own slot 1..N-1. Sibling - /// entries have `isMultiBufferSlot == true` and live in `StorageEntryVec` - /// independently of their primary -- the planner assigns each one its own - /// `bitsOffset` via the same Stage0/Stage2 logic used for normal allocs. - SmallVector relationOtherBuffers; - - /// True if this entry is a multi-buffer sibling (slot >= 1) that should - /// NOT independently write into `buffer2Offsets` -- the primary entry is - /// responsible for emitting all slot offsets in slot order. - bool isMultiBufferSlot{false}; - - /// The number of multibuffer optimization. - /// note: default 1 which means single buffer and does not do multibuffer - /// optimization. - uint32_t multiBufferNum{1}; - - /// Get Bufferlife by vaule - std::shared_ptr GetBufferLifeByValue(const Value v) const; -}; - -struct MemoryBound { - MemoryBound(BufferLifeVec life, uint64_t o, uint64_t e, const StorageEntry *s) - : bufferLifeVec(std::move(life)), offset(o), extent(e), - lastStorageEntry(s) {} - /// collection of buffer plan and free time which use this Memory - BufferLifeVec bufferLifeVec; - /// offset of tagged memory - uint64_t offset; - /// extent of this bound - uint64_t extent; - /// always record storage entry of last plan - const StorageEntry *lastStorageEntry; -}; -using MemBoundList = std::list>; -using MemBoundListConstIter = MemBoundList::const_iterator; - -/// record of buffer plan. used for speculative rollback -struct PlanRecord { - /// speculative level of this plan - int specLevel; - /// child index - int childIdx; - /// if this plan split last memory bound - bool tailed; - /// if this plan has bank offset memory bound - bool headed; - /// split number of entry - size_t splitNums; - /// record the entry for bank info - StorageEntry *entry; - /// the whole extent,add all the split e together - uint64_t allExtent; - /// inserted memory bound node - std::shared_ptr firstMemBound; - /// replaced memory bound node - MemBoundList replaced; - /// When the current PlanRecord is rolled back, it must be rolled back - /// directly. - bool isDirectlyRollback; -}; - -using PlanRecHis = SmallVector; - -struct SpecInfo { - int maxLevel = SPEC_LEVEL_2; - int minLevel = SPEC_LEVEL_0; - int specLevel = SPEC_LEVEL_2; - int childIdx = -1; - int specStartIdx = 0; - int rollbackIdx = -1; - uint64_t rollbackAddr = UINT64_MAX; -}; - -struct OutlineSectionInfo { - OutlineSectionInfo() = default; - OutlineSectionInfo(MemBoundListConstIter &start, MemBoundListConstIter &end, - uint64_t s, bool isDirectlyRollback) - : mem_start(start), mem_end(end), size(s), - isDirectlyRollback(isDirectlyRollback) {} - /// The start of memory plan - MemBoundListConstIter mem_start; - /// The end of memory plan - MemBoundListConstIter mem_end; - /// The size of memory plan - uint64_t size{0}; - /// When the current PlanRecord is rolled back, it must be rolled back - /// directly. - bool isDirectlyRollback; -}; - -/// comparator of buffer life -struct CompareBufferLife { - bool operator()(const std::shared_ptr &lhs, - const std::shared_ptr &rhs) const { - if (lhs->allocTime == rhs->allocTime) { - return lhs->freeTime < rhs->freeTime; - } - return lhs->allocTime < rhs->allocTime; - } -}; - -struct StatusWrapper { - /// Is it enough to roll back - bool hasEnoughRollBackSize; - /// The size required for the buffer - uint64_t alignedConstBits; - /// spec info - SpecInfo *si; - /// current outline info - MemBoundList &outline; - /// current history plan info - PlanRecHis &history; - /// for origin e StorageEntry - StorageEntry *RootE; -}; - -class MemLivenessAnalysis { -public: - MemLivenessAnalysis(func::FuncOp func, MemPlanMode planMode) - : func_(func), planMode(planMode) {} - - void build(); - - /// linear operation info. - SmallVector> linearOperation; - - /// map from buffer value to its buffer information. - std::map bufferInfos; - - /// stable IR order for Values used to keep memory planning deterministic. - StableValueOrderMap stableValueOrder; - - /// map from buffer to its lifetime. - DenseMap> buffer2Life; - - /// map from operation to its gen and kill buffer. - DenseMap genKillMap; - - /// record the map from the buffer to its number of buffer if it does - /// multibuffer optimization. - /// note: the map only record the buffer which do multi buffer - /// optimization and ignore single buffer. - DenseMap buffer2MultiNum; - - /// record inplace pair list. - SmallVector inplacePairList; - - /// record semantic conflict pair list. - SmallVector semanticConflictPairs; - - /// now plan mode is LOCAL_MEM_PLAN. - bool isLocalMemPlan() const; - - /// now plan mode is GLOBAL_WORKSPACE_PLAN. - bool isGlobalWorkSpaceMemPlan() const; - -private: - void RecursionIR(Region *region, Liveness live); - - /// Get the buffer used within the loop and defined outside the loop. - SmallVector GetLiveBuffersInLoop(Operation *loopOp, Liveness live); - - /// Update for Op tensor init args and tensor result args alias info. - void UpdateInitAndResAlias(DestinationStyleOpInterface dstStyleOp); - - /// Recursive operation for. - void RecursiveForOp(scf::ForOp forOp, Liveness live); - - /// Recursive operation for a generic scf.while loop. - void RecursiveWhileOp(scf::WhileOp whileOp, Liveness live); - - /// Update for Op init args region iter args alias info. - void UpdateForOpInitArgsAlias(scf::ForOp forOp); - - /// Update forOp result buffer/region iter arg/yielded buffer args alias info. - void UpdateForOpBufferAlias(scf::ForOp forOp); - - /// Update aliases crossing both regions of an scf.while loop. - void UpdateWhileOpBufferAlias(scf::WhileOp whileOp); - - /// Recursive operation if. - void RecursiveIfOp(scf::IfOp ifOp, Liveness live); - - /// Update buffer alias information for ifop. - void UpdateIfOpBufferAlias(scf::IfOp ifOp, scf::YieldOp yieldOp); - - /// Recursive operation for pto.fusion_region. - void RecursiveFusionRegionOp(pto::FusionRegionOp fusionRegion, Liveness live); - - /// Update buffer alias information for pto.fusion_region results. - void UpdateFusionRegionBufferAlias(pto::FusionRegionOp fusionRegion, - pto::YieldOp yieldOp); - - /// Update and obtain op info information. - OpInfo *UpdateLinearOperation(Operation *op); - - /// Obtain all information about the buffer. - void UpdateOpBufferInfo(Operation *op, const ValueRange &results); - - /// Generate buffer info. - BufferInfo GenerateBufferInfo(Operation *op, Value operand); - - /// Obtain the buffer info of plan operation. - BufferInfo GetBufferInfo(Operation *op, Value operand, - pto::AddressSpace bufferScope); - - /// Process gen buffer based on the result value of op. - void UpdateOpGenInfo(OpInfo *opInfo, const ValueRange &results); - - /// Update normal operand gen information on buffer. - void UpdateOperandGenInfo(OpInfo *opInfo, Value operand); - - /// Update temp buffer for DestinationStyleOpInterface op. - void UpdateOpTempGenInfo(OpInfo *opInfo); - - /// Update the relationship of buffer aliases. - void UpdateBufferAlias(Value buffer, Value aliasBuffer, - bool isIgnoreInplace = false); - - /// Return the union of set1 and set2. - SetVector Union(SetVector set1, SetVector set2); - - /// Get alias buffer information. - SetVector GetAliasBuffers(Value aliasBuffer); - - /// Check whether there is an unknown operation with buffer - /// information. - LogicalResult CheckIfUnknownOpTouchBuffer(Operation *op) const; - - /// Determine whether the current operation can be skipped. - bool isSkippableOp(Operation *op) const; - - /// Kill a buffer handle at the current operation. - void OpKillHandle(OpInfo *opInfo, Liveness live, Block *block); - - /// Process kill buffer based on the result live of op. - void UpdateOpKillInfo(OpInfo *opInfo, Value operand, Liveness live); - - /// Have all alias buffer been killed. - bool AllDeadAfter(Operation *op, SetVector aliasVec, - Liveness live) const; - - /// Determine whether two operation are in the same block. - bool IsInSameBlock(Operation *op1, Operation *op2) const; - - /// Generate buffer's life time. - void GenerateBufferLife(); - - /// initialize the buffers that must be inplaced together - /// namely, tile buffer aliases such as iter args and yields. - void InitializeInplacePairList(); - - /// Record semantic non-reuse pairs for buffers that may be used - /// simultaneously inside one instruction, such as scratch and dst. - void RecordSemanticConflict(Value lhs, Value rhs); - - func::FuncOp func_; - - /// different mode for mem plan. - MemPlanMode planMode; - - /// Gen-kill status corresponding to buffer. - DenseMap buffer2status; - - /// map on buffer alias - DenseMap> buffer2AliasVec; - - int seqIndex{0}; -}; - -/// Pair of StorageEntry. -using StorageEntryPair = std::pair; - -class MemPlan { -public: - MemPlan(MemPlanMode planMode, bool enableGlobalReuse, bool enablePrintMemoryAllocatedSize, - bool restrictInplaceAsISA, bool orderBySize) - : planMode(planMode), enableGlobalReuse(enableGlobalReuse), - enablePrintMemoryAllocatedSize(enablePrintMemoryAllocatedSize), - restrictInplaceAsISA(restrictInplaceAsISA), orderBySize(orderBySize) {} - - LogicalResult plan(); - - /// Get buffer2Offsets - inline DenseMap> GetBuffer2Offsets() { - return buffer2Offsets; - } - - inline void - SetLinearOperation(SmallVector> &linearOp) { - linearOperation = std::move(linearOp); - }; - - inline void - SetBufferInfos(std::map bufsInfo) { - bufferInfos = bufsInfo; - } - - inline void - SetBuffer2Life(DenseMap> buf2Life) { - buffer2Life = buf2Life; - } - - inline void SetGenKillMap(DenseMap gkMap) { - genKillMap = gkMap; - } - - inline void SetBuffer2MultiNum(DenseMap buf2MulBufNum) { - buffer2MultiNum = buf2MulBufNum; - } - - inline void SetInplacePairList(SmallVector inplaceList) { - inplacePairList = inplaceList; - } - - inline void SetSemanticConflictPairs(SmallVector conflictPairs) { - semanticConflictPairs = std::move(conflictPairs); - } - - inline void SetStableValueOrder(StableValueOrderMap valueOrder) { - stableValueOrder = std::move(valueOrder); - } - - /// Setup the device's storage specs - LogicalResult InitMemSpecsFromModule(func::FuncOp funcOp); - - func::FuncOp func_; - -private: - /// different mode for mem plan. - MemPlanMode planMode; - - /// Enable global workspace reuse. - bool enableGlobalReuse; - - /// Enable print memory allocated size. - bool enablePrintMemoryAllocatedSize; - - /// enable PTO op plan memory inplace - bool restrictInplaceAsISA; - - /// Process buffers largest-first (first-fit-decreasing) instead of DMA-first. - bool orderBySize; - - /// StorageEntry generate. - void GenerateStorageEntry(); - - /// Print successful memory alloc. - void PrintSuccessfulAllocatedMaxBits(); - - /// Post-plan sanity check for local memory overflow. - bool RecordOverflowIfAny(); - - /// Prepare the local tile buffer plan. - PlanStatus PlanLocalMemAddress(); - - /// Prepare the global workspace plan. - PlanStatus PlanWorkSpaceMemAddress(); - - /// merge all storage entry to the first storage entry for WorkSpaceArg. - void MergeSameWorkSpaceArgSE(); - - /// Start plan for same work space arg offset. - PlanStatus PlanMemOffsetOfWholeWorkSpace(); - - /// Enable global workspace no reuse. - void GlobalWorkspaceNoReuse(StorageEntry *rootStorageEntry); - - /// Verify that constBits is legal. - void ValidateParameters(std::unique_ptr &e) const; - - /// Expanding the Storage Entry due to the addition of MultiBuffer. - void ExpandMultiBufferStorageEntry(); - - /// merge all storage entry to the first storage entry. - void MergeSameScopeSE(); - - /// merge all storage entry which can be inplaced. - void MergeInplaceSE(); - - /// Start plan. - PlanStatus PlanMemAddressOfWholeLocalBuffer(); - - /// Plan a single local buffer without reuse. - PlanStatus PlanSingleLocalBuffer(StorageEntry *rootStorageEntry, size_t align, - size_t maxBits); - - /// Plan a reusable local buffer scope. - PlanStatus PlanReusableLocalBuffer(StorageEntry *rootStorageEntry, - size_t align, size_t maxBits); - - /// Plan memory only by level0 to report failure info. - void PlanMemAddressForLevel0(StorageEntry *rootStorageEntry); - - /// Determine if the current space is enough to allocate all buffers. - bool IsEnoughForBuffersNoReuse(StorageEntry *rootStorageEntry, - size_t restBufferSize, size_t alignUnit); - - /// Adjust the allocation order of rootStoreEntry to prioritize the allocation - /// of buffers corresponding to DMA. - StorageEntry *GetReorderRootStorageEntry(StorageEntry *rootStorageEntry); - - /// Reorder rootStorageEntry's children largest-first (first-fit-decreasing) - /// across every memory space, keeping ping-pong pairs contiguous. Used when - /// the order-by-size option is enabled. - StorageEntry *GetSizeOrderedRootStorageEntry(StorageEntry *rootStorageEntry); - - /// Assign addresses without reuse. - void PlanBuffersWithoutReuse(StorageEntry *rootStorageEntry, - size_t alignUnit); - - /// Obtain buffer space size and alignment information. - std::pair GetBufferSpaceInfo(pto::AddressSpace &space) const; - - /// Emit buffer applied failure message. - void EmitPlanMemoryFailureInfo(); - - /// Multi level plan strategy. - LogicalResult MultiSpecPlan(SpecInfo &si, MemBoundList &outline, - PlanRecHis &history, StorageEntry *entry); - - /// plan buffer in speculative ways. - LogicalResult SpecAlloc(MemBoundList &outline, PlanRecHis &his, - StorageEntry *e, const SpecInfo &si, int localLevel); - - /// spec_level == SPEC_LEVEL_2, mte2/3 do not reuse with vector. - bool VerifyConflictStage2(PlanRecHis &his, const StorageEntry *e, - int specLevel, MemBoundListConstIter &start, - const MemBoundList &outline); - - /// spec_level == SPEC_LEVEL_1, pure single can reuse with db. - bool VerifyConflictStage1(MemBoundList &outline, PlanRecHis &his, - StorageEntry *e, - const OutlineSectionInfo &outlineInfo, - uint64_t &pongOffset); - - /// check if e1 and e2 has pipe conflict. - bool PipeConflict(const StorageEntry *e1, const StorageEntry *e2, - DenseMap &conflictMap); - - /// spec_level == SPEC_LEVEL_2, MTE2/MTE3 is pipe conflict with all existing - /// allocation. check if current entry has OptDmaPipe-conflict with buffers - /// already allocate at current position. if conflict exists, continue loop - /// until first not-conflict iter is found. Then update start as the first - /// bound right before the not-conflict one. - bool VerifyDmaPipeConflict(const StorageEntry *e, int specLevel, - MemBoundListConstIter &start, - MemBoundListConstIter &end); - - /// Check if it matches the previous rollback result. - bool IsSamePlanAsLastRollBack(uint64_t allocOffset, int curChildIdx, - const SpecInfo &si) const; - - /// spec_level == SPEC_LEVEL_0, life time reuse. - inline bool VerifyConflictStage0(StorageEntry *e, - const std::shared_ptr &last); - - /// Update the outline information and record history - void UpdateOutline(MemBoundList &outline, PlanRecHis &his, StorageEntry *e, - const OutlineSectionInfo &outlineInfo, - int localLevel) const; - - /// plan strategy is achieved through split method. - void AddMemBoundInSectionalWay( - StorageEntry *e, MemBoundListConstIter start, MemBoundListConstIter end, - SmallVector> &splitBound) const; - - /// merge the buffer life between start and end. - inline void MergeBufferLife(MemBoundList::const_iterator start, - MemBoundList::const_iterator end, - BufferLifeVec &newLife) const; - - /// merge buffers in a vector. - void MergeBufferVec(BufferLifeVec &bufferLife) const; - - /// Judge if need to restart plan memory with other strategy after - /// plan failed. - PlanStatus ApplyFailStrategy(StatusWrapper &statusWrapper, - const size_t maxBits); - - void RollBackForAllocFail(StatusWrapper &statusWrapper, const size_t maxBits); - - /// Check if memory plan can be rolled back. - bool ContinueRollBack(const StatusWrapper &statusWrapper) const; - - /// Memory plan fallback information processing. - void RollBackForAllocFailInner(StatusWrapper &statusWrapper, - const size_t maxBits); - - /// Fallback outline plan. - PlanRecord RollbackOutline(PlanRecHis &history, MemBoundList &outline) const; - - /// Update the plan memory address corresponding to mem buffer. - void UpdateBuffer2Offsets(); - - /// Update extra addresses offset caused by multi buffer reuse. - void UpdateMultiBufferReuseExtraOffset(); - - /// generate inplace list by some rules - SmallVector GenerateInplaceList(); - - /// the ptoop that can reuse dst address and src address in limited situation - bool IsReusePTOOp(Operation *op) const; - - /// Get overlap buffer life. - DenseMap - GetOverlapBufferLife(const BufferLifeVec &b1, const BufferLifeVec &b2) const; - - bool HasSemanticConflict(const StorageEntry *entry, - const BufferLifeVec &bufferLives) const; - - /// Reorder and make the storage entries of ping and pong continuous. - void - ReorderContinuousPingPongEntry(SmallVector &storageEntryVec); - - /// Determine if the current buffer life of the Storage Entry conflicts with - /// the memory that has already been allocated in history. - bool IsBufferLifeVecConflict(PlanRecord &r, uint64_t offset, - const StorageEntry *e) const; - - /// Assign pong storage entry's address. - void PlanRelationPongEntryAddress(uint64_t offset, StorageEntry *e); - - /// Processing Pong Storage Entry Information. - void SpecAllocRelationPongEntry(MemBoundList &outline, PlanRecHis &his, - StorageEntry *e, uint64_t offset); - - /// Get relative pong storage entry when the current reuse bound storage entry - /// is of type db. - StorageEntry * - GetMultiRelationPongEntry(const StorageEntry *reuseBoundStorageEntry); - - /// Get the innermost for loop of buffer definition. - LoopLikeOpInterface GetBufferParentLoop(const SmallVector &buffers); - - /// Report all tensors life time info. - void ReportMemLifeDebugInfo(StorageEntry *rootStorageEntry); - - /// Report tensor life time for debug. - void MemLifeDebugInfo(StorageEntry *storageEntry); - - /// Report the allocation root represented by this entry. - void ReportCurEntryDebugInfo(const StorageEntry *curEntry); - - /// Report tensor allocate info. - void ReportAllocatedEntryDebugInfo(StorageEntry *rootStorageEntry); - -private: - /// The buffer corresponding to each operation. - SmallVector> linearOperation; - - /// map from buffer value to its buffer information. - std::map bufferInfos; - - /// map from buffer to its lifetime. - DenseMap> buffer2Life; - - /// record the map from the buffer to its number of buffer if it does - /// multibuffer optimization. - /// note: the map only record the buffer which do multi buffer optimization - /// and ignore single buffer. - DenseMap buffer2MultiNum; - - /// map from operation to its gen and kill buffer. - DenseMap genKillMap; - - /// record all storage entry to be plan address. - SmallVector> StorageEntryVec; - - /// The current status of memory plan. - PlanStatus planStatus{PlanStatus::PLAN_SUCCESS}; - - /// Whether to adopt a split strategy. - bool splitOutline{false}; - - /// Map from local allocation root to planned addresses. - DenseMap> buffer2Offsets; - - /// map from each scope to its root StorageEntry. - DenseMap memscope2rootStorageEntry; - - /// map from workspace arg to its root StorageEntry. - DenseMap workSpaceArg2rootStorageEntry; - - /// map from buffer scope to its required size to plan rest memory without any - /// reuse. - DenseMap bufferScope2RequiredSize; - - /// map from buffer value to its storage entry info - DenseMap buffer2storageEntry; - - /// stable IR order for Values used to keep memory planning deterministic. - StableValueOrderMap stableValueOrder; - - /// Memory dma pipe first plan optimization. - OptMemPlanForDma dmaFirstPipelineOpt; - - /// Map from the storage entry pair to its pipeDma conflict info. - DenseMap pipeDmaConflictMap; - - /// Ping storage entry corresponding to reused additional Pong entry. - DenseMap> - pingEntry2RelationPongEntry; - - SmallVector allocatedEntry; - - /// record inplace pair list. - SmallVector inplacePairList; - - /// record semantic conflict pair list. - SmallVector semanticConflictPairs; - - /// inplace-reuse info for the vf call. - //VFCallInplaceReuseInfo *vfInplaceReuseInfo; - - /// The scope of the buffer applied memory fail and the max bits it applied. - std::map failApplyBufferInfo; - - /// The device's UB storage size - int ubSpaceSize{0}; - - /// The device's L1 storage size - int l1SpaceSize{0}; - - /// The device's L0A storage size - int l0aSpaceSize{0}; - - /// The device's L0B storage size - int l0bSpaceSize{0}; - - /// The device's L0C storage size - int l0cSpaceSize{0}; - - /// The device's UB align size - int ubAlignSize{0}; - - /// The device's L1 align size - int l1AlignSize{0}; - - /// The device's L0C align size - int l0cAlignSize{0}; - - int l0aAlignSize{0}; - - int l0bAlignSize{0}; - - int biasAlignSize{0}; - - int biasSpaceSize{0}; - - /// The device's SCALING align size - int scalingAlignSize{0}; - - /// The device's SCALING storage size - int scalingSpaceSize{0}; -}; -} // namespace pto -} // namespace mlir - -#endif // BISHENG_DIALECT_PTO_TRANSFORMS_PLAN_MEMORY_H diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 9b00833b9d..68310a7af5 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -1858,3 +1858,13 @@ std::unique_ptr mlir::pto::createPlanMemoryModernPass(const PlanMemoryOptions &options) { return std::make_unique(options); } + +// Anchor the generated pass base registration used by pto-test-opt and the +// textual pass pipeline (`-pass-pipeline=...`), replacing the removed legacy +// planner under the same `pto-plan-memory` argument. +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_PLANMEMORY +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir diff --git a/test/lit/pto/alloc_tile_low_precision_valid.pto b/test/lit/pto/alloc_tile_low_precision_valid.pto index f385fb0290..b123142156 100644 --- a/test/lit/pto/alloc_tile_low_precision_valid.pto +++ b/test/lit/pto/alloc_tile_low_precision_valid.pto @@ -18,6 +18,6 @@ module { } // CHECK: func.func @alloc_tile_low_precision_valid() -// CHECK: pto.alloc_tile : !pto.tile_buf -// CHECK: pto.alloc_tile : !pto.tile_buf -// CHECK: pto.alloc_tile : !pto.tile_buf +// CHECK: pto.alloc_tile{{.*}} : !pto.tile_buf +// CHECK: pto.alloc_tile{{.*}} : !pto.tile_buf +// CHECK: pto.alloc_tile{{.*}} : !pto.tile_buf diff --git a/test/lit/pto/compact_left_blayout_parser_a3.pto b/test/lit/pto/compact_left_blayout_parser_a3.pto index b64df93341..2a00822496 100644 --- a/test/lit/pto/compact_left_blayout_parser_a3.pto +++ b/test/lit/pto/compact_left_blayout_parser_a3.pto @@ -8,4 +8,4 @@ module attributes {"pto.target_arch" = "a3"} { } // CHECK-LABEL: func.func @compact_left_blayout_parser_a3() { -// CHECK: pto.alloc_tile : !pto.tile_buf +// CHECK: pto.alloc_tile{{.*}} : !pto.tile_buf diff --git a/test/lit/pto/compact_left_blayout_parser_a5.pto b/test/lit/pto/compact_left_blayout_parser_a5.pto index f34458bd1b..8bcd5a2db4 100644 --- a/test/lit/pto/compact_left_blayout_parser_a5.pto +++ b/test/lit/pto/compact_left_blayout_parser_a5.pto @@ -8,4 +8,4 @@ module attributes {"pto.target_arch" = "a5"} { } // CHECK-LABEL: func.func @compact_left_blayout_parser_a5() { -// CHECK: pto.alloc_tile : !pto.tile_buf +// CHECK: pto.alloc_tile{{.*}} : !pto.tile_buf diff --git a/test/lit/pto/declare_tile_tile_native.pto b/test/lit/pto/declare_tile_tile_native.pto index b365375ee6..660f34e073 100644 --- a/test/lit/pto/declare_tile_tile_native.pto +++ b/test/lit/pto/declare_tile_tile_native.pto @@ -8,7 +8,7 @@ // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE -// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --mlir-print-ir-before=pto-inline-backend-helpers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto b/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto index 4f7aecd75e..3ce707446c 100644 --- a/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto +++ b/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s -// RUN: ptoas --pto-arch=a3 --pto-level=level2 --plan-memory-impl=modern --emit-pto-ir %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s module { func.func @implicit_remaining_tmps(%executed : vector<4xi16>) { diff --git a/test/lit/pto/issue706_comm_staging_buffer_sync.pto b/test/lit/pto/issue706_comm_staging_buffer_sync.pto index af52891765..2ab4a71ff0 100644 --- a/test/lit/pto/issue706_comm_staging_buffer_sync.pto +++ b/test/lit/pto/issue706_comm_staging_buffer_sync.pto @@ -92,11 +92,13 @@ module { // CHECK-LABEL: AICORE void tput_reuses_pong_after_call( // CHECK: pto::comm::TPUT( // CHECK: set_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[TPUT:[0-9]+]]); -// CHECK-NEXT: wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[TPUT]]); +// CHECK: wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[TPUT]]); +// CHECK-NOT: TMOV( // CHECK-NEXT: TMOV( // CHECK-LABEL: AICORE void tbroadcast_reuses_pong_after_call( // CHECK: pto::comm::TBROADCAST( // CHECK: set_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[BCAST:[0-9]+]]); -// CHECK-NEXT: wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[BCAST]]); +// CHECK: wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[BCAST]]); +// CHECK-NOT: TMOV( // CHECK-NEXT: TMOV( diff --git a/test/lit/pto/issue706_treduce_recv_pong_staging_sync.pto b/test/lit/pto/issue706_treduce_recv_pong_staging_sync.pto index aa6891a0bd..c1c225e61e 100644 --- a/test/lit/pto/issue706_treduce_recv_pong_staging_sync.pto +++ b/test/lit/pto/issue706_treduce_recv_pong_staging_sync.pto @@ -57,5 +57,6 @@ module { // CHECK-LABEL: AICORE void treduce_reuses_recv_pong_after_call( // CHECK: pto::comm::TREDUCE( // CHECK: set_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[PONG:[0-9]+]]); -// CHECK-NEXT: wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[PONG]]); +// CHECK-NOT: TMOV( +// CHECK: wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID[[PONG]]); // CHECK-NEXT: TMOV( diff --git a/test/lit/pto/mgather_gm2l1_elem_scratch_reuse_sync_a5.pto b/test/lit/pto/mgather_gm2l1_elem_scratch_reuse_sync_a5.pto index b845b1656a..324b81381b 100644 --- a/test/lit/pto/mgather_gm2l1_elem_scratch_reuse_sync_a5.pto +++ b/test/lit/pto/mgather_gm2l1_elem_scratch_reuse_sync_a5.pto @@ -49,5 +49,6 @@ module { // CHECK-NEXT: set_flag(PIPE_S, PIPE_MTE2 // CHECK-NEXT: set_flag(PIPE_MTE2, PIPE_S // CHECK-NEXT: wait_flag(PIPE_MTE2, PIPE_S +// CHECK-NEXT: pipe_barrier(PIPE_MTE2); // CHECK-NEXT: wait_flag(PIPE_S, PIPE_MTE2 // CHECK: MGATHER diff --git a/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto b/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto index 12dd840d88..32b04a2561 100644 --- a/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto +++ b/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s // Sanity for N == 3 to confirm `ExpandMultiBufferStorageEntry` produces 3 // physical slots and `UpdateBuffer2Offsets` emits them in slot order. diff --git a/test/lit/pto/multi_tile_buf_type_parse_print.pto b/test/lit/pto/multi_tile_buf_type_parse_print.pto index 195c346024..c20b9a0901 100644 --- a/test/lit/pto/multi_tile_buf_type_parse_print.pto +++ b/test/lit/pto/multi_tile_buf_type_parse_print.pto @@ -34,10 +34,10 @@ module { } // CHECK-LABEL: func.func @mtb_compact_n2 -// CHECK: pto.alloc_multi_tile : !pto.multi_tile_buf, count=2> +// CHECK: pto.alloc_multi_tile{{.*}} : !pto.multi_tile_buf, count=2> // CHECK-LABEL: func.func @mtb_compact_n4 -// CHECK: pto.alloc_multi_tile : !pto.multi_tile_buf, count=4> +// CHECK: pto.alloc_multi_tile{{.*}} : !pto.multi_tile_buf, count=4> // CHECK-LABEL: func.func @mtb_verbose -// CHECK: pto.alloc_multi_tile : !pto.multi_tile_buf, count=3> +// CHECK: pto.alloc_multi_tile{{.*}} : !pto.multi_tile_buf, count=3> diff --git a/test/lit/pto/multi_tile_get_const_slot_lowering.pto b/test/lit/pto/multi_tile_get_const_slot_lowering.pto index cb6bd1f8b8..5142918adf 100644 --- a/test/lit/pto/multi_tile_get_const_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_const_slot_lowering.pto @@ -9,8 +9,8 @@ // RUN: ptoas --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT +// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN +// RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT // Verifies the tile-native multi-buffer pipeline for constant slots: // 1. PTOResolveReservedBuffers preserves alloc_multi_tile and multi_tile_get. diff --git a/test/lit/pto/multi_tile_n4_planmem_e2e.pto b/test/lit/pto/multi_tile_n4_planmem_e2e.pto index 13244bc2a6..c061c23b68 100644 --- a/test/lit/pto/multi_tile_n4_planmem_e2e.pto +++ b/test/lit/pto/multi_tile_n4_planmem_e2e.pto @@ -8,8 +8,8 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT +// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN +// RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT // End-to-end test for N == 4: PlanMemory records four physical slot addresses, // then PTOResolveBufferSelect materializes each selected tile handle. diff --git a/test/lit/pto/mx_noalias_plan_memory.pto b/test/lit/pto/mx_noalias_plan_memory.pto index ef08fb4474..bb0baa9e0c 100644 --- a/test/lit/pto/mx_noalias_plan_memory.pto +++ b/test/lit/pto/mx_noalias_plan_memory.pto @@ -9,102 +9,53 @@ // RUN: split-file %s %t // RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir \ // RUN: --mlir-print-ir-after=pto-plan-memory %t/planned.pto -o /dev/null 2>&1 \ -// RUN: | FileCheck %s --check-prefix=LEGACY -// RUN: ptoas --pto-arch=a5 --pto-level=level2 --plan-memory-impl=modern --emit-pto-ir \ -// RUN: --mlir-print-ir-after=pto-plan-memory %t/planned.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=MODERN // RUN: not pto-test-opt --pto-plan-memory %t/explicit_overlap.pto 2>&1 \ // RUN: | FileCheck %s --check-prefix=OVERLAP -// RUN: not ptoas --pto-arch=a5 --pto-level=level3 --plan-memory-impl=modern \ +// RUN: not ptoas --pto-arch=a5 --pto-level=level3 \ // RUN: %t/explicit_overlap.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=OVERLAP // RUN: not ptoas --pto-arch=a5 --pto-level=level3 \ // RUN: %t/src_tmp_overlap.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=SRC-TMP-OVERLAP -// RUN: not ptoas --pto-arch=a5 --pto-level=level3 --plan-memory-impl=modern \ +// RUN: not ptoas --pto-arch=a5 --pto-level=level3 \ // RUN: %t/src_tmp_overlap.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=SRC-TMP-OVERLAP // RUN: not ptoas --pto-arch=a5 --pto-level=level3 \ // RUN: %t/tmp_dst_overlap.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=TMP-DST-OVERLAP -// RUN: not ptoas --pto-arch=a5 --pto-level=level3 --plan-memory-impl=modern \ +// RUN: not ptoas --pto-arch=a5 --pto-level=level3 \ // RUN: %t/tmp_dst_overlap.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=TMP-DST-OVERLAP // RUN: not ptoas --pto-arch=a5 --pto-level=level3 \ // RUN: %t/tquant_output_overlap.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=TQUANT-OUTPUT-OVERLAP -// RUN: not ptoas --pto-arch=a5 --pto-level=level3 --plan-memory-impl=modern \ +// RUN: not ptoas --pto-arch=a5 --pto-level=level3 \ // RUN: %t/tquant_output_overlap.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=TQUANT-OUTPUT-OVERLAP // RUN: not ptoas --pto-arch=a5 --pto-level=level2 %t/view_alias.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=VIEW -// RUN: not ptoas --pto-arch=a5 --pto-level=level2 --plan-memory-impl=modern \ +// RUN: not ptoas --pto-arch=a5 --pto-level=level2 \ // RUN: %t/view_alias.pto -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=VIEW // RUN: not pto-test-opt --pto-plan-memory %t/strided_view_overlap.pto \ // RUN: -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=STRIDED -// RUN: ptoas --pto-arch=a5 --pto-level=level3 --plan-memory-impl=legacy \ -// RUN: %t/plain_control.pto -o /dev/null -// RUN: ptoas --pto-arch=a5 --pto-level=level3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-arch=a5 --pto-level=level3 \ // RUN: %t/plain_control.pto -o /dev/null -// LEGACY-LABEL: func.func @dn_x2zz_src_last_use_dst_first_write -// LEGACY-DAG: %[[DN0:.*]] = arith.constant 0 : i64 -// LEGACY-DAG: %[[DN512:.*]] = arith.constant 512 : i64 -// LEGACY-DAG: %[[DN768:.*]] = arith.constant 768 : i64 -// LEGACY: pto.alloc_tile addr = %[[DN0]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[DN512]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[DN768]] : !pto.tile_buf -// LEGACY-LABEL: func.func @nd_x2zz_src_last_use_dst_first_write -// LEGACY-DAG: %[[ND0:.*]] = arith.constant 0 : i64 -// LEGACY-DAG: %[[ND1024:.*]] = arith.constant 1024 : i64 -// LEGACY-DAG: %[[ND1280:.*]] = arith.constant 1280 : i64 -// LEGACY: pto.alloc_tile addr = %[[ND0]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[ND1024]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[ND1280]] : !pto.tile_buf -// LEGACY-LABEL: func.func @tquant_src_vs_each_output -// LEGACY-DAG: %[[TQ_SRC_C:.*]] = arith.constant 0 : i64 -// LEGACY-DAG: %[[TQ_DST_C:.*]] = arith.constant 2048 : i64 -// LEGACY-DAG: %[[TQ_EXP_C:.*]] = arith.constant 2560 : i64 -// LEGACY-DAG: %[[TQ_MAX_C:.*]] = arith.constant 2816 : i64 -// LEGACY-DAG: %[[TQ_SCALING_C:.*]] = arith.constant 3072 : i64 -// LEGACY: pto.alloc_tile addr = %[[TQ_SRC_C]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[TQ_DST_C]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[TQ_EXP_C]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[TQ_MAX_C]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[TQ_SCALING_C]] : !pto.tile_buf -// LEGACY-LABEL: func.func @tquant_src_vs_exp_zz -// LEGACY-DAG: %[[FUSED_SRC:.*]] = arith.constant 0 : i64 -// LEGACY-DAG: %[[FUSED_DST:.*]] = arith.constant 4096 : i64 -// LEGACY-DAG: %[[FUSED_EXP:.*]] = arith.constant 5120 : i64 -// LEGACY-DAG: %[[FUSED_MAX:.*]] = arith.constant 5376 : i64 -// LEGACY-DAG: %[[FUSED_SCALING:.*]] = arith.constant 5632 : i64 -// LEGACY-DAG: %[[FUSED_EXP_ZZ:.*]] = arith.constant 5888 : i64 -// LEGACY: pto.alloc_tile addr = %[[FUSED_SRC]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[FUSED_DST]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[FUSED_EXP]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[FUSED_MAX]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[FUSED_SCALING]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[FUSED_EXP_ZZ]] : !pto.tile_buf -// LEGACY-LABEL: func.func @plain_tmov_reuse_control -// LEGACY-DAG: %[[PLAIN_C:.*]] = arith.constant 0 : i64 -// LEGACY-DAG: %[[PLAIN_NEXT:.*]] = arith.constant 2048 : i64 -// LEGACY: pto.alloc_tile addr = %[[PLAIN_C]] : !pto.tile_buf -// LEGACY: pto.alloc_tile addr = %[[PLAIN_NEXT]] : !pto.tile_buf - // MODERN-LABEL: func.func @dn_x2zz_src_last_use_dst_first_write // MODERN-DAG: %[[DN0:.*]] = arith.constant 0 : i64 -// MODERN-DAG: %[[DN512:.*]] = arith.constant 512 : i64 +// MODERN-DAG: %[[DN768:.*]] = arith.constant 768 : i64 // MODERN: pto.alloc_tile addr = %[[DN0]] : !pto.tile_buf // MODERN: pto.alloc_tile addr = %{{.*}} : !pto.tile_buf -// MODERN: pto.alloc_tile addr = %[[DN512]] : !pto.tile_buf +// MODERN: pto.alloc_tile addr = %[[DN768]] : !pto.tile_buf // MODERN-LABEL: func.func @nd_x2zz_src_last_use_dst_first_write // MODERN-DAG: %[[ND0:.*]] = arith.constant 0 : i64 -// MODERN-DAG: %[[ND1024:.*]] = arith.constant 1024 : i64 +// MODERN-DAG: %[[ND1280:.*]] = arith.constant 1280 : i64 // MODERN: pto.alloc_tile addr = %[[ND0]] : !pto.tile_buf // MODERN: pto.alloc_tile addr = %{{.*}} : !pto.tile_buf -// MODERN: pto.alloc_tile addr = %[[ND1024]] : !pto.tile_buf +// MODERN: pto.alloc_tile addr = %[[ND1280]] : !pto.tile_buf // MODERN-LABEL: func.func @tquant_src_vs_each_output // MODERN-DAG: %[[TQ_SRC:.*]] = arith.constant 0 : i64 // MODERN-DAG: %[[TQ_DST:.*]] = arith.constant 2048 : i64 diff --git a/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto b/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto index 1f0ddd9ce0..9a1da1af28 100644 --- a/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto +++ b/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_five_gates_phi_family.pto b/test/lit/pto/plan_memory_five_gates_phi_family.pto index 5007e8874f..b694f64974 100644 --- a/test/lit/pto/plan_memory_five_gates_phi_family.pto +++ b/test/lit/pto/plan_memory_five_gates_phi_family.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_five_gates_phi_family_select.pto b/test/lit/pto/plan_memory_five_gates_phi_family_select.pto index 07512b3022..f850e68c25 100644 --- a/test/lit/pto/plan_memory_five_gates_phi_family_select.pto +++ b/test/lit/pto/plan_memory_five_gates_phi_family_select.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_fusion_region_alias.pto b/test/lit/pto/plan_memory_fusion_region_alias.pto index 5a2fa499fd..170ef2e23a 100644 --- a/test/lit/pto/plan_memory_fusion_region_alias.pto +++ b/test/lit/pto/plan_memory_fusion_region_alias.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_impl_invalid_level3.pto b/test/lit/pto/plan_memory_impl_invalid_level3.pto deleted file mode 100644 index 6c66d5d194..0000000000 --- a/test/lit/pto/plan_memory_impl_invalid_level3.pto +++ /dev/null @@ -1,17 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -// RUN: not ptoas --pto-level=level3 --plan-memory-impl=typo %s 2>&1 | FileCheck %s - -module { - func.func @empty() attributes {pto.entry} { - return - } -} - -// CHECK: Error: invalid --plan-memory-impl='typo', expected 'legacy' or 'modern'. diff --git a/test/lit/pto/plan_memory_inplace_forbid_alias.pto b/test/lit/pto/plan_memory_inplace_forbid_alias.pto index 8bfed61303..081a88b679 100644 --- a/test/lit/pto/plan_memory_inplace_forbid_alias.pto +++ b/test/lit/pto/plan_memory_inplace_forbid_alias.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_legacy_scalar_pipe_conflict.pto b/test/lit/pto/plan_memory_legacy_scalar_pipe_conflict.pto deleted file mode 100644 index 5a432fb12b..0000000000 --- a/test/lit/pto/plan_memory_legacy_scalar_pipe_conflict.pto +++ /dev/null @@ -1,30 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --emit-pto-ir \ -// RUN: --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @scalar_pipe_conflict() -> f32 { - %c0 = arith.constant 0 : index - %scalar = pto.alloc_tile : !pto.tile_buf - %value = pto.tgetval ins(%scalar, %c0 : !pto.tile_buf, index) - outs : f32 - - %vector = pto.alloc_tile : !pto.tile_buf - pto.tprint ins(%vector : !pto.tile_buf) - return %value : f32 - } -} - -// A PIPE_S-touched tile must not share an address with a non-scalar tile even -// when their linear liveness intervals do not overlap. -// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0 : i64 -// CHECK-DAG: %[[NEXT:.*]] = arith.constant 1024 : i64 -// CHECK: pto.alloc_tile addr = %[[ZERO]] : !pto.tile_buf -// CHECK: pto.alloc_tile addr = %[[NEXT]] : !pto.tile_buf diff --git a/test/lit/pto/plan_memory_loop_backedge_liveness.pto b/test/lit/pto/plan_memory_loop_backedge_liveness.pto index a1a2ac7184..dc9c9b24a8 100644 --- a/test/lit/pto/plan_memory_loop_backedge_liveness.pto +++ b/test/lit/pto/plan_memory_loop_backedge_liveness.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_modern_bias_capacity_invalid.pto b/test/lit/pto/plan_memory_modern_bias_capacity_invalid.pto index 0d400bb176..43011f28d5 100644 --- a/test/lit/pto/plan_memory_modern_bias_capacity_invalid.pto +++ b/test/lit/pto/plan_memory_modern_bias_capacity_invalid.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: not ptoas --pto-level=level2 --pto-arch=a3 \ -// RUN: --plan-memory-impl=modern --emit-pto-ir %s 2>&1 | FileCheck %s +// RUN: --emit-pto-ir %s 2>&1 | FileCheck %s module { func.func @bias_capacity_overflow() { diff --git a/test/lit/pto/plan_memory_modern_child_module.pto b/test/lit/pto/plan_memory_modern_child_module.pto index 415bfe7f1e..564caa24b5 100644 --- a/test/lit/pto/plan_memory_modern_child_module.pto +++ b/test/lit/pto/plan_memory_modern_child_module.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level2 --pto-arch=a3 --pto-backend=vpto \ -// RUN: --plan-memory-impl=modern \ +// RUN: \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_modern_first_writer_reuse.pto b/test/lit/pto/plan_memory_modern_first_writer_reuse.pto index fb1542f629..f011fa9d80 100644 --- a/test/lit/pto/plan_memory_modern_first_writer_reuse.pto +++ b/test/lit/pto/plan_memory_modern_first_writer_reuse.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_modern_multiblock_invalid.pto b/test/lit/pto/plan_memory_modern_multiblock_invalid.pto index c0c775ad55..cbfd598005 100644 --- a/test/lit/pto/plan_memory_modern_multiblock_invalid.pto +++ b/test/lit/pto/plan_memory_modern_multiblock_invalid.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level2 --pto-arch=a3 \ -// RUN: --plan-memory-impl=modern --emit-pto-ir %s 2>&1 | FileCheck %s +// RUN: --emit-pto-ir %s 2>&1 | FileCheck %s module { func.func @multiblock_cfg() { diff --git a/test/lit/pto/plan_memory_modern_region_branch_alias.pto b/test/lit/pto/plan_memory_modern_region_branch_alias.pto index bc18e607e7..77cb2f8d5c 100644 --- a/test/lit/pto/plan_memory_modern_region_branch_alias.pto +++ b/test/lit/pto/plan_memory_modern_region_branch_alias.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_modern_unmodeled_region_invalid.pto b/test/lit/pto/plan_memory_modern_unmodeled_region_invalid.pto index 4155137a03..9d7af01ccb 100644 --- a/test/lit/pto/plan_memory_modern_unmodeled_region_invalid.pto +++ b/test/lit/pto/plan_memory_modern_unmodeled_region_invalid.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level2 --pto-arch=a3 \ -// RUN: --plan-memory-impl=modern --emit-pto-ir %s 2>&1 | FileCheck %s +// RUN: --emit-pto-ir %s 2>&1 | FileCheck %s module { func.func @unmodeled_execute_region() { diff --git a/test/lit/pto/plan_memory_order_by_size_noreuse.pto b/test/lit/pto/plan_memory_order_by_size_noreuse.pto index 88ef992031..5448e8b47a 100644 --- a/test/lit/pto/plan_memory_order_by_size_noreuse.pto +++ b/test/lit/pto/plan_memory_order_by_size_noreuse.pto @@ -7,10 +7,8 @@ // tile (dst) is placed first at offset 0. This is the regression guard for the // contract that the option means the same thing whether or not reuse kicks in. // -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT -// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE +// RUN: ptoas --pto-arch=a3 --plan-memory-order-by-size=false --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT // RUN: ptoas --pto-arch=a3 --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE -// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE module { func.func @order_by_size_noreuse(%src_ptr: !pto.ptr, %idx_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.kernel} { diff --git a/test/lit/pto/plan_memory_order_by_size_reuse.pto b/test/lit/pto/plan_memory_order_by_size_reuse.pto index b83ff45090..8f7546d17e 100644 --- a/test/lit/pto/plan_memory_order_by_size_reuse.pto +++ b/test/lit/pto/plan_memory_order_by_size_reuse.pto @@ -5,10 +5,8 @@ // allocated last and lands at a high offset. With --plan-memory-order-by-size // (first-fit-decreasing) the largest tile is allocated first and gets offset 0. // -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT -// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE +// RUN: ptoas --pto-arch=a3 --plan-memory-order-by-size=false --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT // RUN: ptoas --pto-arch=a3 --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE -// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE module { func.func @order_by_size_reuse(%src_ptr: !pto.ptr, %idx_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.kernel} { diff --git a/test/lit/pto/plan_memory_pipev_reuse_cost_state_update.pto b/test/lit/pto/plan_memory_pipev_reuse_cost_state_update.pto index eafbda622a..df9b87dc99 100644 --- a/test/lit/pto/plan_memory_pipev_reuse_cost_state_update.pto +++ b/test/lit/pto/plan_memory_pipev_reuse_cost_state_update.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: | FileCheck %s diff --git a/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto b/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto index 5dffba7a56..aadb8b8cd5 100644 --- a/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto +++ b/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto @@ -108,8 +108,8 @@ module attributes {pto.target_arch = "a2a3"} { } // PLAN-LABEL: func.func @plan_memory_reused_tstore_sync_level2( -// PLAN-DAG: %[[REUSED_ADDR:[A-Za-z0-9_]+]] = arith.constant 8448 : i64 -// PLAN-DAG: %[[STORED_ADDR:[A-Za-z0-9_]+]] = arith.constant 8192 : i64 +// PLAN-DAG: %[[REUSED_ADDR:[A-Za-z0-9_]+]] = arith.constant 180224 : i64 +// PLAN-DAG: %[[STORED_ADDR:[A-Za-z0-9_]+]] = arith.constant 0 : i64 // PLAN: %[[STORED:[0-9]+]] = pto.alloc_tile addr = %[[STORED_ADDR]] : !pto.tile_buf // PLAN: pto.tstore ins(%[[STORED]] // PLAN: %[[REUSED:[0-9]+]] = pto.alloc_tile addr = %[[REUSED_ADDR]] : !pto.tile_buf @@ -118,7 +118,6 @@ module attributes {pto.target_arch = "a2a3"} { // SYNC-LABEL: AICORE void plan_memory_reused_tstore_sync_level2( // SYNC: // pto: %stored // SYNC: TSTORE( -// SYNC-NEXT: set_flag(PIPE_MTE3, PIPE_V, EVENT_ID[[MTE3_TO_V:[0-9]+]]); // SYNC-NOT: TMULS( -// SYNC: wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID[[MTE3_TO_V]]); +// SYNC: pipe_barrier(PIPE_V); // SYNC-NEXT: TMULS( diff --git a/test/lit/pto/plan_memory_row_plus_one_footprint.pto b/test/lit/pto/plan_memory_row_plus_one_footprint.pto index 41dfbaec03..857aa8f65e 100644 --- a/test/lit/pto/plan_memory_row_plus_one_footprint.pto +++ b/test/lit/pto/plan_memory_row_plus_one_footprint.pto @@ -9,7 +9,7 @@ // RUN: ptoas --pto-level=level2 --pto-arch=a3 --emit-pto-ir \ // RUN: --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null \ // RUN: | FileCheck %s --check-prefix=LEGACY -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 \ // RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ // RUN: 1>/dev/null | FileCheck %s --check-prefix=MODERN diff --git a/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto b/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto index bdef29fd59..c5241e02c9 100644 --- a/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto +++ b/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern --emit-pto-ir \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 --emit-pto-ir \ // RUN: --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module attributes {"pto.target_arch" = "a3"} { diff --git a/test/lit/pto/plan_memory_spec_level0_reuse.pto b/test/lit/pto/plan_memory_spec_level0_reuse.pto index 33a6db3d08..3879df1087 100644 --- a/test/lit/pto/plan_memory_spec_level0_reuse.pto +++ b/test/lit/pto/plan_memory_spec_level0_reuse.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern --emit-pto-ir \ +// RUN: ptoas --pto-level=level2 --pto-arch=a3 --emit-pto-ir \ // RUN: --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module attributes {"pto.target_arch" = "a3"} { diff --git a/test/lit/pto/tfillpad_plan_memory_inference.pto b/test/lit/pto/tfillpad_plan_memory_inference.pto index 346d6e083f..2bf4bba98e 100644 --- a/test/lit/pto/tfillpad_plan_memory_inference.pto +++ b/test/lit/pto/tfillpad_plan_memory_inference.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --plan-memory-impl=modern %s | FileCheck %s +// RUN: ptoas %s | FileCheck %s module { func.func @tfillpad_plan_memory_alias() { diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto index 2e0c89f927..976a70d997 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto @@ -19,12 +19,10 @@ // CHECK-NOT: pto.tfillpad ins // CHECK: pto.vecscope // CHECK: pto.castptr -// Normal lowering is the conservative fallback for unprovable addresses. It -// must finish copying the valid source region before issuing any pad stores so -// the generated template remains correct when the addresses alias at runtime. -// CHECK-NOT: pto.vdup -// CHECK: pto.vlds -// CHECK: pto.vsts +// The modern planner assigns src and dst the same tile address (in-place +// fillpad), so the template's conservative fallback elides the self-copy of +// the valid source region and only the pad stores remain. +// CHECK-NOT: pto.vlds // CHECK: pto.pxor // CHECK: pto.vdup // CHECK: pto.vsts diff --git a/test/lit/vpto/fold_tile_buf_intrinsics.pto b/test/lit/vpto/fold_tile_buf_intrinsics.pto index 8f872a449c..09e9b07575 100644 --- a/test/lit/vpto/fold_tile_buf_intrinsics.pto +++ b/test/lit/vpto/fold_tile_buf_intrinsics.pto @@ -44,7 +44,8 @@ // - tile_buf_addr has been folded to concrete pto.castptr addresses // - tile-slice addressing is carried by the vlds/vsts offset operand // NORMALIZED-LABEL: func.func @TADD -// NORMALIZED: pto.castptr +// The modern planner reuses the address of %a for %dst (in-place TADD), so +// only two distinct castptr addresses remain. // NORMALIZED: pto.castptr // NORMALIZED: pto.castptr // NORMALIZED-NOT: pto.tile_buf_addr diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 1b20bb68cc..2cc08bc331 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -222,16 +222,10 @@ llvm::cl::opt planMemoryOrderBySize( "plan-memory-order-by-size", llvm::cl::desc("Plan larger local buffers first inside one AddressSpace " "before applying the basic SPEC_LEVEL_0 reuse strategy. " - "Defaults to true when --plan-memory-impl=modern is " - "explicitly selected"), + "Defaults to false to match the previous default pipeline " + "(largest-first ordering is opt-in)"), llvm::cl::init(false)); -llvm::cl::opt planMemoryImpl( - "plan-memory-impl", - llvm::cl::desc("Select local memory planner implementation: legacy or " - "modern"), - llvm::cl::init("legacy")); - llvm::cl::opt enableBufidSync( "enable-bufid_sync", llvm::cl::desc("Enable A5 buffer-id synchronization insertion pass"), diff --git a/tools/ptoas/ptoas_internal.h b/tools/ptoas/ptoas_internal.h index dad269f6a9..424afdc6e4 100644 --- a/tools/ptoas/ptoas_internal.h +++ b/tools/ptoas/ptoas_internal.h @@ -69,7 +69,6 @@ extern llvm::cl::opt enableVexpdifFusion; extern llvm::cl::opt enableShapeInference; extern llvm::cl::opt enableVfSimCostmodelOptimization; extern llvm::cl::opt dumpVfSimUnrollTest; -extern llvm::cl::opt planMemoryImpl; extern llvm::cl::opt planMemoryOrderBySize; extern llvm::cl::opt ptoBuildLevel; extern llvm::cl::opt disableInferLayout; diff --git a/tools/ptoas/ptoas_pipeline.cpp b/tools/ptoas/ptoas_pipeline.cpp index 058039d0fa..3b2ba61c05 100644 --- a/tools/ptoas/ptoas_pipeline.cpp +++ b/tools/ptoas/ptoas_pipeline.cpp @@ -1355,26 +1355,11 @@ static LogicalResult appendFusionFrontendPasses( static LogicalResult appendPlanMemoryPasses(PassManager &pm, PTOBuildLevel effectiveLevel) { - if (planMemoryImpl != "legacy" && planMemoryImpl != "modern") { - llvm::errs() << "Error: invalid --plan-memory-impl='" << planMemoryImpl - << "', expected 'legacy' or 'modern'.\n"; - return failure(); - } - if (effectiveLevel != PTOBuildLevel::Level3) { pto::PlanMemoryOptions planMemoryOptions; planMemoryOptions.memMode = "local"; - bool effectivePlanMemoryOrderBySize = planMemoryOrderBySize; - if (planMemoryImpl == "modern" && - planMemoryOrderBySize.getNumOccurrences() == 0) { - effectivePlanMemoryOrderBySize = true; - } - planMemoryOptions.orderBySize = effectivePlanMemoryOrderBySize; - if (planMemoryImpl == "legacy") { - pm.addPass(pto::createPlanMemoryPass(planMemoryOptions)); - } else { - pm.addPass(pto::createPlanMemoryModernPass(planMemoryOptions)); - } + planMemoryOptions.orderBySize = planMemoryOrderBySize; + pm.addPass(pto::createPlanMemoryModernPass(planMemoryOptions)); } return success(); }