Torch.istft NOLA check causes synchronization and massive slow down

Nabarun_Goswami: <p>torch.istft spends almost all its time in the NOLA check.</p>
<p>Below is a screenshot of the profiling of stft followed by istft of a signal of shape <code>(32, 4, 2, 7 * 44100)</code>. This is just a toy example, but when the istft is part of a larger neural network, it causes everything to synchronize.</p>
<p>I found the cause of the synchronization (see code trail below). But I am wondering if there is any way around this to avoid the synchronization.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/c/2c931594e3c41ccc7243bcefd430edd4c747e49d.png" data-download-href="https://discuss.pytorch.org/uploads/default/2c931594e3c41ccc7243bcefd430edd4c747e49d" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/c/2c931594e3c41ccc7243bcefd430edd4c747e49d_2_690x236.png" alt="image" data-base62-sha1="6mk8oltptpwFnlOsAAA8xiaXIFT" width="690" height="236" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/c/2c931594e3c41ccc7243bcefd430edd4c747e49d_2_690x236.png, https://discuss.pytorch.org/uploads/default/original/3X/2/c/2c931594e3c41ccc7243bcefd430edd4c747e49d.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/c/2c931594e3c41ccc7243bcefd430edd4c747e49d.png 2x" data-dominant-color="ABC6CD"><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">image</span><span class="informations">752×258 54.6 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>From the  <a href="https://github.com/pytorch/pytorch/blob/e790281a85fe3693fc1d38bf0e2c6e874d5e10b0/aten/src/ATen/native/SpectralOps.cpp#L1154" rel="noopener nofollow ugc">ATen/native/SpectralOps.cpp#L1154</a></p>
<pre><code class="lang-auto">if (at::is_scalar_tensor_true(window_envelop_lowest)) {
    std::ostringstream ss;
    REPR(ss) &lt;&lt; "window overlap add min: " &lt;&lt; window_envelop_lowest;
    AT_ERROR(ss.str());
  }

</code></pre>
<p>where, the call to <code>is_scalar_tensor_true</code> calls <code>at::equal</code> function at <a href="https://github.com/pytorch/pytorch/blob/6faa6c68e8b76fb68f3a2b2783685102d0e87c00/aten/src/ATen/TensorSubclassLikeUtils.h#L84" rel="noopener nofollow ugc">ATen/TensorSubclassLikeUtils.h#L84</a></p>
<pre><code class="lang-auto">inline bool is_scalar_tensor_true(const Tensor&amp; t) {
  TORCH_INTERNAL_ASSERT(t.dim() == 0)
  TORCH_INTERNAL_ASSERT(t.scalar_type() == kBool)
  return at::equal(t, t.new_ones({}, t.options()));
}
</code></pre>
<p>which in turn calls  <a href="https://github.com/pytorch/pytorch/blob/e790281a85fe3693fc1d38bf0e2c6e874d5e10b0/aten/src/ATen/native/cuda/Equal.cpp#L29" rel="noopener nofollow ugc">ATen/native/cuda/Equal.cpp#L29</a></p>
<pre><code class="lang-auto">return at::cuda::eq(self, src).all().item().to&lt;bool&gt;();
</code></pre>
<p>The <code>.item()</code> call does a memcopy from device to host and also synchronizes all operations on the GPU.<br>
Is it possible to maybe allow users to bypass the NOLA check if they want?</p>

ptrblck: <p>The bypassing logic might be a good feature request as we’ve also implemented a similar approach for some linalg methods. Could you create this feature request on GitHub so that the code owners could discuss it, please?</p>