diff --git a/CMakeLists.txt b/CMakeLists.txt index 3ad5932..47780ea 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -29,13 +29,9 @@ if(NVINFER STREQUAL "NVINFER-NOTFOUND") /usr/local/cuda/targets/x86_64-linux/lib/) endif() +file(GLOB tkdnn_CUSRC "src/kernels/*.cu") cuda_include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${NVINFER_INCLUDES}) -cuda_add_library(kernels SHARED src/kernels/activation_elu.cu - src/kernels/activation_leaky.cu - src/kernels/activation_logistic.cu - src/kernels/reorg.cu - src/kernels/softmax.cu - src/kernels/convert.cu) +cuda_add_library(kernels SHARED ${tkdnn_CUSRC}) file(GLOB tkdnn_SRC "src/*.cpp") set(tkdnn_LIBS kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} -lcudnn -lnvinfer ${OpenCV_LIBS}) diff --git a/include/kernels.h b/include/kernels.h index 4c38193..632ca2f 100644 --- a/include/kernels.h +++ b/include/kernels.h @@ -13,5 +13,9 @@ void softmaxForward(float *input, int n, int batch, int batch_offset, int groups, int group_offset, int stride, float temp, float *output, cudaStream_t stream = cudaStream_t(0)); +void shortcutForward(dnnType* srcData, dnnType* dstData, int n1, int c1, int h1, int w1, int s1, + int n2, int c2, int h2, int w2, int s2, + cudaStream_t stream = cudaStream_t(0)); + void float2half(float* srcData, __half* dstData, int size, const cudaStream_t stream = cudaStream_t(0)); #endif //KERNELS_H diff --git a/src/kernels/shortcut.cu b/src/kernels/shortcut.cu new file mode 100644 index 0000000..83539a8 --- /dev/null +++ b/src/kernels/shortcut.cu @@ -0,0 +1,47 @@ +#include "kernels.h" +#include "assert.h" + +__global__ void shortcut_kernel(int size, int minw, int minh, int minc, int stride, int sample, int batch, + int w1, int h1, int c1, dnnType *add, + int w2, int h2, int c2, float s1, float s2, dnnType *out) +{ + int id = (blockIdx.x + blockIdx.y*gridDim.x) * blockDim.x + threadIdx.x; + if (id >= size) return; + int i = id % minw; + id /= minw; + int j = id % minh; + id /= minh; + int k = id % minc; + id /= minc; + int b = id % batch; + + int out_index = i*sample + w2*(j*sample + h2*(k + c2*b)); + int add_index = i*stride + w1*(j*stride + h1*(k + c1*b)); + out[out_index] = s1*out[out_index] + s2*add[add_index]; + //out[out_index] += add[add_index]; +} + +void shortcutForward(dnnType* srcData, dnnType* dstData, int n1, int c1, int h1, int w1, int s1, + int n2, int c2, int h2, int w2, int s2, + cudaStream_t stream) +{ + assert(n1 == n2); + int batch = n1; + + int minw = (w1 < w2) ? w1 : w2; + int minh = (h1 < h2) ? h1 : h2; + int minc = (c1 < c2) ? c1 : c2; + + int stride = w1/w2; + int sample = w2/w1; + assert(stride == h1/h2); + assert(sample == h2/h1); + if(stride < 1) stride = 1; + if(sample < 1) sample = 1; + + int size = batch * minw * minh * minc; + int blocks = (size+255)/256; + int threads = 256; + shortcut_kernel<<>>(size, minw, minh, minc, stride, sample, batch, + w1, h1, c1, srcData, w2, h2, c2, s1, s2, dstData); +}