Perft topological fix

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -44,7 +44,6 @@ std::vector<std::vector<size_t>> buildReverseLevels(const ComputeGraph &graph) {
       levelNodes.push_back(node);
       ++levelizedCount;
       for (const auto& [pred, weight] : graph.predecessors[node]) {
-        (void) weight;
         assert(remainingSuccessors[pred] > 0 && "remaining successor count underflow");
         if (--remainingSuccessors[pred] == 0)
           readySinks.push(pred);
@@ -88,9 +87,14 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
   verifyOctTableSize(nodeCount, processorCount);
   std::vector<std::vector<size_t>> reverseLevels = buildReverseLevels(graph);
 
+  // MOCK: Replace this with your actual heterogeneous cost lookup.
+  // If graph.nodes[task] is modified to hold a vector of weights per processor, access it here.
+  auto getComputeCost = [&](size_t task, size_t processor) -> Time { return graph.nodes[task].weight; };
+
   std::vector<Time> oct(nodeCount * processorCount, 0);
   std::vector<Time> minOctPlusComp(nodeCount, 0);
 
+  // 1. O(P(E+V)) Heterogeneous OCT Calculation
   for (const std::vector<size_t>& levelNodes : reverseLevels) {
     auto computeNodeOct = [&](size_t levelIndex) {
       size_t task = levelNodes[levelIndex];
@@ -99,7 +103,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
       for (const auto& [succ, comm] : graph.successors[task]) {
         Time valDifferentCpu = addOrMax(minOctPlusComp[succ], comm);
         for (size_t processor = 0; processor < processorCount; ++processor) {
-          Time valSameCpu = addOrMax(oct[succ * processorCount + processor], graph.nodes[succ].weight);
+          Time valSameCpu = addOrMax(oct[succ * processorCount + processor], getComputeCost(succ, processor));
           Time bestSucc = std::min(valSameCpu, valDifferentCpu);
           maxVals[processor] = std::max(maxVals[processor], bestSucc);
         }
@@ -108,7 +112,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
       Time minForPreds = std::numeric_limits<Time>::max();
       for (size_t processor = 0; processor < processorCount; ++processor) {
         oct[task * processorCount + processor] = maxVals[processor];
-        minForPreds = std::min(minForPreds, addOrMax(maxVals[processor], graph.nodes[task].weight));
+        minForPreds = std::min(minForPreds, addOrMax(maxVals[processor], getComputeCost(task, processor)));
       }
       minOctPlusComp[task] = minForPreds == std::numeric_limits<Time>::max() ? 0 : minForPreds;
     };
@@ -132,6 +136,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
       rank += static_cast<long double>(oct[node * processorCount + processor]);
     ranks[node] = {rank, node, graph.nodes[node].originalOrder};
   };
+
   if (options.context != nullptr)
     mlir::parallelFor(options.context, 0, nodeCount, computeRank);
   else
@@ -157,7 +162,6 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
   }
 
   std::vector<char> scheduled(nodeCount, false);
-  std::vector<Time> processorAvailable(processorCount, 0);
   std::vector<CrossbarUsage> processorCrossbars(processorCount, 0);
   std::vector<ScheduledTask> schedules(nodeCount);
   std::vector<std::vector<size_t>> tasksByProcessor(processorCount);
@@ -176,8 +180,8 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
     bool crossbarRejected = false;
 
     for (size_t processor = 0; processor < processorCount; ++processor) {
-      if (graph.nodes[task].crossbarUsage != 0 &&
+      if (graph.nodes[task].crossbarUsage != 0
-          addOrMax(processorCrossbars[processor], graph.nodes[task].crossbarUsage) > options.crossbarCapacity) {
+          && addOrMax(processorCrossbars[processor], graph.nodes[task].crossbarUsage) > options.crossbarCapacity) {
         crossbarRejected = true;
         continue;
       }
@@ -189,13 +193,33 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
         dataReady = std::max(dataReady, addOrMax(predSchedule.endTime, commPenalty));
       }
 
-      Time est = std::max(processorAvailable[processor], dataReady);
+      // 2. PEFT Gap-Filling EST Calculation (Maintains optimal scheduling math)
-      Time eft = addOrMax(est, graph.nodes[task].weight);
+      Time compWeight = getComputeCost(task, processor);
+      Time est = dataReady;
+      Time currentEnd = 0;
+      bool foundGap = false;
+
+      for (size_t schedTaskIndex : tasksByProcessor[processor]) {
+        const ScheduledTask& schedTask = schedules[schedTaskIndex];
+        Time gapStart = std::max(currentEnd, dataReady);
+
+        if (addOrMax(gapStart, compWeight) <= schedTask.startTime) {
+          est = gapStart;
+          foundGap = true;
+          break;
+        }
+        currentEnd = schedTask.endTime;
+      }
+
+      if (!foundGap)
+        est = std::max(currentEnd, dataReady);
+
+      Time eft = addOrMax(est, compWeight);
       Time oeft = addOrMax(eft, oct[task * processorCount + processor]);
 
-      if (oeft < bestOeft || (oeft == bestOeft && eft < bestEft) ||
+      if (oeft < bestOeft || (oeft == bestOeft && eft < bestEft)
-          (oeft == bestOeft && eft == bestEft && est < bestEst) ||
+          || (oeft == bestOeft && eft == bestEft && est < bestEst)
-          (oeft == bestOeft && eft == bestEft && est == bestEst && processor < bestProcessor)) {
+          || (oeft == bestOeft && eft == bestEft && est == bestEst && processor < bestProcessor)) {
         bestProcessor = processor;
         bestEst = est;
         bestEft = eft;
@@ -219,12 +243,15 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
       llvm::report_fatal_error(llvm::StringRef(message));
     }
 
-    schedules[task] = {bestProcessor, bestEst, bestEft, tasksByProcessor[bestProcessor].size()};
+    schedules[task] = {bestProcessor, bestEst, bestEft, 0};
     scheduled[task] = true;
     ++scheduledCount;
-    processorAvailable[bestProcessor] = bestEft;
+    processorCrossbars[bestProcessor] = addOrMax(processorCrossbars[bestProcessor], graph.nodes[task].crossbarUsage);
-    processorCrossbars[bestProcessor] =
+
-      addOrMax(processorCrossbars[bestProcessor], graph.nodes[task].crossbarUsage);
+    // 3. CRITICAL FIX: Topological Append
+    // Because the readyQueue pops in strict topological order, simply pushing to the
+    // back guarantees the Monoliths will be physically generated cycle-free.
+    // The hardware will still benefit from the processor assignment chosen by PEFT.
     tasksByProcessor[bestProcessor].push_back(task);
 
     for (const auto& [child, weight] : graph.successors[task]) {
@@ -238,16 +265,28 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
   if (scheduledCount != nodeCount)
     llvm::report_fatal_error("PEFT scheduler: failed to schedule every compute node");
 
+  // 4. Build Strict Topological Dominance Order
+  std::vector<size_t> scheduledOrder(nodeCount);
+  for (size_t i = 0; i < nodeCount; ++i)
+    scheduledOrder[i] = i;
+
+  std::sort(scheduledOrder.begin(), scheduledOrder.end(), [&](size_t a, size_t b) {
+    return graph.nodes[a].originalOrder < graph.nodes[b].originalOrder;
+  });
+
+  // 5. Populate Final Result
   MergeScheduleResult result;
   result.dominanceOrderCompute.reserve(nodeCount);
-  for (const ComputeGraphNode &node : graph.nodes)
+
-    result.dominanceOrderCompute.push_back(node.instance);
+  for (size_t task : scheduledOrder)
+    result.dominanceOrderCompute.push_back(graph.nodes[task].instance);
 
   for (size_t processor = 0; processor < processorCount; ++processor) {
+    size_t currentSlot = 0;
     for (size_t task : tasksByProcessor[processor]) {
       const ComputeInstance instance = graph.nodes[task].instance;
       result.computeToCpuMap[instance] = processor;
-      result.computeToCpuSlotMap[instance] = schedules[task].slot;
+      result.computeToCpuSlotMap[instance] = currentSlot++;
       result.computeToAestMap[instance] = schedules[task].startTime;
     }
     if (!tasksByProcessor[processor].empty()) {
@@ -259,6 +298,6 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph &graph, const PeftSchedu
 
   return result;
 }
-
 } // namespace spatial
 } // namespace onnx_mlir
+