CUDA Stream for PyTorch C++/CUDA Custom Extension

username: leimao
cooked: <p>Hello,</p>
<p>I was studying the “<a href="https://pytorch.org/tutorials/advanced/cpp_extension.html#custom-c-and-cuda-extensions" rel="noopener nofollow ugc">CUSTOM C++ AND CUDA EXTENSIONS</a>” from the PyTorch official tutorial. I am not sure how PyTorch was implemented but I would assume PyTorch would use CUDA stream for asynchronous execution. However, it seems that we don’t have access to the CUDA stream when we are programming the custom extension. This means that the execution of the custom extension will always be synchronous. Can someone provide some instructions on it? Thank you very much.</p>
<p>Best,</p>
<p>Lei</p>

username: glaringlee
cooked: <p>Can you paste some example code? I think you should have access to CUDA stream from C++, if you can share some code and your observation. I might be able to help.</p>

username: leimao
cooked: <p><a class="mention" href="/u/glaringlee">@glaringlee</a><br>
Thank you, Xinyun. I am looking at this particular line of code:<br>
<a href="https://github.com/pytorch/extension-cpp/blob/master/cuda/lltm_cuda_kernel.cu#L119" rel="noopener nofollow ugc">https://github.com/pytorch/extension-cpp/blob/master/cuda/lltm_cuda_kernel.cu#L119</a>.<br>
I don’t think CUDA stream was wrapped up there?<br>
Best,<br>
Lei</p>

username: glaringlee
cooked: <p><a class="mention" href="/u/leimao">@leimao</a><br>
Yes, it doesn’t. It use the default stream in this case. But you can use your own stream by tweaking this line right?<br>
‘<a href="https://github.com/pytorch/extension-cpp/blob/master/cuda/lltm_cuda_kernel.cu#L120" rel="noopener nofollow ugc">https://github.com/pytorch/extension-cpp/blob/master/cuda/lltm_cuda_kernel.cu#L120</a>’<br>
You can assign CUDA stream in the bracket as the 4th parameter (3rd parameter is 0).</p>

username: leimao
cooked: <p>Thank you very much for the quick response <a class="mention" href="/u/glaringlee">@glaringlee</a>. I can create a CUDA stream and use it here in the extension. But it will not make sense because the CUDA stream I used for the extension is not the same as the one used in the neural network backbone. Therefore, the execution of the extension will not be truly async. Ideally, it is PyTorch’s responsibility to expose the PyTorch CUDA stream to the user.<br>
I think TensorFlow exposes CUDA stream to the user:<br>
</p><aside class="onebox allowlistedgeneric">
  <header class="source">
      <img src="https://www.gstatic.com/devrel-devsite/prod/vfe8af62599ec445552c3fb43608c37ff46463c9fce3b14d8ee63b2e71edddffd/tensorflow/images/favicon.png" class="site-icon" width="16" height="16">
      <a href="https://www.tensorflow.org/guide/create_op" target="_blank" rel="noopener nofollow ugc">TensorFlow</a>
  </header>
  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:129/72;"><img src="https://www.tensorflow.org/images/tf_logo_social.png" class="thumbnail" width="129" height="72"></div>

<h3><a href="https://www.tensorflow.org/guide/create_op" target="_blank" rel="noopener nofollow ugc">Create an op &nbsp;|&nbsp; TensorFlow Core</a></h3>



  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
Please ctrl + F and search “stream”.

username: leimao
cooked: <p>That is to say. PyTorch should own the CUDA stream. The extension should not.</p>

username: glaringlee
cooked: <p><a class="mention" href="/u/leimao">@leimao</a><br>
No, I mean, you can use libtorch exposed api to get CUDA stream.<br>
I wrote this recently ‘<a href="https://pytorch.org/cppdocs/notes/tensor_cuda_stream.html" rel="noopener nofollow ugc">https://pytorch.org/cppdocs/notes/tensor_cuda_stream.html</a>’ , not finished yet, someone said it is not working well, but no follow up, can you try and let me know if it works? The other thing is that please pay attention to setCurrentStream if you plan to use it, very tricky.</p>

username: leimao
cooked: <p>Thank you very much <a class="mention" href="/u/glaringlee">@glaringlee</a> . I did not know we have this API. I can certainly try it but I am not sure how should I test to see if it is running asynchronously using the CUDA stream correctly?</p>

username: glaringlee
cooked: <p><a class="mention" href="/u/leimao">@leimao</a><br>
I think what you can do is that for eg. get 2 streams from stream pool, pass each one to a separate kernel, and inside the kernel, check the current stream, then you know whether your kernel is running on both streams.</p>

username: James_Huang
cooked: <p>Hi <a class="mention" href="/u/leimao">@leimao</a>, I am recently working on similar problems and eventually end up in this example: <a href="https://github.com/NVIDIA/flownet2-pytorch/blob/2e9e010c98931bc7cef3eb063b195f1e0ab470ba/networks/channelnorm_package/channelnorm_kernel.cu#L113" rel="noopener nofollow ugc">flownet2</a>.<br>
It seems using the <code>at::cuda::getCurrentCUDAStream()</code> as the 4th parameter in your kernel will respect the stream setting in your network backbone (i.e. in pytorch code, guarded by <code>with torch.cuda.stream(s):</code>).</p>
<p>I tested it myself using nvidia profiler and it did work.</p>

username: leimao
cooked: <p>Thank you very much. I will take a look.</p>

username: alanzhai219
cooked: <p>pytorch has provided cuda stream api getCurrentCUDAStream(), please take a look.</p>