8000 [PT2][memory] correct wait tensor output size by xuanzhang816 · Pull Request #153569 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

[PT2][memory] correct wait tensor output size #153569

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.

Sign up for GitHub

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

Closed
wants to merge 4 commits into from
Closed

[PT2][memory] correct wait tensor output size #153569

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
692be3c
Update
xuanzhang816 May 14, 2025
1bfb776
Update on "[PT2][memory] add missing dependencies due to mutations"
xuanzhang816 May 14, 2025
b5a4be8
Update on "[PT2][memory] add missing dependencies due to mutations"
xuanzhang816 May 17, 2025
65d4a4e
Update on "[PT2][memory] add missing dependencies due to mutations"
xuanzhang816 May 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions torch/_inductor/memory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from torch.utils._ordered_set import OrderedSet

from .ir import MultiOutputLayout, NoneLayout
from .utils import get_dtype_size
from .utils import get_dtype_size, is_wait
from .virtualized import V


Expand Down Expand Up @@ -147,8 +147,23 @@ def _compute_and_update_buf_size(
sched_buf: SchedulerBuffer, user_of_MultiOutputLayout: bool = False
) -> int:
if isinstance(sched_buf.node.layout, NoneLayout):
sched_buf_to_size[sched_buf.get_name()] = (0, 0)
return 0
_size = 0
# for a wait tensor op, its schedulerBuffer NoneLayout layout. However,
# the schedulerBuffer is treated as a mutation of the collective output
# so it needs to inherit the size of the collectives
if (
sched_buf.defining_op
and is_wait(sched_buf.defining_op.node)
and sched_buf.get_mutations()
):
mutated_buf_name = sched_buf.get_mutations()[0]
_size = (
sched_buf_to_size[mutated_buf_name][1]
if mutated_buf_name in sched_buf_to_size
else 0
)
sched_buf_to_size[sched_buf.get_name()] = (_size, _size)
return _size
elif isinstance(sched_buf.node.layout, MultiOutputLayout):
size_alloc = 0
for user in sched_buf.users:
Expand Down
Loading
0