From 8cff886ee535eb5ba94f8545bbd48f1b72928d4d Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Thu, 23 Apr 2020 01:18:27 +0200
Subject: [PATCH] Flatten batch to be checked

---
 include/tkDNN/pluginsRT/FlattenConcatRT.h | 14 +++++++++-----
 tests/test_rtinference/rtinference.cpp    |  3 ++-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/include/tkDNN/pluginsRT/FlattenConcatRT.h b/include/tkDNN/pluginsRT/FlattenConcatRT.h
index e6e4eb2..4ebce69 100644
--- a/include/tkDNN/pluginsRT/FlattenConcatRT.h
+++ b/include/tkDNN/pluginsRT/FlattenConcatRT.h
@@ -47,11 +47,15 @@ public:
 	virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
 		dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
 		dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
-		checkCuda( cudaMemcpy(dstData, srcData, rows*cols*sizeof(dnnType), cudaMemcpyDeviceToDevice));
-	
-		float const alpha(1.0);
-		float const beta(0.0);
-		checkERROR( cublasSgeam( handle, CUBLAS_OP_T, CUBLAS_OP_N, rows, cols, &alpha, srcData, cols, &beta, srcData, rows, dstData, rows ));
+		checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*rows*cols*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
+
+		checkERROR( cublasSetStream(handle, stream) );	
+		for(int i=0; i<batchSize; i++) {
+			float const alpha(1.0);
+			float const beta(0.0);
+			int offset = i*rows*cols;
+			checkERROR( cublasSgeam( handle, CUBLAS_OP_T, CUBLAS_OP_N, rows, cols, &alpha, srcData + offset, cols, &beta, srcData + offset, rows, dstData + offset, rows ));
+		}
 		return 0;
 	}
 
diff --git a/tests/test_rtinference/rtinference.cpp b/tests/test_rtinference/rtinference.cpp
index 2b521ad..1ca9763 100644
--- a/tests/test_rtinference/rtinference.cpp
+++ b/tests/test_rtinference/rtinference.cpp
@@ -46,7 +46,8 @@ int main(int argc, char *argv[]) {
         TIMER_STOP
 
         // control output
-        for(int o=1; o<netRT.getBuffersN(); o++) {
+        std::cout<<"Output Buffers: "<<netRT.getBuffersN()-1<<"\n";
+	for(int o=1; o<netRT.getBuffersN(); o++) {
             for(int b=1; b<BATCH_SIZE; b++) {
                 dnnType *out_d = (dnnType*) netRT.buffersRT[o];
                 dnnType *out0_d = out_d;