fix(vpto): insert mem_bar for V tile hazards#906
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the PTOInsertVMemBar pass to automatically insert VPTO V-pipeline memory barriers (pto.mem_bar) for overlapping UB tile hazards (RAW, WAR, WAW) before TileLang expansion. Feedback on the implementation suggests optimizing the barrier insertion by performing it on-the-fly to prevent redundant or overly conservative barriers, using operand indices instead of raw OpOperand * pointers to improve robustness, and adding overflow checks in the address range overlap calculation.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| SmallVector<std::pair<Operation *, MemBarKind>, 16> insertions; | ||
| bool sawUnknownRange = false; | ||
|
|
||
| for (Block &block : func.getBody()) { | ||
| SmallVector<VTileOpInfo, 16> vOps; | ||
| for (Operation &op : block) { | ||
| if (!isStructuredVTileOp(&op)) | ||
| continue; | ||
| auto accesses = collectTileAccesses(&op); | ||
| if (accesses.empty()) | ||
| continue; | ||
| sawUnknownRange |= hasUnknownRange(accesses); | ||
|
|
||
| HazardSet uncovered; | ||
| for (const VTileOpInfo &previous : vOps) { | ||
| HazardSet hazards = classifyHazards(previous.accesses, accesses); | ||
| if (hazards.empty()) | ||
| continue; | ||
| if (!hasCoveringMemBarBetween(previous.op, &op, hazards)) | ||
| uncovered.merge(hazards); | ||
| } | ||
|
|
||
| if (!uncovered.empty()) | ||
| insertions.push_back({&op, getMemBarKindFor(uncovered)}); | ||
| vOps.push_back(VTileOpInfo{&op, std::move(accesses)}); | ||
| } | ||
| } | ||
|
|
||
| OpBuilder builder(func.getContext()); | ||
| for (auto [op, kind] : insertions) { | ||
| builder.setInsertionPoint(op); | ||
| builder.create<pto::MemBarOp>(op->getLoc(), | ||
| pto::MemBarAttr::get(func.getContext(), | ||
| kind)); | ||
| } |
There was a problem hiding this comment.
We can optimize the memory barrier insertion by performing it on-the-fly instead of deferring it to the end of the block traversal.
Currently, because insertions are deferred, hasCoveringMemBarBetween cannot see the barriers that will be inserted for earlier hazards in the same block. This leads to redundant barriers (e.g., multiple VST_VLD barriers for the same RAW hazard on subsequent reads) and overly conservative barriers (e.g., merging multiple hazards into a heavy VV_ALL barrier when a lighter barrier inserted earlier would have already covered some of them).
By inserting the barriers on-the-fly, subsequent hazard checks will correctly detect the already-inserted barriers, resulting in minimal and optimal barrier insertion.
OpBuilder builder(func.getContext());
bool sawUnknownRange = false;
for (Block &block : func.getBody()) {
SmallVector<VTileOpInfo, 16> vOps;
for (Operation &op : block) {
if (!isStructuredVTileOp(&op))
continue;
auto accesses = collectTileAccesses(&op);
if (accesses.empty())
continue;
sawUnknownRange |= hasUnknownRange(accesses);
HazardSet uncovered;
for (const VTileOpInfo &previous : vOps) {
HazardSet hazards = classifyHazards(previous.accesses, accesses);
if (hazards.empty())
continue;
if (!hasCoveringMemBarBetween(previous.op, &op, hazards))
uncovered.merge(hazards);
}
if (!uncovered.empty()) {
builder.setInsertionPoint(&op);
builder.create<pto::MemBarOp>(
op.getLoc(),
pto::MemBarAttr::get(func.getContext(),
getMemBarKindFor(uncovered)));
}
vOps.push_back(VTileOpInfo{&op, std::move(accesses)});
}
}| SmallPtrSet<OpOperand *, 4> writeOperands; | ||
| for (OpOperand &init : dps.getDpsInitsMutable()) { | ||
| if (isa<TileBufType>(init.get().getType())) | ||
| writeOperands.insert(&init); | ||
| } | ||
|
|
||
| for (OpOperand &operand : op->getOpOperands()) { | ||
| if (!isa<TileBufType>(operand.get().getType())) | ||
| continue; | ||
| accesses.push_back( | ||
| TileAccess{resolveTileAddrRange(operand.get()), | ||
| writeOperands.contains(&operand)}); | ||
| } |
There was a problem hiding this comment.
In collectTileAccesses, storing and comparing OpOperand * pointers from dps.getDpsInitsMutable() and op->getOpOperands() can be fragile and relies on the iterator implementation details of MutableOperandRange.
A more robust and standard approach in MLIR is to store the operand indices (using OpOperand::getOperandNumber()) and check if the operand index is in the set of write operand indices.
| SmallPtrSet<OpOperand *, 4> writeOperands; | |
| for (OpOperand &init : dps.getDpsInitsMutable()) { | |
| if (isa<TileBufType>(init.get().getType())) | |
| writeOperands.insert(&init); | |
| } | |
| for (OpOperand &operand : op->getOpOperands()) { | |
| if (!isa<TileBufType>(operand.get().getType())) | |
| continue; | |
| accesses.push_back( | |
| TileAccess{resolveTileAddrRange(operand.get()), | |
| writeOperands.contains(&operand)}); | |
| } | |
| llvm::SmallDenseSet<unsigned, 4> writeOperandIndices; | |
| for (OpOperand &init : dps.getDpsInitsMutable()) { | |
| if (isa<TileBufType>(init.get().getType())) | |
| writeOperandIndices.insert(init.getOperandNumber()); | |
| } | |
| for (OpOperand &operand : op->getOpOperands()) { | |
| if (!isa<TileBufType>(operand.get().getType())) | |
| continue; | |
| accesses.push_back( | |
| TileAccess{resolveTileAddrRange(operand.get()), | |
| writeOperandIndices.contains(operand.getOperandNumber())}); | |
| } |
| static bool overlaps(const AddrRange &lhs, const AddrRange &rhs) { | ||
| if (!lhs.valid || !rhs.valid) | ||
| return true; | ||
| return lhs.base < rhs.base + rhs.size && rhs.base < lhs.base + lhs.size; | ||
| } |
There was a problem hiding this comment.
The overlaps function performs addition on uint64_t values (rhs.base + rhs.size and lhs.base + lhs.size) without checking for potential integer overflow. While tile sizes and bases are typically small, it is safer and more robust to check for overflow to prevent undefined behavior or incorrect overlap detection.
static bool overlaps(const AddrRange &lhs, const AddrRange &rhs) {
if (!lhs.valid || !rhs.valid)
return true;
uint64_t lhs_end = lhs.base + lhs.size;
uint64_t rhs_end = rhs.base + rhs.size;
if (lhs_end < lhs.base || rhs_end < rhs.base)
return true;
return lhs.base < rhs_end && rhs.base < lhs_end;
}
Codex Review该评论由 review 机器人自动更新。
SummaryReview failed at stage Findings未生成结构化 findings，因为 review 过程提前失败。 Log Tail |
2 participants
Summary
Fixes #866.
This PR adds an A5 VPTO
pto-insert-v-membarpass to insertpto.mem_barfor V-pipeline UB hazards before TileLang expansion. It handles RAW, WAR, and WAW dependencies between structured PIPE_V tile ops using DPS read/write roles andalloc_tile addrbyte ranges.The pass runs before VPTO fusion planning, so inserted
pto.mem_baroperations act as tile-native scheduling boundaries instead of breaking already-created fusion regions.Key behavior
VST_VLD.VLD_VST.VST_VST.VV_ALL.pto.mem_baronly suppresses insertion when its kind covers the hazard.Validation
ninja -C build-llvm21 PTOTransforms ptoasllvm-lit -v build-llvm21/test/lit/vpto/vmembar_tile_hazards.pto build-llvm21/test/lit/vpto/membar_barrier_types_vpto_llvm.ptogit diff --check