[Stream] Fix O(N²) compile-time scaling in partitionRegionConcurrency by isa-Lai · Pull Request #24958 · iree-org/iree

isa-Lai · 2026-09-22T14:11:10Z

When analyzing execution regions where many operations share a common resource operand, partitionRegionConcurrencyReference previously iterated over operand.getUsers() and called tiedOp.hasAnyTiedUses(), causing a nested linear scan across all use sites.

Switch to iterating over operand.getUses() directly and check tiedOp.isOperandTied(use.getOperandNumber()) in O(1) time. Deduplicate accumulated operations using a SmallPtrSet so operands referenced across multiple slots on the same operation are only processed once.

Fixes #24928

Signed-off-by: @isa-Lai
Assisted-by: Gemini CLI

…Reference Fix a compile-time regression in Stream partitioning analysis (`partitionRegionConcurrencyReference`) when processing execution regions where many operations share a common resource operand (e.g., unrolled or sliced dispatches sharing an input tensor). In `partitionRegionConcurrencyReference`, for each operation and resource operand, all users of the operand were inspected. For each user op, the pass invoked `tiedOp.hasAnyTiedUses(operand)`. Because `hasAnyTiedUses` scanned linearly over all use sites of `operand` inside an outer loop over users, the check scaled cubically/quadratically with the number of dispatches sharing that resource: O(ops * users * uses). - Switch from `operand.getUsers()` to `operand.getUses()` in `ReferencePartitioning.cpp` to directly inspect each use site with its `OpOperand`. - Replace `hasAnyTiedUses(operand)` (O(uses)) with `isOperandTied(use.getOperandNumber())` (O(1)), eliminating the inner scan. - Add `seenTiedUsers` (`llvm::SmallPtrSet<Operation *, 4>`) to deduplicate hazard accumulation when an operation references the operand across multiple argument positions. Signed-off-by: isa-lai <isabellai1004@gmail.com>

github-actions · 2026-09-22T14:11:29Z

Hello @isa-Lai 👋

Thank you for submitting a Pull Request to IREE! It looks like this is your first one. We have one ask, and you can also find some general tips below.

Action required: acknowledge IREE project policies

IREE is a Linux Foundation project. All participants are expected to follow the LF Projects Code of Conduct.

All contributions to IREE must follow our IREE AI Tool Use Policy. In particular:

Contributors must fully understand, and vouch for, all submitted changes and the intent behind them.
Substantial use of LLM/generative AI tools must be noted in the PR description, e.g. via Assisted-by: tool-name or Co-authored-by: tool-name tool@email trailers.
Contributors must write PR descriptions themselves. There must always be a human in the loop: contributors must respond to reviews and questions by themselves.
If a response includes LLM-assisted segments (e.g. reproducers, LLM agent analysis excerpts), the segment should be clearly marked as "assisted", same as for PR contents.
GitHub issues labeled as "Good first issue" are explicitly designated as learning opportunities for newcomers to the project. With exceptions for boilerplate edits, AI tool usage for resolutions to such issues is forbidden.

We kindly ask you to reply to this message and confirm that you understand and accept the cited policies, particularly the AI Tool Use Policy.

General guidance

Our general Contributing guide contains information and links to detailed guides on code quality, testing, commit summaries and our CI system.

A common point for new PRs: if a DCO signing check fails for you, check out the section on Developer Certificate of Origin.
In these cases, it should suffice to amend your commit signature(s) per the guide and force-push the PR branch.

If you have any questions, feel free to leave a comment here, or ask away on IREE Discord.

Thank you,
The IREE Community

isa-Lai · 2026-09-22T14:23:35Z

Hello @isa-Lai 👋

Thank you for submitting a Pull Request to IREE! It looks like this is your first one. We have one ask, and you can also find some general tips below.

Action required: acknowledge IREE project policies

IREE is a Linux Foundation project. All participants are expected to follow the LF Projects Code of Conduct.

All contributions to IREE must follow our IREE AI Tool Use Policy. In particular:

Contributors must fully understand, and vouch for, all submitted changes and the intent behind them.

Substantial use of LLM/generative AI tools must be noted in the PR description, e.g. via Assisted-by: tool-name or Co-authored-by: tool-name tool@email trailers.

Contributors must write PR descriptions themselves. There must always be a human in the loop: contributors must respond to reviews and questions by themselves.
If a response includes LLM-assisted segments (e.g. reproducers, LLM agent analysis excerpts), the segment should be clearly marked as "assisted", same as for PR contents.

GitHub issues labeled as "Good first issue" are explicitly designated as learning opportunities for newcomers to the project. With exceptions for boilerplate edits, AI tool usage for resolutions to such issues is forbidden.

We kindly ask you to reply to this message and confirm that you understand and accept the cited policies, particularly the AI Tool Use Policy.

General guidance

Our general Contributing guide contains information and links to detailed guides on code quality, testing, commit summaries and our CI system.

A common point for new PRs: if a DCO signing check fails for you, check out the section on Developer Certificate of Origin. In these cases, it should suffice to amend your commit signature(s) per the guide and force-push the PR branch.

