Skip to content

[AOTI] Fix a small buffer mutation issue#169347

Closed
desertfire wants to merge 4 commits intogh/desertfire/614/basefrom
gh/desertfire/614/head
Closed

[AOTI] Fix a small buffer mutation issue#169347
desertfire wants to merge 4 commits intogh/desertfire/614/basefrom
gh/desertfire/614/head

Conversation

@desertfire
Copy link
Contributor

@desertfire desertfire commented Dec 2, 2025

Stack from ghstack (oldest at bottom):

Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo

Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented Dec 2, 2025

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/169347

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 6c7a272 with merge base c6ae757 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

desertfire added a commit that referenced this pull request Dec 2, 2025
Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

ghstack-source-id: 9840de9
Pull Request resolved: #169347
Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo

[ghstack-poisoned]
desertfire added a commit that referenced this pull request Dec 2, 2025
Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

ghstack-source-id: 842e21c
Pull Request resolved: #169347
Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo

[ghstack-poisoned]
desertfire added a commit that referenced this pull request Dec 2, 2025
Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

ghstack-source-id: bedd5e1
Pull Request resolved: #169347
self.named_buffers: dict[str, torch.Tensor] = (
const_module.named_buffers if const_module else {}
)
self.mutated_named_buffers: OrderedSet[torch.Tensor] = gm.meta.get(
Copy link
Contributor

@yushangdi yushangdi Dec 3, 2025

Choose a reason for hiding this comment

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

can or should we re-use self.mutated_buffers?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

self.mutated_buffers is only filled as we iterate through the graph, so by the point we check target in self.mutated_buffers, self.mutated_buffers may not have accurate enough mutation information.

Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo

[ghstack-poisoned]
@desertfire
Copy link
Contributor Author

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Dec 4, 2025
@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

umechand-amd pushed a commit to ROCm/pytorch that referenced this pull request Dec 8, 2025
Summary: Fix pytorch#169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

Pull Request resolved: pytorch#169347
Approved by: https://github.com/yushangdi
JacobSzwejbka pushed a commit that referenced this pull request Dec 8, 2025
Summary: Fix #169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

Pull Request resolved: #169347
Approved by: https://github.com/yushangdi
tiendatngcs pushed a commit to tiendatngcs/pytorch-Dec25 that referenced this pull request Dec 10, 2025
Summary: Fix pytorch/pytorch#169118. A small named buffer can be optimized as inlined constant, but if that named buffer is mutated later on, AOTI can generate wrong code if not respecting the mutation information. This PR fixes the issue by recording mutated_named_buffers when unlifting the input graph and pass to Inductor graph lowering.

ghstack-source-id: 6eec08b
Pull Request resolved: pytorch/pytorch#169347
@github-actions github-actions bot deleted the gh/desertfire/614/head branch January 4, 2026 02:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants