Enhance the CachingMemoryManager to possibly handle high memory pressure

Enhance the CachingMemoryManager to possibly handle high memory
pressure by adding 2 options.
Note: the default behavior is unchanged.
Side: remove some glog includes failing compilation if glog-dev not
installed locally.

apply-format

@@ -0,0 +1,13 @@
+#! /bin/bash
+# Apply clang-format using the FL .clang-format on the modified/added cpp/h/hpp files of the current commit:
+clang-format --version || exit 1
+if [ ! -f .clang-format ]; then
+  echo "Cannot find .clang-fornmat in the current dir" 
+  exit 1
+fi
+for file in `git show --name-only -1 --pretty="format:" | grep -E '\.(cpp|h|hpp)$'`; do 
+  echo "Formatting $file ..."; 
+  clang-format -style=file -i $file; 
+done
+echo done
+

ext/CMakeLists.txt

@@ -7,6 +7,7 @@ include(${CMAKE_CURRENT_LIST_DIR}/common/CMakeLists.txt)
 
 # --------------------------- Tests ---------------------------
 
+# The SeqBuilder test needs some contribs modules
 if (FL_BUILD_TESTS)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/test)
 endif ()

ext/test/CMakeLists.txt

@@ -3,8 +3,10 @@ cmake_minimum_required(VERSION 3.5.1)
 set(DIR ${CMAKE_CURRENT_LIST_DIR})
 set(LIBS flashlight)
 
