Add nvtx.range() context manager (#42925)
Summary:
Small quality-of-life improvement to NVTX Python bindings, that we're using internally and that would be useful to other folks using NVTX annotations via PyTorch. (And my first potential PyTorch contribution.)
Instead of needing to be careful with try/finally to make sure all your range_push'es are range_pop'ed:
```
nvtx.range_push("Some event")
try:
# Code here...
finally:
nvtx.range_pop()
```
you can simply do:
```
with nvtx.range("Some event"):
# Code here...
```
or even use it as a decorator:
```
class MyModel(nn.Module):
# Other methods here...
nvtx.range("MyModel.forward()")
def forward(self, *input):
# Forward pass code here...
```
A couple small open questions:
1. I also added the ability to call `msg.format()` inside `range()`, with the intention that, if there is nothing listening to NVTX events, we should skip the string formatting, to lower the overhead in that case. If you like that idea, I could add the actual "skip string formatting if nobody is listening to events" parts. We can also just leave it as is. Or I can remove that if you folks don't like it. (In the first two cases, should we add that to `range_push()` and `mark()` too?) Just let me know which one it is, and I'll update the pull request.
2. I don't think there are many places for bugs to hide in that function, but I can certainly add a quick test, if you folks want.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/42925
Reviewed By: gchanan
Differential Revision: D24476977
Pulled By: ezyang
fbshipit-source-id: 874882818d958e167e624052e42d52fae3c4abf1
