NiccoloN
127 lines
4.0 KiB
C++
127 lines
4.0 KiB
C++
#include <cassert>
|
|
#include <cstdlib>
|
|
#include <iostream>
|
|
|
|
#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 {
|
|
|
|
LocalAllocInterval makeInterval(size_t id, size_t size, uint64_t start, uint64_t end) {
|
|
LocalAllocInterval interval;
|
|
interval.id = id;
|
|
interval.size = size;
|
|
interval.start = start;
|
|
interval.end = end;
|
|
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].slotPlanIndex == intervals[1].slotPlanIndex);
|
|
}
|
|
|
|
int testSameSizeNonOverlap() {
|
|
std::cout << "testSameSizeNonOverlap:" << std::endl;
|
|
assertSingleSlotCase(makeInterval(0, 64, 0, 10), makeInterval(1, 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);
|
|
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);
|
|
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);
|
|
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);
|
|
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;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, char *argv[]) {
|
|
(void) argc;
|
|
(void) argv;
|
|
|
|
int failures = 0;
|
|
failures += testSameSizeNonOverlap();
|
|
failures += testLargerFirst();
|
|
failures += testSmallerFirst();
|
|
failures += testOverlapNeedsTwoSlots();
|
|
failures += testReuseChain();
|
|
failures += testPartialAddressReuse();
|
|
failures += testAddressAssignment();
|
|
if (failures != 0) {
|
|
std::cerr << failures << " test failures\n";
|
|
return EXIT_FAILURE;
|
|
}
|
|
return EXIT_SUCCESS;
|
|
}
|