blazingly fasterer

test/PIM/PimMemoryLivenessPlannerTest.cpp

@@ -2,9 +2,11 @@
 #include <cstdlib>
 #include <iostream>
 
-#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
+#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;
 
 namespace {
@@ -18,12 +20,19 @@ LocalAllocInterval makeInterval(size_t id, size_t size, uint64_t start, uint64_t
   return interval;
 }
 
+size_t getPeak(llvm::ArrayRef<onnx_mlir::PlannedPhysicalSlot> slots) {
+  size_t peak = 0;
+  for (const auto& slot : slots)
+    peak = std::max(peak, slot.address + slot.requiredSize);
+  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].physicalSlotId == intervals[1].physicalSlotId);
+  assert(intervals[0].slotPlanIndex == intervals[1].slotPlanIndex);
 }
 
 int testSameSizeNonOverlap() {
@@ -34,13 +43,21 @@ int testSameSizeNonOverlap() {
 
 int testLargerFirst() {
   std::cout << "testLargerFirst:" << std::endl;
-  assertSingleSlotCase(makeInterval(0, 100, 0, 10), makeInterval(1, 40, 11, 20), 100);
+  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);
   return 0;
 }
 
 int testSmallerFirst() {
   std::cout << "testSmallerFirst:" << std::endl;
-  assertSingleSlotCase(makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), 100);
+  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);
   return 0;
 }
 
@@ -50,7 +67,7 @@ int testOverlapNeedsTwoSlots() {
     makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)};
   auto slots = planPhysicalSlots(intervals);
   assert(slots.size() == 2);
-  assert(intervals[0].physicalSlotId != intervals[1].physicalSlotId);
+  assert(intervals[0].slotPlanIndex != intervals[1].slotPlanIndex);
   return 0;
 }
 
@@ -59,10 +76,31 @@ int testReuseChain() {
   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() == 1);
-  assert(slots.front().requiredSize == 100);
-  assert(intervals[0].physicalSlotId == intervals[1].physicalSlotId);
-  assert(intervals[1].physicalSlotId == intervals[2].physicalSlotId);
+  assert(slots.size() == 3);
+  assert(getPeak(slots) == 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);
+  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);
   return 0;
 }
 
@@ -78,6 +116,8 @@ int main(int argc, char *argv[]) {
   failures += testSmallerFirst();
   failures += testOverlapNeedsTwoSlots();
   failures += testReuseChain();
+  failures += testPartialAddressReuse();
+  failures += testAddressAssignment();
   if (failures != 0) {
     std::cerr << failures << " test failures\n";
     return EXIT_FAILURE;