diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-09-25 22:57:54 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-25 23:02:36 -0700 |
| commit | c63d21b0bfc534b6377b332e9d2ba2abbdb7e0eb (patch) | |
| tree | 6fb44dc9eb26de3ff6e25f57e4799ac2da066d95 /tensorflow/core/common_runtime | |
| parent | f2b17b22e12bd743b66945070f338f70b5fa3332 (diff) | |
Adds a build flag to enable MKL (mkl_enabled=true).
PiperOrigin-RevId: 214557082
Diffstat (limited to 'tensorflow/core/common_runtime')
3 files changed, 10 insertions, 7 deletions
diff --git a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc index 2ed4f69f90..efd6185f8b 100644 --- a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc +++ b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc @@ -108,7 +108,7 @@ TEST(DirectSessionWithTrackingAllocTest, CostModelTest) { EXPECT_EQ(2, shape.dim(0).size()); EXPECT_EQ(1, shape.dim(1).size()); if (node->name() == y->name()) { -#ifdef INTEL_MKL +#if defined(INTEL_MKL) && defined(ENABLE_MKL) // if MKL is used, it goes through various additional // graph rewrite pass. In TF, everytime a graph pass // happens, "constant" nodes are allocated @@ -120,13 +120,13 @@ TEST(DirectSessionWithTrackingAllocTest, CostModelTest) { EXPECT_EQ(29, cm->AllocationId(node, 0)); #else EXPECT_EQ(21, cm->AllocationId(node, 0)); -#endif +#endif // INTEL_MKL && ENABLE_MKL } else { -#ifdef INTEL_MKL +#if defined(INTEL_MKL) && defined(ENABLE_MKL) EXPECT_EQ(30, cm->AllocationId(node, 0)); #else EXPECT_EQ(22, cm->AllocationId(node, 0)); -#endif +#endif // INTEL_MKL && ENABLE_MKL } } EXPECT_LE(0, cm->MaxExecutionTime(node)); diff --git a/tensorflow/core/common_runtime/mkl_cpu_allocator_test.cc b/tensorflow/core/common_runtime/mkl_cpu_allocator_test.cc index a67411cd2e..e08ab57638 100644 --- a/tensorflow/core/common_runtime/mkl_cpu_allocator_test.cc +++ b/tensorflow/core/common_runtime/mkl_cpu_allocator_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifdef INTEL_MKL +#if defined(INTEL_MKL) && defined(ENABLE_MKL) #include "tensorflow/core/common_runtime/mkl_cpu_allocator.h" @@ -50,4 +50,4 @@ TEST(MKLBFCAllocatorTest, TestMaxLimit) { } // namespace tensorflow -#endif // INTEL_MKL +#endif // INTEL_MKL && ENABLE_MKL diff --git a/tensorflow/core/common_runtime/threadpool_device.cc b/tensorflow/core/common_runtime/threadpool_device.cc index 0fbc20b34b..8587d1783a 100644 --- a/tensorflow/core/common_runtime/threadpool_device.cc +++ b/tensorflow/core/common_runtime/threadpool_device.cc @@ -113,8 +113,11 @@ class MklCPUAllocatorFactory : public AllocatorFactory { } }; +#ifdef ENABLE_MKL REGISTER_MEM_ALLOCATOR("MklCPUAllocator", 200, MklCPUAllocatorFactory); +#endif // ENABLE_MKL + } // namespace -#endif +#endif // INTEL_MKL } // namespace tensorflow |