diff options
Diffstat (limited to 'tensorflow/core/common_runtime/threadpool_device.cc')
| -rw-r--r-- | tensorflow/core/common_runtime/threadpool_device.cc | 9 |
1 files changed, 9 insertions, 0 deletions
diff --git a/tensorflow/core/common_runtime/threadpool_device.cc b/tensorflow/core/common_runtime/threadpool_device.cc index 686bc6885e..ca6ba7970f 100644 --- a/tensorflow/core/common_runtime/threadpool_device.cc +++ b/tensorflow/core/common_runtime/threadpool_device.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/local_device.h" #include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/allocator_registry.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.pb_text.h" @@ -27,6 +28,10 @@ limitations under the License. #include "tensorflow/core/platform/types.h" #include "tensorflow/core/public/session_options.h" +#ifdef INTEL_MKL +#include "tensorflow/core/common_runtime/mkl_cpu_allocator.h" +#endif + namespace tensorflow { ThreadPoolDevice::ThreadPoolDevice(const SessionOptions& options, @@ -70,4 +75,8 @@ Status ThreadPoolDevice::MakeTensorFromProto( ProtoDebugString(tensor_proto)); } +#ifdef INTEL_MKL +REGISTER_MEM_ALLOCATOR("MklCPUAllocator", 200, MklCPUAllocator); +#endif + } // namespace tensorflow |