Why PyTorch calls so many cudaStreamCreateWithPriority and cudaStreamIsCapturing?

youkaichao1: <p>Here is my simple program:</p>
<pre data-code-wrap="python"><code class="lang-python"># test.py
import torch
from torch.cuda import nvtx

with nvtx.range("init"):
    stream1 = torch.cuda.Stream()
    stream2 = torch.cuda.Stream()
    a = torch.zeros(233, device='cuda')

with torch.cuda.stream(stream1):
    with nvtx.range("first add"):
        a += 1

with torch.cuda.stream(stream2):
    with nvtx.range("second add"):
        a += 1

with nvtx.range("print"):
    torch.cuda.synchronize()
    average = a.mean().item()
    print(f'Average of the elements in tensor a: {average}')
</code></pre>
<p>I’m profiling it using <code>nsys profile -o test python test.py</code>, and during the <code>init</code> stage, in nsight system, I see hundreds of <code>cudaStreamCreateWithPriority</code> and <code>cudaStreamIsCapturing</code> call. But I only created two cuda streams. What happened?</p>

ptrblck: <p>You can use <code>cuda-gdb</code> and set a breakpoint at <code>cudaStreamIsCapturing</code> to check the stacktraces.<br>
This will show that e.g. the <code>CUDAGenerator</code> calls this method during the initialization e.g. to set its <code>seed</code> <a href="https://github.com/pytorch/pytorch/blob/d60135e91548783a6cbcfc69ea0fbe4976251b16/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp#L269">here</a>:</p>
<pre><code class="lang-auto">#0  0x00007fffec4775e4 in cudaStreamIsCapturing () from /usr/local/cuda/targets/x86_64-linux/lib/libcudart.so.12
#1  0x00007fff8894508f in at::cuda::assertNotCapturing(std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt;) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cuda.so
#2  0x00007fff8894228a in at::CUDAGeneratorImpl::set_current_seed(unsigned long) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cuda.so
#3  0x00007fff889425d2 in at::CUDAGeneratorImpl::seed() () from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cuda.so
#4  0x00007fff889441fa in at::cuda::detail::getDefaultCUDAGenerator(signed char) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cuda.so
#5  0x00007fffeb882c0e in THCPModule_initExtension(_object*, _object*) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#6  0x00005555556aba3e in ?? ()
#7  0x000055555569ccfa in _PyEval_EvalFrameDefault ()
#8  0x00005555556ae9fc in _PyFunction_Vectorcall ()
#9  0x0000555555687469 in ?? ()
#10 0x000055555578aa48 in PyObject_CallMethod ()
#11 0x00007fffeb80af2f in torch::utils::device_lazy_init(c10::DeviceType) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#12 0x00007fffeb20ecb7 in torch::autograd::THPVariable_zeros(_object*, _object*, _object*) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
</code></pre>
<p>and later the caching allocator will call into <code>cudaMallocMaybeCapturing</code> <a href="https://github.com/pytorch/pytorch/blob/d60135e91548783a6cbcfc69ea0fbe4976251b16/c10/cuda/CUDACachingAllocator.cpp#L782">here</a>:</p>
<pre><code class="lang-auto">#0  0x00007fffec4775e4 in cudaStreamIsCapturing () from /usr/local/cuda/targets/x86_64-linux/lib/libcudart.so.12
#1  0x00007ffff6596aa6 in c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::cudaMallocMaybeCapturing(void**, unsigned long) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libc10_cuda.so
#2  0x00007ffff65a0331 in c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::alloc_block(c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::AllocParams&amp;, bool, std::shared_ptr&lt;c10::GatheredContext&gt; const&amp;, std::unique_lock&lt;std::recursive_mutex&gt;&amp;) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libc10_cuda.so
#3  0x00007ffff65a4d7a in c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::malloc(signed char, unsigned long, CUstream_st*) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libc10_cuda.so
#4  0x00007ffff65a65d5 in c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::malloc(void**, signed char, unsigned long, CUstream_st*) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libc10_cuda.so
#5  0x00007ffff65a6a5b in c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::allocate(unsigned long) const () from /usr/local/lib/python3.10/dist-packages/torch/lib/libc10_cuda.so
</code></pre>
<p>These checks are performed to fail properly when these disallowed operations are called during a CUDA Graphs capture.</p>