mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
92ddb3d3d3
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:
```python
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
<img width="734" alt="image" src="https://github.com/user-attachments/assets/718ff64e-103b-4b11-b66c-c89cfc770b5d" />
Pull Request resolved: https://github.com/pytorch/pytorch/pull/144561
Approved by: https://github.com/manuelcandales
ghstack dependencies: #144559, #144560
|
||
|---|---|---|
| .. | ||
| actions | ||
| ci_commit_pins | ||
| ISSUE_TEMPLATE | ||
| requirements | ||
| scripts | ||
| templates | ||
| workflows | ||
| actionlint.yaml | ||
| auto_request_review.yml | ||
| label_to_label.yml | ||
| labeler.yml | ||
| merge_rules.yaml | ||
| nitpicks.yml | ||
| PULL_REQUEST_TEMPLATE.md | ||
| pytorch-circleci-labels.yml | ||
| pytorch-probot.yml | ||
| regenerate.sh | ||
| requirements-gha-cache.txt | ||