[Inductor] Add DeviceAssert op to enable device-side assertion in torch.compile #160677

karthickai · 2025-08-14T20:46:46Z

This PR introduces a device_assert op to trigger device-side assertions within torch.compile. This implementation is based on the suggestion in this comment.

Changes Included

Implemented device_assert op and overrides has_side_effect to return True to avoid removal by dead code elimination.
Commented out the assert_async_msg_decomp and functional_assert_async_msg_decomp decompositions to disable the default assert decomposition inside Inductor.
Added lowering for torch.ops.aten._assert_async.msg to convert assert calls into the ops_handler.
Implemented the codegen method for the device_assert op. This supports generating C++ and Triton code.
Added test cases to verify both "should throw" and "should not throw" scenarios.

Fixes #147282

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

pytorch-bot · 2025-08-14T20:46:50Z

🔗 Helpful Links

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

📄 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.

✅ No Failures

As of commit cc442eb with merge base c55bdb2 ():
💚 Looks good so far! There are no failures yet. 💚

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

eellison · 2025-08-15T14:45:53Z

torch/_inductor/ir.py

        )


+class DeviceAssert(ExternKernel):


We should fuse the device assertions. Which means it should be an assertion, and we avoid DCEing buffers with the assertion.

Thank you for the feedback. Removed the ExternKernel approach and changed it to use the ops handler method and modified has_side_effect method to return True to avoid DCE. (b42cb7c)

test/inductor/test_device_assert.py

torch/_inductor/decomposition.py

torch/_inductor/dtype_propagation.py

torch/_inductor/lowering.py

test/inductor/test_device_assert.py

torch/_inductor/lowering.py

torch/_inductor/scheduler.py

mlazos

Looks good!

eellison

You need to define this on FusedSchedulerNode as well, to check all of its snodes. (although its unlikely we will fuse multiple nodes then dce them)

eellison · 2025-08-21T21:19:55Z

torch/_inductor/scheduler.py

                    )
        return buffers_store_as_atomic_add

+    def has_side_effects(self) -> bool:


you might as well cache_on_self this, because it doesnt change after instantiation

Now I added has_side_effects method in FusedSchedulerNode and used cache_on_self.

eellison · 2025-08-21T22:13:28Z

torch/_inductor/scheduler.py


+    @cache_on_self
+    def has_side_effects(self) -> bool:
+        if self.snodes is not None:


this should just be any(node.has_side_effects() for node in self.snodes)

thank you, fixed it.

karthickai · 2025-08-22T01:12:14Z

@pytorchbot retest this please

pytorch-bot · 2025-08-22T01:12:17Z

❌ 🤖 pytorchbot command failed:

@pytorchbot: error: argument command: invalid choice: 'retest' (choose from 'merge', 'revert', 'rebase', 'label', 'drci', 'cherry-pick')

usage: @pytorchbot [-h] {merge,revert,rebase,label,drci,cherry-pick} ...

Try @pytorchbot --help for more info.

karthickai · 2025-08-22T05:20:29Z

@pytorchbot merge

pytorchmergebot · 2025-08-22T05:22:21Z

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team

Raised by workflow job

karthickai · 2025-08-22T05:24:20Z

@pytorchbot merge

mlazos · 2025-08-28T01:34:23Z

@karthickai can you resolve https://github.com/pytorch/pytorch/pull/160677/files#r2299347303 before merging again?

torch/_inductor/lowering.py

…edundant test

facebook-github-bot · 2025-08-28T03:37:21Z

@karthickai has imported this pull request. If you are a Meta employee, you can view this in D81197669.

facebook-github-bot · 2025-08-28T17:35:33Z

@pytorchbot merge

(Initiating merge automatically since Phabricator Diff has merged)

pytorchmergebot · 2025-08-28T17:37:27Z

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 · 2025-08-28T17:37:44Z

Merge failed

Reason: 1 jobs have failed, first few of them are: Meta Internal-Only Changes Check

Details for Dev Infra team

Raised by workflow job

karthickai · 2025-08-28T18:55:42Z

@pytorchbot merge -f "landed in fbcode"

pytorchmergebot · 2025-08-28T18:57:21Z

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

