Avoid folding mutated buffer storages (#185077)
AOTInductor packages unlift exported module buffers into runtime-owned constants. For buffer mutation, export re-inserts a copy_ from the functional mutation output back into that buffer. The uniform-value constant folder could fold a singleton get_attr buffer into a compile-time full(...) value before that copy_, so generated packages wrote the same compile-time value into the buffer on every call. Loaded packages therefore returned 2, 2, 2 instead of preserving the counter state as 2, 3, 4. Track storages written by mutable operator schema arguments and skip uniform-value folding for any node whose fake tensor storage aliases one of those writes. This keeps mutable buffers and their views as runtime reads while preserving uniform folding for immutable values. The alternative would be to special-case AOTI buffer constants, but storage-based alias tracking matches the actual hazard and also covers views of mutated buffers. The regression test compiles and loads an AOTI package, then calls it repeatedly for assignment, add_, slice assignment, shape-(2,) add_, and indexed mutation with both always_keep_tensor_constants settings. Fixes #165444 Generated by my agent Test Plan: - python test/inductor/test_aot_inductor_package.py TestAOTInductorPackage_cpu.test_buffer_mutations_persist_across_package_calls - python -m pytest test/inductor/test_aot_inductor_package.py::TestAOTInductorPackage_cpu::test_buffer_mutations_persist_across_package_calls -q - python -m py_compile torch/_inductor/fx_passes/joint_graph.py test/inductor/test_aot_inductor_package.py - lintrunner -a torch/_inductor/fx_passes/joint_graph.py test/inductor/test_aot_inductor_package.py Pull Request resolved: https://github.com/pytorch/pytorch/pull/185077 Approved by: https://github.com/eellison
J
Jason Ansel committed
96c5e78ddd4918fab2bb4c795bde7fbcb1241d76
Parent: 66eef30
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 8/24/2026, 2:05:28 AM