fix scheduling reasoning on physical cores directly

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp

@@ -6,7 +6,6 @@
 #include "DeferredCommunicationRealization.hpp"
 #include "DeferredCommunicationScheduling.hpp"
 #include "DeferredTransferPlanning.hpp"
-#include "Scheduling/PeftScheduler.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
@@ -34,76 +33,6 @@ static LogicalResult verifyNoEscapingRegionValues(Operation* owner, StringRef ph
                               << escapingUser->getName() << " at " << escapingUser->getLoc();
 }
 
-static LogicalResult placeLogicalProcessorsOnPhysicalCores(
-    DeferredTransferPlan& plan, const SchedulingTarget& target,
-    size_t pipelineStages) {
-  std::vector<Cost> logicalTrafficFlits(target.processorCount * target.processorCount, 0);
-  for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
-    for (const ExternalTransferFamily& transfer : exchange->external) {
-      auto fragmentType = dyn_cast<ShapedType>(transfer.requirement->publicationFragmentType);
-      if (!fragmentType || !fragmentType.hasStaticShape())
-        return exchange->deferred.emitOpError("physical core placement requires a static transfer fragment");
-      auto fragmentBytes = pim::getCheckedShapedTypeSizeInBytes(
-        fragmentType, exchange->deferred, "physical core placement transfer fragment");
-      if (failed(fragmentBytes))
-        return failure();
-      Cost flits = static_cast<Cost>(*fragmentBytes) / target.transferWidthBytes
-                 + (*fragmentBytes % target.transferWidthBytes != 0);
-      for (size_t index = 0; index < transfer.sourceCores.size(); ++index) {
-        size_t sourceLogicalProcessor = static_cast<size_t>(transfer.sourceCores.valueAt(index));
-        size_t targetLogicalProcessor = static_cast<size_t>(transfer.targetCores.valueAt(index));
-        Cost& traffic = logicalTrafficFlits[sourceLogicalProcessor * target.processorCount + targetLogicalProcessor];
-        traffic = checkedAdd(traffic, flits);
-      }
-    }
-
-  std::vector<size_t> placementGroups;
-  if (pipelineStages > 1) {
-    if (plan.processorStages.size() != target.processorCount)
-      return failure();
-    placementGroups = plan.processorStages;
-  }
-  std::vector<size_t> physicalCoreForLogicalProcessor =
-    mapLogicalProcessorsToPhysicalCores(
-      logicalTrafficFlits, target, placementGroups);
-  auto getPhysicalCore = [&](int64_t logicalProcessor) {
-    assert(logicalProcessor >= 0 && static_cast<size_t>(logicalProcessor) < physicalCoreForLogicalProcessor.size()
-           && "logical processor is outside the scheduling target");
-    return static_cast<int64_t>(physicalCoreForLogicalProcessor[logicalProcessor]);
-  };
-  auto remap = [&](StaticIntSequence& logicalProcessors) {
-    SmallVector<int64_t> physicalCores;
-    physicalCores.reserve(logicalProcessors.size());
-    for (size_t index = 0; index < logicalProcessors.size(); ++index)
-      physicalCores.push_back(getPhysicalCore(logicalProcessors.valueAt(index)));
-    logicalProcessors = StaticIntSequence::fromValues(physicalCores);
-  };
-
-  for (ScheduledInfo& scheduled : plan.scheduled) {
-    for (int64_t& logicalProcessor : scheduled.cores)
-      logicalProcessor = getPhysicalCore(logicalProcessor);
-    if (isa<SpatScheduledCompute>(scheduled.op)) {
-      scheduled.op->setAttr(
-        kCoreIdAttrName, IntegerAttr::get(IntegerType::get(scheduled.op->getContext(), 32), scheduled.cores.front()));
-    }
-    else {
-      SmallVector<int32_t> physicalCores;
-      physicalCores.reserve(scheduled.cores.size());
-      for (int64_t physicalCore : scheduled.cores)
-        physicalCores.push_back(static_cast<int32_t>(physicalCore));
-      scheduled.op->setAttr(kCoreIdsAttrName, DenseI32ArrayAttr::get(scheduled.op->getContext(), physicalCores));
-    }
-  }
-  for (const std::unique_ptr<ProducedValue>& produced : plan.producedStorage)
-    produced->core = getPhysicalCore(produced->core);
-  for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
-    for (ExternalTransferFamily& transfer : exchange->external) {
-      remap(transfer.sourceCores);
-      remap(transfer.targetCores);
-    }
-  return success();
-}
-
 static LogicalResult replaceFinalGraphPublications(func::FuncOp funcOp, DeferredTransferPlan& plan) {
   for (Operation& op : funcOp.getOps()) {
     if (!isa<SpatGraphCompute, SpatGraphComputeBatch>(op))
@@ -227,9 +156,6 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
     funcOp, materialization, pipelineStages, target.processorCount);
   if (failed(transfers))
     return funcOp.emitOpError("phase 2 failed to build symbolic transfer families");
-  if (failed(placeLogicalProcessorsOnPhysicalCores(
-        *transfers, target, pipelineStages)))
-    return failure();
   if (transfers->pipelineHostBufferBytes != 0) {
     auto bytes = pim::checkedCast<int64_t>(
       transfers->pipelineHostBufferBytes, funcOp,

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp

@@ -285,7 +285,7 @@ FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp
     enqueue(stream);
   auto advance = [&] {
     bool changed = false;
-    while (!advanceable.empty()) {
+    while (!advanceable.empty() && ready.empty()) {
       unsigned stream = advanceable.front();
       advanceable.pop();
       StreamProgress& progress = streams[stream];
@@ -305,6 +305,8 @@ FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp
   ScheduledCommunicationPlan result;
   unsigned finishedGroups = 0;
   while (finishedGroups != groups.size()) {
+    while (!ready.empty() && groups[ready.top()].scheduled)
+      ready.pop();
     bool progressed = advance();
     std::optional<unsigned> chosen;
     unsigned bestExtension = 0;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp

@@ -111,8 +111,16 @@ void dumpScheduledComputeReport(ModuleOp moduleOp, func::FuncOp funcOp, const Me
      << "  materialized homogeneous runs: " << materializedRuns << "\n"
      << "  largest run: " << largestRun << "\n"
      << "  instances compacted: " << instancesCompacted << "\n"
-     << "  compatible runs rejected: 0\n\n"
-     << "Materialized scheduled ops\n";
+     << "  compatible runs rejected: 0\n";
+  llvm::MapVector<size_t, SmallVector<size_t>> stageCores;
+  for (auto [core, stage] : llvm::enumerate(schedule.processorStages))
+    stageCores[stage].push_back(core);
+  for (const auto &[stage, cores] : stageCores) {
+    os << "  pipeline stage " << stage << " cores=";
+    printIndexedList(os, ArrayRef<size_t>(cores));
+    os << "\n";
+  }
+  os << "\nMaterialized scheduled ops\n";
 
   for (const ScheduledMaterializationRecord &record : records) {
     bool batch = isa<SpatScheduledComputeBatch>(record.scheduledOp);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp

@@ -10,8 +10,6 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Passes/PIMPasses.h"
 
-#include <limits>
-
 using namespace mlir;
 
 namespace onnx_mlir {
@@ -34,38 +32,6 @@ static FailureOr<func::FuncOp> requireEntry(ModuleOp moduleOp) {
   return *entry;
 }
 
-static SchedulingTarget getPipelineSchedulingTarget(
-    const SchedulingTarget& physicalTarget, size_t pipelineStages) {
-  if (pipelineStages == 1)
-    return physicalTarget;
-
-  PipelineCoreLayout layout(physicalTarget.processorCount, pipelineStages);
-  SchedulingTarget schedulingTarget = physicalTarget;
-  schedulingTarget.processorCount = layout.getLogicalProcessorCount();
-  schedulingTarget.residentWeightCapacity = checkedMultiply(
-    physicalTarget.residentWeightCapacity, pipelineStages);
-  schedulingTarget.interProcessorLatencyNs.assign(
-    schedulingTarget.processorCount * schedulingTarget.processorCount, 0);
-  Cost latencySum = 0;
-  size_t pairCount = 0;
-  for (size_t source = 0; source < schedulingTarget.processorCount; ++source)
-    for (size_t destination = 0;
-         destination < schedulingTarget.processorCount; ++destination) {
-      Cost latency = physicalTarget.getInterProcessorLatencyNs(
-        source, destination);
-      schedulingTarget.interProcessorLatencyNs[
-        source * schedulingTarget.processorCount + destination] = latency;
-      if (source != destination) {
-        latencySum = checkedAdd(latencySum, latency);
-        ++pairCount;
-      }
-    }
-  schedulingTarget.averageInterProcessorLatencyNs = pairCount == 0
-    ? 0
-    : (latencySum + pairCount - 1) / pairCount;
-  return schedulingTarget;
-}
-
 struct ScheduleAndRealizeSpatialPass final
     : PassWrapper<ScheduleAndRealizeSpatialPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleAndRealizeSpatialPass)
@@ -92,9 +58,7 @@ struct ScheduleAndRealizeSpatialPass final
     PipelineCoreLayout pipelineLayout(target.processorCount, pipelineStages);
     if (!pipelineLayout.isValid()
         || (pipelineStages > 1
-            && target.synchronizationRegisterCount == 0)
-        || target.residentWeightCapacity
-             > std::numeric_limits<size_t>::max() / pipelineStages) {
+            && target.synchronizationRegisterCount == 0)) {
       moduleOp.emitError("ScheduleAndRealizeSpatial requires valid pipeline stages and resource counts");
       signalPassFailure();
       return;
@@ -106,21 +70,15 @@ struct ScheduleAndRealizeSpatialPass final
     }
     func::FuncOp entryFunc = *entry;
 
-    SchedulingTarget schedulingTarget = getPipelineSchedulingTarget(
-      target, pipelineStages);
     ComputeGraph scheduledGraph;
     MergeScheduleResult schedule;
     for (;;) {
-      MergeSchedulingAnalysis analysis(
-        entryFunc, schedulingTarget,
-        pipelineStages > 1 ? target.processorCount : 0);
-      scheduledGraph = analysis.getGraph();
-      schedule = std::move(analysis.getResult());
+      scheduledGraph = buildComputeGraph(entryFunc, target);
       std::string pipelineError;
       if (pipelineStages > 1) {
         FailureOr<PipelineWorkloadPreparation> preparation =
           preparePipelineWorkload(
-            scheduledGraph, schedule, pipelineStages, target, pipelineError);
+            scheduledGraph, pipelineStages, target, pipelineError);
         if (failed(preparation)) {
           moduleOp.emitError() << pipelineError;
           signalPassFailure();
@@ -129,6 +87,9 @@ struct ScheduleAndRealizeSpatialPass final
         if (*preparation == PipelineWorkloadPreparation::Changed)
           continue;
       }
+      MergeSchedulingAnalysis analysis(
+        scheduledGraph, target, entryFunc.getContext());
+      schedule = std::move(analysis.getResult());
       if (succeeded(applyPipelineScheduling(
             scheduledGraph, schedule, pipelineStages, target, pipelineError)))
         break;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp

@@ -88,22 +88,21 @@ void verifySchedule(const ComputeGraph& graph,
 
 } // namespace
 
-MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op,
+MergeSchedulingAnalysis::MergeSchedulingAnalysis(const ComputeGraph& computeGraph,
                                                  const SchedulingTarget& schedulingTarget,
-                                                 size_t partitionCount)
-: entryOp(op), target(schedulingTarget), computePartitionCount(partitionCount) {
+                                                 mlir::MLIRContext* context)
+: context(context), target(schedulingTarget), graph(computeGraph) {
   result = run();
 }
 
 MergeScheduleResult MergeSchedulingAnalysis::run() {
-  graph = buildComputeGraph(entryOp, target, computePartitionCount);
   if (!verifyAcyclic(graph))
     llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
 
   MergeScheduleResult schedule = runPeftScheduler(
     graph, PeftScheduleOptions {
       target,
-      entryOp->getContext()});
+      context});
   verifySchedule(graph, schedule, target);
   return schedule;
 }

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp

@@ -11,17 +11,15 @@ namespace spatial {
 
 class MergeSchedulingAnalysis {
 public:
-  MergeSchedulingAnalysis(mlir::Operation* op,
+  MergeSchedulingAnalysis(const ComputeGraph& graph,
                           const SchedulingTarget& target,
-                          size_t computePartitionCount = 0);
+                          mlir::MLIRContext* context);
   MergeScheduleResult& getResult() { return result; }
-  const ComputeGraph& getGraph() const { return graph; }
 
 private:
-  mlir::Operation* entryOp = nullptr;
+  mlir::MLIRContext* context = nullptr;
   const SchedulingTarget& target;
-  size_t computePartitionCount = 0;
-  ComputeGraph graph;
+  const ComputeGraph& graph;
   MergeScheduleResult result;
 
   MergeScheduleResult run();

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

@@ -5,7 +5,6 @@
 #include "llvm/Support/FormatVariadic.h"
 
 #include <limits>
-#include <numeric>
 #include <optional>
 #include <queue>
 #include <tuple>
@@ -42,11 +41,6 @@ struct TopologyModel {
   }
 };
 
-Time getAverageTransferTime(const TransferCost& transferCost, const SchedulingTarget& target) {
-  return checkedAdd(transferCost.fixed,
-                    checkedMultiply(transferCost.networkFlits, target.averageInterProcessorLatencyNs));
-}
-
 std::vector<std::vector<size_t>> buildReverseLevels(const ComputeGraph& graph) {
   std::vector<size_t> remainingSuccessors(graph.nodes.size(), 0);
   std::queue<size_t> readySinks;
@@ -83,24 +77,6 @@ std::vector<std::vector<size_t>> buildReverseLevels(const ComputeGraph& graph) {
   return reverseLevels;
 }
 
-void verifyOctTableSize(size_t nodeCount, size_t processorCount) {
-  constexpr size_t kMaxOctTableBytes = 1ull << 35;
-  if (nodeCount == 0 || processorCount == 0)
-    return;
-  if (processorCount > std::numeric_limits<size_t>::max() / sizeof(Time))
-    llvm::report_fatal_error("PEFT scheduler: OCT table size overflow");
-  size_t rowBytes = processorCount * sizeof(Time);
-  if (nodeCount > std::numeric_limits<size_t>::max() / rowBytes)
-    llvm::report_fatal_error("PEFT scheduler: OCT table size overflow");
-  size_t totalBytes = nodeCount * rowBytes;
-  if (totalBytes > kMaxOctTableBytes) {
-    std::string message = llvm::formatv("PEFT scheduler: OCT table would require {0} MiB, exceeding the 1024 MiB guard",
-                                        totalBytes / (1024 * 1024))
-                            .str();
-    llvm::report_fatal_error(llvm::StringRef(message));
-  }
-}
-
 bool hasHighResidentWeightPressure(const ComputeGraph& graph, size_t processorCount, size_t residentWeightCapacity) {
   if (residentWeightCapacity > std::numeric_limits<size_t>::max() / processorCount)
     return false;
@@ -243,79 +219,6 @@ FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeB
 
 } // namespace
 
-std::vector<size_t> mapLogicalProcessorsToPhysicalCores(ArrayRef<Cost> logicalTrafficFlits,
-                                                        const SchedulingTarget& target,
-                                                        ArrayRef<size_t> placementGroups) {
-  const size_t processorCount = target.processorCount;
-  assert(logicalTrafficFlits.size() == processorCount * processorCount
-         && "logical traffic matrix must cover every processor pair");
-  assert((placementGroups.empty() || placementGroups.size() == processorCount)
-         && "physical placement groups must cover every processor");
-  std::vector<size_t> physicalCoreForLogicalProcessor(processorCount);
-  std::iota(physicalCoreForLogicalProcessor.begin(), physicalCoreForLogicalProcessor.end(), 0);
-
-  auto transferCost = [&](size_t sourceLogicalProcessor,
-                          size_t targetLogicalProcessor,
-                          size_t sourcePhysicalCore,
-                          size_t targetPhysicalCore) {
-    Cost traffic = logicalTrafficFlits[sourceLogicalProcessor * processorCount + targetLogicalProcessor];
-    return checkedMultiply(traffic, target.getInterProcessorLatencyNs(sourcePhysicalCore, targetPhysicalCore));
-  };
-
-  for (size_t logicalProcessor = 0; logicalProcessor < processorCount; ++logicalProcessor) {
-    size_t bestPeerLogicalProcessor = logicalProcessor;
-    Cost bestSaving = 0;
-    for (size_t peerLogicalProcessor = 0; peerLogicalProcessor < processorCount; ++peerLogicalProcessor) {
-      if (peerLogicalProcessor == logicalProcessor)
-        continue;
-      if (!placementGroups.empty()
-          && placementGroups[peerLogicalProcessor]
-               != placementGroups[logicalProcessor])
-        continue;
-      size_t physicalCore = physicalCoreForLogicalProcessor[logicalProcessor];
-      size_t peerPhysicalCore = physicalCoreForLogicalProcessor[peerLogicalProcessor];
-      Cost currentCost = 0;
-      Cost swappedCost = 0;
-      for (size_t otherLogicalProcessor = 0; otherLogicalProcessor < processorCount; ++otherLogicalProcessor) {
-        if (otherLogicalProcessor == logicalProcessor || otherLogicalProcessor == peerLogicalProcessor)
-          continue;
-        size_t otherPhysicalCore = physicalCoreForLogicalProcessor[otherLogicalProcessor];
-        currentCost = checkedAdd(
-          currentCost, transferCost(logicalProcessor, otherLogicalProcessor, physicalCore, otherPhysicalCore));
-        currentCost = checkedAdd(
-          currentCost, transferCost(otherLogicalProcessor, logicalProcessor, otherPhysicalCore, physicalCore));
-        currentCost = checkedAdd(
-          currentCost, transferCost(peerLogicalProcessor, otherLogicalProcessor, peerPhysicalCore, otherPhysicalCore));
-        currentCost = checkedAdd(
-          currentCost, transferCost(otherLogicalProcessor, peerLogicalProcessor, otherPhysicalCore, peerPhysicalCore));
-        swappedCost = checkedAdd(
-          swappedCost, transferCost(logicalProcessor, otherLogicalProcessor, peerPhysicalCore, otherPhysicalCore));
-        swappedCost = checkedAdd(
-          swappedCost, transferCost(otherLogicalProcessor, logicalProcessor, otherPhysicalCore, peerPhysicalCore));
-        swappedCost = checkedAdd(
-          swappedCost, transferCost(peerLogicalProcessor, otherLogicalProcessor, physicalCore, otherPhysicalCore));
-        swappedCost = checkedAdd(
-          swappedCost, transferCost(otherLogicalProcessor, peerLogicalProcessor, otherPhysicalCore, physicalCore));
-      }
-      currentCost =
-        checkedAdd(currentCost, transferCost(logicalProcessor, peerLogicalProcessor, physicalCore, peerPhysicalCore));
-      currentCost =
-        checkedAdd(currentCost, transferCost(peerLogicalProcessor, logicalProcessor, peerPhysicalCore, physicalCore));
-      swappedCost =
-        checkedAdd(swappedCost, transferCost(logicalProcessor, peerLogicalProcessor, peerPhysicalCore, physicalCore));
-      swappedCost =
-        checkedAdd(swappedCost, transferCost(peerLogicalProcessor, logicalProcessor, physicalCore, peerPhysicalCore));
-      if (currentCost > swappedCost && currentCost - swappedCost > bestSaving) {
-        bestSaving = currentCost - swappedCost;
-        bestPeerLogicalProcessor = peerLogicalProcessor;
-      }
-    }
-    std::swap(physicalCoreForLogicalProcessor[logicalProcessor],
-              physicalCoreForLogicalProcessor[bestPeerLogicalProcessor]);
-  }
-  return physicalCoreForLogicalProcessor;
-}
-
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options) {
   const size_t nodeCount = graph.nodes.size();
   const size_t processorCount = options.target.processorCount;
@@ -327,36 +230,17 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
   std::vector<ResidentWeightSet> capacityReservations = planResidentWeightReservations(
     graph, processorCount, options.target.residentWeightCapacity, topology, preferWeightReuse);
 
-  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 costs per processor, access it here.
-  auto getComputeCost = [&](size_t task, size_t processor) -> Time { return graph.nodes[task].cost; };
-  std::vector<Time> oct(nodeCount * processorCount, 0);
-  std::vector<Time> minOctPlusComp(nodeCount, 0);
-
-  // 1. O(P(E+V)) Heterogeneous OCT Calculation
+  // Compute costs are identical on every core. The optimistic successor can
+  // stay on the same core for zero transfer cost, so every physical OCT column
+  // is identical even with nonuniform links. Store that exact lower bound once.
+  std::vector<Time> oct(nodeCount, 0);
   for (const std::vector<size_t>& levelNodes : reverseLevels) {
     auto computeNodeOct = [&](size_t levelIndex) {
       size_t task = levelNodes[levelIndex];
-      std::vector<Time> maxVals(processorCount, 0);
-
-      for (const auto& [succ, comm] : graph.successors[task]) {
-        Time valDifferentCpu = addOrMax(minOctPlusComp[succ], getAverageTransferTime(comm, options.target));
-        for (size_t processor = 0; processor < processorCount; ++processor) {
-          Time valSameCpu = addOrMax(oct[succ * processorCount + processor], getComputeCost(succ, processor));
-          Time bestSucc = std::min(valSameCpu, valDifferentCpu);
-          maxVals[processor] = std::max(maxVals[processor], bestSucc);
-        }
-      }
-
-      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], getComputeCost(task, processor)));
-      }
-      minOctPlusComp[task] = minForPreds == std::numeric_limits<Time>::max() ? 0 : minForPreds;
+      for (const auto& [succ, comm] : graph.successors[task])
+        oct[task] = std::max(oct[task], addOrMax(oct[succ], graph.nodes[succ].cost));
     };
 
     if (options.context != nullptr)
@@ -373,9 +257,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
   };
   std::vector<RankEntry> ranks(nodeCount);
   auto computeRank = [&](size_t node) {
-    long double rank = 0.0L;
-    for (size_t processor = 0; processor < processorCount; ++processor)
-      rank += static_cast<long double>(oct[node * processorCount + processor]);
+    long double rank = static_cast<long double>(oct[node]);
     ranks[node] = {rank, node, graph.nodes[node].originalOrder};
   };
 
@@ -405,7 +287,6 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     }
 
     std::vector<char> scheduled(nodeCount, false);
-    std::vector<ResidentWeightSet> reservations = capacityReservations;
     std::vector<ResidentWeightSet> processorResidentWeights(processorCount);
     std::vector<ScheduledTask> schedules(nodeCount);
     std::vector<std::vector<size_t>> tasksByProcessor(processorCount);
@@ -435,9 +316,9 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
         unsigned int overlapCount =
           countResidentWeightOverlap(processorResidentWeights[processor], graph.nodes[task].residentWeights);
         size_t residentWeightUnion =
-          getResidentWeightUnionSize(reservations[processor], graph.nodes[task].residentWeights);
+          getResidentWeightUnionSize(capacityReservations[processor], graph.nodes[task].residentWeights);
         smallestResidentWeightUnion = std::min(smallestResidentWeightUnion, residentWeightUnion);
-        if (!graph.nodes[task].residentWeights.empty() && residentWeightUnion > options.target.residentWeightCapacity) {
+        if (!graph.nodes[task].residentWeights.empty() && residentWeightUnion > capacityReservations[processor].size()) {
           residentWeightRejected = true;
           continue;
         }
@@ -448,7 +329,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
           dataReady = std::max(dataReady, addOrMax(predSchedule.endTime, commPenalty));
         }
 
-        Time computeCost = getComputeCost(task, processor);
+        Time computeCost = graph.nodes[task].cost;
         Time est = dataReady;
         Time currentEnd = 0;
         bool foundGap = false;
@@ -467,7 +348,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
           est = std::max(currentEnd, dataReady);
 
         Time eft = addOrMax(est, computeCost);
-        Time oeft = addOrMax(eft, oct[task * processorCount + processor]);
+        Time oeft = addOrMax(eft, oct[task]);
         size_t centerDistance = topology.getCenterDistance(processor);
         size_t taskCount = tasksByProcessor[processor].size();
         bool betterResidentWeightChoice =
@@ -519,7 +400,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
           const ComputeInstance& instance = graph.nodes[task].instance;
           std::string message =
             llvm::formatv("PEFT scheduler: no valid processor for task {0} (lanes {1}..{2}, {3} distinct weights); "
-                          "smallest processor union is {4}, exceeding resident-weight capacity {5}",
+                          "no physical reservation contains the task weights (smallest union {4}, core capacity {5})",
                           graph.nodes[task].originalOrder,
                           instance.laneStart,
                           instance.laneStart + instance.laneCount,
@@ -539,7 +420,6 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
       schedules[task] = {bestProcessor, bestEst, bestEft};
       scheduled[task] = true;
       ++scheduledCount;
-      insertResidentWeights(reservations[bestProcessor], graph.nodes[task].residentWeights);
       insertResidentWeights(processorResidentWeights[bestProcessor], graph.nodes[task].residentWeights);
 
       auto& timeline = timelineByProcessor[bestProcessor];

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp

@@ -21,16 +21,11 @@ inline Time getPeftTransferTime(const TransferCost& transferCost,
   if (sourceProcessor == targetProcessor)
     return 0;
   return checkedAdd(transferCost.fixed,
-                    checkedMultiply(transferCost.networkFlits, target.averageInterProcessorLatencyNs));
+                    checkedMultiply(transferCost.networkFlits, target.getInterProcessorLatencyNs(sourceProcessor, targetProcessor)));
 }
 
-// PEFT assigns logical processors. Physical core IDs are chosen only after
-// materialization exposes the exact transfer traffic.
+// PEFT assigns physical cores using the injected target topology.
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options);
 
-std::vector<size_t> mapLogicalProcessorsToPhysicalCores(llvm::ArrayRef<Cost> logicalTrafficFlits,
-                                                        const SchedulingTarget& target,
-                                                        llvm::ArrayRef<size_t> placementGroups = {});
-
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp

@@ -49,21 +49,8 @@ struct PipelineTaskModel {
   std::vector<TaskList> predecessors;
 };
 
-static bool fits(const ComputeGraph &graph,
-                 const TaskCosts &taskCosts,
-                 const TaskList &tasks,
-                 Cost maximumCost,
-                 size_t residentWeightCapacity,
-                 size_t maximumPacks);
-static Cost findMaximumPackCost(const ComputeGraph &graph,
-                                const TaskCosts &taskCosts,
-                                const TaskList &tasks,
-                                size_t residentWeightCapacity,
-                                size_t pipelineStages);
-
 static BlueprintAssemblyInfo getBlueprintAssemblyInfo(
     const ComputeGraph &graph, const ComputeGraphNode &node,
-    size_t consumerCpu, const MergeScheduleResult &schedule,
     const SchedulingTarget &target) {
   BlueprintAssemblyInfo info;
   if (!node.instance.op)
@@ -112,12 +99,7 @@ static BlueprintAssemblyInfo getBlueprintAssemblyInfo(
           elements, static_cast<Cost>(fragmentSizes[fragment * rank + dimension]));
       Cost bytes = (checkedMultiply(elements, target.computeBitwidth) + 7) / 8;
       TransferCost transferCost = getTransferCostFromBytes(bytes, target);
-      size_t producerCpu = schedule.computeToCpuMap.lookup(producerInstance);
-      Cost transferCycles = (producerCpu == consumerCpu
-                               ? transferCost.fixed
-                               : getPeftTransferTime(
-                                   transferCost, producerCpu, consumerCpu, target))
-                          / target.processorPeriodNs;
+      Cost transferCycles = transferCost.fixed / target.processorPeriodNs;
       info.schedulingCost = checkedAdd(
         info.schedulingCost, checkedAdd(elements, transferCycles));
       info.work = checkedAdd(info.work, elements);
@@ -127,8 +109,7 @@ static BlueprintAssemblyInfo getBlueprintAssemblyInfo(
 }
 
 static PipelineTaskModel getPipelineTaskModel(
-    const ComputeGraph &graph, const MergeScheduleResult &schedule,
-    const SchedulingTarget &target) {
+    const ComputeGraph &graph, const SchedulingTarget &target) {
   PipelineTaskModel model;
   model.schedulingCosts.reserve(graph.nodes.size());
   model.assemblyCosts.reserve(graph.nodes.size());
@@ -143,9 +124,8 @@ static PipelineTaskModel getPipelineTaskModel(
     }
 
   for (size_t task = 0; task < graph.nodes.size(); ++task) {
-    size_t cpu = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
     BlueprintAssemblyInfo blueprint = getBlueprintAssemblyInfo(
-      graph, graph.nodes[task], cpu, schedule, target);
+      graph, graph.nodes[task], target);
     model.schedulingCosts[task] = checkedAdd(
       model.schedulingCosts[task], blueprint.schedulingCost);
     model.assemblyCosts[task] = checkedAdd(
@@ -156,25 +136,11 @@ static PipelineTaskModel getPipelineTaskModel(
     for (const auto &[predecessor, transferCost] : graph.predecessors[task]) {
       if (!llvm::is_contained(model.predecessors[task], predecessor))
         model.predecessors[task].push_back(predecessor);
-      size_t predecessorCpu =
-        schedule.computeToCpuMap.lookup(graph.nodes[predecessor].instance);
-      if (predecessorCpu == cpu)
-        continue;
-      model.schedulingCosts[task] = checkedAdd(
-        model.schedulingCosts[task], getPeftTransferTime(
-                                       transferCost, predecessorCpu, cpu, target)
-                                     / target.processorPeriodNs);
     }
   }
   return model;
 }
 
-static TaskCosts getPipelineTaskCosts(const ComputeGraph &graph,
-                                      const MergeScheduleResult &schedule,
-                                      const SchedulingTarget &target) {
-  return getPipelineTaskModel(graph, schedule, target).schedulingCosts;
-}
-
 static void copyPipelineComputeAttrs(SpatGraphComputeBatch source,
                                      SpatGraphComputeBatch destination) {
   for (NamedAttribute attr : source->getAttrs()) {
@@ -609,130 +575,23 @@ static LogicalResult splitBatchCompute(SpatGraphComputeBatch batch,
 }
 
 static FailureOr<PipelineWorkloadPreparation> preparePipelineWorkloadImpl(
-    const ComputeGraph &graph, const MergeScheduleResult &schedule,
-    size_t pipelineStages, const SchedulingTarget &physicalTarget,
-    std::string &error) {
-  size_t groupSize = schedule.processorCount;
-  std::vector<TaskList> tasksByCpu(groupSize);
-  for (size_t task = 0; task < graph.nodes.size(); ++task) {
-    auto cpu = schedule.computeToCpuMap.find(graph.nodes[task].instance);
-    if (cpu == schedule.computeToCpuMap.end() || cpu->second >= groupSize) {
-      error = "pipeline split received an incomplete PEFT schedule for task "
-            + std::to_string(task) + " (cpu "
-            + (cpu == schedule.computeToCpuMap.end()
-                 ? std::string("missing")
-                 : std::to_string(cpu->second))
-            + ", schedule processors " + std::to_string(groupSize) + ")";
-      return failure();
-    }
-    tasksByCpu[cpu->second].push_back(task);
-  }
-  for (TaskList &tasks : tasksByCpu)
-    llvm::sort(tasks, [&](size_t lhs, size_t rhs) {
-      return schedule.computeToCpuSlotMap.lookup(graph.nodes[lhs].instance)
-           < schedule.computeToCpuSlotMap.lookup(graph.nodes[rhs].instance);
-    });
-  TaskCosts taskCosts = getPipelineTaskCosts(graph, schedule, physicalTarget);
-
-  for (const TaskList &tasks : tasksByCpu) {
-    if (tasks.empty())
+    const ComputeGraph &graph, size_t pipelineStages,
+    const SchedulingTarget &physicalTarget, std::string &error) {
+  for (const ComputeGraphNode &node : graph.nodes) {
+    if (node.residentWeights.size() <= physicalTarget.residentWeightCapacity)
       continue;
-    Cost maximumCost = findMaximumPackCost(
-      graph, taskCosts, tasks, physicalTarget.residentWeightCapacity,
-      pipelineStages);
-    if (fits(graph, taskCosts, tasks, maximumCost,
-             physicalTarget.residentWeightCapacity, pipelineStages))
-      continue;
-
-    SmallVector<size_t, 8> candidates(tasks.begin(), tasks.end());
-    llvm::stable_sort(candidates, [&](size_t lhs, size_t rhs) {
-      if (graph.nodes[lhs].residentWeights.size()
-          != graph.nodes[rhs].residentWeights.size())
-        return graph.nodes[lhs].residentWeights.size()
-             > graph.nodes[rhs].residentWeights.size();
-      return taskCosts[lhs] > taskCosts[rhs];
-    });
-    std::string candidateError;
-    for (size_t task : candidates) {
-      auto batch = dyn_cast_or_null<SpatGraphComputeBatch>(
-        graph.nodes[task].instance.op);
-      if (!batch || batch->hasAttr("pipeline.split"))
-        continue;
-      std::string currentError;
-      if (succeeded(splitBatchCompute(
-            batch, pipelineStages, physicalTarget, currentError)))
-        return PipelineWorkloadPreparation::Changed;
-      if (!currentError.empty())
-        candidateError = currentError;
-    }
-    if (llvm::any_of(tasks, [&](size_t task) {
-          return graph.nodes[task].residentWeights.size()
-               > physicalTarget.residentWeightCapacity;
-        })) {
-      error = candidateError.empty()
-        ? "pipeline scheduling cannot split a compute instance to fit one "
-          "physical core's crossbars"
-        : candidateError;
-      return failure();
-    }
+    auto batch = dyn_cast_or_null<SpatGraphComputeBatch>(node.instance.op);
+    if (batch && succeeded(splitBatchCompute(
+                   batch, pipelineStages, physicalTarget, error)))
+      return PipelineWorkloadPreparation::Changed;
+    if (error.empty())
+      error = "pipeline scheduling cannot split a compute instance to fit one "
+              "physical core's crossbars";
+    return failure();
   }
-
   return PipelineWorkloadPreparation::Ready;
 }
 
-bool fits(const ComputeGraph& graph,
-          const TaskCosts& taskCosts,
-          const TaskList& tasks,
-          Cost maximumCost,
-          size_t residentWeightCapacity,
-          size_t maximumPacks) {
-  size_t packs = 1;
-  Cost cost = 0;
-  ResidentWeightSet weights;
-  bool packEmpty = true;
-  for (size_t task : tasks) {
-    const ComputeGraphNode& node = graph.nodes[task];
-    Cost taskCost = taskCosts[task];
-    if (node.residentWeights.size() > residentWeightCapacity)
-      return false;
-    bool startsNewPack = !packEmpty
-                      && (cost > maximumCost - taskCost
-                          || getResidentWeightUnionSize(weights, node.residentWeights) > residentWeightCapacity);
-    if (startsNewPack) {
-      if (++packs > maximumPacks)
-        return false;
-      cost = 0;
-      weights.clear();
-      packEmpty = true;
-    }
-    cost = checkedAdd(cost, taskCost);
-    insertResidentWeights(weights, node.residentWeights);
-    packEmpty = false;
-  }
-  return true;
-}
-
-Cost findMaximumPackCost(const ComputeGraph& graph,
-                         const TaskCosts& taskCosts,
-                         const TaskList& tasks,
-                         size_t residentWeightCapacity,
-                         size_t pipelineStages) {
-  Cost low = 0;
-  Cost high = 0;
-  for (size_t task : tasks) {
-    low = std::max(low, taskCosts[task]);
-    high = checkedAdd(high, taskCosts[task]);
-  }
-  while (low < high) {
-    Cost middle = low + (high - low) / 2;
-    if (fits(graph, taskCosts, tasks, middle, residentWeightCapacity, pipelineStages))
-      high = middle;
-    else
-      low = middle + 1;
-  }
-  return low;
-}
-
 static Cost findMaximumIndexedPackCost(
     const TaskCosts &taskCosts,
     const std::vector<TaskList> &taskWeightIds, size_t weightCount,
@@ -1146,11 +1005,12 @@ static bool packPipelineStage(
     std::vector<TaskList> &tasksByCpu, const PipelineCoreLayout &layout,
     ArrayRef<size_t> topologicalPosition, size_t stage,
     size_t residentWeightCapacity, const SchedulingTarget &target,
-    std::vector<size_t> &taskCpus, bool prioritizeWeightReuse = false) {
+    std::vector<size_t> &taskCpus, std::vector<Time> &taskEndTimes,
+    bool prioritizeWeightReuse = false) {
   PipelineStageRange range = layout.getStageRange(stage);
   TaskList tasks;
   for (size_t cpu = range.begin; cpu < range.begin + range.size; ++cpu)
-    llvm::append_range(tasks, tasksByCpu[cpu]);
+    llvm::append_range(tasks, tasksByCpu[layout.getPhysicalCore(cpu)]);
   llvm::sort(tasks, [&](size_t lhs, size_t rhs) {
     return topologicalPosition[lhs] < topologicalPosition[rhs];
   });
@@ -1158,8 +1018,10 @@ static bool packPipelineStage(
   std::vector<ResidentWeightSet> weights(range.size);
   TaskCosts loads(range.size);
   TaskCosts assemblyLoads(range.size);
+  std::vector<Time> coreReady(range.size);
   for (size_t task : tasks) {
     std::optional<size_t> bestCore;
+    Time bestEndTime = 0;
     using PackScore = std::tuple<Cost, Time, Cost, Cost, size_t>;
     std::optional<PackScore> bestScore;
     for (size_t core = 0; core < range.size; ++core) {
@@ -1172,23 +1034,29 @@ static bool packPipelineStage(
         assemblyLoads[core], assemblyCosts[task]);
       Cost schedulingLoad = checkedAdd(
         loads[core], schedulingCosts[task]);
-      Time transferTime = 0;
-      size_t candidateCpu = range.begin + core;
+      Time startTime = coreReady[core];
+      size_t candidateCpu = layout.getPhysicalCore(range.begin + core);
+      // Only same-stage edges use direct sends/receives. Pipeline boundaries
+      // use host buffers and belong to different samples in steady state.
       for (const auto &[predecessor, transferCost] :
            graph.predecessors[task])
-        if (taskCpus[predecessor] < target.processorCount)
-          transferTime = checkedAdd(
-            transferTime, getPeftTransferTime(
-                            transferCost, taskCpus[predecessor],
-                            candidateCpu, target));
+        if (taskCpus[predecessor] < target.processorCount
+            && layout.getStageForCore(taskCpus[predecessor]) == stage) {
+          Time transfer = getPeftTransferTime(
+            transferCost, taskCpus[predecessor], candidateCpu, target);
+          startTime = std::max(
+            startTime, addOrMax(taskEndTimes[predecessor], transfer));
+        }
+      Time endTime = addOrMax(startTime, schedulingCosts[task]);
       PackScore score = prioritizeWeightReuse
         ? PackScore {
-            addedWeights, assemblyLoad, transferTime, schedulingLoad, core}
+            addedWeights, endTime, assemblyLoad, schedulingLoad, core}
         : PackScore {
-            assemblyLoad, transferTime, schedulingLoad, addedWeights, core};
+            endTime, assemblyLoad, schedulingLoad, addedWeights, core};
       if (!bestScore || score < *bestScore) {
         bestCore = core;
         bestScore = score;
+        bestEndTime = endTime;
       }
     }
     if (!bestCore)
@@ -1199,10 +1067,12 @@ static bool packPipelineStage(
     loads[*bestCore] = checkedAdd(loads[*bestCore], schedulingCosts[task]);
     assemblyLoads[*bestCore] = checkedAdd(
       assemblyLoads[*bestCore], assemblyCosts[task]);
-    taskCpus[task] = range.begin + *bestCore;
+    taskCpus[task] = layout.getPhysicalCore(range.begin + *bestCore);
+    taskEndTimes[task] = bestEndTime;
+    coreReady[*bestCore] = taskEndTimes[task];
   }
   for (size_t core = 0; core < range.size; ++core)
-    tasksByCpu[range.begin + core] = std::move(packed[core]);
+    tasksByCpu[layout.getPhysicalCore(range.begin + core)] = std::move(packed[core]);
   return true;
 }
 
@@ -1244,15 +1114,16 @@ static LogicalResult packPipelineStages(
 
   const TaskCosts &balanceCosts = getPipelineBalanceCosts(graph, model);
   std::vector<size_t> taskCpus(graph.nodes.size(), target.processorCount);
+  std::vector<Time> taskEndTimes(graph.nodes.size());
   for (size_t stage = 0; stage < pipelineStages; ++stage)
     if (!packPipelineStage(
           graph, model.schedulingCosts, balanceCosts, tasksByCpu,
           layout, topologicalPosition, stage, residentWeightCapacity,
-          target, taskCpus)
+          target, taskCpus, taskEndTimes)
         && !packPipelineStage(
           graph, model.schedulingCosts, balanceCosts, tasksByCpu,
           layout, topologicalPosition, stage, residentWeightCapacity,
-          target, taskCpus, /*prioritizeWeightReuse=*/true)) {
+          target, taskCpus, taskEndTimes, /*prioritizeWeightReuse=*/true)) {
       failedStage = stage;
       error = "pipeline scheduling cannot pack dependency-monotone stage "
             + std::to_string(stage)
@@ -1330,7 +1201,7 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
   PipelineCoreLayout balancedLayout(
     physicalTarget.processorCount, pipelineStages);
   if (!balancedLayout.isValid()
-      || groupSize != balancedLayout.getLogicalProcessorCount()) {
+      || groupSize != physicalTarget.processorCount) {
     error = "pipeline scheduling received an incompatible physical core layout";
     return mlir::failure();
   }
@@ -1345,7 +1216,7 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
     }
   }
   PipelineTaskModel taskModel = getPipelineTaskModel(
-    graph, schedule, physicalTarget);
+    graph, physicalTarget);
   FailureOr<PipelineStageAssignment> assignment = assignPipelineStages(
     graph, taskModel, balancedLayout,
     physicalTarget.residentWeightCapacity, error);
@@ -1356,13 +1227,13 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
   std::string packingError;
   bool packed = false;
   for (size_t attempt = 0; attempt < physicalTarget.processorCount; ++attempt) {
-    PipelineCoreLayout candidateLayout(assignment->stageSizes);
+    PipelineCoreLayout candidateLayout(assignment->stageSizes, physicalTarget);
     for (TaskList &tasks : tasksByPhysicalCpu)
       tasks.clear();
     for (size_t task = 0; task < graph.nodes.size(); ++task) {
       PipelineStageRange range =
         candidateLayout.getStageRange(assignment->taskStages[task]);
-      tasksByPhysicalCpu[range.begin].push_back(task);
+      tasksByPhysicalCpu[candidateLayout.getPhysicalCore(range.begin)].push_back(task);
     }
     size_t failedStage = 0;
     if (succeeded(packPipelineStages(
@@ -1393,7 +1264,7 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
     error = packingError;
     return failure();
   }
-  PipelineCoreLayout pipelineLayout(assignment->stageSizes);
+  PipelineCoreLayout pipelineLayout(assignment->stageSizes, physicalTarget);
   if (failed(verifyPipelineStageAssignment(
         graph, taskModel, tasksByPhysicalCpu, pipelineLayout, error)))
     return failure();
@@ -1403,8 +1274,8 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
   schedule.processorStages.resize(physicalTarget.processorCount);
   for (size_t stage = 0; stage < pipelineLayout.getStageCount(); ++stage) {
     PipelineStageRange range = pipelineLayout.getStageRange(stage);
-    std::fill_n(
-      schedule.processorStages.begin() + range.begin, range.size, stage);
+    for (size_t index = range.begin; index < range.begin + range.size; ++index)
+      schedule.processorStages[pipelineLayout.getPhysicalCore(index)] = stage;
   }
   schedule.computeToCpuSlotMap.clear();
   schedule.computeToAestMap.clear();
@@ -1478,6 +1349,48 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
 
 } // namespace
 
+PipelineCoreLayout::PipelineCoreLayout(ArrayRef<size_t> sizes,
+                                       const SchedulingTarget &target)
+    : PipelineCoreLayout(sizes) {
+  assert(processorCount == target.processorCount);
+  std::vector<bool> assigned(processorCount);
+  std::vector<Cost> remainingDistance(processorCount);
+  auto distanceBetween = [&](size_t source, size_t destination) {
+    return checkedAdd(target.getInterProcessorLatencyNs(source, destination),
+                      target.getInterProcessorLatencyNs(destination, source));
+  };
+  for (size_t core = 0; core < processorCount; ++core)
+    for (size_t other = 0; other < processorCount; ++other)
+      remainingDistance[core] = checkedAdd(
+        remainingDistance[core], distanceBetween(core, other));
+  physicalCores.reserve(processorCount);
+  physicalCoreStages.resize(processorCount);
+  // ponytail: greedy O(P^2) clustering; use a graph partitioner only if measured
+  // intra-stage communication warrants a more expensive placement search.
+  for (auto [stage, size] : llvm::enumerate(stageSizes)) {
+    std::vector<Cost> distance(processorCount);
+    for (size_t index = 0; index < size; ++index) {
+      size_t best = processorCount;
+      for (size_t core = 0; core < processorCount; ++core)
+        if (!assigned[core]
+            && (best == processorCount
+                || (index == 0 ? remainingDistance[core] > remainingDistance[best]
+                               : distance[core] < distance[best])))
+          best = core;
+      assert(best != processorCount);
+      assigned[best] = true;
+      physicalCores.push_back(best);
+      physicalCoreStages[best] = stage;
+      for (size_t core = 0; core < processorCount; ++core)
+        if (!assigned[core]) {
+          Cost link = distanceBetween(best, core);
+          distance[core] = checkedAdd(distance[core], link);
+          remainingDistance[core] -= link;
+        }
+    }
+  }
+}
+
 mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
                                             MergeScheduleResult& schedule,
                                             size_t pipelineStages,
@@ -1489,7 +1402,7 @@ mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
     physicalTarget.processorCount, pipelineStages);
   if (!pipelineLayout.isValid() || schedule.processorCount == 0
       || schedule.processorCount
-           != pipelineLayout.getLogicalProcessorCount()) {
+           != physicalTarget.processorCount) {
     error = "pipeline scheduling requires a valid balanced physical core layout";
     return mlir::failure();
   }
@@ -1497,11 +1410,10 @@ mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
 }
 
 mlir::FailureOr<PipelineWorkloadPreparation> preparePipelineWorkload(
-    const ComputeGraph &graph, const MergeScheduleResult &schedule,
-    size_t pipelineStages, const SchedulingTarget &physicalTarget,
-    std::string &error) {
+    const ComputeGraph &graph, size_t pipelineStages,
+    const SchedulingTarget &physicalTarget, std::string &error) {
   return preparePipelineWorkloadImpl(
-    graph, schedule, pipelineStages, physicalTarget, error);
+    graph, pipelineStages, physicalTarget, error);
 }
 
 } // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp

@@ -40,23 +40,24 @@ public:
           stageSizes.begin(), stageSizes.end(), size_t {0})),
         stageSizes(stageSizes.begin(), stageSizes.end()) {}
 
+  PipelineCoreLayout(llvm::ArrayRef<size_t> stageSizes,
+                     const SchedulingTarget &target);
+
   bool isValid() const {
     return !stageSizes.empty()
         && llvm::none_of(stageSizes, [](size_t size) { return size == 0; });
   }
 
-  size_t getLogicalProcessorCount() const {
-    return isValid()
-      ? *std::min_element(stageSizes.begin(), stageSizes.end())
-      : 0;
-  }
-
   size_t getStageCount() const { return stageSizes.size(); }
 
   size_t getProcessorCount() const { return processorCount; }
 
   llvm::ArrayRef<size_t> getStageSizes() const { return stageSizes; }
 
+  size_t getPhysicalCore(size_t position) const {
+    return physicalCores.empty() ? position : physicalCores[position];
+  }
+
   PipelineStageRange getStageRange(size_t stage) const {
     return {std::accumulate(
               stageSizes.begin(), stageSizes.begin() + stage, size_t {0}),
@@ -66,6 +67,8 @@ public:
   std::optional<size_t> getStageForCore(size_t core) const {
     if (!isValid() || core >= processorCount)
       return std::nullopt;
+    if (!physicalCoreStages.empty())
+      return physicalCoreStages[core];
     size_t end = 0;
     for (auto [stage, size] : llvm::enumerate(stageSizes)) {
       end += size;
@@ -78,6 +81,8 @@ public:
 private:
   size_t processorCount;
   std::vector<size_t> stageSizes;
+  std::vector<size_t> physicalCores;
+  std::vector<size_t> physicalCoreStages;
 };
 
 mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
@@ -92,8 +97,8 @@ enum class PipelineWorkloadPreparation {
 };
 
 mlir::FailureOr<PipelineWorkloadPreparation> preparePipelineWorkload(
-    const ComputeGraph& graph, const MergeScheduleResult& schedule,
-    size_t pipelineStages, const SchedulingTarget& physicalTarget,
+    const ComputeGraph& graph, size_t pipelineStages,
+    const SchedulingTarget& physicalTarget,
     std::string& error);
 
 } // namespace onnx_mlir::spatial