|
|
#include "../cuda_utils.h" |
|
|
#include "aggregation_cuda_kernel.h" |
|
|
|
|
|
|
|
|
__global__ void aggregation_forward_cuda_kernel(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output) { |
|
|
// input: input: (n, c), position: (n, nsample, c), weight: (n, nsample, w_c), idx: (n, nsample), output: (n, c) |
|
|
int index = blockIdx.x * blockDim.x + threadIdx.x; |
|
|
if (index >= n * c) return; |
|
|
const int c_idx = index % c; |
|
|
const int n_idx = index / c; |
|
|
const int w_c_idx = c_idx % w_c; |
|
|
for (int nsample_idx = 0; nsample_idx < nsample; nsample_idx++) |
|
|
{ |
|
|
int idx_idx = n_idx * nsample + nsample_idx; |
|
|
int input_idx = idx[idx_idx] * c + c_idx; |
|
|
int position_idx = n_idx * nsample * c + nsample_idx * c + c_idx; |
|
|
int weight_idx = n_idx * nsample * w_c + nsample_idx * w_c + w_c_idx; |
|
|
output[index] += (input[input_idx] + position[position_idx]) * weight[weight_idx]; |
|
|
} |
|
|
} |
|
|
|
|
|
__global__ void aggregation_backward_cuda_kernel(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight) { |
|
|
// input: grad_output: (n, c), output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight: (n, nsample, w_c) |
|
|
int index = blockIdx.x * blockDim.x + threadIdx.x; |
|
|
if (index >= n * c) return; |
|
|
const int c_idx = index % c; |
|
|
const int n_idx = index / c; |
|
|
const int w_c_idx = c_idx % w_c; |
|
|
for (int nsample_idx = 0; nsample_idx < nsample; nsample_idx++) |
|
|
{ |
|
|
int idx_idx = n_idx * nsample + nsample_idx; |
|
|
int input_idx = idx[idx_idx] * c + c_idx; |
|
|
int position_idx = n_idx * nsample * c + nsample_idx * c + c_idx; |
|
|
int weight_idx = n_idx * nsample * w_c + nsample_idx * w_c + w_c_idx; |
|
|
atomicAdd(grad_input + input_idx, grad_output[index] * weight[weight_idx]); |
|
|
grad_position[position_idx] = grad_output[index] * weight[weight_idx]; |
|
|
atomicAdd(grad_weight + weight_idx, grad_output[index] * (input[input_idx] + position[position_idx])); |
|
|
} |
|
|
} |
|
|
|
|
|
void aggregation_forward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output) { |
|
|
// input: input: (n, c), position: (n, nsample, c), weight: (n, nsample, w_c), idx: (n, nsample), output: (n, c) |
|
|
dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); |
|
|
dim3 threads(THREADS_PER_BLOCK); |
|
|
aggregation_forward_cuda_kernel<<<blocks, threads, 0>>>(n, nsample, c, w_c, input, position, weight, idx, output); |
|
|
} |
|
|
|
|
|
void aggregation_backward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight) { |
|
|
// input: grad_output: (n, c), output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight: (n, nsample, w_c) |
|
|
dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); |
|
|
dim3 threads(THREADS_PER_BLOCK); |
|
|
aggregation_backward_cuda_kernel<<<blocks, threads, 0>>>(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight); |
|
|
} |
|
|
|
