[AOTI] Support module buffer mutation

[desertfire]

Stack from ghstack (oldest at bottom):

• -> [AOTI] Support module buffer mutation #123164

Summary: Fixes #120424. Because in a forward pass module buffers may be mutated, we need to allow that in AOTI. In addition, this will be a necessary step if we want to extend AOTI to training.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @amjames @chauhang

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit 4faec2d with merge base bcb6e5a ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / linux-docs / build-docs-python-false (gh)
  Process completed with exit code 2.

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• pull / linux-jammy-py3.8-gcc11 / test (docs_test, 1, 1, linux.2xlarge, unstable) (gh)
  Process completed with exit code 2.

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

[desertfire]

Reland #122824 because there was a merge conflict.

[digantdesai]

Stamping

[malfet]

LGTM (provided that Mac CI is green)

[desertfire]

@pytorchbot merge

[pytorchmergebot]

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

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

• pull / linux-docs / build-docs-python-false

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[desertfire]

@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 2 checks: pull / linux-docs / build-docs-python-false, pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 3, 5, linux.4xlarge.nvidia.gpu)

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

[pytorchmergebot]

Merge failed

Reason: Command git -C /home/runner/work/pytorch/pytorch rebase origin/main returned non-zero exit code 1

Rebasing (1/1)
Auto-merging torch/_inductor/compile_fx.py
CONFLICT (content): Merge conflict in torch/_inductor/compile_fx.py
Auto-merging torch/_inductor/lowering.py
Auto-merging torch/_inductor/utils.py
error: could not apply 37dc4037b85... [AOTI] Support module buffer mutation (#123164)
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
Could not apply 37dc4037b85... [AOTI] Support module buffer mutation (#123164)

Details for Dev Infra team

Raised by workflow job

[desertfire]

@pytorchbot merge -f "Trunk test passed previously; fixed a minor merge conflict"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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

[jithunnair-amd]

@desertfire This is the second recent AOTI-related PR that has broken ROCm CI due to nosignal; the other one was 122882. Since the only ROCm CI jobs that run by default on AOTI PRs are the ones for ciflow/trunk (pre-merge checks) and ciflow/inductor (auto-label by bot), and neither of those currently run test_aot_inductor.py, there is no CI signal on the PR indicating a ROCm breakage (we can't run the entire list of test suites on PRs due to CI capacity constraints). To remedy this, I propose adding inductor/test_aot_inductor to the inductor workflow here:

pytorch/.ci/pytorch/test.sh

Line 318 in 9288b27

python test/run_test.py --include inductor/test_torchinductor inductor/test_torchinductor_opinfo --verbose

It seems like it would add about an hour more to the inductor workflow test jobs for both CUDA and ROCm. Would that be acceptable?

Testing this out in #123340

cc @huydhn @malfet @atalman

[atalman]

@desertfire Please provide forward fix for this issue this is a failing test on the trunk:

024-04-03T07:37:22.3706672Z =========================== short test summary info ============================
2024-04-03T07:37:22.3707311Z FAILED [6.3134s] inductor/test_aot_inductor.py::AOTInductorTestABICompatibleCuda::test_buffer_mutation_3_abi_compatible_cuda - AssertionError: False is not true
2024-04-03T07:37:22.3707315Z 
2024-04-03T07:37:22.3707505Z To execute this test, run the following from the base repo dir:
2024-04-03T07:37:22.3707879Z     PYTORCH_TEST_WITH_ROCM=1 python test_aot_inductor.py -k test_buffer_mutation_3_abi_compatible_cuda
2024-04-03T07:37:22.3707883Z