8000 [MPS] Expose `MPSProfiler::start/stopCapture` to Python by malfet · Pull Request #144561 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

[MPS] Expose MPSProfiler::start/stopCapture to Python #144561

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 5 commits into from
Closed

[MPS] Expose MPSProfiler::start/stopCapture to Python #144561

wants to merge 5 commits into from

Conversation

malfet
Copy link
Contributor
@malfet malfet commented Jan 10, 2025 •
edited
Loading

Stack from ghstack (oldest at bottom):

I.e. when MTL_CAPTURE_ENABLED environment variable is set to 1, one should be able to invoke wrap the code with torch.mps.profiler.capture_metal to generate gputrace for shaders invoked inside the context manager.

For example, code below:

import torch
import os

def foo(x):
   return x[:,::2].sin() + x[:, 1::2].cos()


if __name__ == "__main__":
    os.environ["MTL_CAPTURE_ENABLED"] = "1"
    x = torch.rand(32, 1024, device="mps")

    with torch.mps.profiler.metal_capture("compiled_shader"):
        torch.compile(foo)(x)

should capture the execution of a torch.compile generated shader
image

@malfet
Update
7a1d5e6 
[ghstack-poisoned]
This was referenced Jan 10, 2025
Closed
Closed
@pytorch-bot PyTorch Bot
Copy link
pytorch-bot bot commented Jan 10, 2025 •
edited
Loading

🔗 Helpful Links

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

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

⏳ 13 Pending, 1 Unrelated Failure

As of commit 5bf42fd with merge base a222029 (image):

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

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

@malfet malfet marked this pull request as draft January 10, 2025 16:12
@malfet
Update
24d5d70 
[ghstack-poisoned]
@pytorch-bot pytorch-bot bot added ciflow/mps Run MPS tests (subset of trunk) release notes: releng release notes category labels Jan 10, 2025
@malfet
Update
9bb3086 
[ghstack-poisoned]
@malfet malfet added topic: improvements topic category release notes: mps Release notes category and removed release notes: releng release notes category labels Jan 10, 2025
@malfet malfet requested a review from manuelcandales January 10, 2025 20:33
@malfet malfet marked this pull request as ready for review January 10, 2025 20:33
@malfet malfet requested review from kulinseth and a team as code owners January 10, 2025 20:33
manuelcandales
manuelcandales reviewed Jan 10, 2025
View reviewed changes
torch/mps/profiler.py
torch._C._mps_startCapture(fname)
yield
finally:
torch.mps.synchronize()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this call to synchronize distort times in any way?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the contrary, it ensures that everything that is submitted asynchronously will be captured

@manuelcandales
Copy link
Contributor

Could you add documentation for this new feature?
Actually, is the existing MPS profiler documented? I found this https://pytorch.org/docs/stable/mps.html#mps-profiler, but looks like automatically generated docs.

malfet added 2 commits January 10, 2025 14:43
@malfet
Update
e30ad98 
[ghstack-poisoned]
@malfet
Update
5bf42fd 
[ghstack-poisoned]
@malfet malfet changed the title [MPS] Expose MetalCapture primitive to Python [MPS] Expose MPSProfiler::start/stopCapture to Python Jan 11, 2025
@malfet
Copy link
Contributor Author
malfet commented Jan 11, 2025

Could you add documentation for this new feature? Actually, is the existing MPS profiler documented?

Generates docs are in: https://docs-preview.pytorch.org/pytorch/pytorch/144561/mps.html#mps-profiler Will submit a followup to improve the docs later on

@malfet
Copy link
Contributor Author
malfet commented Jan 11, 2025

@pytorchbot merge -f "Lint + MPS is green"

@pytorchmergebot
Copy link
Collaborator

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

@github-actions github-actions bot deleted the gh/malfet/117/head branch February 12, 2025 02:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@manuelcandales manuelcandales manuelcandales approved these changes

@kulinseth kulinseth Awaiting requested review from kulinseth

Assignees
No one assigned
Labels
ciflow/mps Run MPS tests (subset of trunk) Merged release notes: mps Release notes category topic: improvements topic category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@malfet @manuelcandales @pytorchmergebot
0