If you have any questions, feel free to leave a comment here, or ask away on IREE Discord.

Thank you, The IREE Community

I understand and accept the cited policies.

Manewing · 2026-09-22T16:33:59Z

+      llvm::SmallPtrSet<Operation *, 4> seenTiedUsers;
+      for (auto &use : operand.getUses()) {
+        Operation *user = use.getOwner();
        if (user == &op || user->getBlock() != block ||
            user->isBeforeInBlock(&op)) {
          continue;
        }
        auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(user);
-        if (!tiedOp || !tiedOp.hasAnyTiedUses(operand)) {
+        if (!tiedOp || !tiedOp.isOperandTied(use.getOperandNumber())) {
+          continue;
+        }
+        if (!seenTiedUsers.insert(user).second) {


isOperandTied changes semantics we now only check a single use instead of checking all uses of the operand globally.

However, given that hasAnyTiedUses is a static function there is no need at all to have this called in the function itself. We can check it before the loop and skip the entire user loop.

Suggested change

llvm::SmallPtrSet<Operation *, 4> seenTiedUsers;

for (auto &use : operand.getUses()) {

Operation *user = use.getOwner();

if (user == &op || user->getBlock() != block ||

user->isBeforeInBlock(&op)) {

continue;

}

auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(user);

if (!tiedOp || !tiedOp.hasAnyTiedUses(operand)) {

if (!tiedOp || !tiedOp.isOperandTied(use.getOperandNumber())) {

continue;

}

if (!seenTiedUsers.insert(user).second) {

if (!IREE::Util::TiedOpInterface::hasAnyTiedUses(operand)) {

continue;

}

llvm::SmallPtrSet<Operation *, 4> seenTiedUsers;

for (auto &use : operand.getUses()) {

Operation *user = use.getOwner();

if (user == &op || user->getBlock() != block ||

user->isBeforeInBlock(&op)) {

continue;

}

if (!isa<IREE::Util::TiedOpInterface>(user)) {

continue;

}

if (!seenTiedUsers.insert(user).second) {

The formatting got a bit messed up here, I hope it is clear what I am suggesting 😅

Manewing

Thank you for your submission! :)

I think we can improve even further by skipping the entire loop since the tied uses check is using a static function, no need at all to call this within the loop.

Also would you mind trimming down your PR description a bit it will be the commit message for this change :)

isa-Lai · 2026-09-24T00:11:01Z

Hi @Manewing. Thanks for the review. But I think we still need to check isOperandTied in the uses loop (please correct me if I am wrong) because of this case:

%buf1
%out0: use %buf1
%out1: tied %buf1. 
%out2: use %buf1 and tied %buf2

Assuming we execute your suggested code on %out0 now, hasAnyTiedUses will return true because %out1 tied, and we go into the for (auto &use : operand.getUses()) loops.

We check %out1 and it is tied %buf1. Your code will behave correctly as it cannot run concurrently with %out1.
We check %out2. In your code isa<IREE::Util::TiedOpInterface>(user) return true so the code think this cannot run concurrently, but it actually can, because it does not mutate buf1.
That is why I still check if (!tiedOp || !tiedOp.isOperandTied(use.getOperandNumber())), if this is tiedOp but does not tied the same operand, it can still be valid. This may help with the final performance.

In terms of the early stop using hasAnyTiedUses outside the loop, I can add that, but might gain very little performance improvement if we are also checking isOperandTied inside the loop, so I am keeping my PR as it.

I also shorten the descriptor a bit.

isa-Lai requested review from Manewing, benvanik, devtbi and ziereis as code owners September 22, 2026 14:11

Manewing reviewed Sep 22, 2026

View reviewed changes

Navigation Menu

[Stream] Fix O(N²) compile-time scaling in partitionRegionConcurrency - #24958

[Stream] Fix O(N²) compile-time scaling in partitionRegionConcurrency#24958
isa-Lai wants to merge 1 commit into
iree-org:mainfrom
isa-Lai:bug/use-compile

isa-Lai commented Sep 22, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 22, 2026

Uh oh!

isa-Lai commented Sep 22, 2026 •

edited

Loading

Action required: acknowledge IREE project policies

General guidance

Uh oh!

Manewing Sep 22, 2026 •

edited

Loading

Uh oh!

Manewing Sep 22, 2026

Uh oh!

Manewing left a comment •

edited

Loading

Uh oh!

isa-Lai commented Sep 24, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

isa-Lai commented Sep 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Sep 22, 2026

Action required: acknowledge IREE project policies

General guidance

Uh oh!

isa-Lai commented Sep 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Action required: acknowledge IREE project policies

General guidance

Uh oh!

Manewing Sep 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Manewing Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

Manewing left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

isa-Lai commented Sep 24, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

isa-Lai commented Sep 22, 2026 •

edited

Loading

isa-Lai commented Sep 22, 2026 •

edited

Loading

Manewing Sep 22, 2026 •

edited

Loading

Manewing left a comment •

edited

Loading