This commit is contained in:
@@ -237,12 +237,19 @@ getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const Operati
|
||||
}
|
||||
|
||||
static MemoryTouchInterval
|
||||
computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ordering, uint64_t fallbackEnd) {
|
||||
computeMemoryTouchInterval(memref::AllocOp allocOp,
|
||||
const OperationOrdering& ordering,
|
||||
uint64_t fallbackEnd,
|
||||
bool includeAliasDescriptions) {
|
||||
MemoryTouchInterval interval;
|
||||
interval.start = ordering.position.lookup(allocOp);
|
||||
interval.end = interval.start;
|
||||
interval.startOp = allocOp;
|
||||
interval.endOp = allocOp;
|
||||
auto recordAlias = [&](mlir::Value value) {
|
||||
if (includeAliasDescriptions)
|
||||
appendAliasDescription(interval.aliasesFollowed, value);
|
||||
};
|
||||
|
||||
SmallPtrSet<mlir::Value, 16> visitedValues;
|
||||
SmallPtrSet<Operation*, 32> visitedUsers;
|
||||
@@ -262,7 +269,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
|
||||
if (isSupportedAliasOp(user)) {
|
||||
for (mlir::Value result : user->getResults()) {
|
||||
pendingValues.push_back(result);
|
||||
appendAliasDescription(interval.aliasesFollowed, result);
|
||||
recordAlias(result);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -272,7 +279,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
|
||||
if (!tiedOperand || tiedOperand->get() != value)
|
||||
continue;
|
||||
pendingValues.push_back(result);
|
||||
appendAliasDescription(interval.aliasesFollowed, result);
|
||||
recordAlias(result);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -282,8 +289,8 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
|
||||
continue;
|
||||
pendingValues.push_back(forOp.getRegionIterArgs()[index]);
|
||||
pendingValues.push_back(forOp.getResult(index));
|
||||
appendAliasDescription(interval.aliasesFollowed, forOp.getRegionIterArgs()[index]);
|
||||
appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index));
|
||||
recordAlias(forOp.getRegionIterArgs()[index]);
|
||||
recordAlias(forOp.getResult(index));
|
||||
if (parentLoop && forOp != parentLoop)
|
||||
interval.escapesLoop = true;
|
||||
}
|
||||
@@ -298,7 +305,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
|
||||
if (operand != value)
|
||||
continue;
|
||||
pendingValues.push_back(ifOp.getResult(index));
|
||||
appendAliasDescription(interval.aliasesFollowed, ifOp.getResult(index));
|
||||
recordAlias(ifOp.getResult(index));
|
||||
}
|
||||
}
|
||||
else if (indexSwitch) {
|
||||
@@ -306,8 +313,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
|
||||
if (operand != value)
|
||||
continue;
|
||||
pendingValues.push_back(indexSwitch.getResult(index));
|
||||
appendAliasDescription(interval.aliasesFollowed,
|
||||
indexSwitch.getResult(index));
|
||||
recordAlias(indexSwitch.getResult(index));
|
||||
}
|
||||
}
|
||||
else if (!forOp) {
|
||||
@@ -318,7 +324,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
|
||||
if (operand != value)
|
||||
continue;
|
||||
pendingValues.push_back(forOp.getResult(index));
|
||||
appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index));
|
||||
recordAlias(forOp.getResult(index));
|
||||
if (parentLoop && forOp == parentLoop)
|
||||
interval.escapesLoop = true;
|
||||
}
|
||||
@@ -411,7 +417,8 @@ static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, ArrayRef<Lo
|
||||
} // namespace
|
||||
|
||||
SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp,
|
||||
std::optional<unsigned> lane) {
|
||||
std::optional<unsigned> lane,
|
||||
bool includeAliasDescriptions) {
|
||||
SmallVector<LocalAllocInterval, 0> intervals;
|
||||
OperationOrdering ordering = buildOperationOrdering(coreLikeOp);
|
||||
if (ordering.position.empty())
|
||||
@@ -428,7 +435,8 @@ SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation
|
||||
llvm_unreachable("Failed to compute local allocation size");
|
||||
}
|
||||
|
||||
MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd);
|
||||
MemoryTouchInterval touchInterval =
|
||||
computeMemoryTouchInterval(allocOp, ordering, fallbackEnd, includeAliasDescriptions);
|
||||
LocalAllocInterval interval;
|
||||
interval.id = nextIntervalId++;
|
||||
interval.alloc = allocOp;
|
||||
|
||||
Reference in New Issue
Block a user