[profiler] Nvtx support (#61634)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61634
Legacy profiler supported Nvtx and that was used by emit_nvtx, this PR
adds support for Nvtx in the new compiler, to prepare for the eventual
deprecation of the legacy profiler
Test Plan:
Verified that the profiles produced with nvprof are the same
```
import torch
import torchvision.models as models
from torch.autograd.profiler import emit_nvtx, load_nvprof
model = models.resnet18().cuda()
inputs = torch.randn(5, 3, 224, 224).cuda()
with emit_nvtx(record_shapes=True):
model(inputs)
```
/usr/local/cuda/bin/nvprof -o test_trace2.prof -f -- python test_emit_nvtx.py
```
evt = load_nvprof("/home/iliacher/local/pytorch/test_trace.prof")
```
Imported from OSS
Reviewed By: kimishpatel, gdankel
Differential Revision: D29691316
fbshipit-source-id: 1e186cc072368f3e3987a2da0bfd90ed328817c5