eellison · 2025-09-16T23:02:21Z

test/inductor/test_device_assert.py

+            except Exception:
+                return False
+
+        bisect_result = CompilerBisector.do_bisect(test_fn)


Sorry - getting to this late.

I would avoid using CompilerBisector in this manner. If the intent is to test that different backends all threw, you can parameterize the test by backend so that we get a more explicit failure on which backend is not throwing. This would also allow you to test that "should throw" is in the error msg.

I agree with your point, I will update the tests.

…ch.compile (pytorch#160677) This PR introduces a device_assert op to trigger device-side assertions within torch.compile. This implementation is based on the suggestion in [this comment](pytorch#147282 (comment)). Changes Included - Implemented device_assert op and overrides has_side_effect to return True to avoid removal by dead code elimination. - Commented out the assert_async_msg_decomp and functional_assert_async_msg_decomp decompositions to disable the default assert decomposition inside Inductor. - Added lowering for torch.ops.aten._assert_async.msg to convert assert calls into the ops_handler. - Implemented the codegen method for the device_assert op. This supports generating C++ and Triton code. - Added test cases to verify both "should throw" and "should not throw" scenarios. Fixes pytorch#147282 Pull Request resolved: pytorch#160677 Approved by: https://github.com/mlazos

…n in torch.compile (pytorch#160677)" This reverts commit cddcaa1. Reverted pytorch#160677 on behalf of https://github.com/karthickai due to This is breaking tests on Rocm ([comment](pytorch#160677 (comment)))

…ch.compile (pytorch#160677) This PR introduces a device_assert op to trigger device-side assertions within torch.compile. This implementation is based on the suggestion in [this comment](pytorch#147282 (comment)). Changes Included - Implemented device_assert op and overrides has_side_effect to return True to avoid removal by dead code elimination. - Commented out the assert_async_msg_decomp and functional_assert_async_msg_decomp decompositions to disable the default assert decomposition inside Inductor. - Added lowering for torch.ops.aten._assert_async.msg to convert assert calls into the ops_handler. - Implemented the codegen method for the device_assert op. This supports generating C++ and Triton code. - Added test cases to verify both "should throw" and "should not throw" scenarios. Fixes pytorch#147282 Pull Request resolved: pytorch#160677 Approved by: https://github.com/mlazos

…n in torch.compile (pytorch#160677)" This reverts commit 378edb0. Reverted pytorch#160677 on behalf of https://github.com/atalman due to new test is failing internally ([comment](pytorch#160677 (comment)))

…ch.compile (pytorch#160677) This PR introduces a device_assert op to trigger device-side assertions within torch.compile. This implementation is based on the suggestion in [this comment](pytorch#147282 (comment)). Changes Included - Implemented device_assert op and overrides has_side_effect to return True to avoid removal by dead code elimination. - Commented out the assert_async_msg_decomp and functional_assert_async_msg_decomp decompositions to disable the default assert decomposition inside Inductor. - Added lowering for torch.ops.aten._assert_async.msg to convert assert calls into the ops_handler. - Implemented the codegen method for the device_assert op. This supports generating C++ and Triton code. - Added test cases to verify both "should throw" and "should not throw" scenarios. Fixes pytorch#147282 Pull Request resolved: pytorch#160677 Approved by: https://github.com/mlazos, https://github.com/atalman

Updated the DeviceAssert operation to match the behavior of Store, it will fixes the issue mentioned in [this PR](#163023) and updated testcases as Elias [suggested](#160677 (comment)). cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben mlazos choijon5 [ghstack-poisoned]

Updated the DeviceAssert operation to match the behavior of Store, it will fixes the issue mentioned in [this PR](#163023) and updated testcases as Elias [suggested](#160677 (comment)). Pull Request resolved: #163696 Approved by: https://github.com/mlazos

pytorch-bot bot added ciflow/inductor module: inductor labels Aug 14, 2025

karthickai added the module: decompositions Topics related to decomposition (excluding PrimTorch) label Aug 14, 2025

eellison reviewed Aug 15, 2025

View reviewed changes

karthickai marked this pull request as ready for review August 18, 2025 21:06

karthickai requested a review from mlazos August 18, 2025 21:31