From 6bc7709aa4bf8c4b6df01a72b55bfd5484b193d8 Mon Sep 17 00:00:00 2001
From: ilgeco <gabriele.magnani@mail.polimi.it>
Date: Tue, 21 Jul 2026 15:43:22 +0200
Subject: [PATCH] slighlty faster

---
 .../Scheduling/ComputeGraph.cpp               |  9 ++-
 .../Scheduling/PeftScheduler.cpp              | 60 ++++++++++++++++++-
 2 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
index 25fbefd..93f80c4 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
@@ -195,9 +195,9 @@ struct PimsimSchedulerCostModel {
 };
 
 std::optional<uint64_t> getStaticTripCount(scf::ForOp loop);
-[[maybe_unused]] Cost getOperationCost(Operation& op);
+Cost getOperationCost(Operation& op);
 
-[[maybe_unused]] Cost getRegionCost(Region& body) {
+Cost getRegionCost(Region& body) {
   Cost cost = 0;
   for (Block& block : body)
     for (Operation& op : block)
@@ -205,7 +205,7 @@ std::optional<uint64_t> getStaticTripCount(scf::ForOp loop);
   return cost;
 }
 
-[[maybe_unused]] Cost getOperationCost(Operation& op) {
+Cost getOperationCost(Operation& op) {
   if (auto loop = dyn_cast<scf::ForOp>(&op)) {
     std::optional<uint64_t> tripCount = getStaticTripCount(loop);
     if (!tripCount)
@@ -271,8 +271,7 @@ std::optional<uint64_t> getStaticTripCount(scf::ForOp loop) {
 }
 
 Cost getComputeBodyCost(Region& body) {
-  constexpr Cost kOperationCost = 100;
-  return checkedMultiply(static_cast<Cost>(countComputeBodyOperationInstances(body)), kOperationCost);
+  return getRegionCost(body);
 }
 
 uint64_t countOperationInstances(Operation& op) {
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
index 6011cec..ef381d8 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
@@ -6,7 +6,9 @@
 
 #include <cmath>
 #include <limits>
+#include <optional>
 #include <queue>
+#include <tuple>
 #include <vector>
 
 #include "PeftScheduler.hpp"
@@ -133,6 +135,58 @@ void verifyOctTableSize(size_t nodeCount, size_t processorCount) {
   }
 }
 
+std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
+                                                    size_t processorCount,
+                                                    size_t crossbarCapacity,
+                                                    const MeshModel& mesh) {
+  std::vector<size_t> weightedTasks;
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    if (!graph.nodes[task].crossbarUsage.empty())
+      weightedTasks.push_back(task);
+  llvm::sort(weightedTasks, [&](size_t lhs, size_t rhs) {
+    if (graph.nodes[lhs].crossbarUsage.size() != graph.nodes[rhs].crossbarUsage.size())
+      return graph.nodes[lhs].crossbarUsage.size() > graph.nodes[rhs].crossbarUsage.size();
+    return graph.nodes[lhs].originalOrder < graph.nodes[rhs].originalOrder;
+  });
+
+  std::vector<CrossbarUsage> reservations(processorCount);
+  std::vector<Time> reservedLoad(processorCount, 0);
+  for (size_t task : weightedTasks) {
+    size_t bestProcessor = std::numeric_limits<size_t>::max();
+    using ReservationScore = std::tuple<Time, size_t, size_t, size_t>;
+    std::optional<ReservationScore> bestScore;
+    for (size_t processor = 0; processor < processorCount; ++processor) {
+      size_t crossbarUnion =
+        getCrossbarUnionSize(reservations[processor], graph.nodes[task].crossbarUsage);
+      if (crossbarUnion > crossbarCapacity)
+        continue;
+      size_t addedCrossbars = crossbarUnion - reservations[processor].size();
+      ReservationScore score {reservedLoad[processor],
+                              addedCrossbars,
+                              mesh.getCenterDistance(processor),
+                              processor};
+      if (!bestScore || score < *bestScore) {
+        bestProcessor = processor;
+        bestScore = score;
+      }
+    }
+    if (bestProcessor == std::numeric_limits<size_t>::max()) {
+      std::string message =
+        llvm::formatv("PEFT reservation planner: cannot place task {0} with {1} distinct weights in {2} "
+                      "processors of capacity {3}",
+                      graph.nodes[task].originalOrder,
+                      graph.nodes[task].crossbarUsage.size(),
+                      processorCount,
+                      crossbarCapacity)
+          .str();
+      llvm::report_fatal_error(llvm::StringRef(message));
+    }
+    insertCrossbarWeights(reservations[bestProcessor], graph.nodes[task].crossbarUsage);
+    reservedLoad[bestProcessor] = addOrMax(reservedLoad[bestProcessor], graph.nodes[task].cost);
+  }
+  return reservations;
+}
+
 } // namespace
 
 Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount) {
@@ -145,6 +199,8 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
   if (processorCount == 0)
     llvm::report_fatal_error("PEFT scheduler: processor count must be positive");
   MeshModel mesh = MeshModel::infer(processorCount);
+  std::vector<CrossbarUsage> capacityReservations =
+    planCrossbarReservations(graph, processorCount, options.crossbarCapacity, mesh);
 
   verifyOctTableSize(nodeCount, processorCount);
   std::vector<std::vector<size_t>> reverseLevels = buildReverseLevels(graph);
@@ -245,7 +301,8 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
 
     for (size_t processor = 0; processor < processorCount; ++processor) {
       unsigned int overlapCount = countCrossbarOverlap(processorCrossbars[processor], graph.nodes[task].crossbarUsage);
-      size_t crossbarUnion = getCrossbarUnionSize(processorCrossbars[processor], graph.nodes[task].crossbarUsage);
+      size_t crossbarUnion =
+        getCrossbarUnionSize(capacityReservations[processor], graph.nodes[task].crossbarUsage);
       smallestCrossbarUnion = std::min(smallestCrossbarUnion, crossbarUnion);
       if (!graph.nodes[task].crossbarUsage.empty() && crossbarUnion > options.crossbarCapacity) {
         crossbarRejected = true;
@@ -337,6 +394,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     schedules[task] = {bestProcessor, bestEst, bestEft};
     scheduled[task] = true;
     ++scheduledCount;
+    insertCrossbarWeights(capacityReservations[bestProcessor], graph.nodes[task].crossbarUsage);
     insertCrossbarWeights(processorCrossbars[bestProcessor], graph.nodes[task].crossbarUsage);
 
     // 3. CRITICAL FIX: Topological Append