diff --git a/include/kernels.h b/include/kernels.h index 088e235..1f4df50 100644 --- a/include/kernels.h +++ b/include/kernels.h @@ -1,10 +1,15 @@ +#ifndef KERNELS_H +#define KERNELS_H + #include "utils.h" -void activationELUForward(dnnType* srcData, dnnType* dstData, int size); -void activationLEAKYForward(dnnType* srcData, dnnType* dstData, int size); -void activationLOGISTICForward(dnnType* srcData, dnnType* dstData, int size); +void activationELUForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream = cudaStream_t(0)); +void activationLEAKYForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream = cudaStream_t(0)); +void activationLOGISTICForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream = cudaStream_t(0)); void reorgForward( dnnType* srcData, dnnType* dstData, - int n, int c, int h, int w, int stride); + int n, int c, int h, int w, int stride, cudaStream_t stream = cudaStream_t(0)); void softmaxForward(float *input, int n, int batch, int batch_offset, - int groups, int group_offset, int stride, float temp, float *output); + int groups, int group_offset, int stride, float temp, float *output, cudaStream_t stream = cudaStream_t(0)); + +#endif //KERNELS_H diff --git a/src/kernels/activation_elu.cu b/src/kernels/activation_elu.cu index 5d3ef1d..219541d 100644 --- a/src/kernels/activation_elu.cu +++ b/src/kernels/activation_elu.cu @@ -28,10 +28,10 @@ void activation_elu(dnnType *input, dnnType *output, int size) { /** ELU activation function */ -void activationELUForward(dnnType* srcData, dnnType* dstData, int size) +void activationELUForward(dnnType* srcData, dnnType* dstData, int size, const cudaStream_t stream) { int blocks = (size+255)/256; int threads = 256; - activation_elu<<>>(srcData, dstData, size); + activation_elu<<>>(srcData, dstData, size); } diff --git a/src/kernels/activation_leaky.cu b/src/kernels/activation_leaky.cu index 4b0d3e2..a9029ad 100644 --- a/src/kernels/activation_leaky.cu +++ b/src/kernels/activation_leaky.cu @@ -17,12 +17,12 @@ void activation_leaky(dnnType *input, dnnType *output, int size) { /** ELU activation function */ -void activationLEAKYForward(dnnType* srcData, dnnType* dstData, int size) +void activationLEAKYForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream) { int blocks = (size+255)/256; int threads = 256; - activation_leaky<<>>(srcData, dstData, size); + activation_leaky<<>>(srcData, dstData, size); } diff --git a/src/kernels/activation_logistic.cu b/src/kernels/activation_logistic.cu index bd2a1e9..a1ca770 100644 --- a/src/kernels/activation_logistic.cu +++ b/src/kernels/activation_logistic.cu @@ -14,12 +14,12 @@ void activation_logistic(dnnType *input, dnnType *output, int size) { /** LOGISTIC activation function */ -void activationLOGISTICForward(dnnType* srcData, dnnType* dstData, int size) +void activationLOGISTICForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream) { int blocks = (size+255)/256; int threads = 256; - activation_logistic<<>>(srcData, dstData, size); + activation_logistic<<>>(srcData, dstData, size); } diff --git a/src/kernels/reorg.cu b/src/kernels/reorg.cu index 7afdc2a..dc98d0d 100644 --- a/src/kernels/reorg.cu +++ b/src/kernels/reorg.cu @@ -36,14 +36,14 @@ __global__ void reorg_kernel(int N, float *x, int w, int h, int c, int batch, in reorg function function */ void reorgForward(dnnType* srcData, dnnType* dstData, - int n, int c, int h, int w, int stride) { + int n, int c, int h, int w, int stride, cudaStream_t stream) { int size = n*c*h*w; int blocks = (size+255)/256; int threads = 256; - reorg_kernel<<>>(size, srcData, w, h, c, n, stride, false, dstData); + reorg_kernel<<>>(size, srcData, w, h, c, n, stride, false, dstData); } diff --git a/src/kernels/softmax.cu b/src/kernels/softmax.cu index c746b24..fb5fcd9 100644 --- a/src/kernels/softmax.cu +++ b/src/kernels/softmax.cu @@ -32,11 +32,11 @@ __global__ void softmax_kernel(float *input, int n, int batch, int batch_offset, softmax function */ void softmaxForward(float *input, int n, int batch, int batch_offset, - int groups, int group_offset, int stride, float temp, float *output) + int groups, int group_offset, int stride, float temp, float *output, cudaStream_t stream) { int size = groups*batch; int blocks = (size+255)/256; int threads = 256; - softmax_kernel<<>>(input, n, batch, batch_offset, groups, group_offset, stride, temp, output); + softmax_kernel<<>>(input, n, batch, batch_offset, groups, group_offset, stride, temp, output); } diff --git a/src/pluginsRT/ActivationLeakyRT.cpp b/src/pluginsRT/ActivationLeakyRT.cpp index c693730..651eb36 100644 --- a/src/pluginsRT/ActivationLeakyRT.cpp +++ b/src/pluginsRT/ActivationLeakyRT.cpp @@ -42,7 +42,7 @@ public: virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override { activationLEAKYForward((dnnType*)reinterpret_cast(inputs[0]), - reinterpret_cast(outputs[0]), size); + reinterpret_cast(outputs[0]), size, stream); return 0; } diff --git a/src/pluginsRT/RegionRT.cpp b/src/pluginsRT/RegionRT.cpp index 844199e..eec4771 100644 --- a/src/pluginsRT/RegionRT.cpp +++ b/src/pluginsRT/RegionRT.cpp @@ -52,10 +52,10 @@ public: for (int b = 0; b < batchSize; ++b){ for(int n = 0; n < num; ++n){ int index = entry_index(b, n*w*h, 0, batchSize); - activationLOGISTICForward(srcData + index, dstData + index, 2*w*h); + activationLOGISTICForward(srcData + index, dstData + index, 2*w*h, stream); index = entry_index(b, n*w*h, coords, batchSize); - activationLOGISTICForward(srcData + index, dstData + index, w*h); + activationLOGISTICForward(srcData + index, dstData + index, w*h, stream); } } @@ -63,7 +63,7 @@ public: int index = entry_index(0, 0, coords + 1, batchSize); softmaxForward( srcData + index, classes, batchSize*num, (batchSize*c*h*w)/num, - w*h, 1, w*h, 1, dstData + index); + w*h, 1, w*h, 1, dstData + index, stream); return 0; } diff --git a/src/pluginsRT/ReorgRT.cpp b/src/pluginsRT/ReorgRT.cpp index 766c33a..698101b 100644 --- a/src/pluginsRT/ReorgRT.cpp +++ b/src/pluginsRT/ReorgRT.cpp @@ -42,7 +42,7 @@ public: reorgForward((dnnType*)reinterpret_cast(inputs[0]), reinterpret_cast(outputs[0]), - batchSize, c, h, w, stride); + batchSize, c, h, w, stride, stream); return 0; }