stream in TRT plugin

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

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<<<blocks, threads>>>(srcData, dstData, size);
+    activation_elu<<<blocks, threads, 0, stream>>>(srcData, dstData, size);
 }

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<<<blocks, threads>>>(srcData, dstData, size);
+    activation_leaky<<<blocks, threads, 0, stream>>>(srcData, dstData, size);
 }
 
 

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<<<blocks, threads>>>(srcData, dstData, size);
+    activation_logistic<<<blocks, threads, 0, stream>>>(srcData, dstData, size);
 }
 
 

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<<<blocks, threads>>>(size, srcData, w, h, c, n, stride, false, dstData);
+    reorg_kernel<<<blocks, threads, 0, stream>>>(size, srcData, w, h, c, n, stride, false, dstData);
 }
 
 

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<<<blocks, threads>>>(input, n, batch, batch_offset, groups, group_offset, stride, temp, output);
+    softmax_kernel<<<blocks, threads, 0, stream>>>(input, n, batch, batch_offset, groups, group_offset, stride, temp, output);
 }

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<const dnnType*>(inputs[0]), 
-											reinterpret_cast<dnnType*>(outputs[0]), size);
+											reinterpret_cast<dnnType*>(outputs[0]), size, stream);
 		return 0;
 	}
 

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;
 	}

src/pluginsRT/ReorgRT.cpp

@@ -42,7 +42,7 @@ public:
 
 		reorgForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]), 
 					  reinterpret_cast<dnnType*>(outputs[0]), 
-					  batchSize, c, h, w, stride);
+					  batchSize, c, h, w, stride, stream);
 		return 0;
 	}