Would you tell me why you did not consider cuda shared memory size?

coincheung: <p>Hi,</p>
<p>I am learning from the source code, and I got to know this line:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/8e2841781eb569a43e3f0248ba7686f9cae19b9d/aten/src/ATen/native/cuda/SoftMax.cu#L128" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/8e2841781eb569a43e3f0248ba7686f9cae19b9d/aten/src/ATen/native/cuda/SoftMax.cu#L128" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/8e2841781eb569a43e3f0248ba7686f9cae19b9d/aten/src/ATen/native/cuda/SoftMax.cu#L128</a></h4>
<pre class="onebox"><code class="lang-cu"><ol class="start lines" start="118" style="counter-reset: li-counter 117 ;">
<li>}</li>
<li>
</li><li>
</li><li>template&lt;typename accscalar_t, typename Kernel&gt;</li>
<li>void SpatialSoftMax_getLaunchSizes(</li>
<li>    Kernel k,</li>
<li>    uint64_t outer_size, uint64_t dim_size, uint64_t inner_size,</li>
<li>    dim3&amp; grid, dim3&amp; block, uint32_t&amp; smem_size) {</li>
<li>  block = SpatialSoftMax_getBlockSize(outer_size, dim_size, inner_size);</li>
<li>  uint32_t block_threads = block.x * block.y;</li>
<li class="selected">  smem_size = block.x == 1 ? 0 : block_threads * sizeof(accscalar_t);</li>
<li>  int max_active_blocks;</li>
<li>#if defined(__HIP_PLATFORM_HCC__) &amp;&amp; HIP_VERSION &lt; 305</li>
<li>  // HIP function signature is not compatible yet.</li>
<li>  uint32_t max_blocks;</li>
<li>  cudaOccupancyMaxActiveBlocksPerMultiprocessor(&amp;max_blocks,</li>
<li>                                                k, block_threads, smem_size);</li>
<li>  max_active_blocks = max_blocks;</li>
<li>#else</li>
<li>  cudaOccupancyMaxActiveBlocksPerMultiprocessor(&amp;max_active_blocks,</li>
<li>                                                k, block_threads, smem_size);</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Here we allocate as many shared memory according to the number of parallelization. However, I also learned that for most gpus, the max shm of each CUDA SM is around 48k. If there are 48 blocks, each block would have at most 1k shared memory. Would you please tell me why the pytorch code does not consider that in the implementation of operators?</p>

tom: <p>So the maximum number of threads in the block is 1024 (per <code>SpatialSoftMax_getBlockSize</code>), the accscalar will be 32 or 64 bits, so we’re using 4k-8k shmem per block.  <code>SpatialSoftMax_getLaunchSizes</code> now uses the occupancy API to determine the maximum number of active blocks.<br>
Which part is not considered?</p>
<p>Best regards</p>
<p>Thomas</p>

coincheung: <p>Thanks for replying !!!</p>
<p>Will I be ok if each block does not exceed 48k, but the whole SM exceed 48k ?</p>
<p>By the way, did the number of active blocks is determined by the cuda <code>cudaOccupancyMaxPotentialBlockSize</code> method ? This method outputs both recommended block size and grid sizes, and only the output grid size is used as active number of blocks with the block size not used. Is that the way it works ?</p>