Repository: pytorch
Issue: Make torch.cuda.graphs.is_current_stream_capturing() available in TorchScript
Labels: oncall: jit,module: cuda graphs

### 🚀 The feature, motivation and pitch

I would like to be able to detect if CUDA graph capturing is underway from inside a TorchScript model. Currently this example fails to script:

```python

import torch
import torch.nn as nn
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, x):
        is_capturing = (
            x.is_cuda
            and torch.cuda.graphs.is_current_stream_capturing()
        )
        if not is_capturing:
            torch.cuda.synchronize()
        y = x*2
        return y

x = torch.ones(1).to("cuda").requires_grad_(True)
model = MyModel()
model = torch.jit.script(model)
stream = torch.cuda.Stream()
with torch.cuda.stream(stream):
    #Warmup
    y = model(x)
    model = torch.cuda.make_graphed_callables(model, (x,), allow_unused_input=True)
    y = model(x)

```
With the following error:

```shell
Traceback (most recent call last):
  File "/home/raulp/ex.py", line 20, in <module>
    model = torch.jit.script(model)
            ^^^^^^^^^^^^^^^^^^^^^^^
  File "/shared/raul/mambaforge/envs/torchmd-net-latest/lib/python3.11/site-packages/torch/jit/_script.py", line 1324, in script
    return torch.jit._recursive.create_script_module(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/shared/raul/mambaforge/envs/torchmd-net-latest/lib/python3.11/site-packages/torch/jit/_recursive.py", line 559, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/shared/raul/mambaforge/envs/torchmd-net-latest/lib/python3.11/site-packages/torch/jit/_recursive.py", line 636, in create_script_module_impl
    create_methods_and_properties_from_stubs(
  File "/shared/raul/mambaforge/envs/torchmd-net-latest/lib/python3.11/site-packages/torch/jit/_recursive.py", line 469, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(
  File "/shared/raul/mambaforge/envs/torchmd-net-latest/lib/python3.11/site-packages/torch/jit/_recursive.py", line 1010, in try_compile_fn
    return torch.jit.script(fn, _rcb=rcb)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/shared/raul/mambaforge/envs/torchmd-net-latest/lib/python3.11/site-packages/torch/jit/_script.py", line 1381, in script
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: 
Python builtin <built-in function _cuda_isCurrentStreamCapturing> is currently not supported in Torchscript:
  File "/shared/raul/mambaforge/envs/torchmd-net-latest/lib/python3.11/site-packages/torch/cuda/graphs.py", line 32
    If a CUDA context does not exist on the current device, returns False without initializing the context.
    """
    return _cuda_isCurrentStreamCapturing()
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
'is_current_stream_capturing' is being compiled since it was called from 'MyModel.forward'
  File "/home/raulp/ex.py", line 10
    def forward(self, x):
        is_capturing = (
            x.is_cuda
            ~~~~~~~~~ <--- HERE
            and torch.cuda.graphs.is_current_stream_capturing()
        )
```


### Alternatives

Using this inline extension does the job:
```python
import torch
from torch.utils.cpp_extension import load_inline
import torch.nn as nn

cpp_source = '''
#include <torch/script.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime_api.h>

bool is_stream_capturing() {
  at::cuda::CUDAStream current_stream = at::cuda::getCurrentCUDAStream();
  cudaStream_t cuda_stream = current_stream.stream();
  cudaStreamCaptureStatus capture_status;
  cudaError_t err = cudaStreamGetCaptureInfo(cuda_stream, &capture_status, nullptr);

  if (err != cudaSuccess) {
    throw std::runtime_error(cudaGetErrorString(err));
  }

  return capture_status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive;
}

static auto registry =
  torch::RegisterOperators()
    .op("torch_extension::is_stream_capturing", &is_stream_capturing);
'''

# Create an inline extension
torch_extension = load_inline(
    "is_stream_capturing",
    cpp_sources=cpp_source,
    functions=["is_stream_capturing"],
    with_cuda=True,
    verbose=True,
)


@torch.jit.script
def check_stream_capturing():
    return torch.ops.torch_extension.is_stream_capturing()


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, x):
        is_capturing = (
            x.is_cuda
            and check_stream_capturing()
        )
        if not is_capturing:
            torch.cuda.synchronize()
        y = x*2
        return y


x = torch.ones(1).to("cuda").requires_grad_(True)
model = MyModel()
model = torch.jit.script(model)
stream = torch.cuda.Stream()
with torch.cuda.stream(stream):
    #Warmup
    y = model(x)
    model = torch.cuda.make_graphed_callables(model, (x,), allow_unused_input=True)
    y = model(x)
```
but has some downsides of its own and poses the question of how to ship the extension itself.

### Additional context

_No response_

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel @mcarilli @ezyang