unify memory opimization passes

better reports
cleanups

test/PIM/PimMemoryLivenessPlannerTest.cpp

@@ -4,103 +4,97 @@
 
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
 
-using onnx_mlir::LocalAllocInterval;
-using onnx_mlir::CoreMemoryPlan;
-using onnx_mlir::assignPhysicalSlotAddresses;
-using onnx_mlir::planPhysicalSlots;
+using onnx_mlir::LocalMemoryPlacement;
+using onnx_mlir::pim::LocalMemoryInterval;
+using onnx_mlir::planLocalMemoryPlacements;
 
 namespace {
 
-LocalAllocInterval makeInterval(size_t id, size_t size, uint64_t start, uint64_t end) {
-  LocalAllocInterval interval;
-  interval.id = id;
+LocalMemoryInterval makeInterval(size_t size, uint64_t start, uint64_t end) {
+  LocalMemoryInterval interval;
   interval.size = size;
   interval.start = start;
   interval.end = end;
   return interval;
 }
 
-size_t getPeak(llvm::ArrayRef<onnx_mlir::PlannedPhysicalSlot> slots) {
+size_t getPeak(llvm::ArrayRef<LocalMemoryPlacement> placements) {
   size_t peak = 0;
-  for (const auto& slot : slots)
-    peak = std::max(peak, slot.address + slot.requiredSize);
+  for (const auto& placement : placements)
+    peak = std::max(peak, placement.address + placement.size);
   return peak;
 }
 
-void assertSingleSlotCase(LocalAllocInterval a, LocalAllocInterval b, size_t expectedSlotSize) {
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {a, b};
-  auto slots = planPhysicalSlots(intervals);
-  assert(slots.size() == 1);
-  assert(slots.front().requiredSize == expectedSlotSize);
-  assert(intervals[0].slotPlanIndex == intervals[1].slotPlanIndex);
+void assertSinglePlacementCase(LocalMemoryInterval a, LocalMemoryInterval b, size_t expectedSize) {
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {a, b};
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
+  assert(mlir::succeeded(placements));
+  assert(placements->size() == 1);
+  assert(placements->front().size == expectedSize);
+  assert(placements->front().intervalIndices.size() == 2);
 }
 
 int testSameSizeNonOverlap() {
   std::cout << "testSameSizeNonOverlap:" << std::endl;
-  assertSingleSlotCase(makeInterval(0, 64, 0, 10), makeInterval(1, 64, 11, 20), 64);
+  assertSinglePlacementCase(makeInterval(64, 0, 10), makeInterval(64, 11, 20), 64);
   return 0;
 }
 
 int testLargerFirst() {
   std::cout << "testLargerFirst:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
-    makeInterval(0, 100, 0, 10), makeInterval(1, 40, 11, 20)};
-  auto slots = planPhysicalSlots(intervals);
-  assert(slots.size() == 2);
-  assert(getPeak(slots) == 100);
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
+    makeInterval(100, 0, 10), makeInterval(40, 11, 20)};
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
+  assert(mlir::succeeded(placements) && placements->size() == 2);
+  assert(getPeak(*placements) == 100);
   return 0;
 }
 
 int testSmallerFirst() {
   std::cout << "testSmallerFirst:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
-    makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20)};
-  auto slots = planPhysicalSlots(intervals);
-  assert(slots.size() == 2);
-  assert(getPeak(slots) == 100);
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
+    makeInterval(40, 0, 10), makeInterval(100, 11, 20)};
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
+  assert(mlir::succeeded(placements) && placements->size() == 2);
+  assert(getPeak(*placements) == 100);
   return 0;
 }
 
 int testOverlapNeedsTwoSlots() {
   std::cout << "testOverlapNeedsTwoSlots:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
-    makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)};
-  auto slots = planPhysicalSlots(intervals);
-  assert(slots.size() == 2);
-  assert(intervals[0].slotPlanIndex != intervals[1].slotPlanIndex);
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
+    makeInterval(100, 0, 20), makeInterval(40, 10, 30)};
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
+  assert(mlir::succeeded(placements) && placements->size() == 2);
+  assert((*placements)[0].address != (*placements)[1].address);
   return 0;
 }
 
 int testReuseChain() {
   std::cout << "testReuseChain:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
-    makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), makeInterval(2, 20, 21, 30)};
-  auto slots = planPhysicalSlots(intervals);
-  assert(slots.size() == 3);
-  assert(getPeak(slots) == 100);
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
+    makeInterval(40, 0, 10), makeInterval(100, 11, 20), makeInterval(20, 21, 30)};
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
+  assert(mlir::succeeded(placements) && placements->size() == 3);
+  assert(getPeak(*placements) == 100);
   return 0;
 }
 
 int testPartialAddressReuse() {
   std::cout << "testPartialAddressReuse:" << std::endl;
-  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
-    makeInterval(0, 100, 0, 10), makeInterval(1, 60, 11, 20), makeInterval(2, 40, 11, 20)};
-  auto slots = planPhysicalSlots(intervals);
-  assert(getPeak(slots) == 100);
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
+    makeInterval(100, 0, 10), makeInterval(60, 11, 20), makeInterval(40, 11, 20)};
+  auto placements = planLocalMemoryPlacements(intervals, 1024);
+  assert(mlir::succeeded(placements));
+  assert(getPeak(*placements) == 100);
   return 0;
 }
 
-int testAddressAssignment() {
-  std::cout << "testAddressAssignment:" << std::endl;
-  CoreMemoryPlan plan;
-  plan.intervals = {makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)};
-  plan.slots = planPhysicalSlots(plan.intervals);
-  assert(mlir::succeeded(assignPhysicalSlotAddresses(plan, 1024)));
-  assert(plan.slots.size() == 2);
-  const auto& firstSlot = plan.slots[plan.intervals[0].slotPlanIndex];
-  const auto& secondSlot = plan.slots[plan.intervals[1].slotPlanIndex];
-  assert(firstSlot.address == 0 && firstSlot.requiredSize == 100);
-  assert(secondSlot.address == 100 && secondSlot.requiredSize == 40);
+int testAddressLimit() {
+  std::cout << "testAddressLimit:" << std::endl;
+  llvm::SmallVector<LocalMemoryInterval, 4> intervals = {
+    makeInterval(100, 0, 20), makeInterval(40, 10, 30)};
+  assert(mlir::failed(planLocalMemoryPlacements(intervals, 128)));
   return 0;
 }
 
@@ -117,7 +111,7 @@ int main(int argc, char *argv[]) {
   failures += testOverlapNeedsTwoSlots();
   failures += testReuseChain();
   failures += testPartialAddressReuse();
-  failures += testAddressAssignment();
+  failures += testAddressLimit();
   if (failures != 0) {
     std::cerr << failures << " test failures\n";
     return EXIT_FAILURE;