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: #42925

Reviewed By: gchanan

Differential Revision: D24476977

Pulled By: ezyang

fbshipit-source-id: 874882818d958e167e624052e42d52fae3c4abf1

torch/cuda/nvtx.py

@@ -1,3 +1,5 @@
+from contextlib import contextmanager
+
 try:
     from torch._C import _nvtx
 except ImportError:
@@ -12,7 +14,7 @@ def _fail(*args, **kwargs):
 
     _nvtx = _NVTXStub()  # type: ignore[assignment]
 
-__all__ = ['range_push', 'range_pop', 'mark']
+__all__ = ['range_push', 'range_pop', 'mark', 'range']
 
 
 def range_push(msg):
@@ -42,3 +44,18 @@ def mark(msg):
         msg (string): ASCII message to associate with the event.
     """
     return _nvtx.markA(msg)
+
+
+@contextmanager
+def range(msg, *args, **kwargs):
+    """
+    Context manager / decorator that pushes an NVTX 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().
+
+    Arguments:
+        msg (string): message to associate with the range
+    """
+    range_push(msg.format(*args, **kwargs))
+    yield
+    range_pop()