Custom C++-CUDA Layer save_for_backward_equivalent

k_k: <p>Hello, i have implemented a custom layer for pytorch and want to speed it up by using CUDA.<br>
I was following the tutorial given here <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="nofollow noopener">https://pytorch.org/tutorials/advanced/cpp_extension.html</a>.<br>
But it is not mentioned how to save tensors in the forward pass ctx.save_for_backward in CUDA.<br>
I would really appreciate if someone could explain how to do this if possible or give me suggestion on how to do it in another way.<br>
Here is the code for python implementation:</p>
<pre><code class="lang-auto">class NeighborFill(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, y, z, neighbors):
        dt = TimestampMillisec64()
        nx = torch.zeros((1, x.shape[1], neighbors.shape[1] * neighbors.shape[2])).cuda()
        ny = torch.zeros((1, y.shape[1], neighbors.shape[1] * neighbors.shape[2])).cuda()
        nz = torch.zeros((1, z.shape[1], neighbors.shape[1] * neighbors.shape[2])).cuda()
        for i in range(0, neighbors.shape[1]):
            for j in range(0, neighbors.shape[2]):
                index = neighbors[0][i][j] - 1
                if (index &gt;= 0):
                    nx[0, :, i * neighbors.shape[2] + j] = x[0, :, index]
                    ny[0, :, i * neighbors.shape[2] + j] = y[0, :, index]
                    nz[0, :, i * neighbors.shape[2] + j] = z[0, :, index]
        ctx.save_for_backward(x, y, z, neighbors)
        ctx.mark_non_differentiable(neighbors)
        return nx, ny, nz
    
    @staticmethod
    def backward(ctx, gradoutx, gradouty, gradoutz):
        dt = TimestampMillisec64()
        x, y, z, neighbors = ctx.saved_tensors
        gradinpx = torch.zeros((1, x.shape[1], x.shape[2])).cuda()
        gradinpy = torch.zeros((1, y.shape[1], y.shape[2])).cuda()
        gradinpz = torch.zeros((1, z.shape[1], z.shape[2])).cuda()    
        for i in range(0, neighbors.shape[1]):
            for j in range(0, neighbors.shape[2]):
                index = neighbors[0][i][j] - 1
                if (index &gt;= 0):
                    gradinpx[0, :, i] += gradoutx[0, :, i * neighbors.shape[2] + j]
                    gradinpy[0, :, i] += gradouty[0, :, i * neighbors.shape[2] + j]
                    gradinpz[0, :, i] += gradoutz[0, :, i * neighbors.shape[2] + j]
        
        return gradinpx, gradinpy, gradinpz, neighbors
</code></pre>

k_k: <p>And here is the code for CUDA. I have only the forward for now:</p>
<pre><code class="lang-auto">std::vector&lt;torch::Tensor&gt; neighborfill_forward(torch::Tensor x, torch::Tensor y, torch::Tensor z, torch::Tensor neighbors){
    CHECK_INPUT(x);
    CHECK_INPUT(y);
    CHECK_INPUT(z);
    CHECK_CUDA(neighbors);
    
    return neighborfill_cuda_forward(x, y, z, neighbors)
}

std::vector&lt;torch::Tensor&gt; neighborfill_cuda_forward(torch::Tensor x, torch::Tensor y, torch::Tensor z, torch::Tensor neighbors){
    auto batch_size = x.size(0);
    auto state_size = x.size(2);
    int threads = 1024;
    dim3 blocks((state_size + threads - 1) / threads, batch_size);
    torch::Tensor nx = torch::zeros((1, x.size(1), neighbors.size(1) * neighbors.size(2)));
    torch::Tensor nx = torch::zeros((1, y.size(1), neighbors.size(1) * neighbors.size(2)));
    torch::Tensor nx = torch::zeros((1, z.size(1), neighbors.size(1) * neighbors.size(2)));
    
    neighborfill_cuda_forward_kernel&lt;x.type(), neighbors.type()&gt;&lt;&lt;&lt;blocks, threads&gt;&gt;&gt;(
    nx.packed_accessor32&lt;nx.type(), 3, torch::RestrictPtrTraits&gt;(), 
    ny.packed_accessor32&lt;ny.type(), 3, torch::RestrictPtrTraits&gt;(), 
    nz.packed_accessor32&lt;nz.type(), 3, torch::RestrictPtrTraits&gt;(), 
    x.packed_accessor32&lt;nx.type(), 3, torch::RestrictPtrTraits&gt;(), 
    y.packed_accessor32&lt;ny.type(), 3, torch::RestrictPtrTraits&gt;(), 
    z.packed_accessor32&lt;nz.type(), 3, torch::RestrictPtrTraits&gt;(), 
    neighbors.packed_accessor32&lt;neighbors.type(), 3, torch::RestrictPtrTraits&gt;(), 
    x.size(2), neighbors.size(2), x.size(1)
    );
    
    return {nx, ny, nz}
}

template &lt;typename x_t, typename neighbors_t&gt;
__global__ void fill_neighbors_kernel(torch::PackedTensorAccessor32&lt;x_t,3,torch::RestrictPtrTraits&gt; nx, torch::PackedTensorAccessor32&lt;x_t,3,torch::RestrictPtrTraits&gt; ny, torch::PackedTensorAccessor32&lt;x_t,3,torch::RestrictPtrTraits&gt; nz, torch::PackedTensorAccessor32&lt;x_t,3,torch::RestrictPtrTraits&gt; x, torch::PackedTensorAccessor32&lt;x_t,3,torch::RestrictPtrTraits&gt; y, torch::PackedTensorAccessor32&lt;x_t,3,torch::RestrictPtrTraits&gt; z, torch::PackedTensorAccessor32&lt;neighbor_t,3,torch::RestrictPtrTraits&gt; neighbors, int tsize, int nsize, int fmaps){
    int index = (blockIdx.x + blockIdx.y * gridDim.x) * blockDim.x + threadIdx.x;
    if(index &lt; tsize){ 
        for(int i = 0; i &lt; nsize; i++){
            int nind = neighbors[0][index][i] - 1;   
            if(nind &gt;= 0){
                for(int j = 0 ; j &lt; fmaps; j++){
                    nx[0][j][index * nsize + i] = x[0][j][nind];//TODO INDEXING
                    ny[0][j][index * nsize + i]  = y[0][j][nind];
                    nz[0][j][index * nsize + i]  = z[0][j][nind];
                }
            }
        }
    }
</code></pre>

albanD: <p>In the tutorial you link, the cpp function is then used inside an <code>autograd.Function</code>. So you can save them in the autograd.Function that wraps your cpp implementation as you do in your current implementation.</p>

k_k: <p>Thank you. I didn’t realize i needed to write a wrapper.</p>