Add more size classes

velox/common/memory/Memory.cpp

@@ -49,6 +49,7 @@ std::shared_ptr<MemoryAllocator> createAllocator(
   if (options.useMmapAllocator) {
     MmapAllocator::Options mmapOptions;
     mmapOptions.capacity = options.allocatorCapacity;
+    mmapOptions.largestSizeClass = options.largestSizeClassPages;
     mmapOptions.useMmapArena = options.useMmapArena;
     mmapOptions.mmapArenaCapacityRatio = options.mmapArenaCapacityRatio;
     return std::make_shared<MmapAllocator>(mmapOptions);

velox/common/memory/Memory.h

@@ -92,6 +92,9 @@ struct MemoryManagerOptions {
   /// std::malloc.
   bool useMmapAllocator{false};
 
+  // Number of pages in largest size class in MmapAllocator.
+  int32_t largestSizeClassPages{256};
+
   /// If true, allocations larger than largest size class size will be delegated
   /// to ManagedMmapArena. Otherwise a system mmap call will be issued for each
   /// such allocation.

velox/common/memory/MemoryAllocator.cpp

@@ -29,6 +29,18 @@ DECLARE_bool(velox_memory_use_hugepages);
 
 namespace facebook::velox::memory {
 
+// static
+std::vector<MachinePageCount> MemoryAllocator::makeSizeClassSizes(
+    MachinePageCount largest) {
+  VELOX_CHECK_LE(256, largest);
+  VELOX_CHECK_EQ(largest, bits::nextPowerOfTwo(largest));
+  std::vector<MachinePageCount> sizes;
+  for (auto size = 1; size <= largest; size *= 2) {
+    sizes.push_back(size);
+  }
+  return sizes;
+}
+
 namespace {
 std::string& cacheFailureMessage() {
   thread_local std::string message;

velox/common/memory/MemoryAllocator.h

@@ -406,7 +406,11 @@ class MemoryAllocator : public std::enable_shared_from_this<MemoryAllocator> {
       std::function<int64_t()> ioVolume = nullptr);
 
  protected:
-  explicit MemoryAllocator() = default;
+  MemoryAllocator(MachinePageCount largestSizeClassPages = 256)
+      : sizeClassSizes_(makeSizeClassSizes(largestSizeClassPages)) {}
+
+  static std::vector<MachinePageCount> makeSizeClassSizes(
+      MachinePageCount largest);
 
   /// Represents a mix of blocks of different sizes for covering a single
   /// allocation.

velox/common/memory/MmapAllocator.cpp

@@ -25,7 +25,8 @@
 
 namespace facebook::velox::memory {
 MmapAllocator::MmapAllocator(const Options& options)
-    : kind_(MemoryAllocator::Kind::kMmap),
+    : MemoryAllocator(options.largestSizeClass),
+      kind_(MemoryAllocator::Kind::kMmap),
       useMmapArena_(options.useMmapArena),
       maxMallocBytes_(options.maxMallocBytes),
       mallocReservedBytes_(

velox/common/memory/MmapAllocator.h

@@ -55,6 +55,8 @@ class MmapAllocator : public MemoryAllocator {
     ///  Capacity in bytes, default unlimited.
     uint64_t capacity{kMaxMemory};
 
+    int32_t largestSizeClass{256};
+
     /// If set true, allocations larger than largest size class size will be
     /// delegated to ManagedMmapArena. Otherwise a system mmap call will be
     /// issued for each such allocation.

velox/common/memory/tests/MemoryAllocatorTest.cpp

@@ -1943,4 +1943,50 @@ VELOX_INSTANTIATE_TEST_SUITE_P(
     MemoryAllocatorTestSuite,
     MemoryAllocatorTest,
     testing::ValuesIn({0, 1, 2}));
+
+class MmapConfigTest : public testing::Test {
+ public:
+ protected:
+  void setupAllocator() {
+    constexpr int64_t kCapacityBytes = 900LL << 20; // 900MB.
+    MemoryManagerOptions options;
+    options.useMmapAllocator = true;
+    options.allocatorCapacity = kCapacityBytes;
+    options.largestSizeClassPages = 4096;
+    options.arbitratorCapacity = kCapacityBytes;
+    options.arbitratorReservedCapacity = 128 << 20;
+    options.memoryPoolReservedCapacity = 1 << 20;
+    options.smallAllocationReservePct = 4;
+    options.maxMallocBytes = 3 * 1024;
+    memoryManager_ = std::make_unique<MemoryManager>(options);
+    allocator_ = memoryManager_->allocator();
+    ASSERT_EQ(
+        AllocationTraits::numPages(memoryManager_->allocator()->capacity()),
+        bits::roundUp(
+            kCapacityBytes * (100 - options.smallAllocationReservePct) / 100 /
+                AllocationTraits::kPageSize,
+            64 * allocator_->sizeClasses().back()));
+  }
+
+  std::unique_ptr<MemoryManager> memoryManager_;
+  MemoryAllocator* allocator_;
+};
+
+TEST_F(MmapConfigTest, sizeClasses) {
+  setupAllocator();
+  Allocation result;
+  ASSERT_TRUE(
+      allocator_->allocateNonContiguous(2 * 4096 - 1, result, nullptr, 0));
+  auto g = folly::makeGuard([&]() { allocator_->freeNonContiguous(result); });
+  // Check that the allocation has one page of each size class, largest to
+  // smallest.
+  EXPECT_EQ(4096 * 2 - 1, result.numPages());
+  EXPECT_EQ(13, result.numRuns());
+  int32_t runPages = 4096;
+  for (auto i = 0; i < result.numRuns(); ++i) {
+    EXPECT_EQ(runPages, result.runAt(i).numPages());
+    runPages = runPages / 2;
+  }
+}
+
 } // namespace facebook::velox::memory