Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[stream] Dedeplicate stream.yield values from stream.async.execute #20056

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

rsuderman
Copy link
Contributor

If during cleanup values are replicated in stream.async.execute cases it can result in double allocation of those blocks (this occurs as ScheduleAllocation assumes all returns are unique). This double allocation also results in bad numerical behavior as the separate returns may be used by different dispatches and fetch incorrect values.

@rsuderman rsuderman requested a review from benvanik as a code owner February 21, 2025 04:14
@rsuderman rsuderman force-pushed the dedupe_async branch 2 times, most recently from a80c055 to 7e6abff Compare February 21, 2025 04:17
Copy link
Collaborator

@benvanik benvanik left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is something we want in optimizeClosureLikeOp instead - that already deduplicates operands, removes unused results, and handles tied operands correctly - I'll see if I can get it in real quick

auto yield =
cast<IREE::Stream::YieldOp>(op.getBody().front().getTerminator());
int64_t oldYieldCount = yield.getResourceOperands().size();
for (int i = 0, s = oldYieldCount; i < s; ++i) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

style/warnings: use consistent types - oldYieldCount is int64_t so you'd use int64_t here - but really, size() is size_t so oldYieldCount and i/s here should be size_t
(some compilers care about this more than others especially if signedness changes, but it also helps readability)

auto newExecuteOp = rewriter.create<IREE::Stream::AsyncExecuteOp>(
op.getLoc(), newTypes, newResultSizes, op.getAwaitTimepoint(),
op.getResourceOperands(), op.getResourceOperandSizes(),
llvm::SmallVector<int64_t>());
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we'll need to preserve tied operands - this is stripping them all

(always useful to throw in a tied operand to ensure they are not stripped, like by adding a ..., %tied) -> %tied to the end so that it's clear the indices are updated properly - updating index 0 tied operands is usually too easy :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants