better throughput scheduling with steady state cost model (WIP)

test/PIM/PimPipelineSelectionTest.cpp

@@ -0,0 +1,80 @@
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
+#include "llvm/Support/JSON.h"
+#include <cassert>
+#include <limits>
+using namespace onnx_mlir;
+using namespace onnx_mlir::spatial;
+int main() {
+  mlir::MLIRContext context;
+  context.allowUnregisteredDialects();
+  mlir::OwningOpRef<mlir::ModuleOp> checkpoint = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  mlir::OperationState state(checkpoint->getLoc(), "test.compute");
+  auto *op = mlir::Operation::create(state);
+  checkpoint->getBody()->push_back(op);
+  spatial::ComputeInstance instance{op, 3, 2};
+  spatial::MergeScheduleResult schedule;
+  schedule.dominanceOrderCompute.push_back(instance);
+  schedule.computeToCpuMap[instance] = 7;
+  schedule.computeToCpuSlotMap[instance] = 4;
+  schedule.computeToAestMap[instance] = 42;
+  schedule.isLastComputeOfCpu.insert(instance);
+  schedule.cpuToLastComputeMap[7] = instance;
+  for (int i = 0; i < 3; ++i) {
+    mlir::IRMapping mapping;
+    mlir::OwningOpRef<mlir::ModuleOp> trial = mlir::cast<mlir::ModuleOp>((*checkpoint)->clone(mapping));
+    auto remapped = spatial::remapSpatialSchedule(schedule, mapping);
+    auto copy = remapped.dominanceOrderCompute.front();
+    assert(copy.op != op && copy.laneStart == 3 && copy.laneCount == 2);
+    assert(remapped.computeToCpuMap.lookup(copy) == 7 && remapped.computeToCpuSlotMap.lookup(copy) == 4);
+    assert(remapped.computeToAestMap.lookup(copy) == 42 && remapped.isLastComputeOfCpu.contains(copy));
+    assert(remapped.cpuToLastComputeMap.lookup(7) == copy);
+    copy.op->erase();
+    assert(schedule.dominanceOrderCompute.front().op == op && op->getBlock() == checkpoint->getBody());
+  }
+
+  auto score = [](double value, std::string id) { return WholeScheduleCost{value, value, {}, id, 0.021}; };
+  std::string reason;
+  auto original = score(100, "0"), faster = score(90, "1");
+  auto select = [&](std::initializer_list<WholeScheduleCost> values) {
+    return selectPipelineCandidate(llvm::ArrayRef<WholeScheduleCost>(values), reason);
+  };
+  assert(select({original, faster}) == 1);
+  assert(select({original, score(97, "2")}) == 0);
+  assert(select({original, score(105, "2")}) == 0);
+  auto sensitive = faster; sensitive.alternative = 101;
+  assert(select({original, sensitive}) == 0);
+  assert(select({original, faster, score(91, "2")}) == 1);
+  assert(select({original, faster, score(90, "1")}) == 1);
+  auto unavailable = WholeScheduleCost{};
+  assert(select({unavailable, faster}) == 0);
+  assert(select({original, unavailable}) == 0);
+  auto unknownUncertainty = faster;
+  unknownUncertainty.relativeUncertainty.reset();
+  assert(select({original, unknownUncertainty}) == 0);
+  for (double value : {0., -1., std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()})
+    assert(select({original, score(value, "2")}) == 0);
+  for (int i = 0; i < 8; ++i) assert(select({original, faster}) == 1);
+
+  // The generic query accepts target-lowered IR without a target dialect API.
+  for (double interval : {90., 97.}) {
+    SchedulingTarget target;
+    assert(!target.prepareScheduleCost);
+    target.prepareScheduleCost = [interval](mlir::ModuleOp module,
+        llvm::json::Object &) -> mlir::FailureOr<ScheduleCostQuery> {
+      assert(module.getBody()->getOperations().size() == 1);
+      ScheduleCostQuery query;
+      query.identity = "static-program";
+      query.evaluate = [interval](llvm::json::Object &) {
+        return WholeScheduleCost{interval, interval, {}, "static-program", 0.021};
+      };
+      return query;
+    };
+    llvm::json::Object summary;
+    auto query = target.prepareScheduleCost(*checkpoint, summary);
+    assert(mlir::succeeded(query) && query->identity == "static-program");
+    auto cost = query->evaluate(summary);
+    assert(cost.eligible() && select({original, cost}) == (interval == 90. ? 1 : 0));
+    assert(checkpoint->getBody()->front().getName().getStringRef() == "test.compute");
+  }
+}