mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

stream in TRT plugin

Browse Source
This commit is contained in:
Francesco Gatti
2017-08-10 19:21:15 +02:00
parent 9a6058ac4a
commit 3124f86878
9 changed files with 25 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernels/reorg.cu
+2 -2
View File
@@ -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);
}