Hugging Face's logo

Spaces:
ziqima
/
Find3D
Running on Zero

App Files Community
1
Find3D
File size: 2,074 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
 
#include "../cuda_utils.h"
#include "interpolation_cuda_kernel.h"


__global__ void interpolation_forward_cuda_kernel(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output)
{
    // input: input: (m, c), idx: (n, k), weight: (n, k), output: output (n, c)
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= n * c) return;
    int c_idx = index % c;
    int n_idx = index / c;
    for (int i = 0; i < k; i++)
    {
        int idx_idx = n_idx * k + i;
        int input_idx = idx[idx_idx] * c + c_idx;
        output[index] += input[input_idx] * weight[idx_idx];
    }
}

__global__ void interpolation_backward_cuda_kernel(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input)
{
    // input: grad_output: (n, c), idx: (n, k), weight: (n, k), output: grad_input (m, c)
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= n * c) return;
    int c_idx = index % c;
    int n_idx = index / c;
    for (int i = 0; i < k; i++)
    {
        int idx_idx = n_idx * k + i;
        int input_idx = idx[idx_idx] * c + c_idx;
        atomicAdd(grad_input + input_idx, grad_output[index] * weight[idx_idx]);
    }
}

void interpolation_forward_cuda_launcher(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output) {
    // input: input: (m, c), idx: (n, k), weight: (n, k), output: output (n, c)
    dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK));
    dim3 threads(THREADS_PER_BLOCK);
    interpolation_forward_cuda_kernel<<<blocks, threads, 0>>>(n, c, k, input, idx, weight, output);
}

void interpolation_backward_cuda_launcher(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input) {
    // input: grad_output: (n, c), idx: (n, k), weight: (n, k), output: grad_input (m, c)
    dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK));
    dim3 threads(THREADS_PER_BLOCK);
    interpolation_backward_cuda_kernel<<<blocks, threads, 0>>>(n, c, k, grad_output, idx, weight, grad_input);
}