How to have FFT using multiple GPU streams?

eduardo4jesus: <p>Comparing the FFT execution time with the one of Conv2d, I noticed that FFT operation is way slower than expected.</p>
<p>Doing some investigation, I noticed the Conv2d uses the FFT method, and somehow is able to launch several stream.<br>
While when I called <code>nn.funcitonal.fft.fft2(inputs)</code> myself, I was only able to launch one stream.</p>
<p>Is there any way to improve the FFT calls? I guess it is safe to say it should be possible. <code>cudnn_convolution</code> is somehow doing it.</p>
<p>How does the number of stream affects speed?</p>
<hr>
<p>This is the profiling of <code>nn.functional.conv2d</code></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/b/8b72c725cf31cbb92987ea91752c9dd643c8219c.png" data-download-href="https://discuss.pytorch.org/uploads/default/8b72c725cf31cbb92987ea91752c9dd643c8219c" title="Screen Shot 2021-09-24 at 1.10.23 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/b/8b72c725cf31cbb92987ea91752c9dd643c8219c_2_690x390.png" alt="Screen Shot 2021-09-24 at 1.10.23 PM" data-base62-sha1="jTChiioeHS1X5gIvQBsyk1eOAVe" width="690" height="390" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/b/8b72c725cf31cbb92987ea91752c9dd643c8219c_2_690x390.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/b/8b72c725cf31cbb92987ea91752c9dd643c8219c_2_1035x585.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/b/8b72c725cf31cbb92987ea91752c9dd643c8219c_2_1380x780.png 2x" data-dominant-color="E8E3EC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2021-09-24 at 1.10.23 PM</span><span class="informations">1988×1125 344 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>This is the profiling of <code>nn.functional.fft.fft2</code></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/7/c73cee9d16e983b004037606813c6e20700f0ac9.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/c73cee9d16e983b004037606813c6e20700f0ac9" title="Screen Shot 2021-09-24 at 1.39.39 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c73cee9d16e983b004037606813c6e20700f0ac9_2_690x396.jpeg" alt="Screen Shot 2021-09-24 at 1.39.39 PM" data-base62-sha1="sqxyUcE6q6njtIFPVus2Azdj8o1" width="690" height="396" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c73cee9d16e983b004037606813c6e20700f0ac9_2_690x396.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c73cee9d16e983b004037606813c6e20700f0ac9_2_1035x594.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c73cee9d16e983b004037606813c6e20700f0ac9_2_1380x792.jpeg 2x" data-dominant-color="F3F1F3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2021-09-24 at 1.39.39 PM</span><span class="informations">2032×1169 190 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>You can use <a href="https://pytorch.org/docs/stable/generated/torch.cuda.Stream.html"><code>torch.cuda.Stream</code></a> to create custom streams. Check the docs carefully (as well as the streams part <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-semantics">here</a>) as I would consider it an advanced mechanism since you can easily create race conditions in your code.</p>

eduardo4jesus: <aside class="quote no-group" data-username="eduardo4jesus" data-post="1" data-topic="132738">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/eduardo4jesus/48/20203_2.png" class="avatar"> eduardo4jesus:</div>
<blockquote>
<p>How does the number of stream affects speed?</p>
</blockquote>
</aside>
<p>Thank you so much.</p>
<p>How does the number of stream affects speed? I am also trying to understand which streams have the potential to benefit from stream.</p>
<p>Would the following mindset be accurate?</p>
<p>A grid of blocks is dispatched for a task A on the Main stream. For certain sizes of grids and blocks it could happen that towards the end of the computation some SM become available, while others are still busy. For those scenarios, if multiple if a second stream (S) is used to process a second task (B), at that point, those idle SM would be able to be allocated to start processing the operations for B, given they would not depend on the result of A.</p>

ptrblck: <p>The mindset is generally right and yes, enough resources must be available to execute concurrently.<br>
<a href="https://developer.download.nvidia.com/CUDA/training/StreamsAndConcurrencyWebinar.pdf">This presentation</a> as well as <a href="https://on-demand.gputechconf.com/gtc/2014/presentations/S4158-cuda-streams-best-practices-common-pitfalls.pdf">this one</a> could give you a good idea.</p>

eduardo4jesus: <p>Thank you so much for all feedbacks and great resources.</p>