Brutex
/
cavis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mkldnn concat call cases correction (#471) 

* - disable mkldnn concat when number of input arrays > 3072

Signed-off-by: Yurii <iuriish@yahoo.com>

* - get rid of loop in calculating of input arrays number

Signed-off-by: Yurii <iuriish@yahoo.com>
master
Yurii Shyrma 2020-05-19 21:56:41 +03:00 committed by GitHub
parent ec757f654d
commit 0bc9785508
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
libnd4j/include/ops/declarable/platform/mkldnn/concat.cpp
View File

@ -178,7 +178,11 @@ PLATFORM_CHECK(concat, ENGINE_CPU) {
const auto zType = z->dataType(); const auto zType = z->dataType();
return z->rankOf() < 7 && (zType==DataType::FLOAT32 || zType==DataType::HALF || zType==DataType::BFLOAT16 || zType==DataType::UINT8 || zType==DataType::INT8); const bool isAxisInLastArr = block.getBArguments()->size() == 0 ? false : B_ARG(0);
const int numOfInArrs = isAxisInLastArr ? block.width() - 1 : block.width();
return z->rankOf() < 7 && numOfInArrs <= 3072
&& (zType==DataType::FLOAT32 || zType==DataType::HALF || zType==DataType::BFLOAT16 || zType==DataType::UINT8 || zType==DataType::INT8);
} }
} }