Spaces:
Runtime error
Runtime error
File size: 4,304 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
#include <vector>
#include <torch/serialize/tensor.h>
#include <ATen/cuda/CUDAContext.h>
#include "attention_cuda_kernel.h"
void attention_relation_step_forward_cuda(int m, int g, int c,
at::Tensor query_tensor, at::Tensor key_tensor, at::Tensor weight_tensor,
at::Tensor index_target_tensor, at::Tensor index_refer_tensor,
at::Tensor output_tensor)
{
const float *query = query_tensor.data_ptr<float>();
const float *key = key_tensor.data_ptr<float>();
const float *weight = weight_tensor.data_ptr<float>();
const int *index_target = index_target_tensor.data_ptr<int>();
const int *index_refer = index_refer_tensor.data_ptr<int>();
float *output = output_tensor.data_ptr<float>();
attention_relation_step_forward_cuda_launcher(m, g, c, query, key, weight, index_target, index_refer, output);
}
void attention_relation_step_backward_cuda(int m, int g, int c,
at::Tensor query_tensor, at::Tensor grad_query_tensor,
at::Tensor key_tensor, at::Tensor grad_key_tensor,
at::Tensor weight_tensor, at::Tensor grad_weight_tensor,
at::Tensor index_target_tensor, at::Tensor index_refer_tensor,
at::Tensor grad_output_tensor)
{
const float *query = query_tensor.data_ptr<float>();
float *grad_query = grad_query_tensor.data_ptr<float>();
const float *key = key_tensor.data_ptr<float>();
float *grad_key = grad_key_tensor.data_ptr<float>();
const float *weight = weight_tensor.data_ptr<float>();
float *grad_weight = grad_weight_tensor.data_ptr<float>();
const int *index_target = index_target_tensor.data_ptr<int>();
const int *index_refer = index_refer_tensor.data_ptr<int>();
const float *grad_output = grad_output_tensor.data_ptr<float>();
attention_relation_step_backward_cuda_launcher(m, g, c,
query, grad_query,
key, grad_key,
weight, grad_weight,
index_target, index_refer, grad_output);
}
void attention_fusion_step_forward_cuda(int m, int g, int c,
at::Tensor weight_tensor, at::Tensor value_tensor,
at::Tensor index_target_tensor, at::Tensor index_refer_tensor,
at::Tensor output_tensor)
{
const float *weight = weight_tensor.data_ptr<float>();
const float *value = value_tensor.data_ptr<float>();
const int *index_target = index_target_tensor.data_ptr<int>();
const int *index_refer = index_refer_tensor.data_ptr<int>();
float *output = output_tensor.data_ptr<float>();
attention_fusion_step_forward_cuda_launcher(m, g, c, weight, value, index_target, index_refer, output);
}
void attention_fusion_step_backward_cuda(int m, int g, int c,
at::Tensor weight_tensor, at::Tensor grad_weight_tensor,
at::Tensor value_tensor, at::Tensor grad_value_tensor,
at::Tensor index_target_tensor, at::Tensor index_refer_tensor,
at::Tensor grad_output_tensor)
{
const float *weight = weight_tensor.data_ptr<float>();
float *grad_weight = grad_weight_tensor.data_ptr<float>();
const float *value = value_tensor.data_ptr<float>();
float *grad_value = grad_value_tensor.data_ptr<float>();
const int *index_target = index_target_tensor.data_ptr<int>();
const int *index_refer = index_refer_tensor.data_ptr<int>();
const float *grad_output = grad_output_tensor.data_ptr<float>();
attention_fusion_step_backward_cuda_launcher(m, g, c,
weight, grad_weight,
value, grad_value,
index_target, index_refer, grad_output);
}
|