mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
06b173780d
**Background**: "TorchDynamo Cache Lookup" events appear in traces to indicate a dynamo cache lookup; it's useful to check when cache lookups are taking a long time. To add a profiler event, one can use the `torch.profiler.record_function` context manager, or the C++ equivalent. Previously, the python version was used; first, when the profiler was enabled, callbacks for record_function_enter and record_function_exit were registered; then those would be called before and after every cache lookup.
**This PR**: Instead of calling the python bindings for `torch.profiler.record_function`, directly call the C++ implementation. This simplifies a lot of the code for binding C/C++. It also improves performance; previously there was a lot of overhead in the "TorchDynamo Cache Lookup" event, making the event artificially take a long time. After this change the events now appear shorter, because there's less overhead in starting/stopping the event: in other words, the profiler no longer distorts the results as much.
**Performance results**:
I ran using the script below on a cpu-only 1.6GHz machine. I report the median time (from 100 measurements) of a "TorchDynamo Cache Lookup" event before and after this PR. I think it is reasonable to consider the difference to be due to a reduction in overhead.
<details>
<summary>Benchmarking script</summary>
```python
def fn(x, y):
return (x * y).relu()
a, b = [torch.rand((4, 4), requires_grad=True) for _ in range(2)]
opt_fn = torch.compile(fn)
opt_fn(a, b)
opt_fn(a, b)
with torch.profiler.profile() as prof:
opt_fn(a, b)
```
</details>
Median before PR: 198-228 us (median of 100, measured 5 times)
Median after PR: 27us
Pull Request resolved: https://github.com/pytorch/pytorch/pull/108436
Approved by: https://github.com/anijain2305, https://github.com/jansel
23 lines
551 B
C++
23 lines
551 B
C++
#include <torch/csrc/dynamo/cpp_shim.h>
|
|
|
|
#include <ATen/record_function.h>
|
|
|
|
struct _PytorchRecordFunctionState {
|
|
at::RecordFunction guard;
|
|
|
|
_PytorchRecordFunctionState() : guard(at::RecordScope::FUNCTION) {}
|
|
};
|
|
|
|
_PytorchRecordFunctionState* _pytorch_record_function_enter(const char* name) {
|
|
_PytorchRecordFunctionState* state = new _PytorchRecordFunctionState();
|
|
state->guard.before(name);
|
|
return state;
|
|
}
|
|
|
|
void _pytorch_record_function_exit(_PytorchRecordFunctionState* state) {
|
|
if (state == nullptr) {
|
|
return;
|
|
}
|
|
delete state;
|
|
}
|