Is there any non-blocking version of torch.nonzero / aten::nonzero?

username: thuwzt
cooked: <p>I’m creating my own <code>torch.autograd.Function</code> with torch API. One of the API used is <code>torch.nonzero</code> and I found it becoming the speed bottleneck since it causes host-device synchronization blocking my CPU stream, noted in <a href="https://pytorch.org/docs/stable/generated/torch.nonzero.html#torch-nonzero" rel="noopener nofollow ugc">torch.nonzero()</a>. I tried to use <code>torch.masked_select</code> or indexing with <code>foo[bar]</code> instead, yet I found all of them will call <code>aten::nonzero</code> and cause a sync. The source code of <code>aten::nonzero</code> is <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cuda/Nonzero.cu#L73" rel="noopener nofollow ugc">here</a>, with the line 73 causing a sync when copying the number of nonzero items from device to host, calling <code>at::cuda::memcpy_and_sync</code>. I’m confused about the necessity of this sync and wondering is there any non-blocking version of  <code>nonzero</code>?</p>
<p>Thanks! : )</p>

username: tom
cooked: <p>No: the output size (needed on the CPU) depends on the input data (on the GPU) and this forces the sync. Depending your application and the sizes involved, it might be better to work with a mask or similar.</p>
<p>Best regards</p>
<p>Thomas</p>

username: netanel-haber
cooked: <p>There is now in a way, <code>torch.nonzero_static</code>.</p>
<p>The semantics are different, but it might provide what you need depending on your use case.</p>
<aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/pytorch/pull/141838">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/pull/141838" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row" data-github-private-repo="false">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/pytorch/pull/141838" target="_blank" rel="noopener nofollow ugc">Implements nonzero_static on cuda</a>
      </h4>

    <div class="branches">
      <code>main</code> ← <code>ngimel/nonzero_static</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2024-12-01" data-time="20:40:11" data-timezone="UTC">08:40PM - 01 Dec 24 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/ngimel" target="_blank" rel="noopener nofollow ugc">
            <img alt="" src="https://avatars.githubusercontent.com/u/15841449?v=4" class="onebox-avatar-inline" width="20" height="20">
            ngimel
          </a>
        </div>

        <div class="lines" title="17 commits changed 5 files with 448 additions and 166 deletions">
          <a href="https://github.com/pytorch/pytorch/pull/141838/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+448</span>
            <span class="removed">-166</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">using blockwide cub primitives. 
This adds CUDA functionality for nonzero_stati<span class="show-more-container"><a href="https://github.com/pytorch/pytorch/pull/141838" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">c, which was missing in https://github.com/pytorch/pytorch/pull/97417.
cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @galv 
For `size` approx equal to number of nonzeros, the perf is very close to the regular version, for larger sizes filling in padding  indices takes additional time. 
Disabled for cuda &lt;=11.4</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<aside class="onebox allowlistedgeneric" data-onebox-src="https://ianbarber.blog/2024/12/18/nonzero_static-in-pytorch/">
  <header class="source">
      <img src="https://ianbarber.blog/wp-content/uploads/2024/11/cropped-logo-removebg-preview.png?w=32" class="site-icon" alt="" width="32" height="32">

      <a href="https://ianbarber.blog/2024/12/18/nonzero_static-in-pytorch/" target="_blank" rel="noopener nofollow ugc" title="08:12PM - 18 December 2024">Ian’s Blog – 18 Dec 24</a>
  </header>

  <article class="onebox-body">
    <img width="200" height="200" src="https://ianbarber.blog/wp-content/uploads/2024/11/cropped-logo-removebg-preview.png?w=200" class="thumbnail onebox-avatar" alt="">

<h3><a href="https://ianbarber.blog/2024/12/18/nonzero_static-in-pytorch/" target="_blank" rel="noopener nofollow ugc">nonzero_static in PyTorch</a></h3>

  <p>Natalia Gimelshein added cuda support for nonzero_static to PyTorch the other day — its a feature Jax has had for a while that lets you avoid a bit of data dependent annoyingness. I most ofte…</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>