1 files changed, 61 insertions, 0 deletions
diff --git a/tensorflow/core/framework/allocator_test.cc b/tensorflow/core/framework/allocator_test.cc
new file mode 100644
index 0000000000..6b1e52cfc4
--- /dev/null
+++ b/tensorflow/core/framework/allocator_test.cc
@@ -0,0 +1,61 @@
+#include "tensorflow/core/framework/allocator.h"
+#include <algorithm>
+#include "tensorflow/core/platform/logging.h"
+#include <gtest/gtest.h>
+namespace tensorflow {
+
+TEST(CPUAllocatorTest, Simple) {
+  Allocator* a = cpu_allocator();
+  std::vector<void*> ptrs;
+  for (int s = 1; s < 1024; s++) {
+    void* raw = a->AllocateRaw(1, s);
+    ptrs.push_back(raw);
+  }
+  std::sort(ptrs.begin(), ptrs.end());
+  for (size_t i = 0; i < ptrs.size(); i++) {
+    if (i > 0) {
+      CHECK_NE(ptrs[i], ptrs[i - 1]);  // No dups
+    }
+    a->DeallocateRaw(ptrs[i]);
+  }
+  float* t1 = a->Allocate<float>(1024);
+  double* t2 = a->Allocate<double>(1048576);
+  a->Deallocate(t1);
+  a->Deallocate(t2);
+}
+
+// Define a struct that we will use to observe behavior in the unit tests
+struct TestStruct {
+  int x;  // not used just want to make sure sizeof(TestStruct) > 1
+};
+
+TEST(CPUAllocatorTest, CheckStructSize) { CHECK_GT(sizeof(TestStruct), 1); }
+
+TEST(CPUAllocatorTest, AllocateOverflowMaxSizeT) {
+  Allocator* a = cpu_allocator();
+
+  // The maximum size_t value will definitely overflow.
+  size_t count_to_allocate = std::numeric_limits<size_t>::max();
+  TestStruct* const test_pointer = a->Allocate<TestStruct>(count_to_allocate);
+
+  CHECK_EQ(test_pointer, reinterpret_cast<TestStruct*>(NULL));
+}
+
+TEST(CPUAllocatorTest, AllocateOverflowSmallest) {
+  Allocator* a = cpu_allocator();
+
+  // count_to_allocate is the smallest count that will cause overflow.
+  const size_t count_to_allocate =
+      (std::numeric_limits<size_t>::max() / sizeof(TestStruct)) + 1;
+  TestStruct* const test_pointer = a->Allocate<TestStruct>(count_to_allocate);
+
+  CHECK_EQ(test_pointer, reinterpret_cast<TestStruct*>(NULL));
+}
+
+TEST(CPUAllocatorTest, Sizes) {
+  Allocator* a = cpu_allocator();
+
+  EXPECT_EQ(false, a->TracksAllocationSizes());
+}
+
+}  // namespace tensorflow