better reports cleanups
This commit is contained in:
NiccoloN
@@ -5,4 +5,8 @@ add_pim_library(OMPimLocalMemoryLifetimeAnalysis
|
||||
|
||||
INCLUDE_DIRS PUBLIC
|
||||
${PIM_PUBLIC_INCLUDE_DIRS}
|
||||
|
||||
LINK_LIBS PUBLIC
|
||||
OMPimCommon
|
||||
PimOps
|
||||
)
|
||||
|
||||
@@ -2,70 +2,204 @@
|
||||
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||||
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
|
||||
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
#include "llvm/ADT/SmallPtrSet.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace pim {
|
||||
namespace onnx_mlir::pim {
|
||||
namespace {
|
||||
|
||||
struct OperationOrdering {
|
||||
DenseMap<Operation*, uint64_t> position;
|
||||
DenseMap<Operation*, uint64_t> subtreeEnd;
|
||||
uint64_t next = 0;
|
||||
};
|
||||
|
||||
static void orderOperation(Operation* op, OperationOrdering& ordering) {
|
||||
uint64_t end = ordering.next;
|
||||
ordering.position[op] = ordering.next++;
|
||||
for (Region& region : op->getRegions())
|
||||
for (Block& block : region)
|
||||
for (Operation& nested : block) {
|
||||
orderOperation(&nested, ordering);
|
||||
end = std::max(end, ordering.subtreeEnd.lookup(&nested));
|
||||
}
|
||||
ordering.subtreeEnd[op] = end;
|
||||
}
|
||||
|
||||
static OperationOrdering buildOrdering(Operation* coreLikeOp) {
|
||||
OperationOrdering ordering;
|
||||
if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
|
||||
return ordering;
|
||||
for (Operation& op : coreLikeOp->getRegion(0).front())
|
||||
orderOperation(&op, ordering);
|
||||
return ordering;
|
||||
}
|
||||
|
||||
static bool isRuntimeMemoryTouch(Operation* op) {
|
||||
return isa<PimMemCopyHostToDevOp,
|
||||
PimMemCopyDevToHostOp,
|
||||
PimMemCopyOp,
|
||||
PimReceiveOp,
|
||||
PimSendOp,
|
||||
PimConcatOp,
|
||||
PimVMMOp,
|
||||
PimTransposeOp,
|
||||
PimVVAddOp,
|
||||
PimVVSubOp,
|
||||
PimVVMulOp,
|
||||
PimVVMaxOp,
|
||||
PimVVDMulOp,
|
||||
PimVAvgOp,
|
||||
PimVReluOp,
|
||||
PimVTanhOp,
|
||||
PimVSigmOp,
|
||||
PimVSoftmaxOp>(op);
|
||||
}
|
||||
|
||||
static bool isIgnoredUser(Operation* op) {
|
||||
return isLocalMemoryAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op)
|
||||
|| isCoreStaticAddressOp(op);
|
||||
}
|
||||
|
||||
static bool isWithin(Value value, Region* region) {
|
||||
if (auto argument = dyn_cast<BlockArgument>(value))
|
||||
return argument.getOwner()->getParent() == region;
|
||||
if (Operation* definingOp = value.getDefiningOp())
|
||||
return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion());
|
||||
return false;
|
||||
}
|
||||
|
||||
static std::pair<uint64_t, uint64_t> getTouchRange(Value definingValue, Operation* user, const OperationOrdering& ordering) {
|
||||
uint64_t start = ordering.position.lookup(user);
|
||||
uint64_t end = start;
|
||||
for (Operation* current = user; current; current = current->getParentOp())
|
||||
if (auto forOp = dyn_cast<scf::ForOp>(current); forOp && !isWithin(definingValue, &forOp.getRegion())) {
|
||||
start = std::min(start, ordering.position.lookup(forOp));
|
||||
end = std::max(end, ordering.subtreeEnd.lookup(forOp));
|
||||
}
|
||||
return {start, end};
|
||||
}
|
||||
|
||||
static FailureOr<size_t> getAllocationSize(memref::AllocOp allocation) {
|
||||
auto type = dyn_cast<ShapedType>(allocation.getType());
|
||||
if (!type)
|
||||
return failure();
|
||||
auto bytes = getCheckedShapedTypeSizeInBytes(type, allocation, "memory allocation byte size");
|
||||
if (failed(bytes))
|
||||
return failure();
|
||||
return checkedSize(*bytes, allocation, "memory allocation byte size");
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
bool isLocalMemoryAliasOp(Operation* op) {
|
||||
return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
|
||||
}
|
||||
|
||||
LogicalResult walkLocalMemoryUses(Value root,
|
||||
llvm::function_ref<LogicalResult(Value, Operation*)> visitUser) {
|
||||
LogicalResult walkLocalMemoryUses(Value root, LocalMemoryUseCallback callback) {
|
||||
llvm::SmallPtrSet<Value, 16> visitedValues;
|
||||
llvm::SmallPtrSet<Operation*, 32> visitedUsers;
|
||||
llvm::SmallVector<Value> pendingValues {root};
|
||||
auto addAlias = [&](Value value) { pendingValues.push_back(value); };
|
||||
|
||||
while (!pendingValues.empty()) {
|
||||
Value value = pendingValues.pop_back_val();
|
||||
llvm::SmallPtrSet<OpOperand*, 32> visitedUses;
|
||||
SmallVector<Value> pending = {root};
|
||||
while (!pending.empty()) {
|
||||
Value value = pending.pop_back_val();
|
||||
if (!visitedValues.insert(value).second)
|
||||
continue;
|
||||
for (Operation* user : value.getUsers()) {
|
||||
if (!visitedUsers.insert(user).second)
|
||||
continue;
|
||||
if (failed(visitUser(value, user)))
|
||||
for (OpOperand& use : value.getUses()) {
|
||||
if (!visitedUses.insert(&use).second || failed(callback(use)))
|
||||
return failure();
|
||||
Operation* user = use.getOwner();
|
||||
|
||||
if (isLocalMemoryAliasOp(user))
|
||||
for (Value result : user->getResults())
|
||||
addAlias(result);
|
||||
if (isLocalMemoryAliasOp(user) && use.getOperandNumber() == 0) {
|
||||
if (user->getNumResults() != 1)
|
||||
return failure();
|
||||
pending.push_back(user->getResult(0));
|
||||
}
|
||||
|
||||
if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user))
|
||||
for (OpResult result : user->getResults())
|
||||
if (OpOperand* tied = dpsOp.getTiedOpOperand(result); tied && tied->get() == value)
|
||||
addAlias(result);
|
||||
if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user); dpsOp && dpsOp.isDpsInit(&use))
|
||||
pending.push_back(dpsOp.getTiedOpResult(&use));
|
||||
|
||||
if (auto forOp = dyn_cast<scf::ForOp>(user))
|
||||
for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs()))
|
||||
if (initArg == value) {
|
||||
addAlias(forOp.getRegionIterArgs()[index]);
|
||||
addAlias(forOp.getResult(index));
|
||||
for (auto [index, init] : llvm::enumerate(forOp.getInitArgsMutable()))
|
||||
if (&init == &use) {
|
||||
pending.push_back(forOp.getRegionIterArgs()[index]);
|
||||
pending.push_back(forOp.getResult(index));
|
||||
}
|
||||
|
||||
auto yieldOp = dyn_cast<scf::YieldOp>(user);
|
||||
if (!yieldOp)
|
||||
continue;
|
||||
for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
|
||||
if (operand != value)
|
||||
continue;
|
||||
if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp()))
|
||||
addAlias(forOp.getResult(index));
|
||||
else if (auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp()))
|
||||
addAlias(ifOp.getResult(index));
|
||||
else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp()))
|
||||
addAlias(switchOp.getResult(index));
|
||||
if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
|
||||
unsigned index = use.getOperandNumber();
|
||||
Operation* parent = yieldOp->getParentOp();
|
||||
if (auto forOp = dyn_cast<scf::ForOp>(parent))
|
||||
pending.push_back(forOp.getResult(index));
|
||||
else if (auto ifOp = dyn_cast<scf::IfOp>(parent))
|
||||
pending.push_back(ifOp.getResult(index));
|
||||
else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(parent))
|
||||
pending.push_back(switchOp.getResult(index));
|
||||
else
|
||||
return failure();
|
||||
}
|
||||
}
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
} // namespace pim
|
||||
} // namespace onnx_mlir
|
||||
FailureOr<SmallVector<LocalMemoryInterval, 0>> analyzeLocalMemoryLifetimes(Operation* coreLikeOp) {
|
||||
SmallVector<LocalMemoryInterval, 0> intervals;
|
||||
OperationOrdering ordering = buildOrdering(coreLikeOp);
|
||||
if (ordering.position.empty())
|
||||
return intervals;
|
||||
uint64_t fallbackEnd = ordering.next - 1;
|
||||
bool failedSize = false;
|
||||
coreLikeOp->walk([&](memref::AllocOp allocation) {
|
||||
auto size = getAllocationSize(allocation);
|
||||
if (failed(size)) {
|
||||
failedSize = true;
|
||||
return;
|
||||
}
|
||||
|
||||
LocalMemoryInterval interval {allocation, ordering.position.lookup(allocation), ordering.position.lookup(allocation), *size};
|
||||
bool hasRuntimeTouch = false;
|
||||
bool needsFallback = false;
|
||||
if (failed(walkLocalMemoryUses(allocation.getResult(), [&](OpOperand& use) {
|
||||
Operation* user = use.getOwner();
|
||||
if (isRuntimeMemoryTouch(user)) {
|
||||
auto [start, end] = getTouchRange(allocation.getResult(), user, ordering);
|
||||
if (!hasRuntimeTouch) {
|
||||
interval.start = start;
|
||||
interval.end = end;
|
||||
hasRuntimeTouch = true;
|
||||
}
|
||||
else {
|
||||
interval.start = std::min(interval.start, start);
|
||||
interval.end = std::max(interval.end, end);
|
||||
}
|
||||
}
|
||||
else if (!isIgnoredUser(user)) {
|
||||
needsFallback = true;
|
||||
}
|
||||
return success();
|
||||
})))
|
||||
needsFallback = true;
|
||||
if (!hasRuntimeTouch || needsFallback) {
|
||||
interval.start = std::min(interval.start, ordering.position.lookup(allocation));
|
||||
interval.end = fallbackEnd;
|
||||
}
|
||||
intervals.push_back(interval);
|
||||
});
|
||||
if (failedSize)
|
||||
return failure();
|
||||
return intervals;
|
||||
}
|
||||
|
||||
bool localMemoryLifetimesOverlap(const LocalMemoryInterval& lhs, const LocalMemoryInterval& rhs) {
|
||||
return !(lhs.end < rhs.start || rhs.end < lhs.start);
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir::pim
|
||||
|
||||
@@ -1,17 +1,25 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/IR/Operation.h"
|
||||
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
||||
|
||||
#include "llvm/ADT/STLFunctionalExtras.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace pim {
|
||||
namespace onnx_mlir::pim {
|
||||
|
||||
struct LocalMemoryInterval {
|
||||
mlir::memref::AllocOp allocation;
|
||||
uint64_t start = 0;
|
||||
uint64_t end = 0;
|
||||
size_t size = 0;
|
||||
};
|
||||
|
||||
using LocalMemoryUseCallback = llvm::function_ref<mlir::LogicalResult(mlir::OpOperand&)>;
|
||||
|
||||
bool isLocalMemoryAliasOp(mlir::Operation* op);
|
||||
mlir::LogicalResult walkLocalMemoryUses(mlir::Value root, LocalMemoryUseCallback callback);
|
||||
mlir::FailureOr<llvm::SmallVector<LocalMemoryInterval, 0>>
|
||||
analyzeLocalMemoryLifetimes(mlir::Operation* coreLikeOp);
|
||||
bool localMemoryLifetimesOverlap(const LocalMemoryInterval& lhs, const LocalMemoryInterval& rhs);
|
||||
|
||||
mlir::LogicalResult walkLocalMemoryUses(
|
||||
mlir::Value root,
|
||||
llvm::function_ref<mlir::LogicalResult(mlir::Value, mlir::Operation*)> visitUser);
|
||||
|
||||
} // namespace pim
|
||||
} // namespace onnx_mlir
|
||||
} // namespace onnx_mlir::pim
|
||||
|
||||
@@ -4,7 +4,6 @@ add_onnx_mlir_dialect_doc(pim Pim.td)
|
||||
add_subdirectory(Analysis)
|
||||
add_subdirectory(Transforms/Bufferization)
|
||||
add_subdirectory(Transforms/HostConstantFolding)
|
||||
add_subdirectory(Transforms/MemoryCoalescing)
|
||||
add_subdirectory(Transforms/LocalMemoryPlanning)
|
||||
add_subdirectory(Transforms/Verification)
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ add_pim_library(OMPimLocalMemoryPlanning
|
||||
|
||||
LINK_LIBS PUBLIC
|
||||
OMPimCommon
|
||||
OMPimCompilerOptions
|
||||
OMPimLocalMemoryLifetimeAnalysis
|
||||
PimOps
|
||||
)
|
||||
|
||||
@@ -1,668 +1,169 @@
|
||||
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
|
||||
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
||||
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||||
#include "mlir/Pass/Pass.h"
|
||||
#include "mlir/IR/Value.h"
|
||||
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
#include "llvm/Support/MathExtras.h"
|
||||
#include "llvm/Support/FileSystem.h"
|
||||
#include "llvm/Support/raw_os_ostream.h"
|
||||
#include "llvm/Support/raw_ostream.h"
|
||||
|
||||
#include <fstream>
|
||||
#include <memory>
|
||||
#include <numeric>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
|
||||
#include "Common/Support/CheckedArithmetic.hpp"
|
||||
#include "Common/Support/ReportUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
|
||||
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
|
||||
|
||||
using namespace llvm;
|
||||
using namespace mlir;
|
||||
using namespace onnx_mlir;
|
||||
|
||||
namespace {
|
||||
|
||||
struct MemoryTouchInterval {
|
||||
uint64_t start = 0;
|
||||
uint64_t end = 0;
|
||||
Operation* firstTouchOp = nullptr;
|
||||
Operation* lastTouchOp = nullptr;
|
||||
uint64_t firstTouchPosition = 0;
|
||||
uint64_t lastTouchPosition = 0;
|
||||
bool hasRuntimeUse = false;
|
||||
bool startUsedAllocFallback = false;
|
||||
bool endUsedFallback = false;
|
||||
bool escapesLoop = false;
|
||||
std::string fallbackReason;
|
||||
};
|
||||
|
||||
struct OperationOrdering {
|
||||
llvm::DenseMap<Operation*, uint64_t> position;
|
||||
llvm::DenseMap<Operation*, uint64_t> subtreeEnd;
|
||||
uint64_t nextPosition = 0;
|
||||
};
|
||||
|
||||
static std::string abbreviate(StringRef text, size_t maxLen) {
|
||||
if (text.size() <= maxLen)
|
||||
return text.str();
|
||||
return (text.take_front(maxLen - 3) + "...").str();
|
||||
}
|
||||
|
||||
static std::string summarizeValue(mlir::Value value, size_t maxLen = 72) {
|
||||
std::string text;
|
||||
llvm::raw_string_ostream os(text);
|
||||
if (auto result = dyn_cast<OpResult>(value))
|
||||
os << result.getOwner()->getName() << '#' << result.getResultNumber();
|
||||
else if (auto blockArg = dyn_cast<BlockArgument>(value))
|
||||
os << "block_arg#" << blockArg.getArgNumber();
|
||||
else
|
||||
os << "<unknown value>";
|
||||
os << " : " << value.getType();
|
||||
os.flush();
|
||||
return abbreviate(text, maxLen);
|
||||
}
|
||||
|
||||
static std::string summarizeOperation(Operation* op, size_t maxLen = 96) {
|
||||
if (!op)
|
||||
return "<none>";
|
||||
return abbreviate(op->getName().getStringRef(), maxLen);
|
||||
}
|
||||
|
||||
static void assignOperationOrdering(Operation* op, OperationOrdering& ordering) {
|
||||
uint64_t position = ordering.nextPosition++;
|
||||
ordering.position[op] = position;
|
||||
uint64_t end = position;
|
||||
for (Region& region : op->getRegions())
|
||||
for (Block& block : region)
|
||||
for (Operation& nestedOp : block) {
|
||||
assignOperationOrdering(&nestedOp, ordering);
|
||||
end = std::max(end, ordering.subtreeEnd.lookup(&nestedOp));
|
||||
}
|
||||
ordering.subtreeEnd[op] = end;
|
||||
}
|
||||
|
||||
static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) {
|
||||
OperationOrdering ordering;
|
||||
if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
|
||||
return ordering;
|
||||
|
||||
for (Operation& op : coreLikeOp->getRegion(0).front())
|
||||
assignOperationOrdering(&op, ordering);
|
||||
return ordering;
|
||||
}
|
||||
|
||||
static bool isRuntimeMemoryTouchOp(Operation* op) {
|
||||
return isa<pim::PimMemCopyHostToDevOp,
|
||||
pim::PimMemCopyDevToHostOp,
|
||||
pim::PimMemCopyOp,
|
||||
pim::PimReceiveOp,
|
||||
pim::PimSendOp,
|
||||
pim::PimConcatOp,
|
||||
pim::PimVMMOp,
|
||||
pim::PimTransposeOp,
|
||||
pim::PimVVAddOp,
|
||||
pim::PimVVSubOp,
|
||||
pim::PimVVMulOp,
|
||||
pim::PimVVMaxOp,
|
||||
pim::PimVVDMulOp,
|
||||
pim::PimVAvgOp,
|
||||
pim::PimVReluOp,
|
||||
pim::PimVTanhOp,
|
||||
pim::PimVSigmOp,
|
||||
pim::PimVSoftmaxOp>(op);
|
||||
}
|
||||
|
||||
static bool isIgnoredLivenessUser(Operation* op) {
|
||||
return pim::isLocalMemoryAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op)
|
||||
|| isCoreStaticAddressOp(op);
|
||||
}
|
||||
|
||||
static bool isWithin(mlir::Value value, Region* region) {
|
||||
if (!region)
|
||||
return false;
|
||||
if (auto blockArg = dyn_cast<BlockArgument>(value))
|
||||
return blockArg.getOwner()->getParent() == region;
|
||||
if (Operation* definingOp = value.getDefiningOp())
|
||||
return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion());
|
||||
return false;
|
||||
}
|
||||
|
||||
static bool isNestedAllocation(Operation* coreLikeOp, memref::AllocOp allocOp) {
|
||||
if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
|
||||
return false;
|
||||
return allocOp->getBlock() != &coreLikeOp->getRegion(0).front();
|
||||
}
|
||||
|
||||
static void addFallbackReason(std::string& reason, StringRef newReason) {
|
||||
if (newReason.empty())
|
||||
return;
|
||||
if (!reason.empty())
|
||||
reason += "; ";
|
||||
reason += newReason.str();
|
||||
}
|
||||
|
||||
struct OrderedTouchRange {
|
||||
uint64_t start = 0;
|
||||
uint64_t end = 0;
|
||||
bool escapedLoop = false;
|
||||
};
|
||||
|
||||
static OrderedTouchRange
|
||||
getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const OperationOrdering& ordering) {
|
||||
OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), false};
|
||||
for (Operation* current = user; current; current = current->getParentOp()) {
|
||||
auto forOp = dyn_cast<scf::ForOp>(current);
|
||||
if (!forOp || isWithin(definingValue, &forOp.getRegion()))
|
||||
continue;
|
||||
range.start = std::min(range.start, ordering.position.lookup(forOp));
|
||||
range.end = std::max(range.end, ordering.subtreeEnd.lookup(forOp));
|
||||
range.escapedLoop = true;
|
||||
}
|
||||
return range;
|
||||
}
|
||||
|
||||
static MemoryTouchInterval computeMemoryTouchInterval(memref::AllocOp allocOp,
|
||||
const OperationOrdering& ordering,
|
||||
uint64_t fallbackEnd) {
|
||||
MemoryTouchInterval interval;
|
||||
interval.start = ordering.position.lookup(allocOp);
|
||||
interval.end = interval.start;
|
||||
|
||||
auto parentLoop = allocOp->getParentOfType<scf::ForOp>();
|
||||
(void) pim::walkLocalMemoryUses(
|
||||
allocOp.getResult(),
|
||||
[&](mlir::Value value, Operation* user) {
|
||||
if (auto forOp = dyn_cast<scf::ForOp>(user);
|
||||
forOp && parentLoop && forOp != parentLoop && llvm::is_contained(forOp.getInitArgs(), value))
|
||||
interval.escapesLoop = true;
|
||||
if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
|
||||
auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
|
||||
auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp());
|
||||
auto indexSwitch = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp());
|
||||
if (!forOp && !ifOp && !indexSwitch)
|
||||
addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
|
||||
else if (forOp && parentLoop && forOp == parentLoop && llvm::is_contained(yieldOp.getOperands(), value))
|
||||
interval.escapesLoop = true;
|
||||
}
|
||||
|
||||
if (isRuntimeMemoryTouchOp(user)) {
|
||||
uint64_t touchPosition = ordering.position.lookup(user);
|
||||
if (!interval.hasRuntimeUse || touchPosition < interval.firstTouchPosition) {
|
||||
interval.firstTouchPosition = touchPosition;
|
||||
interval.firstTouchOp = user;
|
||||
}
|
||||
if (!interval.hasRuntimeUse || touchPosition > interval.lastTouchPosition) {
|
||||
interval.lastTouchPosition = touchPosition;
|
||||
interval.lastTouchOp = user;
|
||||
}
|
||||
|
||||
OrderedTouchRange range = getEffectiveTouchRange(allocOp.getResult(), user, ordering);
|
||||
interval.escapesLoop |= range.escapedLoop;
|
||||
if (!interval.hasRuntimeUse) {
|
||||
interval.start = range.start;
|
||||
interval.end = range.end;
|
||||
interval.hasRuntimeUse = true;
|
||||
}
|
||||
else {
|
||||
if (range.start < interval.start)
|
||||
interval.start = range.start;
|
||||
if (range.end > interval.end)
|
||||
interval.end = range.end;
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
if (isIgnoredLivenessUser(user))
|
||||
return success();
|
||||
|
||||
addFallbackReason(interval.fallbackReason, "unhandled user op");
|
||||
interval.endUsedFallback = true;
|
||||
return success();
|
||||
});
|
||||
|
||||
if (!interval.hasRuntimeUse) {
|
||||
interval.startUsedAllocFallback = true;
|
||||
interval.endUsedFallback = true;
|
||||
interval.start = ordering.position.lookup(allocOp);
|
||||
interval.end = fallbackEnd;
|
||||
interval.firstTouchPosition = interval.start;
|
||||
interval.lastTouchPosition = interval.end;
|
||||
addFallbackReason(interval.fallbackReason, "no runtime memory touch");
|
||||
return interval;
|
||||
}
|
||||
|
||||
if (interval.endUsedFallback)
|
||||
interval.end = std::max(interval.end, fallbackEnd);
|
||||
|
||||
return interval;
|
||||
}
|
||||
|
||||
static FailureOr<size_t> getAllocSizeBytes(memref::AllocOp allocOp) {
|
||||
auto type = dyn_cast<ShapedType>(allocOp.getType());
|
||||
if (!type)
|
||||
return failure();
|
||||
auto checkedBytes = pim::getCheckedShapedTypeSizeInBytes(type, allocOp, "memory allocation byte size");
|
||||
if (failed(checkedBytes))
|
||||
return failure();
|
||||
return pim::checkedSize(*checkedBytes, allocOp, "memory allocation byte size");
|
||||
}
|
||||
|
||||
static bool intervalsOverlap(const LocalAllocInterval& lhs, const LocalAllocInterval& rhs) {
|
||||
return !(lhs.end < rhs.start || rhs.end < lhs.start);
|
||||
}
|
||||
|
||||
static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, ArrayRef<LocalAllocInterval> intervals) {
|
||||
uint64_t slotLogicalBytes = 0;
|
||||
for (size_t intervalIndex : slot.intervalIndices)
|
||||
slotLogicalBytes += intervals[intervalIndex].size;
|
||||
return slotLogicalBytes;
|
||||
}
|
||||
|
||||
static bool placementsOverlap(const PlannedPhysicalSlot& lhs, const PlannedPhysicalSlot& rhs) {
|
||||
return lhs.address < rhs.address + rhs.requiredSize && rhs.address < lhs.address + lhs.requiredSize;
|
||||
}
|
||||
|
||||
static bool hasAddressReuse(size_t slotIndex, ArrayRef<PlannedPhysicalSlot> slots) {
|
||||
if (slots[slotIndex].intervalIndices.size() > 1)
|
||||
return true;
|
||||
return llvm::any_of(llvm::enumerate(slots), [&](auto indexedSlot) {
|
||||
return indexedSlot.index() != slotIndex && placementsOverlap(slots[slotIndex], indexedSlot.value());
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp) {
|
||||
SmallVector<LocalAllocInterval, 0> intervals;
|
||||
OperationOrdering ordering = buildOperationOrdering(coreLikeOp);
|
||||
if (ordering.position.empty())
|
||||
return intervals;
|
||||
|
||||
uint64_t fallbackEnd = ordering.nextPosition == 0 ? 0 : ordering.nextPosition - 1;
|
||||
size_t nextIntervalId = 0;
|
||||
coreLikeOp->walk([&](memref::AllocOp allocOp) {
|
||||
auto checkedSize = getAllocSizeBytes(allocOp);
|
||||
if (failed(checkedSize)) {
|
||||
llvm::errs() << "Failed to compute local allocation size for value: ";
|
||||
allocOp.getResult().print(llvm::errs());
|
||||
llvm::errs() << "\n";
|
||||
llvm_unreachable("Failed to compute local allocation size");
|
||||
}
|
||||
|
||||
MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd);
|
||||
LocalAllocInterval interval;
|
||||
interval.id = nextIntervalId++;
|
||||
interval.alloc = allocOp;
|
||||
interval.start = touchInterval.start;
|
||||
interval.end = touchInterval.end;
|
||||
interval.size = *checkedSize;
|
||||
interval.firstTouchOp = touchInterval.firstTouchOp;
|
||||
interval.lastTouchOp = touchInterval.lastTouchOp;
|
||||
interval.firstTouchPosition = touchInterval.firstTouchPosition;
|
||||
interval.lastTouchPosition = touchInterval.lastTouchPosition;
|
||||
interval.startUsedAllocFallback = touchInterval.startUsedAllocFallback;
|
||||
interval.endUsedFallback = touchInterval.endUsedFallback;
|
||||
interval.hasRuntimeUse = touchInterval.hasRuntimeUse;
|
||||
interval.insideNestedRegion = isNestedAllocation(coreLikeOp, allocOp);
|
||||
interval.escapesLoop = touchInterval.escapesLoop;
|
||||
interval.fallbackReason = std::move(touchInterval.fallbackReason);
|
||||
interval.valueSummary = summarizeValue(allocOp.getResult(), 88);
|
||||
interval.firstTouchSummary = summarizeOperation(touchInterval.firstTouchOp);
|
||||
interval.lastTouchSummary = summarizeOperation(touchInterval.lastTouchOp);
|
||||
intervals.push_back(std::move(interval));
|
||||
});
|
||||
|
||||
return intervals;
|
||||
}
|
||||
|
||||
SmallVector<PlannedPhysicalSlot, 0> onnx_mlir::planPhysicalSlots(MutableArrayRef<LocalAllocInterval> intervals) {
|
||||
SmallVector<PlannedPhysicalSlot, 0> slots;
|
||||
SmallVector<size_t> intervalOrder(intervals.size());
|
||||
std::iota(intervalOrder.begin(), intervalOrder.end(), 0);
|
||||
llvm::stable_sort(intervalOrder, [&](size_t lhsIndex, size_t rhsIndex) {
|
||||
const LocalAllocInterval& lhs = intervals[lhsIndex];
|
||||
const LocalAllocInterval& rhs = intervals[rhsIndex];
|
||||
if (lhs.size != rhs.size)
|
||||
return lhs.size > rhs.size;
|
||||
if (lhs.start != rhs.start)
|
||||
return lhs.start < rhs.start;
|
||||
if (lhs.end != rhs.end)
|
||||
return lhs.end < rhs.end;
|
||||
return lhs.id < rhs.id;
|
||||
});
|
||||
|
||||
for (size_t intervalIndex : intervalOrder) {
|
||||
LocalAllocInterval& interval = intervals[intervalIndex];
|
||||
SmallVector<const PlannedPhysicalSlot*, 16> conflictingSlots;
|
||||
SmallVector<size_t, 16> candidateAddresses = {0};
|
||||
size_t currentPeak = 0;
|
||||
for (const PlannedPhysicalSlot& slot : slots) {
|
||||
currentPeak = std::max(currentPeak, slot.address + slot.requiredSize);
|
||||
if (llvm::any_of(slot.intervalIndices, [&](size_t otherIndex) {
|
||||
return intervalsOverlap(interval, intervals[otherIndex]);
|
||||
})) {
|
||||
conflictingSlots.push_back(&slot);
|
||||
size_t candidate = llvm::alignTo(slot.address + slot.requiredSize, 4);
|
||||
if (!llvm::is_contained(candidateAddresses, candidate))
|
||||
candidateAddresses.push_back(candidate);
|
||||
}
|
||||
}
|
||||
|
||||
size_t bestAddress = std::numeric_limits<size_t>::max();
|
||||
auto bestKey = std::tuple<size_t, size_t>(std::numeric_limits<size_t>::max(),
|
||||
std::numeric_limits<size_t>::max());
|
||||
for (size_t candidate : candidateAddresses) {
|
||||
bool overlaps = llvm::any_of(conflictingSlots, [&](const PlannedPhysicalSlot* slot) {
|
||||
return candidate < slot->address + slot->requiredSize && slot->address < candidate + interval.size;
|
||||
});
|
||||
if (overlaps)
|
||||
continue;
|
||||
auto candidateKey = std::tuple<size_t, size_t>(std::max(currentPeak, candidate + interval.size), candidate);
|
||||
if (candidateKey < bestKey) {
|
||||
bestKey = candidateKey;
|
||||
bestAddress = candidate;
|
||||
}
|
||||
}
|
||||
assert(bestAddress != std::numeric_limits<size_t>::max() && "address after all conflicts must be available");
|
||||
|
||||
auto reusable = llvm::find_if(slots, [&](const PlannedPhysicalSlot& slot) {
|
||||
return slot.address == bestAddress && slot.requiredSize == interval.size;
|
||||
});
|
||||
if (reusable != slots.end()) {
|
||||
reusable->intervalIndices.push_back(intervalIndex);
|
||||
interval.slotPlanIndex = static_cast<size_t>(reusable - slots.begin());
|
||||
}
|
||||
else {
|
||||
slots.push_back({slots.size(), interval.size, bestAddress, {intervalIndex}});
|
||||
interval.slotPlanIndex = slots.size() - 1;
|
||||
}
|
||||
}
|
||||
|
||||
return slots;
|
||||
}
|
||||
|
||||
CoreMemoryPlan onnx_mlir::buildCoreMemoryPlan(Operation* coreLikeOp) {
|
||||
CoreMemoryPlan plan;
|
||||
plan.coreLikeOp = coreLikeOp;
|
||||
plan.intervals = buildLocalAllocIntervals(coreLikeOp);
|
||||
plan.slots = planPhysicalSlots(plan.intervals);
|
||||
return plan;
|
||||
}
|
||||
|
||||
LogicalResult onnx_mlir::assignPhysicalSlotAddresses(CoreMemoryPlan& plan, size_t addressLimit) {
|
||||
SmallVector<size_t> slotOrder(plan.slots.size());
|
||||
std::iota(slotOrder.begin(), slotOrder.end(), 0);
|
||||
llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) {
|
||||
const PlannedPhysicalSlot& lhs = plan.slots[lhsIndex];
|
||||
const PlannedPhysicalSlot& rhs = plan.slots[rhsIndex];
|
||||
if (lhs.requiredSize != rhs.requiredSize)
|
||||
return lhs.requiredSize > rhs.requiredSize;
|
||||
return lhs.id < rhs.id;
|
||||
});
|
||||
|
||||
size_t nextSlotId = 0;
|
||||
for (size_t slotIndex : slotOrder) {
|
||||
PlannedPhysicalSlot& slot = plan.slots[slotIndex];
|
||||
if (slot.address > addressLimit || slot.requiredSize > addressLimit - slot.address) {
|
||||
plan.coreLikeOp->emitError() << "PIM local memory plan exceeds the signed int32 address range while assigning "
|
||||
"a "
|
||||
<< slot.requiredSize << " byte physical placement at address " << slot.address;
|
||||
return failure();
|
||||
}
|
||||
|
||||
slot.id = nextSlotId++;
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
std::string onnx_mlir::buildMemoryPlanReport(Operation* coreLikeOp,
|
||||
ArrayRef<LocalAllocInterval> intervals,
|
||||
ArrayRef<PlannedPhysicalSlot> slots,
|
||||
size_t addressLimit,
|
||||
PimMemoryReportLevel reportLevel) {
|
||||
std::string artifacts;
|
||||
|
||||
uint64_t totalLogicalBytes = 0;
|
||||
uint64_t totalPhysicalBytes = 0;
|
||||
uint64_t fallbackIntervals = 0;
|
||||
uint64_t noRuntimeTouchIntervals = 0;
|
||||
uint64_t reusedAllocations = 0;
|
||||
uint64_t nestedIntervals = 0;
|
||||
uint64_t loopEscapingIntervals = 0;
|
||||
size_t largestLogicalAllocation = 0;
|
||||
size_t largestPhysicalSlot = 0;
|
||||
size_t maximumAssignedAddress = 0;
|
||||
|
||||
for (const LocalAllocInterval& interval : intervals) {
|
||||
totalLogicalBytes += interval.size;
|
||||
largestLogicalAllocation = std::max(largestLogicalAllocation, interval.size);
|
||||
if (interval.startUsedAllocFallback || interval.endUsedFallback)
|
||||
++fallbackIntervals;
|
||||
if (!interval.hasRuntimeUse)
|
||||
++noRuntimeTouchIntervals;
|
||||
if (interval.insideNestedRegion)
|
||||
++nestedIntervals;
|
||||
if (interval.escapesLoop)
|
||||
++loopEscapingIntervals;
|
||||
}
|
||||
for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
|
||||
const PlannedPhysicalSlot& slot = slots[slotIndex];
|
||||
largestPhysicalSlot = std::max(largestPhysicalSlot, slot.requiredSize);
|
||||
maximumAssignedAddress = std::max(maximumAssignedAddress, slot.address + slot.requiredSize);
|
||||
if (hasAddressReuse(slotIndex, slots))
|
||||
reusedAllocations += slot.intervalIndices.size();
|
||||
}
|
||||
totalPhysicalBytes = maximumAssignedAddress;
|
||||
|
||||
uint64_t savedBytes = totalLogicalBytes >= totalPhysicalBytes ? totalLogicalBytes - totalPhysicalBytes : 0;
|
||||
double savedPercent =
|
||||
totalLogicalBytes == 0 ? 0.0 : 100.0 * static_cast<double>(savedBytes) / static_cast<double>(totalLogicalBytes);
|
||||
|
||||
raw_string_ostream os(artifacts);
|
||||
os << "=== PIM Memory Liveness Report ===\n";
|
||||
os << "Op: " << coreLikeOp->getName() << "\n";
|
||||
os << "Summary:\n";
|
||||
os << " logical allocation bytes: " << formatReportMemory(totalLogicalBytes) << " (" << totalLogicalBytes << ")\n";
|
||||
os << " physical allocation bytes: " << formatReportMemory(totalPhysicalBytes) << " (" << totalPhysicalBytes
|
||||
<< ")\n";
|
||||
os << " saved bytes: " << formatReportMemory(savedBytes) << " (" << savedBytes << ")\n";
|
||||
os << " saved percent: " << format("%.2f%%", savedPercent) << "\n";
|
||||
os << " intervals: " << intervals.size() << "\n";
|
||||
os << " address placements: " << slots.size() << "\n";
|
||||
os << " allocations sharing address space: " << reusedAllocations << "\n";
|
||||
os << " fallback intervals: " << fallbackIntervals << "\n";
|
||||
os << " intervals with no runtime memory touch: " << noRuntimeTouchIntervals << "\n";
|
||||
os << " nested allocations: " << nestedIntervals << "\n";
|
||||
os << " loop-escaping allocations: " << loopEscapingIntervals << "\n";
|
||||
os << " largest logical allocation: " << largestLogicalAllocation << "\n";
|
||||
os << " largest address placement: " << largestPhysicalSlot << "\n";
|
||||
os << " address limit: " << addressLimit << "\n";
|
||||
os << " peak physical memory: " << formatReportMemory(maximumAssignedAddress) << " (" << maximumAssignedAddress
|
||||
<< ")\n";
|
||||
os << " maximum assigned address: " << maximumAssignedAddress << "\n";
|
||||
|
||||
SmallVector<const PlannedPhysicalSlot*> reusedSlots;
|
||||
SmallVector<const PlannedPhysicalSlot*> singleUseSlots;
|
||||
for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
|
||||
if (hasAddressReuse(slotIndex, slots))
|
||||
reusedSlots.push_back(&slots[slotIndex]);
|
||||
else
|
||||
singleUseSlots.push_back(&slots[slotIndex]);
|
||||
}
|
||||
|
||||
llvm::stable_sort(reusedSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
|
||||
uint64_t lhsLogicalBytes = getSlotLogicalBytes(*lhs, intervals);
|
||||
uint64_t rhsLogicalBytes = getSlotLogicalBytes(*rhs, intervals);
|
||||
if (lhs->intervalIndices.size() != rhs->intervalIndices.size())
|
||||
return lhs->intervalIndices.size() > rhs->intervalIndices.size();
|
||||
if (lhsLogicalBytes != rhsLogicalBytes)
|
||||
return lhsLogicalBytes > rhsLogicalBytes;
|
||||
if (lhs->requiredSize != rhs->requiredSize)
|
||||
return lhs->requiredSize > rhs->requiredSize;
|
||||
return lhs->id < rhs->id;
|
||||
});
|
||||
llvm::stable_sort(singleUseSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
|
||||
if (lhs->requiredSize != rhs->requiredSize)
|
||||
return lhs->requiredSize > rhs->requiredSize;
|
||||
return lhs->id < rhs->id;
|
||||
});
|
||||
|
||||
constexpr size_t kSummaryReuseLimit = 6;
|
||||
constexpr size_t kSummaryOffenderLimit = 10;
|
||||
|
||||
os << "\nBest Address Reuse:\n";
|
||||
if (reusedSlots.empty()) {
|
||||
os << " no address ranges were reused\n";
|
||||
}
|
||||
else {
|
||||
for (const PlannedPhysicalSlot* slot : ArrayRef(reusedSlots).take_front(kSummaryReuseLimit)) {
|
||||
uint64_t slotLogicalBytes = getSlotLogicalBytes(*slot, intervals);
|
||||
os << " slot #" << slot->id << " addr=" << slot->address << " size=" << formatReportMemory(slot->requiredSize)
|
||||
<< " intervals=" << slot->intervalIndices.size() << " logical_sum=" << formatReportMemory(slotLogicalBytes)
|
||||
<< "\n";
|
||||
}
|
||||
}
|
||||
|
||||
os << "\nTop Offenders:\n";
|
||||
for (const PlannedPhysicalSlot* slot : ArrayRef(singleUseSlots).take_front(kSummaryOffenderLimit)) {
|
||||
const LocalAllocInterval& interval = intervals[slot->intervalIndices.front()];
|
||||
os << " slot #" << slot->id << " is single-use"
|
||||
<< " size=" << formatReportMemory(slot->requiredSize) << " interval=#" << interval.id
|
||||
<< " value=" << abbreviate(interval.valueSummary, 56) << "\n";
|
||||
os << " first=" << abbreviate(interval.firstTouchSummary, 40)
|
||||
<< " last=" << abbreviate(interval.lastTouchSummary, 40)
|
||||
<< " nested=" << (interval.insideNestedRegion ? "yes" : "no")
|
||||
<< " escapes_loop=" << (interval.escapesLoop ? "yes" : "no") << "\n";
|
||||
}
|
||||
if (singleUseSlots.empty())
|
||||
os << " no obvious blockers detected in this core\n";
|
||||
|
||||
if (reportLevel == PimMemoryReportFull) {
|
||||
os << "\nSlot Reuse:\n";
|
||||
for (const PlannedPhysicalSlot& slot : slots) {
|
||||
uint64_t slotLogicalBytes = getSlotLogicalBytes(slot, intervals);
|
||||
os << " slot #" << slot.id << " addr=" << slot.address << " size=" << formatReportMemory(slot.requiredSize)
|
||||
<< " (" << slot.requiredSize << ")"
|
||||
<< " intervals=" << slot.intervalIndices.size() << " logical_sum=" << formatReportMemory(slotLogicalBytes)
|
||||
<< "\n";
|
||||
for (size_t intervalIndex : slot.intervalIndices) {
|
||||
const LocalAllocInterval& interval = intervals[intervalIndex];
|
||||
os << " [" << interval.start << "," << interval.end << "]"
|
||||
<< " #" << interval.id << " logical=" << formatReportMemory(interval.size) << " (" << interval.size
|
||||
<< ")"
|
||||
<< " nested=" << (interval.insideNestedRegion ? "yes" : "no")
|
||||
<< " escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
|
||||
<< " first=" << abbreviate(interval.firstTouchSummary, 48)
|
||||
<< " last=" << abbreviate(interval.lastTouchSummary, 48) << "\n";
|
||||
os << " value=" << abbreviate(interval.valueSummary, 72) << "\n";
|
||||
if (!interval.fallbackReason.empty())
|
||||
os << " fallback_reason=" << interval.fallbackReason << "\n";
|
||||
}
|
||||
}
|
||||
}
|
||||
os.flush();
|
||||
|
||||
return artifacts;
|
||||
}
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace {
|
||||
|
||||
static bool rangesOverlap(size_t lhsAddress, size_t lhsSize, size_t rhsAddress, size_t rhsSize) {
|
||||
return lhsAddress < rhsAddress + rhsSize && rhsAddress < lhsAddress + lhsSize;
|
||||
}
|
||||
|
||||
static FailureOr<size_t> alignedEnd(size_t address, size_t size, size_t limit) {
|
||||
if (address > limit || size > limit - address)
|
||||
return failure();
|
||||
size_t end = address + size;
|
||||
size_t padding = (4 - end % 4) % 4;
|
||||
if (padding > limit - end)
|
||||
return failure();
|
||||
return end + padding;
|
||||
}
|
||||
|
||||
struct PimLocalMemoryPlanningPass : PassWrapper<PimLocalMemoryPlanningPass, OperationPass<ModuleOp>> {
|
||||
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimLocalMemoryPlanningPass)
|
||||
|
||||
StringRef getArgument() const override { return "pim-local-memory-planning"; }
|
||||
StringRef getDescription() const override {
|
||||
return "Plan liveness-based physical slots and addresses for PIM core-local memory";
|
||||
return "Plan liveness-based addresses for PIM core-local memory";
|
||||
}
|
||||
|
||||
void runOnOperation() override {
|
||||
std::fstream reportFile;
|
||||
std::unique_ptr<raw_os_ostream> report;
|
||||
if (pimMemoryReport != PimMemoryReportNone) {
|
||||
reportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt");
|
||||
if (reportFile.is_open())
|
||||
report = std::make_unique<raw_os_ostream>(reportFile);
|
||||
}
|
||||
else if (std::string outputRoot = getOutputDir(); !outputRoot.empty()) {
|
||||
sys::fs::remove(outputRoot + "/reports/pim_memory_liveness_report.txt");
|
||||
}
|
||||
|
||||
Builder builder(&getContext());
|
||||
uint64_t nextBatchId = 0;
|
||||
bool failedPlanning = false;
|
||||
bool hasFailure = false;
|
||||
getOperation().walk([&](Operation* op) {
|
||||
if (failedPlanning || !isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
|
||||
if (hasFailure || !isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
|
||||
return;
|
||||
|
||||
CoreMemoryPlan plan = buildCoreMemoryPlan(op);
|
||||
if (failed(assignPhysicalSlotAddresses(plan, kPimLocalMemoryAddressLimit))) {
|
||||
failedPlanning = true;
|
||||
return;
|
||||
}
|
||||
|
||||
size_t arenaSize = 0;
|
||||
for (const PlannedPhysicalSlot& placement : plan.slots)
|
||||
arenaSize = std::max(arenaSize, placement.address + placement.requiredSize);
|
||||
for (const LocalAllocInterval& interval : plan.intervals) {
|
||||
const PlannedPhysicalSlot& slot = plan.slots[interval.slotPlanIndex];
|
||||
interval.alloc->setAttr(kLocalMemoryAddressAttrName, builder.getI64IntegerAttr(slot.address));
|
||||
interval.alloc->setAttr(kLocalMemorySlotAttrName, builder.getI64IntegerAttr(0));
|
||||
interval.alloc->setAttr(kLocalMemorySlotSizeAttrName, builder.getI64IntegerAttr(arenaSize));
|
||||
}
|
||||
uint64_t fallbackIntervals = llvm::count_if(plan.intervals, [](const LocalAllocInterval& interval) {
|
||||
return interval.startUsedAllocFallback || interval.endUsedFallback;
|
||||
op->removeAttr(kLocalMemorySizeAttrName);
|
||||
for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
|
||||
op->removeAttr(name);
|
||||
op->walk([&](memref::AllocOp allocation) {
|
||||
allocation->removeAttr(kLocalMemoryAddressAttrName);
|
||||
for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
|
||||
allocation->removeAttr(name);
|
||||
});
|
||||
uint64_t nestedSingleUseIntervals = llvm::count_if(plan.intervals, [&](const LocalAllocInterval& interval) {
|
||||
return interval.insideNestedRegion && !hasAddressReuse(interval.slotPlanIndex, plan.slots);
|
||||
});
|
||||
op->setAttr(kLocalMemoryFallbackCountAttrName, builder.getI64IntegerAttr(fallbackIntervals));
|
||||
op->setAttr(kLocalMemoryNestedSingleUseCountAttrName, builder.getI64IntegerAttr(nestedSingleUseIntervals));
|
||||
|
||||
if (!report)
|
||||
return;
|
||||
std::string planReport = buildMemoryPlanReport(
|
||||
op, plan.intervals, plan.slots, kPimLocalMemoryAddressLimit, pimMemoryReport);
|
||||
if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
|
||||
*report << "Core " << coreOp.getCoreId() << ":\n";
|
||||
*report << planReport;
|
||||
auto plan = buildCoreMemoryPlan(op);
|
||||
if (failed(plan)) {
|
||||
hasFailure = true;
|
||||
return;
|
||||
}
|
||||
|
||||
SmallVector<int32_t> coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
|
||||
llvm::sort(coreIds);
|
||||
coreIds.erase(std::unique(coreIds.begin(), coreIds.end()), coreIds.end());
|
||||
uint64_t batchId = nextBatchId++;
|
||||
for (int32_t coreId : coreIds)
|
||||
*report << "Batch " << batchId << " core " << coreId << ":\n" << planReport;
|
||||
op->setAttr(kLocalMemorySizeAttrName, builder.getI64IntegerAttr(plan->arenaSize));
|
||||
for (const LocalMemoryPlacement& placement : plan->placements)
|
||||
for (size_t intervalIndex : placement.intervalIndices)
|
||||
plan->intervals[intervalIndex].allocation->setAttr(
|
||||
kLocalMemoryAddressAttrName, builder.getI64IntegerAttr(placement.address));
|
||||
});
|
||||
|
||||
if (report) {
|
||||
report->flush();
|
||||
reportFile.close();
|
||||
}
|
||||
if (failedPlanning) {
|
||||
if (hasFailure) {
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
dumpModule(getOperation(), "pim4_memory_planned");
|
||||
dumpModule(getOperation(), "pim3_memory_planned");
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
FailureOr<SmallVector<LocalMemoryPlacement, 0>>
|
||||
planLocalMemoryPlacements(ArrayRef<pim::LocalMemoryInterval> intervals, size_t addressLimit) {
|
||||
SmallVector<LocalMemoryPlacement, 0> placements;
|
||||
SmallVector<size_t> order(intervals.size());
|
||||
std::iota(order.begin(), order.end(), 0);
|
||||
llvm::stable_sort(order, [&](size_t lhsIndex, size_t rhsIndex) {
|
||||
const auto& lhs = intervals[lhsIndex];
|
||||
const auto& rhs = intervals[rhsIndex];
|
||||
if (lhs.size != rhs.size)
|
||||
return lhs.size > rhs.size;
|
||||
if (lhs.start != rhs.start)
|
||||
return lhs.start < rhs.start;
|
||||
if (lhs.end != rhs.end)
|
||||
return lhs.end < rhs.end;
|
||||
return lhsIndex < rhsIndex;
|
||||
});
|
||||
|
||||
for (size_t intervalIndex : order) {
|
||||
const auto& interval = intervals[intervalIndex];
|
||||
SmallVector<const LocalMemoryPlacement*, 16> conflicts;
|
||||
SmallVector<size_t, 16> candidates = {0};
|
||||
size_t currentPeak = 0;
|
||||
for (const LocalMemoryPlacement& placement : placements) {
|
||||
auto end = alignedEnd(placement.address, placement.size, addressLimit);
|
||||
if (failed(end))
|
||||
return failure();
|
||||
currentPeak = std::max(currentPeak, *end);
|
||||
if (!llvm::any_of(placement.intervalIndices, [&](size_t otherIndex) {
|
||||
return pim::localMemoryLifetimesOverlap(interval, intervals[otherIndex]);
|
||||
}))
|
||||
continue;
|
||||
conflicts.push_back(&placement);
|
||||
if (!llvm::is_contained(candidates, *end))
|
||||
candidates.push_back(*end);
|
||||
}
|
||||
|
||||
size_t bestAddress = std::numeric_limits<size_t>::max();
|
||||
auto bestKey = std::tuple(std::numeric_limits<size_t>::max(), bestAddress);
|
||||
for (size_t candidate : candidates) {
|
||||
auto end = alignedEnd(candidate, interval.size, addressLimit);
|
||||
if (failed(end) || llvm::any_of(conflicts, [&](const LocalMemoryPlacement* placement) {
|
||||
return rangesOverlap(candidate, interval.size, placement->address, placement->size);
|
||||
}))
|
||||
continue;
|
||||
auto key = std::tuple(std::max(currentPeak, *end), candidate);
|
||||
if (key < bestKey) {
|
||||
bestKey = key;
|
||||
bestAddress = candidate;
|
||||
}
|
||||
}
|
||||
if (bestAddress == std::numeric_limits<size_t>::max())
|
||||
return failure();
|
||||
|
||||
auto reusable = llvm::find_if(placements, [&](const LocalMemoryPlacement& placement) {
|
||||
return placement.address == bestAddress && placement.size == interval.size;
|
||||
});
|
||||
if (reusable != placements.end())
|
||||
reusable->intervalIndices.push_back(intervalIndex);
|
||||
else
|
||||
placements.push_back({bestAddress, interval.size, {intervalIndex}});
|
||||
}
|
||||
return placements;
|
||||
}
|
||||
|
||||
FailureOr<CoreMemoryPlan> buildCoreMemoryPlan(Operation* coreLikeOp) {
|
||||
CoreMemoryPlan plan;
|
||||
plan.coreLikeOp = coreLikeOp;
|
||||
auto intervals = pim::analyzeLocalMemoryLifetimes(coreLikeOp);
|
||||
if (failed(intervals))
|
||||
return failure();
|
||||
plan.intervals = std::move(*intervals);
|
||||
auto placements = planLocalMemoryPlacements(plan.intervals, kPimLocalMemoryAddressLimit);
|
||||
if (failed(placements)) {
|
||||
coreLikeOp->emitError("PIM local-memory plan exceeds the signed int32 address range");
|
||||
return failure();
|
||||
}
|
||||
plan.placements = std::move(*placements);
|
||||
for (const LocalMemoryPlacement& placement : plan.placements) {
|
||||
auto end = alignedEnd(placement.address, placement.size, kPimLocalMemoryAddressLimit);
|
||||
if (failed(end)) {
|
||||
coreLikeOp->emitError("PIM local-memory plan has invalid address arithmetic");
|
||||
return failure();
|
||||
}
|
||||
plan.arenaSize = std::max(plan.arenaSize, *end);
|
||||
}
|
||||
return plan;
|
||||
}
|
||||
|
||||
std::unique_ptr<Pass> createPimLocalMemoryPlanningPass() {
|
||||
return std::make_unique<PimLocalMemoryPlanningPass>();
|
||||
}
|
||||
|
||||
@@ -1,64 +1,24 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
||||
|
||||
#include "llvm/ADT/ArrayRef.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
|
||||
#include <limits>
|
||||
#include <string>
|
||||
|
||||
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
struct LocalAllocInterval {
|
||||
size_t id = 0;
|
||||
mlir::memref::AllocOp alloc;
|
||||
uint64_t start = 0;
|
||||
uint64_t end = 0;
|
||||
size_t size = 0;
|
||||
mlir::Operation* firstTouchOp = nullptr;
|
||||
mlir::Operation* lastTouchOp = nullptr;
|
||||
uint64_t firstTouchPosition = 0;
|
||||
uint64_t lastTouchPosition = 0;
|
||||
bool startUsedAllocFallback = false;
|
||||
bool endUsedFallback = false;
|
||||
bool hasRuntimeUse = false;
|
||||
bool insideNestedRegion = false;
|
||||
bool escapesLoop = false;
|
||||
std::string fallbackReason;
|
||||
std::string valueSummary;
|
||||
std::string firstTouchSummary;
|
||||
std::string lastTouchSummary;
|
||||
size_t slotPlanIndex = std::numeric_limits<size_t>::max();
|
||||
};
|
||||
|
||||
struct PlannedPhysicalSlot {
|
||||
size_t id = std::numeric_limits<size_t>::max();
|
||||
size_t requiredSize = 0;
|
||||
struct LocalMemoryPlacement {
|
||||
size_t address = 0;
|
||||
size_t size = 0;
|
||||
llvm::SmallVector<size_t, 8> intervalIndices;
|
||||
};
|
||||
|
||||
struct CoreMemoryPlan {
|
||||
mlir::Operation* coreLikeOp = nullptr;
|
||||
llvm::SmallVector<LocalAllocInterval, 0> intervals;
|
||||
llvm::SmallVector<PlannedPhysicalSlot, 0> slots;
|
||||
llvm::SmallVector<pim::LocalMemoryInterval, 0> intervals;
|
||||
llvm::SmallVector<LocalMemoryPlacement, 0> placements;
|
||||
size_t arenaSize = 0;
|
||||
};
|
||||
|
||||
llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp);
|
||||
|
||||
llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);
|
||||
|
||||
CoreMemoryPlan buildCoreMemoryPlan(mlir::Operation* coreLikeOp);
|
||||
|
||||
mlir::LogicalResult assignPhysicalSlotAddresses(CoreMemoryPlan& plan, size_t addressLimit);
|
||||
|
||||
std::string buildMemoryPlanReport(mlir::Operation* coreLikeOp,
|
||||
llvm::ArrayRef<LocalAllocInterval> intervals,
|
||||
llvm::ArrayRef<PlannedPhysicalSlot> slots,
|
||||
size_t addressLimit,
|
||||
PimMemoryReportLevel reportLevel);
|
||||
mlir::FailureOr<llvm::SmallVector<LocalMemoryPlacement, 0>>
|
||||
planLocalMemoryPlacements(llvm::ArrayRef<pim::LocalMemoryInterval> intervals, size_t addressLimit);
|
||||
mlir::FailureOr<CoreMemoryPlan> buildCoreMemoryPlan(mlir::Operation* coreLikeOp);
|
||||
|
||||
} // namespace onnx_mlir
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
add_pim_library(OMPimMemoryCoalescing
|
||||
MemoryCoalescing.cpp
|
||||
MemoryCoalescing.hpp
|
||||
MemoryCoalescingPass.cpp
|
||||
|
||||
EXCLUDE_FROM_OM_LIBS
|
||||
|
||||
INCLUDE_DIRS PUBLIC
|
||||
${PIM_PUBLIC_INCLUDE_DIRS}
|
||||
|
||||
LINK_LIBS PUBLIC
|
||||
OMPimCommon
|
||||
OMPimLocalMemoryLifetimeAnalysis
|
||||
PimOps
|
||||
)
|
||||
@@ -1,194 +0,0 @@
|
||||
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
||||
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
|
||||
#include <limits>
|
||||
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace pim {
|
||||
|
||||
namespace {
|
||||
|
||||
static bool isCandidateAllocType(MemRefType type) {
|
||||
return type && type.hasStaticShape() && type.getLayout().isIdentity()
|
||||
&& hasByteSizedElementType(type.getElementType());
|
||||
}
|
||||
|
||||
static uint64_t getTypeSizeBytes(MemRefType type) {
|
||||
return static_cast<uint64_t>(type.getNumElements() * getElementTypeSizeInBytes(type.getElementType()));
|
||||
}
|
||||
|
||||
static Operation* getTopLevelAncestorInBlock(Operation* op, Block& block) {
|
||||
Operation* current = op;
|
||||
while (current && current->getBlock() != &block)
|
||||
current = current->getParentOp();
|
||||
return current;
|
||||
}
|
||||
|
||||
static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis);
|
||||
|
||||
static FailureOr<uint64_t>
|
||||
getLastUseInstruction(memref::AllocOp allocOp, Block& scopeBlock, const DenseMap<Operation*, uint64_t>& opOrder) {
|
||||
uint64_t endInstruction = opOrder.lookup(allocOp);
|
||||
if (failed(walkLocalMemoryUses(allocOp.getResult(), [&](Value, Operation* user) {
|
||||
if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
|
||||
if (!isa<scf::ForOp>(yieldOp->getParentOp()))
|
||||
return failure();
|
||||
}
|
||||
Operation* orderedUser = getTopLevelAncestorInBlock(user, scopeBlock);
|
||||
if (!orderedUser)
|
||||
return failure();
|
||||
auto order = opOrder.find(orderedUser);
|
||||
if (order == opOrder.end())
|
||||
return failure();
|
||||
endInstruction = std::max(endInstruction, order->second);
|
||||
return success();
|
||||
})))
|
||||
return failure();
|
||||
|
||||
return endInstruction;
|
||||
}
|
||||
|
||||
static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis) {
|
||||
for (Operation& op : block)
|
||||
for (Region& region : op.getRegions())
|
||||
for (Block& nestedBlock : region)
|
||||
analyzeBlock(nestedBlock, analysis);
|
||||
|
||||
DenseMap<Operation*, uint64_t> opOrder;
|
||||
uint64_t nextInstruction = 0;
|
||||
for (Operation& op : block)
|
||||
opOrder.try_emplace(&op, nextInstruction++);
|
||||
|
||||
MemoryCoalescingBlockAnalysis blockAnalysis;
|
||||
blockAnalysis.block = █
|
||||
|
||||
for (Operation& op : block) {
|
||||
auto allocOp = dyn_cast<memref::AllocOp>(&op);
|
||||
if (!allocOp)
|
||||
continue;
|
||||
|
||||
auto allocType = dyn_cast<MemRefType>(allocOp.getType());
|
||||
if (!isCandidateAllocType(allocType)) {
|
||||
++blockAnalysis.skippedAllocations;
|
||||
continue;
|
||||
}
|
||||
|
||||
auto endInstruction = getLastUseInstruction(allocOp, block, opOrder);
|
||||
if (failed(endInstruction)) {
|
||||
++blockAnalysis.skippedAllocations;
|
||||
continue;
|
||||
}
|
||||
|
||||
blockAnalysis.candidates.push_back(
|
||||
AllocationCandidate {allocOp, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)});
|
||||
}
|
||||
|
||||
analysis.skippedAllocations += blockAnalysis.skippedAllocations;
|
||||
if (!blockAnalysis.candidates.empty() || blockAnalysis.skippedAllocations != 0)
|
||||
analysis.blocks.push_back(std::move(blockAnalysis));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
uint64_t MemoryCoalescingAnalysis::getCandidateCount() const {
|
||||
uint64_t total = 0;
|
||||
for (const MemoryCoalescingBlockAnalysis& block : blocks)
|
||||
total += block.candidates.size();
|
||||
return total;
|
||||
}
|
||||
|
||||
uint64_t MemoryCoalescingAnalysis::getCandidateBytes() const {
|
||||
uint64_t total = 0;
|
||||
for (const MemoryCoalescingBlockAnalysis& block : blocks)
|
||||
for (const AllocationCandidate& candidate : block.candidates)
|
||||
total += candidate.sizeBytes;
|
||||
return total;
|
||||
}
|
||||
|
||||
MemoryCoalescingAnalysis analyzeMemoryCoalescingCandidates(Operation* coreLikeOp) {
|
||||
MemoryCoalescingAnalysis analysis;
|
||||
if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
|
||||
return analysis;
|
||||
|
||||
analyzeBlock(coreLikeOp->getRegion(0).front(), analysis);
|
||||
return analysis;
|
||||
}
|
||||
|
||||
MemoryCoalescingStats
|
||||
coalesceMemory(Operation* coreLikeOp, const MemoryCoalescingAnalysis& analysis, RewriterBase& rewriter) {
|
||||
(void) coreLikeOp;
|
||||
|
||||
MemoryCoalescingStats stats;
|
||||
stats.skippedAllocations = analysis.skippedAllocations;
|
||||
|
||||
for (const MemoryCoalescingBlockAnalysis& blockAnalysis : analysis.blocks) {
|
||||
auto candidates = blockAnalysis.candidates;
|
||||
llvm::sort(candidates, [](const AllocationCandidate& lhs, const AllocationCandidate& rhs) {
|
||||
if (lhs.startInstruction != rhs.startInstruction)
|
||||
return lhs.startInstruction < rhs.startInstruction;
|
||||
return lhs.endInstruction < rhs.endInstruction;
|
||||
});
|
||||
|
||||
struct ActiveStorage {
|
||||
memref::AllocOp root;
|
||||
uint64_t endInstruction = 0;
|
||||
};
|
||||
|
||||
SmallVector<ActiveStorage> active;
|
||||
SmallVector<memref::AllocOp> freeList;
|
||||
|
||||
for (AllocationCandidate& candidate : candidates) {
|
||||
for (auto it = active.begin(); it != active.end();) {
|
||||
if (it->endInstruction < candidate.startInstruction) {
|
||||
freeList.push_back(it->root);
|
||||
it = active.erase(it);
|
||||
continue;
|
||||
}
|
||||
++it;
|
||||
}
|
||||
|
||||
auto bestFit = freeList.end();
|
||||
uint64_t bestFitBytes = std::numeric_limits<uint64_t>::max();
|
||||
auto candidateType = cast<MemRefType>(candidate.alloc.getType());
|
||||
for (auto it = freeList.begin(); it != freeList.end(); ++it) {
|
||||
auto freeType = cast<MemRefType>((*it).getType());
|
||||
if (freeType != candidateType)
|
||||
continue;
|
||||
|
||||
uint64_t freeBytes = getTypeSizeBytes(freeType);
|
||||
if (freeBytes < candidate.sizeBytes || freeBytes >= bestFitBytes)
|
||||
continue;
|
||||
|
||||
bestFit = it;
|
||||
bestFitBytes = freeBytes;
|
||||
}
|
||||
|
||||
if (bestFit == freeList.end()) {
|
||||
active.push_back(ActiveStorage {candidate.alloc, candidate.endInstruction});
|
||||
continue;
|
||||
}
|
||||
|
||||
memref::AllocOp root = *bestFit;
|
||||
freeList.erase(bestFit);
|
||||
candidate.alloc.getResult().replaceAllUsesWith(root.getResult());
|
||||
rewriter.eraseOp(candidate.alloc);
|
||||
active.push_back(ActiveStorage {root, candidate.endInstruction});
|
||||
++stats.removedAllocs;
|
||||
stats.savedBytes += candidate.sizeBytes;
|
||||
}
|
||||
}
|
||||
|
||||
return stats;
|
||||
}
|
||||
|
||||
} // namespace pim
|
||||
} // namespace onnx_mlir
|
||||
@@ -1,44 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
||||
#include "mlir/IR/PatternMatch.h"
|
||||
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace pim {
|
||||
|
||||
struct AllocationCandidate {
|
||||
mlir::memref::AllocOp alloc;
|
||||
uint64_t startInstruction = 0;
|
||||
uint64_t endInstruction = 0;
|
||||
uint64_t sizeBytes = 0;
|
||||
};
|
||||
|
||||
struct MemoryCoalescingBlockAnalysis {
|
||||
mlir::Block* block = nullptr;
|
||||
llvm::SmallVector<AllocationCandidate> candidates;
|
||||
uint64_t skippedAllocations = 0;
|
||||
};
|
||||
|
||||
struct MemoryCoalescingAnalysis {
|
||||
llvm::SmallVector<MemoryCoalescingBlockAnalysis> blocks;
|
||||
uint64_t skippedAllocations = 0;
|
||||
|
||||
uint64_t getCandidateCount() const;
|
||||
uint64_t getCandidateBytes() const;
|
||||
};
|
||||
|
||||
struct MemoryCoalescingStats {
|
||||
uint64_t removedAllocs = 0;
|
||||
uint64_t savedBytes = 0;
|
||||
uint64_t skippedAllocations = 0;
|
||||
};
|
||||
|
||||
MemoryCoalescingAnalysis analyzeMemoryCoalescingCandidates(mlir::Operation* coreLikeOp);
|
||||
|
||||
MemoryCoalescingStats
|
||||
coalesceMemory(mlir::Operation* coreLikeOp, const MemoryCoalescingAnalysis& analysis, mlir::RewriterBase& rewriter);
|
||||
|
||||
} // namespace pim
|
||||
} // namespace onnx_mlir
|
||||
@@ -1,101 +0,0 @@
|
||||
#include "mlir/Pass/Pass.h"
|
||||
|
||||
#include "llvm/Support/raw_os_ostream.h"
|
||||
|
||||
#include <fstream>
|
||||
|
||||
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.hpp"
|
||||
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace {
|
||||
|
||||
struct CoalescingReportRow {
|
||||
uint64_t candidates = 0;
|
||||
uint64_t logicalBytes = 0;
|
||||
uint64_t skipped = 0;
|
||||
uint64_t removed = 0;
|
||||
uint64_t savedBytes = 0;
|
||||
};
|
||||
|
||||
struct CoalescingReportEntry {
|
||||
std::string label;
|
||||
uint64_t coreCount = 1;
|
||||
CoalescingReportRow row;
|
||||
};
|
||||
|
||||
static void emitReport(ArrayRef<CoalescingReportEntry> entries) {
|
||||
std::fstream file = openReportFile("memory_coalescing_report");
|
||||
if (!file.is_open())
|
||||
return;
|
||||
llvm::raw_os_ostream os(file);
|
||||
CoalescingReportRow total;
|
||||
for (const CoalescingReportEntry& entry : entries) {
|
||||
total.candidates += entry.row.candidates * entry.coreCount;
|
||||
total.logicalBytes += entry.row.logicalBytes * entry.coreCount;
|
||||
total.skipped += entry.row.skipped * entry.coreCount;
|
||||
total.removed += entry.row.removed * entry.coreCount;
|
||||
total.savedBytes += entry.row.savedBytes * entry.coreCount;
|
||||
}
|
||||
printReportTotalsBlock(os,
|
||||
{{"Number of candidates", std::to_string(total.candidates)},
|
||||
{"Logical candidate bytes", formatReportMemory(total.logicalBytes)},
|
||||
{"Skipped allocations", std::to_string(total.skipped)},
|
||||
{"Removed allocations", std::to_string(total.removed)},
|
||||
{"Saved memory", formatReportMemory(total.savedBytes)}});
|
||||
for (const CoalescingReportEntry& entry : entries) {
|
||||
os << "\n" << entry.label << ":\n";
|
||||
printReportFlatFields(os,
|
||||
{{"Number of candidates", std::to_string(entry.row.candidates)},
|
||||
{"Logical candidate bytes", formatReportMemory(entry.row.logicalBytes)},
|
||||
{"Skipped allocations", std::to_string(entry.row.skipped)},
|
||||
{"Removed allocations", std::to_string(entry.row.removed)},
|
||||
{"Saved memory", formatReportMemory(entry.row.savedBytes)}});
|
||||
}
|
||||
os.flush();
|
||||
}
|
||||
|
||||
struct PimMemoryCoalescingPass : PassWrapper<PimMemoryCoalescingPass, OperationPass<ModuleOp>> {
|
||||
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimMemoryCoalescingPass)
|
||||
|
||||
StringRef getArgument() const override { return "pim-memory-coalescing"; }
|
||||
StringRef getDescription() const override { return "Safely coalesce compatible block-local PIM allocations"; }
|
||||
|
||||
void runOnOperation() override {
|
||||
IRRewriter rewriter(&getContext());
|
||||
SmallVector<CoalescingReportEntry, 32> reportEntries;
|
||||
uint64_t nextBatchId = 0;
|
||||
getOperation().walk([&](Operation* op) {
|
||||
if (!isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
|
||||
return;
|
||||
auto analysis = pim::analyzeMemoryCoalescingCandidates(op);
|
||||
auto stats = pim::coalesceMemory(op, analysis, rewriter);
|
||||
CoalescingReportRow row {analysis.getCandidateCount(),
|
||||
analysis.getCandidateBytes(),
|
||||
stats.skippedAllocations,
|
||||
stats.removedAllocs,
|
||||
stats.savedBytes};
|
||||
if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
|
||||
reportEntries.push_back({"Core " + std::to_string(coreOp.getCoreId()), 1, row});
|
||||
return;
|
||||
}
|
||||
SmallVector<int32_t> coreIds = getBatchCoreIds(cast<pim::PimCoreBatchOp>(op));
|
||||
reportEntries.push_back(
|
||||
{"Batch " + std::to_string(nextBatchId++), std::max<uint64_t>(1, coreIds.size()), row});
|
||||
});
|
||||
emitReport(reportEntries);
|
||||
dumpModule(getOperation(), "pim3_coalesced");
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<Pass> createPimMemoryCoalescingPass() { return std::make_unique<PimMemoryCoalescingPass>(); }
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -7,6 +7,7 @@ add_pim_library(OMPimVerification
|
||||
OMPimCommon
|
||||
OMPimCompilerOptions
|
||||
OMPimBufferization
|
||||
OMPimLocalMemoryLifetimeAnalysis
|
||||
PimOps
|
||||
SpatialOps
|
||||
)
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
#include "llvm/Support/FormatVariadic.h"
|
||||
#include "llvm/Support/raw_ostream.h"
|
||||
|
||||
#include <map>
|
||||
#include <string>
|
||||
|
||||
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
|
||||
@@ -18,6 +19,7 @@
|
||||
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
|
||||
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
|
||||
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
@@ -109,61 +111,127 @@ static bool isCoreWeightBlockArgument(Value value) {
|
||||
return false;
|
||||
}
|
||||
|
||||
struct VerifiedLocalMemorySlot {
|
||||
uint64_t size = 0;
|
||||
};
|
||||
|
||||
static LogicalResult
|
||||
verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diagnostics) {
|
||||
DenseMap<uint64_t, VerifiedLocalMemorySlot> slots;
|
||||
bool hasFailure = false;
|
||||
auto fallbackAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryFallbackCountAttrName);
|
||||
auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
|
||||
if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
|
||||
diagnostics.report(coreLikeOp, [&](Operation* op) {
|
||||
op->emitError("requires complete non-negative PIM local-memory planning summary attributes");
|
||||
for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
|
||||
if (coreLikeOp->hasAttr(name)) {
|
||||
diagnostics.report(coreLikeOp, [name](Operation* op) {
|
||||
op->emitError() << "contains removed PIM local-memory planning attribute '" << name << "'";
|
||||
});
|
||||
hasFailure = true;
|
||||
}
|
||||
|
||||
auto arenaAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemorySizeAttrName);
|
||||
if (!arenaAttr || arenaAttr.getInt() < 0
|
||||
|| static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
|
||||
diagnostics.report(coreLikeOp, [](Operation* op) {
|
||||
op->emitError("requires one valid non-negative pim.local_memory_size attribute");
|
||||
});
|
||||
hasFailure = true;
|
||||
}
|
||||
coreLikeOp->walk([&](memref::AllocOp allocOp) {
|
||||
auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
|
||||
auto slotAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotAttrName);
|
||||
auto slotSizeAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotSizeAttrName);
|
||||
if (!addressAttr || !slotAttr || !slotSizeAttr || addressAttr.getInt() < 0 || slotAttr.getInt() < 0
|
||||
|| slotSizeAttr.getInt() < 0) {
|
||||
diagnostics.report(allocOp, [&](Operation*) {
|
||||
allocOp.emitOpError("requires a complete non-negative PIM local-memory plan");
|
||||
if (hasFailure)
|
||||
return failure();
|
||||
uint64_t arenaSize = static_cast<uint64_t>(arenaAttr.getInt());
|
||||
|
||||
auto analyzed = pim::analyzeLocalMemoryLifetimes(coreLikeOp);
|
||||
if (failed(analyzed)) {
|
||||
diagnostics.report(coreLikeOp, [](Operation* op) {
|
||||
op->emitError("cannot analyze PIM local-memory lifetimes for plan verification");
|
||||
});
|
||||
return failure();
|
||||
}
|
||||
|
||||
struct Event {
|
||||
uint64_t position;
|
||||
bool start;
|
||||
size_t intervalIndex;
|
||||
};
|
||||
SmallVector<Event, 0> events;
|
||||
SmallVector<size_t, 0> addresses(analyzed->size());
|
||||
for (auto indexed : llvm::enumerate(*analyzed)) {
|
||||
size_t index = indexed.index();
|
||||
auto& interval = indexed.value();
|
||||
memref::AllocOp allocation = interval.allocation;
|
||||
for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
|
||||
if (allocation->hasAttr(name)) {
|
||||
diagnostics.report(allocation, [name](Operation* op) {
|
||||
op->emitOpError() << "contains removed PIM local-memory planning attribute '" << name << "'";
|
||||
});
|
||||
hasFailure = true;
|
||||
}
|
||||
auto addressAttr = allocation->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
|
||||
if (!addressAttr || addressAttr.getInt() < 0) {
|
||||
diagnostics.report(allocation, [](Operation* op) {
|
||||
op->emitOpError("requires one non-negative pim.local_memory_address attribute");
|
||||
});
|
||||
hasFailure = true;
|
||||
return;
|
||||
continue;
|
||||
}
|
||||
|
||||
uint64_t address = static_cast<uint64_t>(addressAttr.getInt());
|
||||
uint64_t slotId = static_cast<uint64_t>(slotAttr.getInt());
|
||||
uint64_t slotSize = static_cast<uint64_t>(slotSizeAttr.getInt());
|
||||
auto logicalSize = pim::getCheckedShapedTypeSizeInBytes(
|
||||
cast<ShapedType>(allocOp.getType()), allocOp, "planned local allocation byte size");
|
||||
bool invalidRange = failed(logicalSize) || address > slotSize || *logicalSize > slotSize - address
|
||||
|| slotSize > kPimLocalMemoryAddressLimit || address % 4 != 0 || slotId != 0;
|
||||
if (invalidRange) {
|
||||
diagnostics.report(allocOp, [&](Operation*) {
|
||||
allocOp.emitOpError() << "has an invalid PIM local-memory range: address=" << address
|
||||
<< ", logical size=" << (succeeded(logicalSize) ? *logicalSize : 0)
|
||||
<< ", slot size=" << slotSize;
|
||||
if (address % 4 != 0 || address > arenaSize || interval.size > arenaSize - address) {
|
||||
diagnostics.report(allocation, [&](Operation* op) {
|
||||
op->emitOpError() << "has invalid PIM local-memory range [" << address << ", "
|
||||
<< (address <= arenaSize && interval.size <= arenaSize - address
|
||||
? address + interval.size
|
||||
: arenaSize)
|
||||
<< ") for arena size " << arenaSize;
|
||||
});
|
||||
hasFailure = true;
|
||||
return;
|
||||
continue;
|
||||
}
|
||||
addresses[index] = static_cast<size_t>(address);
|
||||
events.push_back({interval.start, true, index});
|
||||
events.push_back({interval.end, false, index});
|
||||
}
|
||||
if (hasFailure)
|
||||
return failure();
|
||||
|
||||
llvm::sort(events, [](const Event& lhs, const Event& rhs) {
|
||||
return lhs.position != rhs.position ? lhs.position < rhs.position : lhs.start > rhs.start;
|
||||
});
|
||||
std::map<size_t, size_t> active;
|
||||
for (const Event& event : events) {
|
||||
size_t index = event.intervalIndex;
|
||||
size_t address = addresses[index];
|
||||
if (!event.start) {
|
||||
auto it = active.find(address);
|
||||
if (it != active.end() && it->second == index)
|
||||
active.erase(it);
|
||||
continue;
|
||||
}
|
||||
|
||||
auto [slotIt, inserted] = slots.try_emplace(slotId, VerifiedLocalMemorySlot {slotSize});
|
||||
if (!inserted && slotIt->second.size != slotSize) {
|
||||
diagnostics.report(allocOp, [&](Operation*) {
|
||||
allocOp.emitOpError() << "has inconsistent size for PIM local-memory arena " << slotId;
|
||||
const auto& interval = (*analyzed)[index];
|
||||
auto successor = active.lower_bound(address);
|
||||
auto conflict = [&](std::map<size_t, size_t>::iterator it) {
|
||||
if (it == active.end())
|
||||
return false;
|
||||
const auto& other = (*analyzed)[it->second];
|
||||
return address < it->first + other.size && it->first < address + interval.size;
|
||||
};
|
||||
auto predecessor = successor;
|
||||
bool hasPredecessor = predecessor != active.begin();
|
||||
if (hasPredecessor)
|
||||
--predecessor;
|
||||
auto conflicting = conflict(successor) ? successor
|
||||
: hasPredecessor && conflict(predecessor) ? predecessor : active.end();
|
||||
if (conflicting != active.end()) {
|
||||
const auto& other = (*analyzed)[conflicting->second];
|
||||
memref::AllocOp allocation = interval.allocation;
|
||||
memref::AllocOp otherAllocation = other.allocation;
|
||||
diagnostics.report(allocation, [&](Operation*) {
|
||||
auto diagnostic = allocation.emitOpError()
|
||||
<< "PIM local-memory plan assigns simultaneously live allocations to overlapping ranges; first range ["
|
||||
<< conflicting->first << ", " << conflicting->first + other.size << "), second range [" << address
|
||||
<< ", " << address + interval.size << "), live positions overlap at ["
|
||||
<< std::max(interval.start, other.start) << ", " << std::min(interval.end, other.end) << "]";
|
||||
diagnostic.attachNote(otherAllocation.getLoc()) << "first allocation";
|
||||
});
|
||||
hasFailure = true;
|
||||
return failure();
|
||||
}
|
||||
});
|
||||
return success(!hasFailure);
|
||||
active.emplace(address, index);
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
static bool isSupportedCoreInstructionOp(Operation* op) {
|
||||
|
||||
@@ -16,6 +16,11 @@ namespace spatial {
|
||||
|
||||
namespace {
|
||||
|
||||
// Pressure means distinct weights exceed half the fleet's one-copy capacity.
|
||||
// The reserved headroom conservatively avoids greedy capacity fragmentation;
|
||||
// makespan, communication, instruction-path, and local-memory proxies remain stable.
|
||||
constexpr size_t kHighCrossbarPressureCapacityDivisor = 2;
|
||||
|
||||
struct ScheduledTask {
|
||||
size_t processor = std::numeric_limits<size_t>::max();
|
||||
Time startTime = 0;
|
||||
@@ -136,7 +141,7 @@ void verifyOctTableSize(size_t nodeCount, size_t processorCount) {
|
||||
bool hasHighCrossbarPressure(const ComputeGraph& graph, size_t processorCount, size_t crossbarCapacity) {
|
||||
if (crossbarCapacity > std::numeric_limits<size_t>::max() / processorCount)
|
||||
return false;
|
||||
const size_t threshold = processorCount * crossbarCapacity / 2;
|
||||
const size_t threshold = processorCount * crossbarCapacity / kHighCrossbarPressureCapacityDivisor;
|
||||
CrossbarUsage distinctWeights;
|
||||
for (const ComputeGraphNode& node : graph.nodes) {
|
||||
for (const CrossbarWeight& weight : node.crossbarUsage) {
|
||||
|
||||
Reference in New Issue
Block a user