Internal errors with cuda graph (CUBLAS_STATUS_NOT_INITIALIZED and jit failure)
Labels: triaged, module: cuda graphs

### 🐛 Describe the bug

# Cublas bug

Cuda graphs fail if cublas is not initialized. The code
```
import torch
x=torch.empty([32, 32]).cuda()
with torch.cuda.graph(torch.cuda.CUDAGraph()):
    torch.matmul(x,x)
```
fails with:
```
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasCreate(handle)`

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "/usr/local/lib/python3.8/dist-packages/torch/cuda/graphs.py", line 176, in __exit__
    self.cuda_graph.capture_end()
  File "/usr/local/lib/python3.8/dist-packages/torch/cuda/graphs.py", line 82, in capture_end
    super(CUDAGraph, self).capture_end()
RuntimeError: CUDA error: operation failed due to a previous error during capture
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```
The code runs fine if I do any multiplication before creating the graph.

# Jit bug

Compiling jit code while creating a cuda graph breaks jit. The code
```
import torch

@torch.jit.script
def f(x):
    return 2*x+1

x=torch.randn([32, 32]).cuda()
g0 = torch.cuda.CUDAGraph()
with torch.cuda.graph(g0):
    y = f(x)

g1 = torch.cuda.CUDAGraph()
with torch.cuda.graph(g1):
    y = f(x)
```
fails with:
```
a.py:16: UserWarning: FALLBACK path has been taken inside: runCudaFusionGroup. This is an indication that codegen Failed for some reason.
To debug try disable codegen fallback path via setting the env variable `export PYTORCH_NVFUSER_DISABLE=fallback`
 (Triggered internally at /opt/pytorch/pytorch/third_party/nvfuser/csrc/manager.cpp:335.)
  y= f(x)
Traceback (most recent call last):
  File "a.py", line 16, in <module>
    y= f(x)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<string>", line 3, in fallback_cuda_fuser

def mul(a : int, b : Tensor) -> Tensor:
  return b * a
         ~~~~~ <--- HERE
def add(a : int, b : Tensor) -> Tensor:
  return b + a
RuntimeError: status != cudaStreamCaptureStatus::cudaStreamCaptureStatusInvalidated INTERNAL ASSERT FAILED at "/opt/pytorch/pytorch/c10/cuda/CUDACachingAllocator.cpp":1490, please report a bug to PyTorch. 



During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "a.py", line 16, in <module>
    y= f(x)
  File "/usr/local/lib/python3.8/dist-packages/torch/cuda/graphs.py", line 176, in __exit__
    self.cuda_graph.capture_end()
  File "/usr/local/lib/python3.8/dist-packages/torch/cuda/graphs.py", line 82, in capture_end
    super(CUDAGraph, self).capture_end()
RuntimeError: CUDA error: operation failed due to a previous error during capture
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```
The first graph compiles correctly and is perfectly usable, with the crash happens when creating the second graph. Again, the crash can be avoided by first calling the method, though it needs the exact same tensor shape and dtype.

### Versions

[Temporarily unable to run, will update later. Using docker image nvcr.io/nvidia/pytorch:23.03-py3 (pytorch version [2.0.0a0+1767026](https://github.com/pytorch/pytorch/commit/1767026)) and running on a A100-80 GB]


cc @mcarilli @ezyang