[Mlir-commits] [mlir] [mlir][bufferization] Add best-fit algorithm to static memory planner (PR #207403)

[Javed Absar]

================
@@ -0,0 +1,109 @@
+//===- StaticMemoryPlanning.cpp - Memory planning algorithms --------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h"
+#include "llvm/Support/MathExtras.h"
+#include <numeric>
+
+using namespace mlir::bufferization;
+
+/// Align an offset to the specified alignment.
+static int64_t alignOffset(int64_t offset, int64_t alignment) {
+  return llvm::alignTo(offset, alignment);
+}
+
+llvm::SmallVector<int64_t> mlir::bufferization::trivialMemoryPlanner(
+    int64_t arenaAlignment, llvm::ArrayRef<MemoryPlannerAlloc> allocs) {
+  llvm::SmallVector<int64_t> offsets;
+  int64_t currentOffset = 0;
+
+  for (const auto &alloc : allocs) {
+    currentOffset = alignOffset(currentOffset, alloc.alignment);
+    assert((arenaAlignment + currentOffset) % alloc.alignment == 0 &&
+           "invalid alignment");
+    offsets.push_back(currentOffset);
+    currentOffset += alloc.sizeInBytes;
+  }
+
+  return offsets;
+}
+
+llvm::SmallVector<int64_t> mlir::bufferization::bestFitMemoryPlanner(
+    int64_t arenaAlignment, llvm::ArrayRef<MemoryPlannerAlloc> allocs) {
+  struct Placement {
+    int64_t offset;
+    int64_t size;
+    int64_t timeEnd;
+  };
+
+  // Process allocations in order of start time.
+  llvm::SmallVector<unsigned> order(allocs.size());
+  std::iota(order.begin(), order.end(), 0);
+  llvm::sort(order, [&](unsigned a, unsigned b) {
+    return allocs[a].timeStart < allocs[b].timeStart;
+  });
+
+  llvm::SmallVector<Placement> placements;
+  llvm::SmallVector<int64_t> offsets(allocs.size(), 0);
+
+  for (unsigned idx : order) {
+    const MemoryPlannerAlloc &alloc = allocs[idx];
+
+    // Collect intervals that are still live at this allocation's start time.
+    llvm::SmallVector<std::pair<int64_t, int64_t>> occupied;
+    for (const auto &p : placements) {
+      if (p.timeEnd > alloc.timeStart)
+        occupied.push_back({p.offset, p.offset + p.size});
+    }
+    llvm::sort(occupied);
+
+    // Find the best (smallest) gap that fits this allocation.
+    int64_t bestOffset = -1;
+    int64_t bestGapSize = INT64_MAX;
+
+    // Arena high-water mark from all placements so far.
+    int64_t arenaEnd = 0;
+    for (const auto &p : placements)
+      arenaEnd = std::max(arenaEnd, p.offset + p.size);
+
+    int64_t gapStart = 0;
+    for (const auto &[occStart, occEnd] : occupied) {
+      int64_t alignedStart = alignOffset(gapStart, alloc.alignment);
+      int64_t gapEnd = occStart;
+      if (alignedStart + alloc.sizeInBytes <= gapEnd) {
+        int64_t gapSize = gapEnd - alignedStart;
+        if (gapSize < bestGapSize) {
+          bestGapSize = gapSize;
+          bestOffset = alignedStart;
+        }
+      }
+      gapStart = std::max(gapStart, occEnd);
+    }
+
+    // Check the trailing gap (between last occupied and arena end).
+    int64_t alignedTrailing = alignOffset(gapStart, alloc.alignment);
+    if (alignedTrailing + alloc.sizeInBytes <= arenaEnd) {
----------------
javedabsar1 wrote:

explained.

https://github.com/llvm/llvm-project/pull/207403