Libtorch, CUDAGraph, and indexing

jakkes12: <p>Hi,</p>
<p>I am running into issues when capturing a CUDA graph in libtorch. Specifically, around indexing. The example below fails, with an error along the lines of “operation is not permitted when capturing”. Specifically, the indexing fails</p>
<pre><code class="lang-cpp">#include &lt;c10/cuda/CUDAStream.h&gt;
#include &lt;ATen/cuda/CUDAGraph.h&gt;
#include &lt;torch/torch.h&gt;


int main()
{
    torch::StreamGuard stream_guard{at::cuda::getStreamFromPool()};
    at::cuda::CUDAGraph graph{};
    auto options = torch::TensorOptions{}.device(torch::kCUDA);

    auto x = torch::randn({25, 5}, options);
    c10::List&lt;c10::optional&lt;at::Tensor&gt;&gt; index_list{x &lt; 0};
    x = x.index_put(index_list, - x.index(index_list));
    auto y = x.square();

    x.copy_(torch::randn({25, 5}, options));
    graph.capture_begin();
    c10::List&lt;c10::optional&lt;at::Tensor&gt;&gt; index_list{x &lt; 0};
    x = x.index_put(index_list, - x.index(index_list));
    y = x.square();
    graph.capture_end();
    std::cout &lt;&lt; y;
}
</code></pre>
<p>Similarly, in Python, this does not work. The snippet below fails with the seemingly random error</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/jakkes/projects/rl-cpp/test.py", line 16, in &lt;module&gt;
    x[x &lt; 0] = -x[x &lt; 0]
RuntimeError: numel: integer multiplication overflow
</code></pre>
<p>or</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/jakkes/projects/rl-cpp/test.py", line 16, in &lt;module&gt;
    x[x &lt; 0] = -x[x &lt; 0]
RuntimeError: status != cudaStreamCaptureStatus::cudaStreamCaptureStatusInvalidated INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":1220, please report a bug to PyTorch.
</code></pre>
<pre><code class="lang-python">import torch


if __name__ == "__main__":
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        x = torch.randn(25, 5, device="cuda")
        x[x &lt; 0] = -x[x &lt; 0]
        y = x.square()
    torch.cuda.current_stream().wait_stream(s)
    
    x.copy_(torch.randn(25, 5, device="cuda"))
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        x[x &lt; 0] = -x[x &lt; 0]
        y = x.square()

    for _ in range(3):
        x.copy_(torch.randn(25, 5, device="cuda"))
        g.replay()
        print(y)

</code></pre>
<p>In this example, I can use <code>torch::where</code> to go around the indexing issues. However, in the specific use case I am having, I need to make use of the accumulate feature, i.e.</p>
<pre><code class="lang-cpp">x = x.index_put_({someindicies}, some_values, /* accumulate */ true)
</code></pre>
<p>with <code>x</code> being a rather large matrix, I would rather not have to create a big loop of <code>torch::where</code> calls.</p>
<p>Any advice is very much welcomed, thank you.</p>

ptrblck: <p>Indexing with a <code>BoolTensor</code> is disallowed as this operation is synchronizing the code with the host as the result tensor has a data dependency (i.e. the shape depends on the number of <code>True</code> entries).<br>
You would either need to move this operation out of the capture or try to replace it with an indexing operation using indices.</p>

jakkes12: <p>Okay, I see. Thank you.</p>