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

Flatten batch to be checked

Browse Source
This commit is contained in:
Francesco Gatti
2020-04-23 01:18:27 +02:00
parent 488887992c
commit 8cff886ee5
2 changed files with 11 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/pluginsRT/FlattenConcatRT.h
+9 -5
View File
@@ -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;
}
tests/test_rtinference/rtinference.cpp
+2 -1
View File
@@ -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;