From 1baac7862739525351d25202800dc04e8ec3868b Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 8 Feb 2018 10:20:26 -0800 Subject: [PATCH] Make MklCpuAllocator a VisitableAllocator, instead of just an Allocator. This allows it to work more efficiently with RDMA networking. PiperOrigin-RevId: 185013628 --- tensorflow/core/common_runtime/mkl_cpu_allocator.h | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/common_runtime/mkl_cpu_allocator.h b/tensorflow/core/common_runtime/mkl_cpu_allocator.h index c7a2b61..0eb47f4 100644 --- a/tensorflow/core/common_runtime/mkl_cpu_allocator.h +++ b/tensorflow/core/common_runtime/mkl_cpu_allocator.h @@ -46,7 +46,7 @@ class MklSubAllocator : public SubAllocator { /// CPU allocator for MKL that wraps BFC allocator and intercepts /// and redirects memory allocation calls from MKL. -class MklCPUAllocator : public Allocator { +class MklCPUAllocator : public VisitableAllocator { public: // Constructor and other standard functions @@ -119,6 +119,14 @@ class MklCPUAllocator : public Allocator { void ClearStats() override { allocator_->ClearStats(); } + void AddAllocVisitor(Visitor visitor) override { + allocator_->AddAllocVisitor(visitor); + } + + void AddFreeVisitor(Visitor visitor) override { + allocator_->AddFreeVisitor(visitor); + } + private: // Hooks provided by this allocator for memory allocation routines from MKL -- 2.7.4