[compiler][stream] Avoid circular dependency between partitions in execution scheduling #18217
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -44,6 +44,14 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config, | |
| Block *block) { | ||
| PartitionSet partitionSet; | ||
|
|
||
| struct OpInfo { | ||
| // Which partitions the op is contained within. | ||
| llvm::BitVector membership; | ||
| // Which partitions transitively depend on this operation. | ||
| llvm::BitVector hazards; | ||
| }; | ||
| DenseMap<Operation *, OpInfo> opInfos; | ||
|
|
||
| struct PartitionBuilder { | ||
| unsigned ordinal; | ||
| // Affinity of the partition. | ||
|
|
@@ -52,24 +60,76 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config, | |
| SetVector<Operation *> ops; | ||
| // Ops that were cloned and are known not to have their values escape. | ||
| DenseSet<Operation *> clonedOps; | ||
| // Which partitions transitively depend on this partition. | ||
| llvm::BitVector hazards; | ||
| void insert(Operation *op, OpInfo &opInfo) { | ||
| if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) { | ||
| affinity = affinity ? affinity.joinAND(affinityOp.getAffinityAttr()) | ||
| : affinityOp.getAffinityAttr(); | ||
| } | ||
| opInfo.membership.set(ordinal); | ||
| if (opInfo.hazards.size() > ordinal) | ||
| opInfo.hazards.reset(ordinal); | ||
| ops.insert(op); | ||
| hazards |= opInfo.hazards; | ||
| } | ||
| }; | ||
| SmallVector<std::unique_ptr<PartitionBuilder>> builders; | ||
| llvm::BitVector usableBuilders; | ||
|
|
||
| auto willCreateCircularDependencyBetweenPartitions = | ||
| [&](unsigned sourceOrdinal, unsigned targetOrdinal) -> bool { | ||
| // Returns: | ||
| // If we are to make partition with ordinal targetOrdinal to | ||
| // depend on partition with ordinal sourceOrdinal, | ||
| // will this create a circular dependency. | ||
| if (sourceOrdinal == targetOrdinal) | ||
| return false; | ||
| return builders[sourceOrdinal]->hazards.size() > targetOrdinal && | ||
| builders[sourceOrdinal]->hazards[targetOrdinal]; | ||
| }; | ||
|
|
||
| auto canAddOpToPartition = [&](Operation &op, OpInfo &opInfo, | ||
| unsigned partitionOrdinal) { | ||
| auto streamableOp = dyn_cast<IREE::Stream::StreamableOpInterface>(op); | ||
| if (!streamableOp) | ||
| return false; | ||
| IREE::Stream::AffinityAttr affinityAttr; | ||
| if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) | ||
| affinityAttr = affinityOp.getAffinityAttr(); | ||
| if (!IREE::Stream::AffinityAttr::canExecuteTogether( | ||
| affinityAttr, builders[partitionOrdinal]->affinity)) | ||
| return false; | ||
|
|
||
| bool preferCloneToConsumers = streamableOp.preferCloneToConsumers(); | ||
| llvm::BitVector *opHazards = nullptr; | ||
| llvm::BitVector opHazardsInCandidatePartition; | ||
| if (preferCloneToConsumers) { | ||
| // If we are cloning we care only about users that are a part of the | ||
| // candidate partition. | ||
| opHazards = &opHazardsInCandidatePartition; | ||
| for (auto user : op.getUsers()) { | ||
| if (builders[partitionOrdinal]->ops.contains(user)) | ||
| opHazardsInCandidatePartition |= opInfos[user].hazards; | ||
| } | ||
| } else | ||
| opHazards = &opInfo.hazards; | ||
|
|
||
| for (auto opHazardOrdinal : opHazards->set_bits()) { | ||
| if (partitionOrdinal < opHazardOrdinal) | ||
| // Reject partition ordering that would require partition sorting. | ||
|
There was a problem hiding this comment. use |
||
| // TODO: It is probably more optimal to reorder the partitions after | ||
| // their formation based on their dependency graph instead of rejecting | ||
| // here. Since this is considered not a good partitioning algorithm | ||
| // and will probably get removed, we leave it like that. | ||
| return false; | ||
| // Check for formation of circular dependency between partitions. | ||
| if (willCreateCircularDependencyBetweenPartitions(opHazardOrdinal, | ||
| partitionOrdinal)) | ||
| return false; | ||
| } | ||
| return true; | ||
| }; | ||
|
|
||
| auto asmState = getRootAsmState(block); | ||
|
|
||
|
|
@@ -107,11 +167,6 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config, | |
| opInfo.hazards.reserve(builders.size() + 1); | ||
| opInfo.hazards.resize(builders.size(), /*t=*/false); | ||
|
|
||
| LLVM_DEBUG({ | ||
| llvm::dbgs() << "====\nPartitioning op:\n"; | ||
| op.print(llvm::dbgs(), *asmState); | ||
|
|
@@ -149,8 +204,7 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config, | |
|
|
||
| // Prune candidates that do not have a compatible affinity. | ||
| for (auto ordinal : candidates.set_bits()) { | ||
| if (!canAddOpToPartition(op, opInfo, ordinal)) { | ||
| LLVM_DEBUG(llvm::dbgs() | ||
| << "Candidate partition " << ordinal << " incompatible\n"); | ||
| candidates.reset(ordinal); | ||
|
|
@@ -181,19 +235,15 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config, | |
| LLVM_DEBUG(llvm::dbgs() << "Cloning into consumer partition " | ||
| << consumerOrdinal << "\n"); | ||
| auto &consumerBuilder = builders[consumerOrdinal]; | ||
| consumerBuilder->insert(&op, opInfo); | ||
| consumerBuilder->clonedOps.insert(&op); | ||
| } | ||
| } else { | ||
| int consumerOrdinal = consumers.find_last(); | ||
| LLVM_DEBUG(llvm::dbgs() << "Moving into consumer partition " | ||
| << consumerOrdinal << "\n"); | ||
| auto &consumerBuilder = builders[consumerOrdinal]; | ||
| consumerBuilder->insert(&op, opInfo); | ||
| } | ||
| LLVM_DEBUG(llvm::dbgs() << "Handled streamable (continue)\n"); | ||
| continue; | ||
|
|
@@ -204,13 +254,13 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config, | |
| if (firstCandidateOrdinal != -1) { | ||
| LLVM_DEBUG(llvm::dbgs() << "Moving to first candidate partition " | ||
| << firstCandidateOrdinal << " (continue)\n"); | ||
| builders[firstCandidateOrdinal]->insert(&op, opInfo); | ||
| continue; | ||
| } | ||
|
|
||
| // Mark the op as having hazards against all other partitions. | ||
| // Why are we that conservative? | ||
| // Why we don't take the hazards for the users? | ||
| if (!builders.empty()) { | ||
| opInfo.hazards.set(0, builders.size() - 1); | ||
| } | ||
|
There was a problem hiding this comment. @benvanik, could you shed some light on what was your reasoning here when you wrote the algorithm? There was a problem hiding this comment. It's better to be safe than incorrect, especially with our current minimal test coverage. It's not always safe to reorder things - if anything we are unlikely to be conservative enough here - for example, if there's a |
||
|
|
@@ -219,8 +269,7 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config, | |
| opInfo.membership.resize(opInfo.membership.size() + 1, /*t=*/true); | ||
| auto builder = std::make_unique<PartitionBuilder>(); | ||
| builder->ordinal = builders.size(); | ||
| builder->insert(&op, opInfo); | ||
| LLVM_DEBUG(llvm::dbgs() | ||
| << "Created partition " << builder->ordinal << "\n"); | ||
| builders.push_back(std::move(builder)); | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actually, here we would need to walk further down the users if a user is also cloned into the partition. This will be useful if we have a block of cloneable ops. If left like that, other than the inefficiency, I don't think it will produce invalid partitioning.