Version aware fusion plan#945
Conversation
…e fusionPlan pass
![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 version-aware planning to the Tile Fusion planner, allowing it to evaluate and compare fusion groups based on both compute nodes and their selected implementation versions. It refactors template candidate metadata parsing into a shared utility and updates the fusion planner to enumerate versioned states, score them with an extended cost model, and narrow candidate lists accordingly. The review feedback highlights two key improvements: caching the parsed template candidate versions in enumerateVersionedStatesFromSeed to avoid redundant parsing in a nested loop, and replacing a non-standard getAs<IntegerAttr> call on DictionaryAttr with dyn_cast_or_null to ensure MLIR version compatibility.
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.
| static FailureOr<SmallVector<GroupState, 16>> | ||
| enumerateVersionedStatesFromSeed(const PlanningContext &ctx, | ||
| const CostModel &costModel, | ||
| const pto::FusionComputeNode &seed, | ||
| const DenseSet<unsigned> &assignedNodes) { | ||
| FailureOr<SmallVector<GroupState, 4>> seedStates = createSeedStates(seed); | ||
| if (failed(seedStates)) | ||
| return failure(); | ||
|
|
||
| SmallVector<GroupState, 16> validStates; | ||
| SmallVector<GroupState, 16> frontier(seedStates->begin(), seedStates->end()); | ||
| std::set<std::string> seenStates; | ||
| for (const GroupState &state : frontier) | ||
| seenStates.insert(getStateSignature(state)); | ||
|
|
||
| // Explore all reachable dataflow-connected groups for this seed, not just | ||
| // the first greedy path through block order. The state signature includes the | ||
| // selected implementation version for each node, so the same node set can be | ||
| // scored multiple ways when TileOp template candidates differ. | ||
| while (!frontier.empty()) { | ||
| GroupState state = std::move(frontier.pop_back_val()); | ||
| for (const pto::FusionComputeNode &candidate : | ||
| ctx.blockAnalysis.computeNodes) { | ||
| if (assignedNodes.contains(candidate.id) || state.contains(candidate)) | ||
| continue; | ||
|
|
||
| FailureOr<SmallVector<TileOpImplVersion, 4>> versions = | ||
| getLegalImplVersions(candidate); | ||
| if (failed(versions)) | ||
| return failure(); | ||
|
|
||
| for (const TileOpImplVersion &version : *versions) { | ||
| FailureOr<std::optional<GroupState>> maybeNextState = | ||
| tryAppendVersionedCandidate(ctx, costModel, state, candidate, | ||
| version); | ||
| if (failed(maybeNextState)) | ||
| return failure(); | ||
| if (!*maybeNextState) | ||
| continue; | ||
|
|
||
| GroupState nextState = std::move(**maybeNextState); | ||
| std::string signature = getStateSignature(nextState); | ||
| if (!seenStates.insert(signature).second) | ||
| continue; | ||
|
|
||
| if (nextState.members.size() >= 2) | ||
| validStates.push_back(nextState); | ||
| frontier.push_back(std::move(nextState)); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| return validStates; | ||
| } |
There was a problem hiding this comment.
Calling getLegalImplVersions inside the nested loop of enumerateVersionedStatesFromSeed is highly inefficient. It parses the template candidate attributes from the operation every time, which involves dictionary lookups, attribute casting, and dynamic memory allocations for std::string names. Since the candidate nodes in the block do not change during the enumeration, these versions should be pre-parsed and cached at the start of the function to avoid redundant parsing.
static FailureOr<SmallVector<GroupState, 16>>
enumerateVersionedStatesFromSeed(const PlanningContext &ctx,
const CostModel &costModel,
const pto::FusionComputeNode &seed,
const DenseSet<unsigned> &assignedNodes) {
DenseMap<unsigned, SmallVector<TileOpImplVersion, 4>> cachedVersions;
for (const pto::FusionComputeNode &node : ctx.blockAnalysis.computeNodes) {
FailureOr<SmallVector<TileOpImplVersion, 4>> versions = getLegalImplVersions(node);
if (failed(versions))
return failure();
cachedVersions[node.id] = std::move(*versions);
}
auto seedIt = cachedVersions.find(seed.id);
if (seedIt == cachedVersions.end())
return failure();
SmallVector<GroupState, 16> validStates;
SmallVector<GroupState, 16> frontier;
for (const TileOpImplVersion &version : seedIt->second) {
GroupState state;
state.append(PlannedFusionMember{&seed, version});
frontier.push_back(std::move(state));
}
std::set<std::string> seenStates;
for (const GroupState &state : frontier)
seenStates.insert(getStateSignature(state));
// Explore all reachable dataflow-connected groups for this seed, not just
// the first greedy path through block order. The state signature includes the
// selected implementation version for each node, so the same node set can be
// scored multiple ways when TileOp template candidates differ.
while (!frontier.empty()) {
GroupState state = std::move(frontier.pop_back_val());
for (const pto::FusionComputeNode &candidate :
ctx.blockAnalysis.computeNodes) {
if (assignedNodes.contains(candidate.id) || state.contains(candidate))
continue;
auto candIt = cachedVersions.find(candidate.id);
if (candIt == cachedVersions.end())
return failure();
const SmallVector<TileOpImplVersion, 4> &versions = candIt->second;
for (const TileOpImplVersion &version : versions) {
FailureOr<std::optional<GroupState>> maybeNextState =
tryAppendVersionedCandidate(ctx, costModel, state, candidate,
version);
if (failed(maybeNextState))
return failure();
if (!*maybeNextState)
continue;
GroupState nextState = std::move(**maybeNextState);
std::string signature = getStateSignature(nextState);
if (!seenStates.insert(signature).second)
continue;
if (nextState.members.size() >= 2)
validStates.push_back(nextState);
frontier.push_back(std::move(nextState));
}
}
}
return validStates;
}There was a problem hiding this comment.
fixed by adding cache to prevent redundant metadata parsing
| auto id = candidate.getAs<IntegerAttr>(pto::kTemplateCandidateIdAttr); | ||
| if (!id) |
There was a problem hiding this comment.
Using getAs<IntegerAttr> on DictionaryAttr is not standard in MLIR and might lead to compilation errors depending on the exact MLIR version. It is safer and more consistent with the rest of the codebase (e.g., TemplateAttributes.cpp) to use dyn_cast_or_null<IntegerAttr>(candidate.get(...)).
| auto id = candidate.getAs<IntegerAttr>(pto::kTemplateCandidateIdAttr); | |
| if (!id) | |
| auto id = dyn_cast_or_null<IntegerAttr>(candidate.get(pto::kTemplateCandidateIdAttr)); | |
| if (!id) |
1 participant
Summary
Implementation
candidateslist to all equally best-scoring versions without reordering or selecting a finalversion.
pto.fusion.group_idandpto.fusion.ordercontract.Tests
fusion_plan_alternate_versioned_path.ptofusion_plan_version_score_skips_bad_extension.ptofusion_plan_retains_tied_candidates.pto