Slice copy from gpu to cpu tensors

pickandroll3: <p>I’m still new to CUDA programming and thought I could write a specific kernel for my application. I have a tensor <code>src</code> which I want to write some of the elements from and copy those values to their corresponding locations in a tensor <code>tgt</code>. Both tensors are 1d, however <code>src</code> is on a gpu and <code>tgt</code> is on CPU memory. I also have access to both the bytemask (<code>torch.BoolTensor</code>) as well as index locations of the elements I want to write to.</p>
<p>I tried using <code>tgt[idxs].copy_(src[idx])</code>, however this fails silently (data isn’t copied to <code>tgt</code> but instead is copied to the new storage of <code>tgt[idxs]</code>), and <code>tgt[idxs] = src[idxs]</code> fails because of different devices.</p>
<p>I have tried writing an <code>at::parallel_for</code> operation for this but it is much slower than I would want, so I am looking to try and write a GPU kernel for this. I know I can access a device-side host pointer with <code>cudaHostGetDevicePointer</code>, however if I’m reading the documentation correctly, this allocates a new host side buffer whereas I already have the host buffer in <code>tgt</code>. Can anyone help me with this?</p>
<p>I have already tried the following kernel which fails due to the <code>tgt</code> buffer being on CPU memory:</p>
<pre><code class="lang-auto">template&lt;typename scalar_t&gt;
__global__
void
cuda_mask_copy_kernel(scalar_t*         __restrict__ tgt_1d,
                      const bool*       __restrict__ mask_1d,
                      const scalar_t*   __restrict__ src_1d,
                      size_t N)
{
    const int index = blockIdx.x * blockDim.x + threadIdx.x;

    if (index &lt; N &amp;&amp; mask_1d[index])
    {
        tgt_1d[index] = src_1d[index];
    }
}
</code></pre>
<p>For completeness, here is my <code>at::parallel_for</code> implementation, but as I said, it is much slower than I would wish:</p>
<pre><code class="lang-auto">torch::Tensor&amp;
idx_copy(torch::Tensor&amp; tgt,
         const torch::Tensor&amp; idxs,
         const torch::Tensor&amp; src)
{
    TORCH_CHECK(idxs.sizes().size() == 1);
    CHECK_TYPE(idxs, torch::kLong);
    CHECK_SAME_SIZE(tgt, src);
    TORCH_CHECK(tgt.scalar_type() == src.scalar_type());

    torch::Tensor tgt_1d = tgt.view({-1});
    const torch::Tensor idxs_1d = idxs.view({-1});
    const torch::Tensor src_1d = src.view({-1});

    at::parallel_for(0, idxs_1d.size(0), at::internal::GRAIN_SIZE, [&amp;](int64_t start, int64_t end)
        {
            for(int64_t idx = start; idx &lt; end; ++idx)
            {
                tgt_1d[idxs_1d[idx]] = src_1d[idxs_1d[idx]];
            }
        }
    );

    return tgt;
}

</code></pre>

tom: <p>I think the title is slightly inaccurate, but the analysis is spot on. The problem is that the indexing you do (with a list of indices or a mask) necessarily creates a new tensor in new memory. If you had slicing, the indexing would create a view and <code>copy_</code> would work.<br>
<code>copy_</code> is the “only way” to transfer from the GPU to CPU (and <code>.to</code> uses it under the hood), so there is not much to do to avoid the transfer and then work on the CPU.<br>
An alternative <code>tgt[idxs] = src[idxs]</code> could be to use where: <code>tgt = torch.where(idxs, src, tgt)</code>, this creates a new tensor <code>tgt</code> but does not create <code>src[idxs]</code>. It also is relatively friendly to optimizers like the nvFuser if they are part of a longer computation. Personally, I find <code>where</code> a good fit to how I think about these operations, but of course, this may or may not apply to your case.</p>
<p>Best regards</p>
<p>Thomas</p>
<p>Best</p>

pickandroll3: <p>Hey Tom,</p>
<p>Thanks for the response, and I am sorry for the inaccuracy in the title. I would love to use <code>torch.where</code> however I believe that would allocate more gpu memory and I am already stressed for it. I also really want to avoid having to copy the gpu tensor to cpu memory and apply an augmented assignment (<code>tgt[idxs] = src[idxs].cpu()</code>): that would have to be an absolute worst case scenario for me.</p>

AKT_TARAFDER: <p><a class="mention" href="/u/tom">@tom</a> , <a class="mention" href="/u/pickandroll3">@pickandroll3</a><br>
I do have this operation:</p>
<pre><code class="lang-auto">def mlp(rank, backend, x, Wi, bi, Wo, bo, W_ln, b_ln,  wi_shape, wo_shape, expected_output, output, WORLD_SIZE, all_reduce_in_gpu, load_from_cpu):
### ..
    wo = torch.empty(row_splits_num, Wo.shape[1], device=cuda_device)

    wo.copy_(Wo[(LOCAL_RANK * row_splits_num): ((LOCAL_RANK * row_splits_num) + row_splits_num), :])
</code></pre>
<p>the Wo is in CPU, and I want to put ‘wo’ in a particular GPU using LOCAl_RANK to save memory<br>
but it is not saving any, for some reason the full Wo is copied.</p>
<p>cuda_device: 2 cuda devices selected</p>
<pre><code class="lang-auto">wo=(Wo[(LOCAL_RANK * row_splits_num): ((LOCAL_RANK * row_splits_num) + row_splits_num), :]).to(cuda_device)
</code></pre>
<p>This does not work either<br>
but the following saves, but is not generic</p>
<pre><code class="lang-auto">        wi0 = Wi[:,(0 * col_splits_num):( (0 * col_splits_num) + col_splits_num)]
        wi1 = Wi[:,(1 * col_splits_num):( (1 * col_splits_num) + col_splits_num)]
        
        wi = wi0.to('cuda:0')
        wi = wi1.to('cuda:1')
</code></pre>
<p>LOCAL_RANK =0/1</p>
<p>I am using multiprocessing calls:</p>
<pre><code class="lang-auto">    mp.spawn(init_processes,
             args=(backend, x, Wi, bi, Wo, bo, W_ln, b_ln, wi_shape, wo_shape, expected_output, output, world_size, all_reduce_in_gpu),
             nprocs=world_size,
             join=True)
</code></pre>