Switch open-source to use jemalloc for CPU Tensor memory allocation, gRPC, and other places where we call malloc/free.

- Only enabled on Linux for now.
- Added as a ./configure option defaulting to enabled.
Change: 144266237

1 files changed, 3 insertions, 2 deletions

diff --git a/tensorflow/core/framework/tracking_allocator_test.cc b/tensorflow/core/framework/tracking_allocator_test.cc
index 98134392ef..850cdc3909 100644
--- a/tensorflow/core/framework/tracking_allocator_test.cc
+++ b/tensorflow/core/framework/tracking_allocator_test.cc
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/mem.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
@@ -27,7 +28,7 @@ class TestableSizeTrackingAllocator : public Allocator {
  public:
   string Name() override { return "test"; }
   void* AllocateRaw(size_t /*alignment*/, size_t num_bytes) override {
-    void* ptr = malloc(num_bytes);
+    void* ptr = port::Malloc(num_bytes);
     size_map_[ptr] = num_bytes;
     return ptr;
   }
@@ -35,7 +36,7 @@ class TestableSizeTrackingAllocator : public Allocator {
     const auto& iter = size_map_.find(ptr);
     EXPECT_NE(size_map_.end(), iter);
     size_map_.erase(iter);
-    free(ptr);
+    port::Free(ptr);
   }
   bool TracksAllocationSizes() override { return true; }
   size_t RequestedSize(void* ptr) override {