Index_put_ really slow on gpu

Zartris: <p>I am trying to understand why my code is 3x slower on GPU. I’m profiling and debugging my update state function, which should be straightforward.<br>
Everything should be on the GPU, so I do not understand why placing these values would take so long. However, when looking at the profiling, we can see that the index_put_ is running from libtorch_cpu.so, which indicates that most of the function is actually performed on the CPU, even though all the values are on the GPU.</p>
<p>I’m new to using torch in c++, so please help me understand.</p>
<pre><code class="lang-auto">    void AAGV::update_state(double dt) {
        auto deltaTime = torch::scalar_tensor(dt, TOptions(torch::kDouble, device));
        torch::Tensor orientation = getOrientation();

        // TODO:: Very slow on gpu, OPTIMIZE THIS
        _B.index_put_({0, 0}, torch::cos(orientation) * deltaTime);
        _B.index_put_({1, 0}, torch::sin(orientation) * deltaTime);
        _B.index_put_({2, 1}, deltaTime);

//        torch::Tensor B = torch::tensor({{torch::cos(orientation) * deltaTime, 0.},
//                                         {torch::sin(orientation) * deltaTime, 0.},
//                                         {0.,                                  deltaTime}},
//                                        TOptions(torch::kDouble, device));

//        double data[] = {torch::cos(orientation).*dt, 0., torch::sin(orientation) * dt, 0., 0., dt};
//        torch::Tensor B = torch::from_blob(data, {3, 2}, torch::kDouble).to(device);


        torch::Tensor vel = torch::stack({getLinearVelocity(), getAngularVelocity()});

        torch::Tensor new_state = A.matmul(getState()) + _B.matmul(vel);
        setState(new_state);
    }
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/b/8b3185279fa691735b3add037f6932bd882d3230.png" data-download-href="https://discuss.pytorch.org/uploads/default/8b3185279fa691735b3add037f6932bd882d3230" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/b/8b3185279fa691735b3add037f6932bd882d3230_2_690x394.png" alt="image" data-base62-sha1="jRmsOFt2jf6eSrORpKzVpCI0SiY" width="690" height="394" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/b/8b3185279fa691735b3add037f6932bd882d3230_2_690x394.png, https://discuss.pytorch.org/uploads/default/original/3X/8/b/8b3185279fa691735b3add037f6932bd882d3230.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/b/8b3185279fa691735b3add037f6932bd882d3230.png 2x" data-dominant-color="35383D"><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">779×445 78.4 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>Could you describe how <code>getOrientation</code> is defined and how you’ve verified <code>orientation</code> is on the GPU?</p>

Zartris: <p>Yes, getOrientation is just a getter indexing into a state tensor.<br>
I only have this function to make my code more readable.</p>
<pre><code class="lang-auto">torch::Tensor getOrientation() {
            return state.index({2});
        }
torch::Tensor state = torch::zeros({3, 1}, torch::TensorOptions().dtype(torch::kDouble).device(device));
</code></pre>
<p>And I verify it by debuging and printing out the device.:</p>
<pre><code class="lang-auto">Orientation: 
0.01 *
-8.0727
[ CUDADoubleType{1} ]
</code></pre>
<p>Could it be because of my indexing is not on the GPU?<br>
the <code>{0, 0}</code> part of <code>_B.index_put_({0, 0}, torch::cos(orientation) * deltaTime);</code></p>

CBGxd: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Same issue in here. In my case, the index_put_ on Cuda is much slower than CPU when calculating the confusion matrix.</p>
<p>Here is the snip code:</p>
<pre><code class="lang-auto">    def add_batch(self, x, y):
        x_row = x.reshape(-1)
        y_row = y.reshape(-1)

        idxs = torch.stack([x_row, y_row], dim=0)

        if self.ones is None or self.last_scan_size != idxs.shape[-1]:
            self.ones = torch.ones((idxs.shape[-1]), device=x.device, dtype=torch.long)
            self.last_scan_size = idxs.shape[-1]

        self.conf_matrix = self.conf_matrix.index_put_(tuple(idxs), self.ones, accumulate=True)
</code></pre>
<p>With CPU version, it takes around 5 mins <strong>[02:14&lt;03:05, 18.62it/s]</strong>, while the time become 12x times slower on CUDA <strong>[00:16&lt;1:05:09,  1.53it/s]</strong>.</p>
<p>Since translating the tensor from CUDA to CPU results extra cost, any idea to solve this problem will be very helpful for my project.</p>
<p>Thanks.</p>