experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

minor bug fixes for unit tests.

Browse Source
This commit is contained in:
Srinivasan Narayanamoorthy 2020-05-27 14:48:52 -07:00
parent be46769cee
commit 0101b7cad5
2 changed files with 5 additions and 1 deletions
tensorflow/core/util
mkl_threadpool.h
third_party/mkl_dnn
mkldnn_v1.BUILD

3
tensorflow/core/util/mkl_threadpool.h
View File

@ -24,7 +24,6 @@ limitations under the License.
#include <unordered_map>
#include <utility>
#include <vector>
#include "mkldnn.hpp"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/platform/threadpool.h"
@ -114,6 +113,7 @@ class MklDnnThreadPoolWrapper {
return instance_;
}
MklDnnThreadPool* CreateThreadPoolPtr(OpKernelContext* ctx) {
mutex_lock l(m_);
if (threadpool_map_.empty() ||
threadpool_map_.find(ctx->device()) == threadpool_map_.end()) {
auto tp_iface = new MklDnnThreadPool(ctx);
@ -126,6 +126,7 @@ class MklDnnThreadPoolWrapper {
}
private:
mutex m_;
std::unordered_map<DeviceBase*, MklDnnThreadPool*> threadpool_map_;
MklDnnThreadPoolWrapper() {}
MklDnnThreadPoolWrapper(const MklDnnThreadPoolWrapper&) = delete;

3
third_party/mkl_dnn/mkldnn_v1.BUILD vendored
View File

@ -86,6 +86,9 @@ cc_library(
]) + if_mkl_v1_open_source_only([
"-UUSE_MKL",
"-UUSE_CBLAS",
]) + if_mkldnn_threadpool([
"-UUSE_MKL",
"-UUSE_CBLAS",
]) + select({
"@org_tensorflow//tensorflow:linux_x86_64": [
"-fopenmp", # only works with gcc