Faster codegen and merge
Validate Operations / validate-operations (push) Has been cancelled

This commit is contained in:
ilgeco
2026-07-16 11:12:14 +02:00
parent c744f388dc
commit d788178749
5 changed files with 116 additions and 40 deletions
+19 -11
View File
@@ -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;