Add simplest memory reuse algo linear scanner

Author: Alwaysproblem

mlir/optimization/scheduler/include/lab/MemoryPlan.h

@@ -0,0 +1,18 @@
+#pragma once
+#include "mlir/IR/Value.h"
+
+struct Slot {
+  int id;
+  int64_t sizeBytes;
+  int64_t availableAfter;
+};
+
+struct MemorySlotAssignment {
+  int64_t slotId = -1;
+  int64_t offset = 0;
+};
+
+struct MemoryPlanResult {
+  llvm::DenseMap<mlir::Value, MemorySlotAssignment> assignments;
+  int64_t totalPoolBytes = 0;
+};

mlir/optimization/scheduler/include/lab/SchedulePlan.h

@@ -0,0 +1,15 @@
+#pragma once
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+
+struct ScheduleCandidate {
+  llvm::SmallVector<int64_t> tileSizes;
+  bool promote = false;
+  bool fuse = false;
+  bool doubleBuffer = false;
+  int pipelineStages = 1;
+
+  int64_t estimatedTrafficBytes = 0;
+  int64_t estimatedPeakMemoryBytes = 0;
+  double estimatedCost = 0.0;
+};

mlir/optimization/scheduler/lib/BufferAnalysisPass.cpp

@@ -1,13 +1,16 @@
+#include "lab/MemoryPlan.h"
 #include "mlir/Analysis/AliasAnalysis.h"
 #include "mlir/Analysis/Liveness.h"
-#include "mlir/IR/AsmState.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
-#include "llvm/Support/DebugLog.h"
 #include "mlir/Pass/Pass.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/DebugLog.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cstdint>
 #include <sys/types.h>
