[inductor] allocate non-blocking copy destinations in pinned memory (#155121) #158758

v0i0 · 2025-07-21T18:37:55Z

Fixes #155121
Stack from ghstack (oldest at bottom):

[inductor] move all cpu scalars using pinned memory for graph partition (#155360) #158983
-> [inductor] allocate non-blocking copy destinations in pinned memory (#155121) #158758

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

…155121) [ghstack-poisoned]

pytorch-bot · 2025-07-21T18:37:59Z

🔗 Helpful Links

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

📄 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 ada0bdc with merge base dad2a05 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

trunk / linux-jammy-rocm-py3.10 / test (distributed, 1, 1, linux.rocm.gpu.4) (gh) (similar failure)
distributed/test_c10d_nccl.py::ProcessGroupNCCLGroupTest::test_extra_cuda_context

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / linux-jammy-py3_9-clang9-xla / test (xla, 1, 1, lf.linux.12xlarge, unstable) (gh) (#158876)
/var/lib/jenkins/workspace/xla/torch_xla/csrc/runtime/BUILD:476:14: Compiling torch_xla/csrc/runtime/xla_util_test.cpp failed: (Exit 1): gcc failed: error executing CppCompile command (from target //torch_xla/csrc/runtime:xla_util_test) /usr/bin/gcc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wunused-but-set-parameter -Wno-free-nonheap-object -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections ... (remaining 229 arguments skipped)

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

…155121) ghstack-source-id: f865a6f Pull Request resolved: #158758

…ed memory (#155121)" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Lucaskabela [ghstack-poisoned]

…155121) ghstack-source-id: 0a98613 Pull Request resolved: #158758

eellison

Looking good !

Add test both for the cat and non-cat cases ? and a test case when cpp wrapper is turned on ?

torch/_inductor/codegen/cpp_wrapper_cpu_array_ref.py

eellison · 2025-07-21T21:33:34Z

torch/_inductor/ir.py

@@ -3534,6 +3534,7 @@ def __init__(
        size: Sequence[Expr],
        stride: Optional[Sequence[Expr]] = None,
        offset: Expr = Integer(0),
+        is_pinned: bool = False,


You got the important cases, but I think there are a few more places you'll need to thread the is_pinned field through. Grep through FixedLayout/FlexibleLayout construction.

This might be a bit tricky without more context, but the instantiations that happen as views should preserve whether the input layout was pinned.

See, ExpandView::create, PermuteView::create, SqueezeView::create View::create, SliceView::create, convert_to_reinterpret_view.

…ed memory (#155121)" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Lucaskabela [ghstack-poisoned]

…155121) ghstack-source-id: 61c061d Pull Request resolved: #158758

…ed memory (#155121)" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Lucaskabela [ghstack-poisoned]

…155121) ghstack-source-id: 7f0eeb6 Pull Request resolved: #158758

…ed memory (#155121)" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Lucaskabela [ghstack-poisoned]

…155121) ghstack-source-id: b041df3 Pull Request resolved: #158758

test/inductor/test_torchinductor.py

…ed memory (#155121)" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Lucaskabela [ghstack-poisoned]

…155121) ghstack-source-id: 95308ae Pull Request resolved: #158758

test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/csrc/kernel.cpp

…ed memory (#155121)" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Lucaskabela [ghstack-poisoned]

…ed memory (#155121)" Fixes #155121 cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Lucaskabela [ghstack-poisoned]

v0i0 · 2025-08-04T21:14:38Z

@pytorchbot merge

pytorchmergebot · 2025-08-04T21:16:32Z

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

v0i0 · 2025-08-04T21:34:48Z

@yushangdi i tried this, it works on my end

…on (#155360) (#158983) Pull Request resolved: #158983 Approved by: https://github.com/eellison ghstack dependencies: #158758

davidberard98 · 2025-08-04T21:45:14Z

@pytorchbot revert -c ghfirst -m "I need to revert #158462 (it causes device-side asserts), and this PR causes a merge conflict in the test file. Sorry about that!"

pytorchmergebot · 2025-08-04T21:46:58Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

… partition (#155360) (#158983)" This reverts commit 444e238. Reverted #158983 on behalf of https://github.com/davidberard98 due to I need to revert #158462 (it causes device-side asserts), and this PR causes a merge conflict in the test file. Sorry about that! ([comment](#158758 (comment)))

…memory (#155121) (#158758)" This reverts commit 6085bf7. Reverted #158758 on behalf of https://github.com/davidberard98 due to I need to revert #158462 (it causes device-side asserts), and this PR causes a merge conflict in the test file. Sorry about that! ([comment](#158758 (comment)))

pytorchmergebot · 2025-08-04T21:47:18Z

@v0i0 your PR has been successfully reverted.

github-actions · 2025-08-04T21:51:19Z

Attention! PyTorch one of the C-stable API file was changed

You MUST NOT change existing function declarations in this, as this header defines a stable C ABI. If you need to change the signature for a function, introduce a new v2 version of the function and modify code generation to target the new version of the function.

Caused by:

torch/csrc/inductor/aoti_torch/c/shim.h

[ghstack-poisoned]

v0i0 · 2025-08-07T17:04:03Z

@pytorchbot merge

pytorchmergebot · 2025-08-07T17:06:42Z

Starting merge as part of PR stack under #158983

pytorchmergebot · 2025-08-07T17:07:23Z

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

…on (#155360) (#158983) Pull Request resolved: #158983 Approved by: https://github.com/eellison ghstack dependencies: #158758

[inductor] allocate non-blocking copy destinations in pinned memory (#…

2193daf

…155121) [ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo module: inductor release notes: inductor (aoti) labels Jul 21, 2025

v0i0 added a commit that referenced this pull request Jul 21, 2025

[inductor] allocate non-blocking copy destinations in pinned memory (#…

8e1a001

…155121) ghstack-source-id: f865a6f Pull Request resolved: #158758

v0i0 requested a review from eellison July 21, 2025 18:45

v0i0 added a commit that referenced this pull request Jul 21, 2025

[inductor] allocate non-blocking copy destinations in pinned memory (#…

4821efd

…155121) ghstack-source-id: 0a98613 Pull Request resolved: #158758

eellison reviewed Jul 21, 2025

View reviewed changes

v0i0 added a commit that referenced this pull request Jul 22, 2025

[inductor] allocate non-blocking copy destinations in pinned memory (#…

c7fcc1f

…155121) ghstack-source-id: 61c061d Pull Request resolved: #158758

v0i0 added a commit that referenced this pull request Jul 22, 2025

[inductor] allocate non-blocking copy destinations in pinned memory (#…

f198537

…155121) ghstack-source-id: 7f0eeb6 Pull Request resolved: #158758

v0i0 added a commit that referenced this pull request Jul 22, 2025

[inductor] allocate non-blocking copy destinations in pinned memory (#…

53aad55

…155121) ghstack-source-id: b041df3 Pull Request resolved: #158758

guangyey reviewed Jul 23, 2025

View reviewed changes

test/inductor/test_torchinductor.py Outdated Show resolved Hide resolved

guangyey reviewed Jul 23, 2025

View reviewed changes

test/inductor/test_torchinductor.py Outdated Show resolved Hide resolved

v0i0 added a commit that referenced this pull request Jul 23, 2025

[inductor] allocate non-blocking copy destinations in pinned memory (#…

c9aada7

…155121) ghstack-source-id: 95308ae Pull Request resolved: #158758

eellison reviewed Jul 23, 2025

View reviewed changes

test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/csrc/kernel.cpp Outdated Show resolved Hide resolved

v0i0 mentioned this pull request Jul 23, 2025

[inductor] move all cpu scalars using pinned memory for graph partition (#155360) #158983

Closed

EikanWang approved these changes Jul 24, 2025

View reviewed changes

v0i0 added 2 commits August 1, 2025 17:03

pytorchmergebot added the merging label Aug 4, 2025

pytorchmergebot closed this in 6085bf7 Aug 4, 2025

pytorchmergebot added Merged and removed merging labels Aug 4, 2025

pytorchmergebot pushed a commit that referenced this pull request Aug 4, 2025

[inductor] move all cpu scalars using pinned memory for graph partiti…

444e238

…on (#155360) (#158983) Pull Request resolved: #158983 Approved by: https://github.com/eellison ghstack dependencies: #158758

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Aug 4, 2025

pytorchmergebot reopened this Aug 4, 2025

v0i0 added 2 commits August 6, 2025 11:05

Update

4c5bbfe

[ghstack-poisoned]

Update

ada0bdc

[ghstack-poisoned]

pytorchmergebot added the merging label Aug 7, 2025

pytorchmergebot closed this in e167c7d Aug 7, 2025

pytorchmergebot pushed a commit that referenced this pull request Aug 7, 2025

[inductor] move all cpu scalars using pinned memory for graph partiti…

57f738b

…on (#155360) (#158983) Pull Request resolved: #158983 Approved by: https://github.com/eellison ghstack dependencies: #158758

pytorchmergebot removed the merging label Aug 7, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[inductor] allocate non-blocking copy destinations in pinned memory (#155121) #158758

[inductor] allocate non-blocking copy destinations in pinned memory (#155121) #158758

v0i0 commented Jul 21, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jul 21, 2025 •

edited

Loading

Uh oh!

eellison left a comment

Uh oh!

Uh oh!

eellison Jul 21, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

v0i0 commented Aug 4, 2025

Uh oh!

pytorchmergebot commented Aug 4, 2025

Uh oh!

v0i0 commented Aug 4, 2025

Uh oh!

davidberard98 commented Aug 4, 2025

Uh oh!

pytorchmergebot commented Aug 4, 2025

Uh oh!

pytorchmergebot commented Aug 4, 2025

Uh oh!

github-actions bot commented Aug 4, 2025

Uh oh!

v0i0 commented Aug 7, 2025

Uh oh!

pytorchmergebot commented Aug 7, 2025

Uh oh!

pytorchmergebot commented Aug 7, 2025

Uh oh!

Uh oh!

[inductor] allocate non-blocking copy destinations in pinned memory (#155121) #158758

[inductor] allocate non-blocking copy destinations in pinned memory (#155121) #158758

Conversation

v0i0 commented Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ You can merge normally! (2 Unrelated Failures)

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

eellison Jul 21, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

v0i0 commented Aug 4, 2025

Uh oh!

pytorchmergebot commented Aug 4, 2025

Merge started

Uh oh!

v0i0 commented Aug 4, 2025

Uh oh!

davidberard98 commented Aug 4, 2025

Uh oh!

pytorchmergebot commented Aug 4, 2025

Uh oh!

pytorchmergebot commented Aug 4, 2025

Uh oh!

github-actions bot commented Aug 4, 2025

Attention! PyTorch one of the C-stable API file was changed

Uh oh!

v0i0 commented Aug 7, 2025

Uh oh!

pytorchmergebot commented Aug 7, 2025

Uh oh!

pytorchmergebot commented Aug 7, 2025

Merge started

Uh oh!

Uh oh!

v0i0 commented Jul 21, 2025 •

edited

Loading

pytorch-bot bot commented Jul 21, 2025 •

edited

Loading