mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01: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
49 lines
897 B
ReStructuredText
49 lines
897 B
ReStructuredText
torch.mps
|
|
===================================
|
|
.. automodule:: torch.mps
|
|
.. currentmodule:: torch.mps
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
device_count
|
|
synchronize
|
|
get_rng_state
|
|
set_rng_state
|
|
manual_seed
|
|
seed
|
|
empty_cache
|
|
set_per_process_memory_fraction
|
|
current_allocated_memory
|
|
driver_allocated_memory
|
|
recommended_max_memory
|
|
|
|
MPS Profiler
|
|
------------
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
profiler.start
|
|
profiler.stop
|
|
profiler.profile
|
|
|
|
profiler.is_capturing_metal
|
|
profiler.is_metal_capture_enabled
|
|
profiler.metal_capture
|
|
|
|
MPS Event
|
|
------------
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
event.Event
|
|
|
|
|
|
.. This module needs to be documented. Adding here in the meantime
|
|
.. for tracking purposes
|
|
.. py:module:: torch.mps.event
|
|
.. py:module:: torch.mps.profiler
|