Spaces:
Runtime error
Runtime error
File size: 3,336 Bytes
4893ce0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
#include <vector>
#include <torch/serialize/tensor.h>
#include <ATen/cuda/CUDAContext.h>
#include "relative_pos_encoding_cuda_kernel.h"
void dot_prod_with_idx_forward_cuda(int N, int M, int h, int hdim, at::Tensor q_tensor, at::Tensor index_tensor,
at::Tensor table_tensor, at::Tensor rel_idx_tensor, at::Tensor output_tensor)
{
const float *q = q_tensor.data_ptr<float>();
const float *table = table_tensor.data_ptr<float>();
const int *index = index_tensor.data_ptr<int>();
const int *rel_idx = rel_idx_tensor.data_ptr<int>();
float *output = output_tensor.data_ptr<float>();
dot_prod_with_idx_forward_cuda_launcher(N, M, h, hdim, q, index, table, rel_idx, output);
}
void dot_prod_with_idx_backward_cuda(int N, int M, int h, int hdim, at::Tensor grad_out_tensor,
at::Tensor q_tensor, at::Tensor index_tensor, at::Tensor table_tensor, at::Tensor rel_idx_tensor,
at::Tensor grad_q_tensor, at::Tensor grad_table_tensor)
{
const float *grad_out = grad_out_tensor.data_ptr<float>();
const float *q = q_tensor.data_ptr<float>();
const int *index = index_tensor.data_ptr<int>();
const float *table = table_tensor.data_ptr<float>();
const int *rel_idx = rel_idx_tensor.data_ptr<int>();
float *grad_q = grad_q_tensor.data_ptr<float>();
float *grad_table = grad_table_tensor.data_ptr<float>();
dot_prod_with_idx_backward_cuda_launcher(N, M, h, hdim, grad_out, q, index, table, rel_idx, grad_q, grad_table);
}
void attention_step2_with_rel_pos_value_forward_cuda(int N, int M, int h, int hdim, at::Tensor attn_tensor, at::Tensor v_tensor,
at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor table_tensor, at::Tensor rel_idx_tensor, at::Tensor output_tensor)
{
const float *attn = attn_tensor.data_ptr<float>();
const float *v = v_tensor.data_ptr<float>();
const int *index0 = index0_tensor.data_ptr<int>();
const int *index1 = index1_tensor.data_ptr<int>();
const float *table = table_tensor.data_ptr<float>();
const int *rel_idx = rel_idx_tensor.data_ptr<int>();
float *output = output_tensor.data_ptr<float>();
attention_step2_with_rel_pos_value_forward_cuda_launcher(N, M, h, hdim, attn, v, index0, index1, table, rel_idx, output);
}
void attention_step2_with_rel_pos_value_backward_cuda(int N, int M, int h, int hdim, at::Tensor grad_out_tensor,
at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor, at::Tensor v_tensor, at::Tensor table_tensor,
at::Tensor rel_idx_tensor, at::Tensor grad_attn_tensor, at::Tensor grad_v_tensor, at::Tensor grad_table_tensor)
{
const float *grad_out = grad_out_tensor.data_ptr<float>();
const int *index0 = index0_tensor.data_ptr<int>();
const int *index1 = index1_tensor.data_ptr<int>();
const float *attn = attn_tensor.data_ptr<float>();
const float *v = v_tensor.data_ptr<float>();
const float *table = table_tensor.data_ptr<float>();
const int *rel_idx = rel_idx_tensor.data_ptr<int>();
float *grad_attn = grad_attn_tensor.data_ptr<float>();
float *grad_v = grad_v_tensor.data_ptr<float>();
float *grad_table = grad_table_tensor.data_ptr<float>();
attention_step2_with_rel_pos_value_backward_cuda_launcher(N, M, h, hdim, grad_out, index0, index1, attn, v, table, rel_idx, grad_attn, grad_v, grad_table);
}
|