@@ -126,7 +129,6 @@ class LabBufferStatsAnalysis {
         return; // 没有语义使用点？先跳过
       int lastUse = opIndex[indexOp];
 
-
       buffers.push_back(BufferRecord{
           result,
           *sizeBytes,
@@ -179,8 +181,8 @@ class LabBufferStatsAnalysis {
   }
 
   static Operation *findLastSemanticUser(func::FuncOp funcOp,
-                                        mlir::Liveness &liveness,
-                                        Value value) {
+                                         mlir::Liveness &liveness,
+                                         Value value) {
     Operation *lastUser = nullptr;
 
     funcOp.walk([&](Operation *op) {
@@ -216,6 +218,80 @@ class LabBufferStatsAnalysis {
   SmallVector<BufferRecord> buffers;
 };
 
+class MemoryPlaner {
+public:
+  explicit MemoryPlaner(
+      const SmallVector<LabBufferStatsAnalysis::BufferRecord> &buffers)
+      : buffers_(buffers.begin(), buffers.end()) {}
+
+  void plan() {
+    llvm::sort(buffers_, [](const LabBufferStatsAnalysis::BufferRecord &a,
+                            const LabBufferStatsAnalysis::BufferRecord &b) {
+      return a.defIndex < b.defIndex;
+    });
+
+    int nextSlotId = 0;
+    int globalSlotOffset = 0;
+    for (const auto &buf : buffers_) {
+      // Try to find a slot for this buffer
+      int64_t offset = 0;
+      bool assigned = false;
+      for (auto &slot : slots_) {
+        if (slot.availableAfter <= buf.defIndex &&
+            slot.sizeBytes >= buf.sizeBytes) {
+          // This slot is available and large enough, assign it to the buffer
+          result_.assignments[buf.value] =
+              MemorySlotAssignment{slot.id, offset};
+          slot.availableAfter = buf.lastUseIndex + 1; // Update availability
+          assigned = true;
+          break;
+        }
+        offset += slot.sizeBytes; // Next buffer in the same slot will be placed
+                                  // after this one
+      }
+      if (!assigned) {
+        // No existing slot can accommodate this buffer, create a new slot
+        Slot newSlot{nextSlotId++, buf.sizeBytes, buf.lastUseIndex + 1};
+        result_.assignments[buf.value] =
+            MemorySlotAssignment{newSlot.id, globalSlotOffset};
+        globalSlotOffset +=
+            newSlot.sizeBytes; // Update offset for the next buffer
+        slots_.push_back(newSlot);
+      }
+    }
+
+    for (const auto &slot : slots_) {
+      result_.totalPoolBytes += slot.sizeBytes;
+    }
+  };
+  const MemoryPlanResult &getPlan() const { return result_; }
+
+  void printPlan(func::FuncOp func) const {
+    llvm::outs() << "Memory Plan:\n";
+    llvm::outs() << "  num_slots=" << slots_.size() << "\n";
+    for (const auto &slot : slots_) {
+      llvm::outs() << "  slot_id=" << slot.id << " size=" << slot.sizeBytes
+                   << "B available_after_op_index=" << slot.availableAfter
+                   << "\n";
+    }
+    for (const auto &entry : result_.assignments) {
+      std::string valueStr;
+      llvm::raw_string_ostream rso(valueStr);
+      AsmState asmState(func);
+      entry.first.printAsOperand(rso, asmState);
+      llvm::outs() << "  value=" << rso.str()
+                   << " -> slot=" << entry.second.slotId
+                   << " offset=" << entry.second.offset << "\n";
+    }
+    llvm::outs() << "  total_pool_bytes=" << result_.totalPoolBytes << "B\n";
+  }
+
+private:
+  SmallVector<LabBufferStatsAnalysis::BufferRecord> buffers_;
+  MemoryPlanResult result_;
+  llvm::SmallVector<Slot> slots_;
+};
+
 struct LabBufferStats
     : public PassWrapper<LabBufferStats, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LabBufferStats)
@@ -241,14 +317,18 @@ struct LabBufferStats
       llvm::raw_string_ostream rso(valueStr);
       AsmState asmState(func);
       buf.value.printAsOperand(rso, asmState);
-      llvm::outs() << "  value=" << rso.str() << " size=" << buf.sizeBytes << "B"
-                   << " def=#" << buf.defIndex
-                   << " last_use=#" << buf.lastUseIndex
-                   << " lifetime=[" << buf.defIndex
-                   << "," << buf.lastUseIndex << "]\n";
+      llvm::outs() << "  value=" << rso.str() << " size=" << buf.sizeBytes
+                   << "B"
+                   << " def=#" << buf.defIndex << " last_use=#"
+                   << buf.lastUseIndex << " lifetime=[" << buf.defIndex << ","
+                   << buf.lastUseIndex << "]\n";
     }
-    llvm::outs() << "  peak_live_memory=" << analysis.getStates().peak_memory << "B\n";
+    llvm::outs() << "  peak_live_memory=" << analysis.getStates().peak_memory
+                 << "B\n";
 
+    MemoryPlaner planer(buffers);
+    planer.plan();
+    planer.printPlan(func);
   }
 };
 } // namespace

mlir/optimization/scheduler/tests/memory_plan_reuse.mlir

@@ -0,0 +1,29 @@
+// RUN: ../build/scheduler/lab-scheduler %s --pass-pipeline='builtin.module(func.func(lab-buffer-stats))' --mlir-disable-threading
+
+module {
+  func.func @memory_plan_reuse_demo() {
+    %c0 = arith.constant 0.0 : f32
+
+    // Buffer A lives first.
+    %a = memref.alloc() : memref<64x64xf32>
+    linalg.fill ins(%c0 : f32) outs(%a : memref<64x64xf32>)
+
+    // Buffer B overlaps with A during the copy, so it cannot reuse A's slot.
+    %b = memref.alloc() : memref<64x64xf32>
+    memref.copy %a, %b : memref<64x64xf32> to memref<64x64xf32>
+    memref.dealloc %a : memref<64x64xf32>
+
+    // Buffer C starts after A is dead and is smaller, so it is a good reuse candidate.
+    %c = memref.alloc() : memref<32x32xf32>
+    linalg.fill ins(%c0 : f32) outs(%c : memref<32x32xf32>)
+    memref.dealloc %c : memref<32x32xf32>
+
+    // Buffer D overlaps with B during the copy, so it needs another live slot.
+    %d = memref.alloc() : memref<64x64xf32>
+    memref.copy %b, %d : memref<64x64xf32> to memref<64x64xf32>
+
+    memref.dealloc %b : memref<64x64xf32>
+    memref.dealloc %d : memref<64x64xf32>
+    return
+  }
+}