diff options
| author |  Xiaoming (Jason) Cui <xiaoming.cui@intel.com> | 2018-09-25 00:42:42 -0700 |
|---|---|---|
| committer | Xiaoming (Jason) Cui <xiaoming.cui@intel.com> | 2018-09-28 12:23:18 -0700 |
| commit | b5feceb9058e06eac3de86ec45c44f5637054855 (patch) | |
| tree | 08d828f1033dacec9d476a3b9a5b88c3af13622d /tensorflow/core/util | |
| parent | 986193d79e00f1780fb3278ed890a72f7285f66e (diff) | |
Added the feature to disable MKL support of TensorFlow by environmental variable TF_DISABLE_MKL=1
Diffstat (limited to 'tensorflow/core/util')
| -rw-r--r-- | tensorflow/core/util/util.cc | 20 | ||||
| -rw-r--r-- | tensorflow/core/util/util.h | 5 |
2 files changed, 25 insertions, 0 deletions
diff --git a/tensorflow/core/util/util.cc b/tensorflow/core/util/util.cc index 1e5a9c5712..44d5becb9c 100644 --- a/tensorflow/core/util/util.cc +++ b/tensorflow/core/util/util.cc @@ -120,4 +120,24 @@ string SliceDebugString(const TensorShape& shape, const int64 flat) { return result; } +#ifdef INTEL_MKL +bool DisableMKL() { + enum MklStatus { + MKL_DEFAULT = 0, + MKL_ON = 1, + MKL_OFF = 2 + }; + static MklStatus status = MKL_DEFAULT; + if (status == MKL_DEFAULT) { + char* tf_disable_mkl = getenv("TF_DISABLE_MKL"); + if ((tf_disable_mkl != NULL) && (std::stoi(tf_disable_mkl) == 1)) { + VLOG(2) << "TF-MKL: Disabling MKL"; + status = MKL_OFF; + } else { + status = MKL_ON; + } + } + return status == MKL_OFF ? true : false; +} +#endif } // namespace tensorflow diff --git a/tensorflow/core/util/util.h b/tensorflow/core/util/util.h index 93dfd51ab5..ba90ad52c2 100644 --- a/tensorflow/core/util/util.h +++ b/tensorflow/core/util/util.h @@ -56,6 +56,11 @@ string PrintMemory(const char* ptr, size_t n); // "tensor", "tensor[i]", "tensor[i, j]", etc. string SliceDebugString(const TensorShape& shape, const int64 flat); +// disable MKL in runtime +#ifdef INTEL_MKL +bool DisableMKL(); +#endif + } // namespace tensorflow #endif // TENSORFLOW_CORE_UTIL_UTIL_H_ |