Add runtime profiler info for AOTDispatcher prologue #155785

azahed98 · 2025-06-12T06:53:02Z

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @jerryzh168 @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @xmfan

pytorch-bot · 2025-06-12T06:53:06Z

🔗 Helpful Links

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

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

⏳ 1 Pending, 1 Unrelated Failure

As of commit c1de92d with merge base b4442f4 ():

UNSTABLE - The following jobs are marked as unstable, possibly due to flakiness on trunk:

inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_torchbench_cpu_smoketest_perf, 1, 1, linux.24xl.spr-metal) (gh) (#156521)
Process completed with exit code 1.
⏳ inductor-rocm / rocm-py3.10-inductor / test (inductor, 1, 2, linux.rocm.gpu.2, unstable) (gh) (#155917)

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

bdhirsh · 2025-06-12T14:11:17Z

torch/_functorch/_aot_autograd/runtime_wrappers.py

            try:
                if grad_enabled:
                    torch._C._set_grad_enabled(False)
+                record_runtime_wrapper_prologue_exit(cm)


Hmm at first I was going to ask you to use the context manager directly, so we don't need to worry about carefully making sure that the enter/exit line up in all the various branches.

I guess this is actually a pain to do though, since you need different global state to be toggled in the different branches, and you want to capture that inside of the recorded region

bdhirsh · 2025-06-12T14:15:24Z

torch/_functorch/_aot_autograd/runtime_wrappers.py

+    def record_runtime_wrapper_prologue_enter() -> AbstractContextManager[None]:
+        cm: AbstractContextManager[None] = (
+            torch._C._profiler._RecordFunctionFast("AOTDispatcher Runtime Wrapper Prologue")
+            if torch.autograd.profiler._is_profiler_enabled


@anijain2305 I saw that you have an additional config.record_pre_graph_bytecode_in_traces that you branch on in your case. Did you add that because you're worried about the extra overhead even in this case where profiling is off? https://github.com/pytorch/pytorch/pull/154769/files#diff-0bb1756725c4426408938314b0c9d3988ae5bf49994892d7038ad7746e209e9fR617

torch/_functorch/_aot_autograd/runtime_wrappers.py

bdhirsh

nice work!

test/dynamo/test_repros.py

azahed98 · 2025-06-18T20:09:09Z

@pytorchbot merge -r

pytorchmergebot · 2025-06-18T20:10:52Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2025-06-18T20:10:55Z

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/viable/strict pull/155785/head returned non-zero exit code 1

Rebasing (1/53)
Rebasing (2/53)
Rebasing (3/53)
Auto-merging test/dynamo/test_repros.py
CONFLICT (content): Merge conflict in test/dynamo/test_repros.py
error: could not apply 57d296a1085... Add test_aot_autograd_runtime_wrapper_prologue_profiled to test_repros
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".
hint: Disable this message with "git config set advice.mergeConflict false"
Could not apply 57d296a1085... Add test_aot_autograd_runtime_wrapper_prologue_profiled to test_repros

Raised by https://github.com/pytorch/pytorch/actions/runs/15742581401

azahed98 · 2025-06-19T23:49:28Z

@pytorchbot merge -r

pytorchmergebot · 2025-06-19T23:51:04Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2025-06-19T23:51:05Z

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/viable/strict pull/155785/head returned non-zero exit code 1

Rebasing (1/12)
Rebasing (2/12)
Rebasing (3/12)
Auto-merging test/dynamo/test_repros.py
CONFLICT (content): Merge conflict in test/dynamo/test_repros.py
error: could not apply 57d296a1085... Add test_aot_autograd_runtime_wrapper_prologue_profiled to test_repros
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".
hint: Disable this message with "git config set advice.mergeConflict false"
Could not apply 57d296a1085... Add test_aot_autograd_runtime_wrapper_prologue_profiled to test_repros

Raised by https://github.com/pytorch/pytorch/actions/runs/15768102447

pytorchmergebot · 2025-06-20T17:53:13Z

Successfully rebased feat/aotdispatcher_prologue_profiler onto refs/remotes/origin/main, please pull locally before adding more changes (for example, via git checkout feat/aotdispatcher_prologue_profiler && git pull --rebase)

azahed98 · 2025-06-20T17:54:15Z

@pytorchbot merge

pytorchmergebot · 2025-06-20T17:56:20Z

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-06-20T17:56:35Z

Merge failed

Reason: 1 jobs have failed, first few of them are: inductor-rocm / rocm-py3.10-inductor / build

Details for Dev Infra team

Raised by workflow job

azahed98 · 2025-06-20T18:53:10Z

@pytorchbot merge

pytorchmergebot · 2025-06-20T18:54:56Z

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-06-20T18:55:08Z

Merge failed

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

pull / linux-jammy-cpu-py3.10-gcc11-bazel-test / build-and-test (default, 1, 1, linux.4xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

azahed98 · 2025-06-20T20:34:32Z

@pytorchbot merge -r

pytorchmergebot · 2025-06-20T20:36:11Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2025-06-20T20:36:14Z

Successfully rebased feat/aotdispatcher_prologue_profiler onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout feat/aotdispatcher_prologue_profiler && git pull --rebase)

pytorchmergebot · 2025-06-20T20:37:17Z

This PR updates submodules third_party/eigen

If those updates are intentional, please add "submodule" keyword to PR title/description.

azahed98 · 2025-06-20T22:32:37Z

@pytorchbot merge

pytorchmergebot · 2025-06-20T22:34:21Z

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

pytorch-bot bot added the ciflow/inductor label Jun 12, 2025

bdhirsh reviewed Jun 12, 2025

View reviewed changes

pytorch-bot bot added module: dynamo module: inductor labels Jun 12, 2025

bdhirsh reviewed Jun 13, 2025

View reviewed changes

torch/_functorch/_aot_autograd/runtime_wrappers.py Outdated Show resolved Hide resolved

bdhirsh approved these changes Jun 13, 2025

View reviewed changes

bdhirsh reviewed Jun 13, 2025

View reviewed changes

test/dynamo/test_repros.py Outdated Show resolved Hide resolved

azahed98 marked this pull request as ready for review June 13, 2025 21:14

azahed98 added the release notes: aotdispatcher label Jun 13, 2025

azahed98 force-pushed the feat/aotdispatcher_prologue_profiler branch from af4604f to c843fc9 Compare June 17, 2025 20:28

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 18, 2025

azahed98 requested review from eqy, syed-ahmed, mikaylagawarecki, bobrenjc93, laithsakka, avikchaudhuri, tugsbayasgalan, zhxchen17, ydwu4, angelayi and EikanWang as code owners June 20, 2025 00:08

pytorchmergebot force-pushed the feat/aotdispatcher_prologue_profiler branch from 13581a7 to f92f250 Compare June 20, 2025 17:53

pytorchmergebot added the merging label Jun 20, 2025

pytorchmergebot removed the merging label Jun 20, 2025

azahed98 force-pushed the feat/aotdispatcher_prologue_profiler branch from f92f250 to 13581a7 Compare June 20, 2025 18:39

pytorchmergebot added the merging label Jun 20, 2025

pytorchmergebot removed the merging label Jun 20, 2025

pytorchmergebot force-pushed the feat/aotdispatcher_prologue_profiler branch from da90d5a to db38892 Compare June 20, 2025 20:36

pytorch-bot bot added the release notes: inductor (aoti) label Jun 20, 2025

azahed98 force-pushed the feat/aotdispatcher_prologue_profiler branch from db38892 to 13581a7 Compare June 20, 2025 20:43

Add prologue record and test

140d34e

azahed98 force-pushed the feat/aotdispatcher_prologue_profiler branch from 13581a7 to 140d34e Compare June 20, 2025 20:56

Lint

c1de92d

pytorchmergebot added the merging label Jun 20, 2025

pytorchmergebot added the Merged label Jun 21, 2025

pytorchmergebot closed this in c09b054 Jun 21, 2025

pytorchmergebot removed the merging label Jun 21, 2025

Add runtime profiler info for AOTDispatcher prologue #155785

Add runtime profiler info for AOTDispatcher prologue #155785

Uh oh!

Conversation

azahed98 commented Jun 12, 2025 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jun 12, 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/155785

⏳ 1 Pending, 1 Unrelated Failure

Uh oh!

bdhirsh Jun 12, 2025

Choose a reason for hiding this comment

Uh oh!

bdhirsh Jun 12, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bdhirsh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

azahed98 commented Jun 18, 2025

Uh oh!

pytorchmergebot commented Jun 18, 2025

Uh oh!

pytorchmergebot commented Jun 18, 2025

Uh oh!

azahed98 commented Jun 19, 2025

Uh oh!

pytorchmergebot commented Jun 19, 2025

Uh oh!

pytorchmergebot commented Jun 19, 2025

Uh oh!

pytorchmergebot commented Jun 20, 2025

Uh oh!

azahed98 commented Jun 20, 2025

Uh oh!

pytorchmergebot commented Jun 20, 2025

Merge started

Uh oh!

pytorchmergebot commented Jun 20, 2025

Merge failed

Uh oh!

azahed98 commented Jun 20, 2025

Uh oh!

pytorchmergebot commented Jun 20, 2025

Merge started

Uh oh!

pytorchmergebot commented Jun 20, 2025

Merge failed

Uh oh!

azahed98 commented Jun 20, 2025

Uh oh!

pytorchmergebot commented Jun 20, 2025

Uh oh!

pytorchmergebot commented Jun 20, 2025

Uh oh!

pytorchmergebot commented Jun 20, 2025

Uh oh!

azahed98 commented Jun 20, 2025

Uh oh!

pytorchmergebot commented Jun 20, 2025

Merge started

Uh oh!

Uh oh!

azahed98 commented Jun 12, 2025 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jun 12, 2025 •

edited

Loading