Class AttentionStep2WithRelPosValue_v2

pointops2/functions/pointops.py:584–644

class AttentionStep2WithRelPosValue_v2(Function):
    @staticmethod
    def forward(ctx, attn, v, index0_offsets, n_max, index1, table, rel_idx):
        """
        input: attn: (M, h), v: (N, h, hdim), index0_offsets: (M), index1: (M), table: (L, h, hdim, 3), rel_idx: (M, 3)
        output: output: [N, h, hdim]
        """
        assert attn.is_contiguous() and v.is_contiguous() and index0_offsets.is_contiguous() and index1.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous()

        M, h = attn.shape
        N, h, hdim = v.shape
        # N_q = int(index0_offsets.max().item()) + 1

        output = torch.cuda.FloatTensor(N, h, hdim).zero_()
        pointops_cuda.attention_step2_with_rel_pos_value_forward_cuda_v2(N, M, h, hdim, n_max, attn, v, index0_offsets, index1, table, rel_idx, output)

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

        ctx.n_max = n_max
        ctx.save_for_backward(attn, v, index0_offsets, index1, table, rel_idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: (N, h, C//h)
        output: (M, h), (N, h, C//h), None, None, (L, h, hdim, 3), None
        """
        n_max = ctx.n_max
        attn, v, index0_offsets, index1, table, rel_idx = ctx.saved_tensors
        N, h, hdim = grad_output.shape
        N = v.shape[0]
        M = attn.shape[0]
        L = table.shape[0]

        # grad_output = grad_output.contiguous()
        # print("grad_output.is_contiguous(): ", grad_output.is_contiguous())
        assert attn.is_contiguous() and v.is_contiguous() and index0_offsets.is_contiguous() and index1.is_contiguous() and grad_output.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous()

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

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

        grad_attn = torch.cuda.FloatTensor(M, h).zero_()
        grad_v = torch.cuda.FloatTensor(N, h, hdim).zero_()
        grad_table = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()

        # print("attn.shape: {}, grad_attn.shape: {}".format(attn.shape, grad_attn.shape))
        # print("v.shape: {}, grad_v.shape: {}".format(v.shape, grad_v.shape))
        # print("table.shape: {}, grad_table.shape: {}".format(table.shape, grad_table.shape))

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.attention_step2_with_rel_pos_value_backward_cuda_v2(N, M, h, hdim, n_max, grad_output, index0_offsets, index1, attn, v, table, rel_idx, grad_attn, grad_v, grad_table)
        
        # torch.cuda.synchronize()
        # end = time.time()