Class DotProdWithIdx

pointops2/functions/pointops.py:320–368

attention_step2_v2 = AttentionStep2_v2.apply

class DotProdWithIdx(Function):
    @staticmethod
    def forward(ctx, q, index, table, rel_idx):
        """
        input: q: (N, h, hdim), index: (M), table: (L, h, hdim, 3), rel_idx: (M, 3)
        output: output: [M, h]
        """
        assert q.is_contiguous() and index.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous()

        N, h, hdim = q.shape
        M = index.shape[0]

        output = torch.cuda.FloatTensor(M, h).zero_()
        pointops_cuda.dot_prod_with_idx_forward_cuda(N, M, h, hdim, q, index, table, rel_idx, output)
        ctx.save_for_backward(q, index, table, rel_idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: [M, h]
        output: (N, h, hdim), None, (L, h, hdim, 3), None
        """
        q, index, table, rel_idx = ctx.saved_tensors
        M, h = grad_output.shape
        N, _, hdim = q.shape
        L = table.shape[0]
        
        grad_output = grad_output.contiguous()
        assert q.is_contiguous() and index.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous() and grad_output.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_q = torch.cuda.FloatTensor(N, h, hdim).zero_()
        grad_table = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.dot_prod_with_idx_backward_cuda(N, M, h, hdim, grad_output, q, index, table, rel_idx, grad_q, grad_table)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v9: {}".format(end - start))
        # # input()
        
        return grad_q, None, grad_table, None

dot_prod_with_idx = DotProdWithIdx.apply