stream in TRT plugin

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