mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
3c7044728b
More detailed description of benefits can be found at #41001. This is Intel's counterpart of NVidia’s NVTX (https://pytorch.org/docs/stable/autograd.html#torch.autograd.profiler.emit_nvtx). ITT is a functionality for labeling trace data during application execution across different Intel tools. For integrating Intel(R) VTune Profiler into Kineto, ITT needs to be integrated into PyTorch first. It works with both standalone VTune Profiler [(https://www.intel.com/content/www/us/en/developer/tools/oneapi/vtune-profiler.html](https://www.intel.com/content/www/us/en/developer/tools/oneapi/vtune-profiler.html)) and Kineto-integrated VTune functionality in the future. It works for both Intel CPU and Intel XPU devices. Pitch Add VTune Profiler's ITT API function calls to annotate PyTorch ops, as well as developer customized code scopes on CPU, like NVTX for NVidia GPU. This PR rebases the code changes at https://github.com/pytorch/pytorch/pull/61335 to the latest master branch. Usage example: ``` with torch.autograd.profiler.emit_itt(): for i in range(10): torch.itt.range_push('step_{}'.format(i)) model(input) torch.itt.range_pop() ``` cc @ilia-cher @robieta @chaekit @gdankel @bitfort @ngimel @orionr @nbcsm @guotuofeng @guyang3532 @gaoteng-git Pull Request resolved: https://github.com/pytorch/pytorch/pull/63289 Approved by: https://github.com/malfet
57 lines
1.2 KiB
Python
57 lines
1.2 KiB
Python
from contextlib import contextmanager
|
|
|
|
try:
|
|
from torch._C import _itt
|
|
except ImportError:
|
|
class _ITTStub(object):
|
|
@staticmethod
|
|
def _fail(*args, **kwargs):
|
|
raise RuntimeError("ITT functions not installed. Are you sure you have a ITT build?")
|
|
|
|
rangePush = _fail
|
|
rangePop = _fail
|
|
mark = _fail
|
|
|
|
_itt = _ITTStub() # type: ignore[assignment]
|
|
|
|
|
|
__all__ = ['range_push', 'range_pop', 'mark', 'range']
|
|
|
|
|
|
def range_push(msg):
|
|
"""
|
|
Arguments:
|
|
msg (string): ASCII message to associate with range
|
|
"""
|
|
return _itt.rangePush(msg)
|
|
|
|
|
|
def range_pop():
|
|
"""
|
|
"""
|
|
return _itt.rangePop()
|
|
|
|
|
|
def mark(msg):
|
|
"""
|
|
Describe an instantaneous event that occurred at some point.
|
|
Arguments:
|
|
msg (string): ASCII message to associate with the event.
|
|
"""
|
|
return _itt.mark(msg)
|
|
|
|
|
|
@contextmanager
|
|
def range(msg, *args, **kwargs):
|
|
"""
|
|
Context manager / decorator that pushes an ITT range at the beginning
|
|
of its scope, and pops it at the end. If extra arguments are given,
|
|
they are passed as arguments to msg.format().
|
|
|
|
Args:
|
|
msg (string): message to associate with the range
|
|
"""
|
|
range_push(msg.format(*args, **kwargs))
|
|
yield
|
|
range_pop()
|