-build_test(
-  ${DIR}/common/SequentialBuilderTest.cpp
-  ${LIBS}
-  "ARCHDIR=\"${DIR}/common/\""
+if(FL_BUILD_CONTRIB)
+  build_test(
+    ${DIR}/common/SequentialBuilderTest.cpp
+    ${LIBS}
+    "ARCHDIR=\"${DIR}/common/\""
   )
+endif()

flashlight/memory/managers/CachingMemoryManager.cpp

@@ -10,7 +10,9 @@
 
 #include <algorithm>
 #include <iostream>
+#include <limits>
 #include <mutex>
+#include <numeric>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -78,7 +80,9 @@ CachingMemoryManager::DeviceMemoryInfo::DeviceMemoryInfo(int id)
 CachingMemoryManager::CachingMemoryManager(
     int numDevices,
     std::shared_ptr<MemoryManagerDeviceInterface> deviceInterface)
-    : MemoryManagerAdapter(deviceInterface) {
+    : MemoryManagerAdapter(deviceInterface),
+      recyclingSizeLimit(std::numeric_limits<size_t>::max()),
+      splitSizeLimit(std::numeric_limits<size_t>::max()) {
   for (int i = 0; i < numDevices; ++i) {
     deviceMemInfos_.emplace(
         i, fl::cpp::make_unique<CachingMemoryManager::DeviceMemoryInfo>(i));
@@ -87,6 +91,14 @@ CachingMemoryManager::CachingMemoryManager(
 
 void CachingMemoryManager::initialize() {}
 
+void CachingMemoryManager::setRecyclingSizeLimit(size_t limit) {
+  recyclingSizeLimit = limit;
+}
+
+void CachingMemoryManager::setSplitSizeLimit(size_t limit) {
+  splitSizeLimit = limit;
+}
+
 void CachingMemoryManager::shutdown() {
   signalMemoryCleanup();
 }
@@ -126,7 +138,9 @@ void* CachingMemoryManager::alloc(
 
   CachingMemoryManager::Block* block = nullptr;
   auto it = pool.lower_bound(&searchKey);
-  if (it != pool.end()) {
+  // Recycle blocks if any found, and if small alloc or the block size is not
+  // too large:
+  if (it != pool.end() && (isSmallAlloc || (*it)->size_ < recyclingSizeLimit)) {
     block = *it;
     pool.erase(it);
     memoryInfo.stats_.cachedBytes_ -= block->size_;
@@ -144,7 +158,10 @@ void* CachingMemoryManager::alloc(
   // implementation simple.
   CachingMemoryManager::Block* remaining = nullptr;
   size_t diff = block->size_ - size;
-  if (diff >= (isSmallAlloc ? kMinBlockSize : kSmallSize)) {
+  if ((diff >= (isSmallAlloc ? kMinBlockSize : kSmallSize)) &&
+      (block->size_ < splitSizeLimit) // possibly dont split large buffers to
+                                      // minimize risk of fragmentation
+  ) {
     remaining = block;
     block = new Block(size, block->ptr_);
     block->prev_ = remaining->prev_;

flashlight/memory/managers/CachingMemoryManager.h

@@ -7,8 +7,11 @@
 
 #pragma once
 
+#include <atomic>
+#include <limits>
 #include <memory>
 #include <mutex>
+#include <numeric>
 #include <set>
 #include <unordered_map>
 #include <vector>
@@ -49,6 +52,9 @@ class CachingMemoryManager : public MemoryManagerAdapter {
   void addMemoryManagement(int device) override;
   void removeMemoryManagement(int device) override;
 
+  void setRecyclingSizeLimit(size_t);
+  void setSplitSizeLimit(size_t);
+
   // Block denotes a single allocated unit of memory.
   struct Block {
     size_t size_; // size of block in bytes
@@ -133,6 +139,14 @@ class CachingMemoryManager : public MemoryManagerAdapter {
 
   void tryMergeBlocks(Block* dst, Block* src, BlockSet& freeBlocks);
   void freeBlock(Block* block);
+
+ private:
+  // Non-const runtime options in order to fine tune the behavior of this
+  // manager. Prevents to recycle some buffers, to be set by the user if
+  // desired:
+  std::atomic<size_t> recyclingSizeLimit;
+  // Prevents to split big buffers, to be set by the user if desired:
+  std::atomic<size_t> splitSizeLimit;
 };
 
 } // namespace fl

flashlight/test/memory/CachingMemoryManagerTest.cpp

@@ -20,9 +20,8 @@
 class CachingMemoryManagerTest : public ::testing::Test {
  protected:
   virtual void SetUp() override {
-    auto deviceInterface_ =
-        std::make_shared<fl::MemoryManagerDeviceInterface>();
-    auto adapter_ = std::make_shared<fl::CachingMemoryManager>(
+    deviceInterface_ = std::make_shared<fl::MemoryManagerDeviceInterface>();
+    adapter_ = std::make_shared<fl::CachingMemoryManager>(
         af::getDeviceCount(), deviceInterface_);
     installer_ = fl::cpp::make_unique<fl::MemoryManagerInstaller>(adapter_);
     installer_->setAsMemoryManager();
@@ -32,6 +31,8 @@ class CachingMemoryManagerTest : public ::testing::Test {
     af_unset_memory_manager();
   }
 
+  std::shared_ptr<fl::MemoryManagerDeviceInterface> deviceInterface_;
+  std::shared_ptr<fl::CachingMemoryManager> adapter_;
   std::unique_ptr<fl::MemoryManagerInstaller> installer_;
 };
 
@@ -137,6 +138,54 @@ TEST_F(CachingMemoryManagerTest, OOM) {
   }
 }
 
+void testFragmentation(
+    std::shared_ptr<fl::MemoryManagerDeviceInterface> deviceInterface_,
+    std::shared_ptr<fl::CachingMemoryManager> adapter_,
+    bool expectOOM) {
+  af::Backend b = af::getActiveBackend();
+
+  if (b != AF_BACKEND_CUDA) {
+    GTEST_SKIP()
+        << "CachingMemoryManager fragmentation tests require CUDA backend";
+  }
+
+  const auto mms = deviceInterface_->getMaxMemorySize(0);
+  const auto maxNumf32 = mms / sizeof(float); // AF f32 is supposed to be 32b
+  ASSERT_NE(mms, 0);
+  {
+    af::array a1(.5f * maxNumf32);
+    adapter_->printInfo("After creating a1:", 0);
+  } // The a1 buffer will not be freed here, just registered to the cache
+  adapter_->printInfo("After releasing a1:", 0);
+
+  af::array a2(.1f * maxNumf32);
+  adapter_->printInfo("After creating a2:", 0);
+
+  af::array a3;
+  try {
+    a3 = af::array(.5f * maxNumf32);
+  } catch (af::exception& ex) {
+    if (expectOOM) {
+      ASSERT_EQ(ex.err(), AF_ERR_NO_MEM);
+    } else {
+      EXPECT_TRUE(false)
+          << "CachingMemoryManagerTest fragmentaiton not supposed to throw: "
+          << ex.what();
+    }
+  }
+}
+
+TEST_F(CachingMemoryManagerTest, Fragmentation) {
+  testFragmentation(deviceInterface_, adapter_, true); // should OOM
+}
+
+TEST_F(CachingMemoryManagerTest, RecLimit) {
+  constexpr static size_t ONE_GB = 1 << 30;
+  // Fine set the manager in order not to recycle big tensors:
+  adapter_->setRecyclingSizeLimit(2 * ONE_GB);
+  testFragmentation(deviceInterface_, adapter_, false); // should not OOM
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();

lib/test/common/StringTest.cpp

@@ -7,7 +7,6 @@
 
 #include <string>
 
-#include <glog/logging.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 

lib/test/common/SystemTest.cpp

@@ -9,7 +9,6 @@
 #include <string>
 #include <vector>
 
-#include <glog/logging.h>
 #include <gtest/gtest.h>
 
 #include "flashlight/lib/common/System.h"
@@ -121,4 +120,4 @@ int main(int argc, char** argv) {
 #endif
 
   return RUN_ALL_TESTS();
-}
+}

lib/test/text/dictionary/DictionaryTest.cpp

@@ -7,7 +7,6 @@
 
 #include <string>
 
-#include <glog/logging.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>