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

[Matthias Springer]

================
@@ -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;
----------------
matthias-springer wrote:

Can this variable be gradually updated as the outer `for (unsigned idx : order)` loop is running? (Instead of recomputing it every time.)


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