Francesco Gatti
36 lines
1.1 KiB
Plaintext
36 lines
1.1 KiB
Plaintext
#include "kernels.h"
|
|
|
|
__global__ void upsample_kernel(size_t N, dnnType *x, int w, int h, int c, int batch, int stride, int forward, float scale, dnnType *out)
|
|
{
|
|
size_t i = (blockIdx.x + blockIdx.y*gridDim.x) * blockDim.x + threadIdx.x;
|
|
if(i >= N) return;
|
|
int out_index = i;
|
|
int out_w = i%(w*stride);
|
|
i = i/(w*stride);
|
|
int out_h = i%(h*stride);
|
|
i = i/(h*stride);
|
|
int out_c = i%c;
|
|
i = i/c;
|
|
int b = i%batch;
|
|
|
|
int in_w = out_w / stride;
|
|
int in_h = out_h / stride;
|
|
int in_c = out_c;
|
|
|
|
int in_index = b*w*h*c + in_c*w*h + in_h*w + in_w;
|
|
|
|
|
|
if(forward) out[out_index] += scale * x[in_index];
|
|
else atomicAdd(x+in_index, scale * out[out_index]);
|
|
}
|
|
|
|
void upsampleForward(dnnType* srcData, dnnType* dstData,
|
|
int n, int c, int h, int w, int s, int forward, float scale,
|
|
cudaStream_t stream) {
|
|
|
|
int size = w*h*c*n*s*s;
|
|
int blocks = (size+255)/256;
|
|
int threads = 256;
|
|
upsample_kernel<<<blocks, threads, 0, stream>>>(size, srcData, w, h, c, n, s, forward, scale, dstData);
|
|
}
|