finally fast googlenet with correct latency artifacts for fair comparison

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp

@@ -262,28 +262,30 @@ static void collectClosure(Value value, Block &body, const DeferredInputPlan &pl
 
 } // namespace
 
-bool isDeferredFragmentAssemblyInput(Value input) {
+bool isDeferredFragmentAssemblyInput(Value input, size_t processorCount) {
   auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
   if (!blueprint || blueprint.getMode() != "fragment_assembly")
     return false;
   return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) {
-    return getProducerValueRef(fragment, nullptr).has_value();
+    return getProducerValueRef(fragment, nullptr, processorCount).has_value();
   });
 }
 
 LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
-                                    const ComputeInstance &consumerInstance, const MergeScheduleResult &,
+                                    const ComputeInstance &consumerInstance,
+                                    const MergeScheduleResult &schedule,
                                     ValueRange scheduledInputs, Block &block, unsigned firstInputArgument,
                                     const DenseMap<ProducerValueKey, MaterializedProducerRef> &availableValues,
                                     Value graphLane, Value scheduledGraphLane,
                                     DeferredInputPlan &plan) {
   plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}, {}, {}, {}, 1, nullptr};
-  if (isDeferredFragmentAssemblyInput(input)) {
+  if (isDeferredFragmentAssemblyInput(input, schedule.processorCount)) {
     plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
     plan.originalSources = getBlueprintFragments(plan.blueprint);
     return success();
   }
-  auto producer = getProducerValueRef(input, &consumerInstance);
+  auto producer = getProducerValueRef(
+    input, &consumerInstance, schedule.processorCount);
   if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
   ProducerValueKey key {producer->instance, producer->resultIndex};
   auto batch = dyn_cast<SpatComputeBatch>(producer->instance.op);
@@ -304,17 +306,19 @@ LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, Bloc
 
 LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
                                         const ComputeStepTuple &tuple, const PeftClassPlan &,
-                                        const MergeScheduleResult &, ValueRange scheduledInputs, Block &block,
+                                        const MergeScheduleResult &schedule,
+                                        ValueRange scheduledInputs, Block &block,
                                         unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane, Value scheduledLane,
                                         DeferredInputPlan &plan) {
   const ComputeInstance &representative = tuple.instances.front();
   plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane, {}, {}, {}, 1, nullptr};
-  if (isDeferredFragmentAssemblyInput(input)) {
+  if (isDeferredFragmentAssemblyInput(input, schedule.processorCount)) {
     plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
     plan.originalSources = getBlueprintFragments(plan.blueprint);
     return success();
   }
-  auto producer = getProducerValueRef(input, &representative);
+  auto producer = getProducerValueRef(
+    input, &representative, schedule.processorCount);
   if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
   auto inputs = getComputeInstanceInputs(representative);
   auto it = llvm::find(inputs, input);
@@ -323,7 +327,8 @@ LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input,
   for (const ComputeInstance &instance : tuple.instances) {
     auto laneInputs = getComputeInstanceInputs(instance);
     if (inputIndex >= laneInputs.size()) return emitError(loc) << "scheduled batch step input out of range";
-    auto laneProducer = getProducerValueRef(laneInputs[inputIndex], &instance);
+    auto laneProducer = getProducerValueRef(
+      laneInputs[inputIndex], &instance, schedule.processorCount);
     if (!laneProducer) return emitError(loc) << "scheduled batch step mixes host and producer inputs";
     auto source = getOriginalProducerValue(*laneProducer);
     if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value";

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp

@@ -22,7 +22,7 @@ struct DeferredInputPlan {
   Block *scalarizedHoistBlock = nullptr;
 };
 
-bool isDeferredFragmentAssemblyInput(Value input);
+bool isDeferredFragmentAssemblyInput(Value input, size_t processorCount);
 
 LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input,
                                     BlockArgument graphInput,

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

@@ -1,39 +1,97 @@
-#include "DeferredCommunicationRealization.hpp"
+#include "mlir/IR/Dominance.h"
 
 #include "DeferredBoundaryPlanning.hpp"
 #include "DeferredBoundaryRealization.hpp"
 #include "DeferredCommunicationDeadlock.hpp"
+#include "DeferredCommunicationRealization.hpp"
 #include "DeferredCommunicationScheduling.hpp"
 #include "DeferredTransferPlanning.hpp"
-
-#include "mlir/IR/Dominance.h"
+#include "Scheduling/PeftScheduler.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
 namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
 
-static LogicalResult replaceFinalGraphPublications(
-    func::FuncOp funcOp, DeferredTransferPlan &plan) {
-  for (Operation &op : funcOp.getOps()) {
+static LogicalResult placeLogicalProcessorsOnPhysicalCores(DeferredTransferPlan& plan, const SchedulingTarget& target) {
+  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> physicalCoreForLogicalProcessor =
+    mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, target);
+  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))
       continue;
     auto graphId = op.getAttrOfType<IntegerAttr>("scheduled.graph_id");
     if (!graphId)
       continue;
     for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) {
-      SmallVector<OpOperand *> externalUses;
-      for (OpOperand &use : result.getUses()) {
-        Operation *user = use.getOwner();
-        if (isa<SpatGraphCompute, SpatGraphComputeBatch,
-                SpatDeferredCommunicationOp>(user))
+      SmallVector<OpOperand*> externalUses;
+      for (OpOperand& use : result.getUses()) {
+        Operation* user = use.getOwner();
+        if (isa<SpatGraphCompute, SpatGraphComputeBatch, SpatDeferredCommunicationOp>(user))
           continue;
         if (auto blueprint = dyn_cast<SpatBlueprintOp>(user)) {
-          bool blueprintEscapes = llvm::any_of(
-            blueprint.getOutput().getUses(), [](OpOperand &blueprintUse) {
-              return !isa<SpatGraphCompute, SpatGraphComputeBatch,
-                          SpatDeferredCommunicationOp>(
-                blueprintUse.getOwner());
-            });
+          bool blueprintEscapes = llvm::any_of(blueprint.getOutput().getUses(), [](OpOperand& blueprintUse) {
+            return !isa<SpatGraphCompute, SpatGraphComputeBatch, SpatDeferredCommunicationOp>(blueprintUse.getOwner());
+          });
           if (!blueprintEscapes)
             continue;
         }
@@ -42,21 +100,16 @@ static LogicalResult replaceFinalGraphPublications(
       if (externalUses.empty())
         continue;
       SmallVector<Value> exact;
-      for (ProducedValue *produced :
-           plan.producedByGraph.lookup(graphId.getInt()))
-        if (produced->resultIndex == resultIndex
-            && produced->published
-            && produced->published.getType() == result.getType()
-            && !llvm::is_contained(exact, produced->published))
+      for (ProducedValue* produced : plan.producedByGraph.lookup(graphId.getInt()))
+        if (produced->resultIndex == resultIndex && produced->published
+            && produced->published.getType() == result.getType() && !llvm::is_contained(exact, produced->published))
           exact.push_back(produced->published);
       if (exact.size() != 1)
-        return op.emitOpError(
-          "phase 2 final publication ownership changed after planning");
-      for (OpOperand *use : externalUses) {
-        Operation *consumer = use->getOwner();
-        Operation *producer = exact.front().getDefiningOp();
-        if (consumer->getBlock() == producer->getBlock()
-            && consumer->isBeforeInBlock(producer))
+        return op.emitOpError("phase 2 final publication ownership changed after planning");
+      for (OpOperand* use : externalUses) {
+        Operation* consumer = use->getOwner();
+        Operation* producer = exact.front().getDefiningOp();
+        if (consumer->getBlock() == producer->getBlock() && consumer->isBeforeInBlock(producer))
           consumer->moveAfter(producer);
         use->set(exact.front());
       }
@@ -65,13 +118,12 @@ static LogicalResult replaceFinalGraphPublications(
   return success();
 }
 
-static LogicalResult eraseOldGraph(func::FuncOp funcOp,
-                                   IRRewriter &rewriter) {
-  SmallVector<Operation *> oldGraph;
-  for (Operation &op : funcOp.getOps())
+static LogicalResult eraseOldGraph(func::FuncOp funcOp, IRRewriter& rewriter) {
+  SmallVector<Operation*> oldGraph;
+  for (Operation& op : funcOp.getOps())
     if (isa<SpatGraphCompute, SpatGraphComputeBatch, SpatBlueprintOp>(op))
       oldGraph.push_back(&op);
-  for (Operation *op : llvm::reverse(oldGraph)) {
+  for (Operation* op : llvm::reverse(oldGraph)) {
     if (auto blueprint = dyn_cast<SpatBlueprintOp>(op)) {
       if (blueprint.getOutput().use_empty())
         rewriter.eraseOp(blueprint);
@@ -80,10 +132,9 @@ static LogicalResult eraseOldGraph(func::FuncOp funcOp,
     if (!op->use_empty()) {
       for (OpResult result : op->getResults()) {
         if (!result.use_empty()) {
-          Operation *user = result.use_begin()->getOwner();
-          return op->emitOpError()
-                 << "phase 2 cannot erase old graph result "
-                 << result.getResultNumber() << " used by " << user->getName();
+          Operation* user = result.use_begin()->getOwner();
+          return op->emitOpError() << "phase 2 cannot erase old graph result " << result.getResultNumber()
+                                   << " used by " << user->getName();
         }
       }
     }
@@ -92,34 +143,27 @@ static LogicalResult eraseOldGraph(func::FuncOp funcOp,
   return success();
 }
 
-static LogicalResult eraseDeferredSourceSelectors(
-    func::FuncOp funcOp, IRRewriter &rewriter) {
+static LogicalResult eraseDeferredSourceSelectors(func::FuncOp funcOp, IRRewriter& rewriter) {
   SmallVector<SpatDeferredSourceSelectOp> selectors;
-  funcOp.walk([&](SpatDeferredSourceSelectOp selector) {
-    selectors.push_back(selector);
-  });
+  funcOp.walk([&](SpatDeferredSourceSelectOp selector) { selectors.push_back(selector); });
   for (SpatDeferredSourceSelectOp selector : llvm::reverse(selectors)) {
     if (!selector.getOutput().use_empty())
-      return selector.emitOpError(
-        "phase 2 left a live deferred source selection");
+      return selector.emitOpError("phase 2 left a live deferred source selection");
     rewriter.eraseOp(selector);
   }
   return success();
 }
 
-static void eraseUnusedIdentityDeferredCommunications(
-    func::FuncOp funcOp, IRRewriter &rewriter) {
+static void eraseUnusedIdentityDeferredCommunications(func::FuncOp funcOp, IRRewriter& rewriter) {
   SmallVector<SpatDeferredCommunicationOp> unused;
   funcOp.walk([&](SpatDeferredCommunicationOp deferred) {
     if (!deferred.getOutput().use_empty() || !deferred.getBody().hasOneBlock())
       return;
-    Block &body = deferred.getBody().front();
+    Block& body = deferred.getBody().front();
     auto yield = dyn_cast<SpatYieldOp>(body.getTerminator());
-    auto argument = yield && yield.getOutputs().size() == 1
-                      ? dyn_cast<BlockArgument>(yield.getOutputs().front())
-                      : BlockArgument();
-    if (argument && argument.getOwner() == &body
-        && argument.getArgNumber() < deferred.getSources().size())
+    auto argument =
+      yield && yield.getOutputs().size() == 1 ? dyn_cast<BlockArgument>(yield.getOutputs().front()) : BlockArgument();
+    if (argument && argument.getOwner() == &body && argument.getArgNumber() < deferred.getSources().size())
       unused.push_back(deferred);
   });
   for (SpatDeferredCommunicationOp deferred : llvm::reverse(unused))
@@ -128,11 +172,10 @@ static void eraseUnusedIdentityDeferredCommunications(
 
 static LogicalResult verifyDominance(func::FuncOp funcOp) {
   DominanceInfo dominance(funcOp);
-  WalkResult result = funcOp.walk([&](Operation *op) {
+  WalkResult result = funcOp.walk([&](Operation* op) {
     for (auto [index, operand] : llvm::enumerate(op->getOperands()))
       if (!dominance.dominates(operand, op)) {
-        op->emitOpError() << "phase 2 produced non-dominating operand "
-                          << index << ": " << operand;
+        op->emitOpError() << "phase 2 produced non-dominating operand " << index << ": " << operand;
         return WalkResult::interrupt();
       }
     return WalkResult::advance();
@@ -142,26 +185,23 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) {
 
 } // namespace
 
-LogicalResult realizeDeferredCommunication(
-    func::FuncOp funcOp,
-    const ScheduledComputeMaterializationResult &materialization) {
+LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
+                                           const ScheduledComputeMaterializationResult& materialization,
+                                           const SchedulingTarget& target) {
   IRRewriter rewriter(funcOp.getContext());
   eraseUnusedIdentityDeferredCommunications(funcOp, rewriter);
 
   auto transfers = buildDeferredTransferPlan(funcOp, materialization);
   if (failed(transfers))
-    return funcOp.emitOpError(
-      "phase 2 failed to build symbolic transfer families");
+    return funcOp.emitOpError("phase 2 failed to build symbolic transfer families");
+  if (failed(placeLogicalProcessorsOnPhysicalCores(*transfers, target)))
+    return failure();
   auto schedule = scheduleDeferredCommunication(funcOp, *transfers);
-  if (failed(schedule)
-      || failed(verifyPlannedCommunicationDeadlockFree(
-        funcOp, transfers->stepCounts, *schedule)))
-    return funcOp.emitOpError(
-      "phase 2 failed to schedule symbolic communication");
+  if (failed(schedule) || failed(verifyPlannedCommunicationDeadlockFree(funcOp, transfers->stepCounts, *schedule)))
+    return funcOp.emitOpError("phase 2 failed to schedule symbolic communication");
   auto boundaries = buildDeferredBoundaryPlan(*transfers, *schedule);
   if (failed(boundaries))
-    return funcOp.emitOpError(
-      "phase 2 failed to build sparse boundary programs");
+    return funcOp.emitOpError("phase 2 failed to build sparse boundary programs");
 
   if (failed(retargetDeferredPublications(funcOp, *transfers))
       || failed(replaceFinalGraphPublications(funcOp, *transfers)))
@@ -169,28 +209,22 @@ LogicalResult realizeDeferredCommunication(
   ConstantPool constants(funcOp, rewriter);
   DeferredEmissionContext context(rewriter, constants);
   DeferredReplacementMap replacements;
-  if (failed(realizeDeferredBoundaries(
-        boundaries->boundaries, boundaries->results, context, replacements)))
+  if (failed(realizeDeferredBoundaries(boundaries->boundaries, boundaries->results, context, replacements)))
     return failure();
   for (auto [op, replacement] : replacements) {
     if (op->getResult(0) == replacement)
-      return op->emitOpError(
-        "phase 2 cannot replace deferred communication with itself");
+      return op->emitOpError("phase 2 cannot replace deferred communication with itself");
     op->getResult(0).replaceAllUsesWith(replacement);
     if (!op->use_empty())
-      return op->emitOpError(
-        "phase 2 cannot erase deferred communication with live uses");
+      return op->emitOpError("phase 2 cannot erase deferred communication with live uses");
     rewriter.eraseOp(op);
   }
-  if (failed(eraseDeferredSourceSelectors(funcOp, rewriter))
-      || failed(eraseOldGraph(funcOp, rewriter))
-      || failed(verifyDominance(funcOp))
-      || failed(verifyRealizedCommunicationDeadlockFree(funcOp, *schedule)))
+  if (failed(eraseDeferredSourceSelectors(funcOp, rewriter)) || failed(eraseOldGraph(funcOp, rewriter))
+      || failed(verifyDominance(funcOp)) || failed(verifyRealizedCommunicationDeadlockFree(funcOp, *schedule)))
     return failure();
   bool deferredRemains = false;
   funcOp.walk([&](SpatDeferredCommunicationOp deferred) {
-    deferred.emitOpError(
-      "phase 2 left an unrealized deferred communication");
+    deferred.emitOpError("phase 2 left an unrealized deferred communication");
     deferredRemains = true;
   });
   return success(!deferredRemains);

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp

@@ -5,9 +5,10 @@
 namespace onnx_mlir::spatial {
 
 struct ScheduledComputeMaterializationResult;
+struct SchedulingTarget;
 
-mlir::LogicalResult realizeDeferredCommunication(
-    mlir::func::FuncOp funcOp,
-    const ScheduledComputeMaterializationResult &materialization);
+mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp,
+                                                 const ScheduledComputeMaterializationResult& materialization,
+                                                 const SchedulingTarget& target);
 
 } // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp

@@ -1,15 +1,14 @@
+#include "mlir/Pass/Pass.h"
+
+#include "DeferredCommunicationRealization.hpp"
 #include "ScheduledComputeMaterialization.hpp"
 #include "ScheduledComputeReport.hpp"
 #include "ScheduledComputeVerification.hpp"
-#include "SpatialDataflowCsvExporter.hpp"
-#include "DeferredCommunicationRealization.hpp"
-
-#include "mlir/Pass/Pass.h"
-
 #include "Scheduling/MergeSchedulingAnalysis.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "SpatialDataflowCsvExporter.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace mlir;
@@ -21,6 +20,10 @@ namespace {
 struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass)
 
+  MergeComputeNodesPass() = default;
+  explicit MergeComputeNodesPass(const SchedulingTarget& schedulingTarget)
+  : target(schedulingTarget), hasTarget(true) {}
+
   StringRef getArgument() const override { return "pim-merge-compute-nodes"; }
   StringRef getDescription() const override {
     return "Materialize scheduled Spatial compute with deferred communication placeholders.";
@@ -28,6 +31,13 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
 
   void runOnOperation() override {
     ModuleOp moduleOp = getOperation();
+    if (!hasTarget || target.processorCount == 0 || target.residentWeightCapacity == 0 || target.transferWidthBytes == 0
+        || target.interProcessorLatencyNs.size() != target.processorCount * target.processorCount
+        || (target.processorCount > 1 && target.averageInterProcessorLatencyNs == 0)) {
+      moduleOp.emitError("MergeComputeNodes requires an explicit valid Spatial scheduling target");
+      signalPassFailure();
+      return;
+    }
     auto entryFunc = getPimEntryFunc(moduleOp);
     if (failed(entryFunc)) {
       moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes");
@@ -36,10 +46,10 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
     }
 
     func::FuncOp funcOp = *entryFunc;
-    MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult();
+    MergeScheduleResult logicalSchedule = MergeSchedulingAnalysis(funcOp, target).getResult();
     PatternRewriter rewriter(moduleOp.getContext());
     FailureOr<ScheduledComputeMaterializationResult> materialization =
-      materializeScheduledCompute(funcOp, schedule, rewriter);
+      materializeScheduledCompute(funcOp, logicalSchedule, rewriter);
     if (failed(materialization)) {
       signalPassFailure();
       return;
@@ -48,7 +58,7 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
     // payloads must be diagnosed from the producer-owned body.
     dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true);
     if (failed(verifyMaterializedScheduleMapping(funcOp,
-                                                 schedule,
+                                                 logicalSchedule,
                                                  materialization->peftClassPlans,
                                                  materialization->graphComputeToBlockMap,
                                                  materialization->materializedSchedules))) {
@@ -75,18 +85,14 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
     SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
     if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
         && failed(exportSpatialDataflowCsvScheduled(
-          funcOp, materialization->materializedSchedules,
-          "spatial3_scheduled_no_comm", "spatial3"))) {
+          funcOp, materialization->materializedSchedules, "spatial3_scheduled_no_comm", "spatial3"))) {
       signalPassFailure();
       return;
     }
 
-    dumpScheduledComputeReport(moduleOp,
-                               funcOp,
-                               schedule,
-                               materialization->peftClassPlans,
-                               materialization->materializedSchedules);
-    if (failed(realizeDeferredCommunication(funcOp, *materialization))) {
+    dumpScheduledComputeReport(
+      moduleOp, funcOp, logicalSchedule, materialization->peftClassPlans, materialization->materializedSchedules);
+    if (failed(realizeDeferredCommunication(funcOp, *materialization, target))) {
       moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed");
       signalPassFailure();
       return;
@@ -100,11 +106,14 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
     }
     if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial4)
         && failed(exportSpatialDataflowCsvScheduled(
-          funcOp, materialization->materializedSchedules,
-          "spatial4_scheduled", "spatial4"))) {
+          funcOp, materialization->materializedSchedules, "spatial4_scheduled", "spatial4"))) {
       signalPassFailure();
     }
   }
+
+private:
+  SchedulingTarget target;
+  bool hasTarget = false;
 };
 
 } // namespace
@@ -112,4 +121,8 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
 
 std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<spatial::MergeComputeNodesPass>(); }
 
+std::unique_ptr<Pass> createMergeComputeNodesPass(const spatial::SchedulingTarget& target) {
+  return std::make_unique<spatial::MergeComputeNodesPass>(target);
+}
+
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp

@@ -164,7 +164,8 @@ inline size_t getScheduledCpuForComputeInstance(const ComputeInstance &instance,
   auto batch = dyn_cast<SpatComputeBatch>(instance.op);
   assert(batch && instance.laneCount != 0 && "missing scheduled CPU for non-batch compute instance");
   assert(instance.laneStart < static_cast<uint32_t>(batch.getLaneCount()) && "batch lane start out of range");
-  ComputeInstance chunk = getBatchChunkForLane(batch, instance.laneStart);
+  ComputeInstance chunk = getBatchChunkForLane(
+    batch, instance.laneStart, schedule.processorCount);
   auto it = schedule.computeToCpuMap.find(chunk);
   assert(it != schedule.computeToCpuMap.end() && "missing scheduled CPU for batch chunk");
   return it->second;
@@ -184,13 +185,17 @@ inline unsigned getScheduledBatchResultArgBase(SpatScheduledComputeBatch schedul
   return inputArgBase + scheduled.getInputs().size();
 }
 
-inline SmallVector<GraphComputeBlockKey> collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) {
+inline SmallVector<GraphComputeBlockKey> collectExpectedGraphComputeBlockKeys(
+    func::FuncOp funcOp, size_t processorCount) {
   SmallVector<GraphComputeBlockKey> keys;
   for (Operation &op : funcOp.getOps()) {
     if (auto compute = dyn_cast<SpatGraphCompute>(&op))
       keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1}));
     else if (auto batch = dyn_cast<SpatGraphComputeBatch>(&op))
-      for (ComputeInstance chunk : getBatchChunksForRange(batch, 0, static_cast<uint32_t>(batch.getLaneCount())))
+      for (ComputeInstance chunk :
+           getBatchChunksForRange(
+             batch, 0, static_cast<uint32_t>(batch.getLaneCount()),
+             processorCount))
         keys.push_back(getGraphComputeBlockKey(chunk));
   }
   return keys;

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp

@@ -179,7 +179,9 @@ LogicalResult collectPeftClassOperandsAndResults(
       for (Value weight : getComputeInstanceWeights(instance))
         appendUnique(peftClassPlan.weights, weight);
       for (Value input : getComputeInstanceInputs(instance))
-        if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input))
+        if (!getProducerValueRef(input, &instance, schedule.processorCount)
+            && !isDeferredFragmentAssemblyInput(
+              input, schedule.processorCount))
           appendUnique(peftClassPlan.inputs, input);
     }
     return success();
@@ -222,7 +224,9 @@ LogicalResult collectPeftClassOperandsAndResults(
       for (Value weight : getComputeInstanceWeights(instance))
         appendUnique(peftClassPlan.weights, weight);
       for (Value input : getComputeInstanceInputs(instance))
-        if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input))
+        if (!getProducerValueRef(input, &instance, schedule.processorCount)
+            && !isDeferredFragmentAssemblyInput(
+              input, schedule.processorCount))
           appendUnique(peftClassPlan.inputs, input);
     }
   }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp

@@ -48,7 +48,8 @@ LogicalResult verifyMaterializedScheduleMapping(
     }
   }
 
-  for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) {
+  for (GraphComputeBlockKey key :
+       collectExpectedGraphComputeBlockKeys(funcOp, schedule.processorCount)) {
     if (graphComputeToBlockMap.count(key))
       continue;
     diagnostics.report(key.op, [&](Operation *illegalOp) {
@@ -66,10 +67,12 @@ LogicalResult verifyMaterializedScheduleMapping(
     }
   }
 
-  if (graphComputeToBlockMap.size() != collectExpectedGraphComputeBlockKeys(funcOp).size()) {
+  const size_t expectedGraphComputeBlockCount =
+    collectExpectedGraphComputeBlockKeys(funcOp, schedule.processorCount).size();
+  if (graphComputeToBlockMap.size() != expectedGraphComputeBlockCount) {
     diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
       illegalOp->emitOpError() << "phase-check expected "
-                               << collectExpectedGraphComputeBlockKeys(funcOp).size()
+                               << expectedGraphComputeBlockCount
                                << " graph compute block mappings but saw " << graphComputeToBlockMap.size();
     });
   }

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

@@ -12,7 +12,6 @@
 #include "llvm/Support/Casting.h"
 
 #include <algorithm>
-#include <cmath>
 #include <iterator>
 #include <limits>
 #include <optional>
@@ -22,7 +21,6 @@
 
 #include "ComputeGraph.hpp"
 #include "ComputeInstanceUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
@@ -38,28 +36,25 @@ uint64_t countComputeBodyOperationInstances(Region& body);
 
 namespace {
 
-struct PimsimSchedulerCostModel {
-  static constexpr Cost kDefaultBitwidth = 8;
-  static constexpr Cost kCorePeriodNs = 1;
-  static constexpr Cost kLocalMemoryWidthBytes = 64;
-  static constexpr Cost kLocalMemoryLatencyCycles = 1;
-  static constexpr Cost kNetworkBusWidthBytes = 8;
-  static constexpr Cost kNetworkBaseLatencyNs = 2;
-  static constexpr Cost kNetworkPerHopLatencyNs = 1;
-  static constexpr Cost kVectorWidth = 16;
-  static constexpr Cost kVectorLatencyCycles = 4;
-  static constexpr Cost kDacResolutionBits = 1;
-  static constexpr Cost kDacLatencyCycles = 1;
-  static constexpr Cost kDacCount = 128;
-  static constexpr Cost kXbarReadLatencyNs = 30;
-  static constexpr Cost kSampleHoldLatencyCycles = 1;
-  static constexpr Cost kAdcLatencyCycles = 10;
-  static constexpr Cost kAdcCount = 2;
-  static constexpr Cost kShiftAdderLatencyCycles = 1;
-  static constexpr Cost kOutputBufferLatencyCycles = 1;
-  static constexpr Cost kInputBufferLatencyCycles = 0;
+TransferCost addTransferCosts(const TransferCost& lhs,
+                              const TransferCost& rhs) {
+  return {checkedAdd(lhs.fixed, rhs.fixed),
+          checkedAdd(lhs.networkFlits, rhs.networkFlits)};
+}
+
+TransferCost scaleTransferCost(const TransferCost& cost,
+                               Cost numerator,
+                               Cost denominator = 1) {
+  assert(denominator > 0 && "transfer cost denominator must be positive");
+  return {checkedMultiply(cost.fixed, numerator) / denominator,
+          checkedMultiply(cost.networkFlits, numerator) / denominator};
+}
+
+struct SchedulerCostModel {
   static constexpr Cost kFallbackOperationCost = 1;
 
+  const SchedulingTarget& target;
+
   static Cost ceilDiv(Cost numerator, Cost denominator) {
     assert(denominator > 0 && "denominator must be positive");
     return (numerator + denominator - 1) / denominator;
@@ -72,7 +67,7 @@ struct PimsimSchedulerCostModel {
     return static_cast<Cost>(shaped.getNumElements());
   }
 
-  static Cost getBitwidthOrDefault(Type type) {
+  Cost getBitwidthOrDefault(Type type) const {
     if (auto shaped = dyn_cast<ShapedType>(type))
       type = shaped.getElementType();
     if (auto intType = dyn_cast<IntegerType>(type))
@@ -81,14 +76,14 @@ struct PimsimSchedulerCostModel {
       return floatType.getWidth();
     if (isa<IndexType>(type))
       return 64;
-    return kDefaultBitwidth;
+    return target.computeBitwidth;
   }
 
-  static Cost getComputeBitwidth(Type type) {
-    return std::min(getBitwidthOrDefault(type), kDefaultBitwidth);
+  Cost getComputeBitwidth(Type type) const {
+    return std::min(getBitwidthOrDefault(type), target.computeBitwidth);
   }
 
-  static Cost getByteSize(Type type, Cost fallbackBitwidth = kDefaultBitwidth) {
+  Cost getByteSize(Type type, Cost fallbackBitwidth = 0) const {
     auto elementCount = getStaticElementCount(type);
     if (!elementCount)
       return kFallbackOperationCost;
@@ -96,60 +91,40 @@ struct PimsimSchedulerCostModel {
     if (bitwidth <= 0)
       bitwidth = getBitwidthOrDefault(type);
     if (bitwidth <= 0)
-      bitwidth = kDefaultBitwidth;
+      bitwidth = target.computeBitwidth;
     return ceilDiv(checkedMultiply(*elementCount, bitwidth), static_cast<Cost>(8));
   }
 
-  static Cost getVectorReadWriteCost(Cost readBytes, Cost writeBytes) {
-    Cost totalBytes = checkedAdd(readBytes, writeBytes);
-    return checkedMultiply(ceilDiv(totalBytes, kLocalMemoryWidthBytes),
-                           checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs));
+  Cost getVectorReadWriteCost(Cost readBytes, Cost writeBytes) const {
+    Cost reads = checkedMultiply(ceilDiv(readBytes, target.localMemoryWidthBytes),
+                                 target.localMemoryReadLatencyCycles);
+    Cost writes = checkedMultiply(ceilDiv(writeBytes, target.localMemoryWidthBytes),
+                                  target.localMemoryWriteLatencyCycles);
+    return checkedMultiply(checkedAdd(reads, writes), target.processorPeriodNs);
   }
 
-  static Cost getVectorComputeCost(Cost elementCount) {
-    return checkedMultiply(ceilDiv(elementCount, kVectorWidth),
-                           checkedMultiply(kVectorLatencyCycles, kCorePeriodNs));
+  Cost getVectorComputeCost(Cost elementCount) const {
+    return checkedMultiply(ceilDiv(elementCount, target.vectorWidth),
+                           checkedMultiply(target.vectorLatencyCycles, target.processorPeriodNs));
   }
 
-  static Cost getTensorMoveCost(Type type) {
+  Cost getTensorMoveCost(Type type) const {
     return getVectorReadWriteCost(getByteSize(type), 0);
   }
 
-  static std::pair<Cost, Cost> estimateMeshShape() {
-    Cost coreCount = static_cast<Cost>(std::max<long>(1, coresCount.getValue()));
-    Cost rows = static_cast<Cost>(std::sqrt(static_cast<long double>(coreCount)));
-    if (rows == 0)
-      rows = 1;
-    while (rows > 1 && coreCount % rows != 0)
-      --rows;
-    Cost cols = ceilDiv(coreCount, rows);
-    return {rows, cols};
+  TransferCost getTransferCostFromBytes(Cost bytes) const {
+    Cost localRead = checkedMultiply(ceilDiv(bytes, target.localMemoryWidthBytes),
+                                     checkedMultiply(target.localMemoryReadLatencyCycles,
+                                                     target.processorPeriodNs));
+    Cost localWrite = checkedMultiply(ceilDiv(bytes, target.localMemoryWidthBytes),
+                                      checkedMultiply(target.localMemoryWriteLatencyCycles,
+                                                      target.processorPeriodNs));
+    Cost payloadFlits = ceilDiv(bytes, target.transferWidthBytes);
+    return {checkedAdd(localRead, localWrite),
+            checkedAdd(static_cast<Cost>(2), payloadFlits)};
   }
 
-  static Cost getAverageInterCoreLatencyNs() {
-    auto [rows, cols] = estimateMeshShape();
-    auto averageAxisDistance = [](Cost size) -> Cost {
-      if (size <= 1)
-        return 0;
-      return checkedMultiply(size, size) - 1;
-    };
-    Cost avgRow = averageAxisDistance(rows) / (static_cast<Cost>(3) * rows);
-    Cost avgCol = averageAxisDistance(cols) / (static_cast<Cost>(3) * cols);
-    return checkedAdd(kNetworkBaseLatencyNs, checkedMultiply(kNetworkPerHopLatencyNs, checkedAdd(avgRow, avgCol)));
-  }
-
-  static Cost getInterCoreTransferCostFromBytes(Cost bytes) {
-    Cost localRead = checkedMultiply(ceilDiv(bytes, kLocalMemoryWidthBytes),
-                                     checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs));
-    Cost localWrite = checkedMultiply(ceilDiv(bytes, kLocalMemoryWidthBytes),
-                                      checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs));
-    Cost payloadFlits = ceilDiv(bytes, kNetworkBusWidthBytes);
-    Cost averageNoCLatency = getAverageInterCoreLatencyNs();
-    Cost network = checkedMultiply(checkedAdd(static_cast<Cost>(2), payloadFlits), averageNoCLatency);
-    return checkedAdd(checkedAdd(localRead, localWrite), network);
-  }
-
-  static Cost getUnaryVectorCost(Type inputType, Type outputType, bool scalarOutput = false) {
+  Cost getUnaryVectorCost(Type inputType, Type outputType, bool scalarOutput = false) const {
     auto maybeElements = getStaticElementCount(inputType);
     if (!maybeElements)
       return kFallbackOperationCost;
@@ -159,7 +134,7 @@ struct PimsimSchedulerCostModel {
     return checkedAdd(getVectorReadWriteCost(inputBytes, outputBytes), getVectorComputeCost(*maybeElements));
   }
 
-  static Cost getBinaryVectorCost(Type lhsType, Type rhsType, Type outputType, bool scalarOutput = false) {
+  Cost getBinaryVectorCost(Type lhsType, Type rhsType, Type outputType, bool scalarOutput = false) const {
     auto maybeElements = getStaticElementCount(lhsType);
     if (!maybeElements)
       return kFallbackOperationCost;
@@ -170,24 +145,35 @@ struct PimsimSchedulerCostModel {
     return checkedAdd(getVectorReadWriteCost(readBytes, outputBytes), getVectorComputeCost(*maybeElements));
   }
 
-  static Cost getMatrixComputeLatency(Cost inputBitwidth) {
-    Cost xbarDim = static_cast<Cost>(crossbarSize.getValue());
-    Cost inputTimes = ceilDiv(inputBitwidth, kDacResolutionBits);
-    Cost dacTimes = ceilDiv(xbarDim, kDacCount);
-    Cost adcTimes = ceilDiv(xbarDim, kAdcCount);
-    Cost frontStage = kInputBufferLatencyCycles + kDacLatencyCycles + kXbarReadLatencyNs + kSampleHoldLatencyCycles;
-    Cost backPipe = std::max(kAdcLatencyCycles, checkedAdd(kShiftAdderLatencyCycles, kOutputBufferLatencyCycles));
-    Cost backStage = checkedAdd(checkedAdd(kAdcLatencyCycles, kShiftAdderLatencyCycles), kOutputBufferLatencyCycles);
-    backStage = checkedAdd(backStage, checkedMultiply(adcTimes - 1, backPipe));
-    Cost totalTimes = checkedMultiply(inputTimes, dacTimes);
+  Cost getMatrixComputeLatency(Cost inputBitwidth) const {
+    Cost inputTimes = ceilDiv(inputBitwidth, target.matrixInputResolutionBits);
+    Cost inputPasses = ceilDiv(target.matrixRows, target.matrixInputParallelism);
+    Cost outputPasses = ceilDiv(target.matrixColumns, target.matrixOutputParallelism);
+    Cost readCycles = ceilDiv(target.matrixReadLatencyNs, target.matrixPeriodNs);
+    Cost frontStage = target.matrixInputBufferLatencyCycles + target.matrixInputLatencyCycles
+                    + readCycles + target.matrixSampleLatencyCycles;
+    Cost backPipe = std::max(target.matrixOutputLatencyCycles,
+                            checkedAdd(target.matrixShiftLatencyCycles,
+                                       target.matrixBufferLatencyCycles));
+    Cost backStage = checkedAdd(
+      checkedAdd(target.matrixOutputLatencyCycles, target.matrixShiftLatencyCycles),
+      target.matrixBufferLatencyCycles);
+    backStage = checkedAdd(backStage, checkedMultiply(outputPasses - 1, backPipe));
+    Cost totalTimes = checkedMultiply(inputTimes, inputPasses);
+    if (!target.matrixPipeline)
+      return checkedMultiply(
+        checkedMultiply(checkedAdd(frontStage, backStage), totalTimes),
+        target.matrixPeriodNs);
     Cost stagePipe = std::max(frontStage, backStage);
-    return checkedAdd(checkedAdd(frontStage, backStage),
-                      checkedMultiply(totalTimes - 1, stagePipe));
+    return checkedMultiply(
+      checkedAdd(checkedAdd(frontStage, backStage),
+                 checkedMultiply(totalTimes - 1, stagePipe)),
+      target.matrixPeriodNs);
   }
 
-  static Cost getWvmmCost(Type inputType, Type outputType) {
+  Cost getWvmmCost(Type inputType, Type outputType) const {
     Cost inputBitwidth = getComputeBitwidth(inputType);
-    Cost inputBytes = checkedMultiply(static_cast<Cost>(crossbarSize.getValue()),
+    Cost inputBytes = checkedMultiply(target.matrixRows,
                                       ceilDiv(inputBitwidth, static_cast<Cost>(8)));
     inputBytes = checkedMultiply(inputBytes, static_cast<Cost>(8));
     Cost outputBytes = getByteSize(outputType, getComputeBitwidth(outputType));
@@ -196,22 +182,22 @@ struct PimsimSchedulerCostModel {
 };
 
 std::optional<uint64_t> getStaticTripCount(scf::ForOp loop);
-Cost getOperationCost(Operation& op);
+Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel);
 
-Cost getRegionCost(Region& body) {
+Cost getRegionCost(Region& body, const SchedulerCostModel& costModel) {
   Cost cost = 0;
   for (Block& block : body)
     for (Operation& op : block)
-      cost = checkedAdd(cost, getOperationCost(op));
+      cost = checkedAdd(cost, getOperationCost(op, costModel));
   return cost;
 }
 
-Cost getOperationCost(Operation& op) {
+Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel) {
   if (auto loop = dyn_cast<scf::ForOp>(&op)) {
     std::optional<uint64_t> tripCount = getStaticTripCount(loop);
     if (!tripCount)
-      return PimsimSchedulerCostModel::kFallbackOperationCost;
-    return checkedMultiply(getRegionCost(loop.getRegion()), static_cast<Cost>(*tripCount));
+      return SchedulerCostModel::kFallbackOperationCost;
+    return checkedMultiply(getRegionCost(loop.getRegion(), costModel), static_cast<Cost>(*tripCount));
   }
 
   if (isa<SpatYieldOp, SpatInParallelOp, affine::AffineApplyOp, arith::ConstantOp,
@@ -219,42 +205,43 @@ Cost getOperationCost(Operation& op) {
     return 0;
 
   if (auto wvmm = dyn_cast<SpatVMMOp>(&op))
-    return PimsimSchedulerCostModel::getWvmmCost(wvmm.getInput().getType(), wvmm.getOutput().getType());
+    return costModel.getWvmmCost(wvmm.getInput().getType(), wvmm.getOutput().getType());
   if (auto vvdmul = dyn_cast<SpatVVDMulOp>(&op))
-    return PimsimSchedulerCostModel::getBinaryVectorCost(
+    return costModel.getBinaryVectorCost(
       vvdmul.getLhs().getType(), vvdmul.getRhs().getType(), vvdmul.getOutput().getType(), /*scalarOutput=*/true);
   if (auto vadd = dyn_cast<SpatVAddOp>(&op))
-    return PimsimSchedulerCostModel::getBinaryVectorCost(vadd.getLhs().getType(), vadd.getRhs().getType(),
-                                                         vadd.getOutput().getType());
+    return costModel.getBinaryVectorCost(
+      vadd.getLhs().getType(), vadd.getRhs().getType(), vadd.getOutput().getType());
   if (auto vsub = dyn_cast<SpatVSubOp>(&op))
-    return PimsimSchedulerCostModel::getBinaryVectorCost(vsub.getLhs().getType(), vsub.getRhs().getType(),
-                                                         vsub.getOutput().getType());
+    return costModel.getBinaryVectorCost(
+      vsub.getLhs().getType(), vsub.getRhs().getType(), vsub.getOutput().getType());
   if (auto vmul = dyn_cast<SpatVMulOp>(&op))
-    return PimsimSchedulerCostModel::getBinaryVectorCost(vmul.getLhs().getType(), vmul.getRhs().getType(),
-                                                         vmul.getOutput().getType());
+    return costModel.getBinaryVectorCost(
+      vmul.getLhs().getType(), vmul.getRhs().getType(), vmul.getOutput().getType());
   if (auto vmax = dyn_cast<SpatVMaxOp>(&op))
-    return PimsimSchedulerCostModel::getBinaryVectorCost(vmax.getLhs().getType(), vmax.getRhs().getType(),
-                                                         vmax.getOutput().getType());
+    return costModel.getBinaryVectorCost(
+      vmax.getLhs().getType(), vmax.getRhs().getType(), vmax.getOutput().getType());
   if (auto vavg = dyn_cast<SpatVAvgOp>(&op))
-    return PimsimSchedulerCostModel::getUnaryVectorCost(vavg.getInput().getType(), vavg.getOutput().getType(),
-                                                        /*scalarOutput=*/true);
+    return costModel.getUnaryVectorCost(
+      vavg.getInput().getType(), vavg.getOutput().getType(), /*scalarOutput=*/true);
   if (auto relu = dyn_cast<SpatReluOp>(&op))
-    return PimsimSchedulerCostModel::getUnaryVectorCost(relu.getInput().getType(), relu.getOutput().getType());
+    return costModel.getUnaryVectorCost(relu.getInput().getType(), relu.getOutput().getType());
   if (auto sigm = dyn_cast<SpatSigmoidOp>(&op))
-    return PimsimSchedulerCostModel::getUnaryVectorCost(sigm.getInput().getType(), sigm.getOutput().getType());
+    return costModel.getUnaryVectorCost(sigm.getInput().getType(), sigm.getOutput().getType());
   if (auto softmax = dyn_cast<SpatSoftmaxOp>(&op)) {
-    Cost unary = PimsimSchedulerCostModel::getUnaryVectorCost(softmax.getInput().getType(), softmax.getOutput().getType());
+    Cost unary =
+      costModel.getUnaryVectorCost(softmax.getInput().getType(), softmax.getOutput().getType());
     return checkedMultiply(unary, static_cast<Cost>(4));
   }
   if (auto extract = dyn_cast<tensor::ExtractSliceOp>(&op))
-    return PimsimSchedulerCostModel::getTensorMoveCost(extract.getResult().getType());
+    return costModel.getTensorMoveCost(extract.getResult().getType());
   if (auto insert = dyn_cast<tensor::InsertSliceOp>(&op))
-    return PimsimSchedulerCostModel::getTensorMoveCost(insert.getSource().getType());
+    return costModel.getTensorMoveCost(insert.getSource().getType());
 
   Cost nestedCost = 0;
   for (Region& region : op.getRegions())
-    nestedCost = checkedAdd(nestedCost, getRegionCost(region));
-  return checkedAdd(PimsimSchedulerCostModel::kFallbackOperationCost, nestedCost);
+    nestedCost = checkedAdd(nestedCost, getRegionCost(region, costModel));
+  return checkedAdd(SchedulerCostModel::kFallbackOperationCost, nestedCost);
 }
 
 std::optional<uint64_t> getStaticTripCount(scf::ForOp loop) {
@@ -271,8 +258,8 @@ std::optional<uint64_t> getStaticTripCount(scf::ForOp loop) {
   return (distance + stride - 1) / stride;
 }
 
-Cost getComputeBodyCost(Region& body) {
-  return getRegionCost(body);
+Cost getComputeBodyCost(Region& body, const SchedulerCostModel& costModel) {
+  return getRegionCost(body, costModel);
 }
 
 uint64_t countOperationInstances(Operation& op) {
@@ -348,7 +335,9 @@ std::optional<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp extract) {
   return std::nullopt;
 }
 
-std::optional<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, Value input) {
+std::optional<TransferCost> getBatchProjectedInputTransferCost(
+    SpatComputeBatch batch, Value input,
+    const SchedulerCostModel& costModel) {
   auto inputIt = llvm::find(batch.getInputs(), input);
   if (inputIt == batch.getInputs().end())
     return std::nullopt;
@@ -359,7 +348,7 @@ std::optional<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, V
   if (!inputArg || !laneArg)
     return std::nullopt;
 
-  Cost projectedCost = 0;
+  TransferCost projectedCost;
   for (Operation* user : inputArg->getUsers()) {
     auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
     if (!extract || extract.getSource() != *inputArg)
@@ -370,11 +359,13 @@ std::optional<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, V
     auto resultType = dyn_cast<ShapedType>(extract.getResult().getType());
     if (!resultType || !resultType.hasStaticShape())
       return std::nullopt;
-    projectedCost = checkedAdd(
-      projectedCost, PimsimSchedulerCostModel::getInterCoreTransferCostFromBytes(static_cast<Cost>(getSizeInBytes(resultType))));
+    projectedCost = addTransferCosts(
+      projectedCost,
+      costModel.getTransferCostFromBytes(
+        costModel.getByteSize(resultType, costModel.getComputeBitwidth(resultType))));
   }
 
-  if (projectedCost == 0)
+  if (projectedCost.fixed == 0 && projectedCost.networkFlits == 0)
     return std::nullopt;
   return projectedCost;
 }
@@ -382,7 +373,8 @@ std::optional<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, V
 static std::optional<SmallVector<ProducerValueRef, 4>>
 collectProjectedProducerValueRefs(SpatComputeBatch producer,
                                   Value input,
-                                  const ComputeInstance& consumerInstance) {
+                                  const ComputeInstance& consumerInstance,
+                                  size_t processorCount) {
   auto consumer = dyn_cast<SpatComputeBatch>(consumerInstance.op);
   if (!consumer)
     return std::nullopt;
@@ -417,7 +409,8 @@ collectProjectedProducerValueRefs(SpatComputeBatch producer,
         int64_t producerLane = *offset + index * *stride;
         if (producerLane < 0 || producerLane >= producer.getLaneCount())
           return std::nullopt;
-        ComputeInstance instance = getBatchChunkForLane(producer, static_cast<uint32_t>(producerLane));
+        ComputeInstance instance = getBatchChunkForLane(
+          producer, static_cast<uint32_t>(producerLane), processorCount);
         if (llvm::none_of(producers, [&](const ProducerValueRef& ref) { return ref.instance == instance; }))
           producers.push_back({instance, 0});
       }
@@ -426,12 +419,16 @@ collectProjectedProducerValueRefs(SpatComputeBatch producer,
   return producers;
 }
 
-Cost getInputTransferCost(const ComputeInstance& consumerInstance, Value input) {
+TransferCost getInputTransferCost(const ComputeInstance& consumerInstance,
+                                  Value input,
+                                  const SchedulerCostModel& costModel) {
   auto inputType = cast<ShapedType>(input.getType());
   if (auto batch = dyn_cast<SpatComputeBatch>(consumerInstance.op))
-    if (std::optional<Cost> projectedCost = getBatchProjectedInputTransferCost(batch, input))
+    if (std::optional<TransferCost> projectedCost =
+          getBatchProjectedInputTransferCost(batch, input, costModel))
       return *projectedCost;
-  return PimsimSchedulerCostModel::getInterCoreTransferCostFromBytes(static_cast<Cost>(getSizeInBytes(inputType)));
+  return costModel.getTransferCostFromBytes(
+    costModel.getByteSize(inputType, costModel.getComputeBitwidth(inputType)));
 }
 
 uint32_t getLaneOverlapCount(const ComputeInstance& lhs, const ComputeInstance& rhs) {
@@ -442,15 +439,20 @@ uint32_t getLaneOverlapCount(const ComputeInstance& lhs, const ComputeInstance&
          : 0;
 }
 
-Cost scaleTransferCostByLaneCount(Cost totalCost, uint32_t totalLaneCount, uint32_t fragmentLaneCount) {
+TransferCost scaleTransferCostByLaneCount(
+    const TransferCost& totalCost, uint32_t totalLaneCount,
+    uint32_t fragmentLaneCount) {
   assert(totalLaneCount > 0 && "laneCount must be positive");
   assert(fragmentLaneCount > 0 && "fragmentLaneCount must be positive");
   if (fragmentLaneCount >= totalLaneCount)
     return totalCost;
-  return checkedMultiply(totalCost, static_cast<Cost>(fragmentLaneCount)) / static_cast<Cost>(totalLaneCount);
+  return scaleTransferCost(totalCost, static_cast<Cost>(fragmentLaneCount),
+                           static_cast<Cost>(totalLaneCount));
 }
 
-SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const ComputeInstance& consumerInstance) {
+SmallVector<ProducerValueRef, 4> collectProducerValueRefs(
+    Value value, const ComputeInstance& consumerInstance,
+    size_t processorCount) {
   SmallVector<ProducerValueRef, 4> producers;
   Operation* op = value.getDefiningOp();
   if (!op)
@@ -461,13 +463,16 @@ SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const Com
     auto batch = dyn_cast_or_null<SpatComputeBatch>(source.getDefiningOp());
     if (batch && batch.getNumResults() != 0) {
       if (std::optional<uint32_t> lane = getConstantExtractLane(extract)) {
-        ComputeInstance instance = getBatchChunkForLane(batch, *lane);
+        ComputeInstance instance =
+          getBatchChunkForLane(batch, *lane, processorCount);
         producers.push_back({instance, 0});
         return producers;
       }
 
       for (ComputeInstance instance :
-           getBatchChunksForRange(batch, 0, static_cast<uint32_t>(batch.getLaneCount())))
+           getBatchChunksForRange(batch, 0,
+                                  static_cast<uint32_t>(batch.getLaneCount()),
+                                  processorCount))
         producers.push_back({instance, 0});
       return producers;
     }
@@ -488,19 +493,24 @@ SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const Com
 
   if (auto batch = dyn_cast<SpatComputeBatch>(op)) {
     if (batch.getNumResults() != 0) {
-      if (auto projected = collectProjectedProducerValueRefs(batch, value, consumerInstance))
+      if (auto projected = collectProjectedProducerValueRefs(
+            batch, value, consumerInstance, processorCount))
         return *projected;
-      std::optional<ProducerValueRef> producer = getProducerValueRef(value, &consumerInstance);
+      std::optional<ProducerValueRef> producer =
+        getProducerValueRef(value, &consumerInstance, processorCount);
       if (!producer)
         return producers;
       for (ComputeInstance instance :
-           getBatchChunksForRange(batch, producer->instance.laneStart, producer->instance.laneCount))
+           getBatchChunksForRange(batch, producer->instance.laneStart,
+                                  producer->instance.laneCount,
+                                  processorCount))
         producers.push_back({instance, 0});
       return producers;
     }
 
     uint32_t lane = cast<OpResult>(value).getResultNumber();
-    ComputeInstance instance = getBatchChunkForLane(batch, lane);
+    ComputeInstance instance =
+      getBatchChunkForLane(batch, lane, processorCount);
     producers.push_back({instance, lane - instance.laneStart});
     return producers;
   }
@@ -508,18 +518,23 @@ SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const Com
   return producers;
 }
 
-Cost getProducerTransferCost(Value input,
-                             const ComputeInstance& consumerInstance,
-                             const ProducerValueRef& producerRef) {
-  Cost transferCost = getInputTransferCost(consumerInstance, input);
+TransferCost getProducerTransferCost(
+    Value input, const ComputeInstance& consumerInstance,
+    const ProducerValueRef& producerRef,
+    const SchedulerCostModel& costModel) {
+  TransferCost transferCost =
+    getInputTransferCost(consumerInstance, input, costModel);
   auto producerBatch = dyn_cast<SpatComputeBatch>(producerRef.instance.op);
   if (!producerBatch || producerBatch.getNumResults() == 0)
     return transferCost;
 
   if (auto consumerBatch = dyn_cast<SpatComputeBatch>(consumerInstance.op)) {
-    if (std::optional<Cost> projectedCost = getBatchProjectedInputTransferCost(consumerBatch, input)) {
+    if (std::optional<TransferCost> projectedCost =
+          getBatchProjectedInputTransferCost(consumerBatch, input, costModel)) {
       uint32_t overlapLaneCount = getLaneOverlapCount(consumerInstance, producerRef.instance);
-      return checkedMultiply(*projectedCost, static_cast<Cost>(std::max<uint32_t>(1, overlapLaneCount)));
+      return scaleTransferCost(
+        *projectedCost,
+        static_cast<Cost>(std::max<uint32_t>(1, overlapLaneCount)));
     }
   }
 
@@ -527,8 +542,8 @@ Cost getProducerTransferCost(Value input,
     transferCost, static_cast<uint32_t>(producerBatch.getLaneCount()), producerRef.instance.laneCount);
 }
 
-static CrossbarWeight getOpaqueCrossbarWeight(Value value, std::optional<uint32_t> lane) {
-  CrossbarWeight weight;
+static ResidentWeight getOpaqueResidentWeight(Value value, std::optional<uint32_t> lane) {
+  ResidentWeight weight;
   weight.opaqueValue = value;
   weight.opaqueLane = lane.value_or(std::numeric_limits<uint32_t>::max());
   return weight;
@@ -619,7 +634,7 @@ static FailureOr<SmallVector<int64_t, 4>> evaluateIndexList(ArrayRef<OpFoldResul
   return result;
 }
 
-static Value resolveCrossbarWeightRoot(Operation* owner, Value root) {
+static Value resolveResidentWeightRoot(Operation* owner, Value root) {
   if (auto arg = dyn_cast<BlockArgument>(root)) {
     if (auto compute = dyn_cast<SpatCompute>(owner)) {
       for (auto [index, operand] : llvm::enumerate(compute.getWeights()))
@@ -637,11 +652,11 @@ static Value resolveCrossbarWeightRoot(Operation* owner, Value root) {
   return root;
 }
 
-static CrossbarWeight completeCrossbarWeight(Value root,
+static ResidentWeight completeResidentWeight(Value root,
                                              SmallVector<int64_t, 4> offsets,
                                              SmallVector<int64_t, 4> sizes,
                                              SmallVector<int64_t, 4> strides) {
-  CrossbarWeight weight;
+  ResidentWeight weight;
   weight.root = root;
   if (auto constant = root.getDefiningOp<arith::ConstantOp>())
     weight.rootAttr = static_cast<Attribute>(constant.getValue());
@@ -651,14 +666,14 @@ static CrossbarWeight completeCrossbarWeight(Value root,
   return weight;
 }
 
-static FailureOr<CrossbarWeight> getStaticCrossbarWeight(Operation* owner,
+static FailureOr<ResidentWeight> getStaticResidentWeight(Operation* owner,
                                                          Value value,
                                                          const DenseMap<Value, int64_t>& bindings,
                                                          std::optional<uint32_t> lane,
                                                          Value laneArg) {
   if (auto extract = value.getDefiningOp<tensor::ExtractSliceOp>()) {
-    FailureOr<CrossbarWeight> sourceWeight =
-      getStaticCrossbarWeight(owner, extract.getSource(), bindings, lane, laneArg);
+    FailureOr<ResidentWeight> sourceWeight =
+      getStaticResidentWeight(owner, extract.getSource(), bindings, lane, laneArg);
     auto offsets = evaluateIndexList(extract.getMixedOffsets(), bindings, lane, laneArg);
     auto sizes = evaluateIndexList(extract.getMixedSizes(), bindings, lane, laneArg);
     auto strides = evaluateIndexList(extract.getMixedStrides(), bindings, lane, laneArg);
@@ -678,7 +693,7 @@ static FailureOr<CrossbarWeight> getStaticCrossbarWeight(Operation* owner,
     return *sourceWeight;
   }
 
-  Value root = resolveCrossbarWeightRoot(owner, value);
+  Value root = resolveResidentWeightRoot(owner, value);
   auto type = dyn_cast<ShapedType>(root.getType());
   if (!type || !type.hasStaticShape())
     return failure();
@@ -686,18 +701,18 @@ static FailureOr<CrossbarWeight> getStaticCrossbarWeight(Operation* owner,
   SmallVector<int64_t, 4> offsets(type.getRank(), 0);
   SmallVector<int64_t, 4> sizes(type.getShape().begin(), type.getShape().end());
   SmallVector<int64_t, 4> strides(type.getRank(), 1);
-  return completeCrossbarWeight(root, std::move(offsets), std::move(sizes), std::move(strides));
+  return completeResidentWeight(root, std::move(offsets), std::move(sizes), std::move(strides));
 }
 
-static void addCrossbarWeight(CrossbarUsage& usage, CrossbarWeight weight) {
-  if (!containsCrossbarWeight(usage, weight))
+static void addResidentWeight(ResidentWeightSet& usage, ResidentWeight weight) {
+  if (!containsResidentWeight(usage, weight))
     usage.push_back(std::move(weight));
 }
 
-static void collectCrossbarWeightsFromOp(Operation* op,
+static void collectResidentWeightsFromOp(Operation* op,
                                          Operation* owner,
                                          DenseMap<Value, int64_t>& bindings,
-                                         CrossbarUsage& usage,
+                                         ResidentWeightSet& usage,
                                          Value laneArg,
                                          std::optional<uint32_t> lane) {
   if (auto loop = dyn_cast<scf::ForOp>(op)) {
@@ -710,36 +725,37 @@ static void collectCrossbarWeightsFromOp(Operation* op,
     for (int64_t iv = *lb; iv < *ub; iv += *step) {
       bindings[loop.getInductionVar()] = iv;
       for (Operation& nested : loop.getBody()->without_terminator())
-        collectCrossbarWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane);
+        collectResidentWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane);
     }
     bindings.erase(loop.getInductionVar());
     return;
   }
 
   if (auto vmm = dyn_cast<SpatVMMOp>(op)) {
-    FailureOr<CrossbarWeight> weight = getStaticCrossbarWeight(owner, vmm.getWeight(), bindings, lane, laneArg);
+    FailureOr<ResidentWeight> weight = getStaticResidentWeight(owner, vmm.getWeight(), bindings, lane, laneArg);
     if (failed(weight)) {
-      addCrossbarWeight(usage, getOpaqueCrossbarWeight(vmm.getWeight(), lane));
+      addResidentWeight(usage, getOpaqueResidentWeight(vmm.getWeight(), lane));
       return;
     }
-    addCrossbarWeight(usage, *weight);
+    addResidentWeight(usage, *weight);
     return;
   }
 
   for (Region& region : op->getRegions())
     for (Block& block : region)
       for (Operation& nested : block.without_terminator())
-        collectCrossbarWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane);
+        collectResidentWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane);
 }
 
 std::vector<ComputeGraphEdge> aggregateEdges(llvm::ArrayRef<ComputeGraphEdge> edges) {
-  llvm::DenseMap<std::pair<size_t, size_t>, Cost> edgeCosts;
+  llvm::DenseMap<std::pair<size_t, size_t>, TransferCost> edgeCosts;
   for (const ComputeGraphEdge& edge : edges) {
     if (edge.source == edge.target)
       continue;
     auto inserted = edgeCosts.try_emplace({edge.source, edge.target}, edge.transferCost);
     if (!inserted.second)
-      inserted.first->second = checkedAdd(inserted.first->second, edge.transferCost);
+      inserted.first->second =
+        addTransferCosts(inserted.first->second, edge.transferCost);
   }
 
   std::vector<ComputeGraphEdge> aggregatedEdges;
@@ -770,8 +786,8 @@ uint64_t countComputeBodyOperationInstances(Region& body) {
   return instances;
 }
 
-CrossbarUsage collectDistinctCrossbarWeights(Operation* owner, std::optional<uint32_t> lane) {
-  CrossbarUsage usage;
+ResidentWeightSet collectDistinctResidentWeights(Operation* owner, std::optional<uint32_t> lane) {
+  ResidentWeightSet usage;
   DenseMap<Value, int64_t> bindings;
   Value laneArg;
   if (auto batch = dyn_cast<SpatComputeBatch>(owner))
@@ -781,54 +797,87 @@ CrossbarUsage collectDistinctCrossbarWeights(Operation* owner, std::optional<uin
   for (Region& region : owner->getRegions())
     for (Block& block : region)
       for (Operation& op : block.without_terminator())
-        collectCrossbarWeightsFromOp(&op, owner, bindings, usage, laneArg, lane);
+        collectResidentWeightsFromOp(&op, owner, bindings, usage, laneArg, lane);
   return usage;
 }
 
-Cost getComputeInstanceCost(const ComputeInstance& instance) {
+Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target) {
+  SchedulerCostModel costModel {target};
   if (auto spatCompute = dyn_cast<SpatCompute>(instance.op))
-    return getComputeBodyCost(spatCompute.getBody());
+    return getComputeBodyCost(spatCompute.getBody(), costModel);
   auto batch = cast<SpatComputeBatch>(instance.op);
-  return checkedMultiply(getComputeBodyCost(batch.getBody()), static_cast<Cost>(instance.laneCount));
+  return checkedMultiply(
+    getComputeBodyCost(batch.getBody(), costModel), static_cast<Cost>(instance.laneCount));
 }
 
-bool containsCrossbarWeight(ArrayRef<CrossbarWeight> usage, const CrossbarWeight& weight) {
+bool containsResidentWeight(ArrayRef<ResidentWeight> usage, const ResidentWeight& weight) {
   return llvm::is_contained(usage, weight);
 }
 
-unsigned countCrossbarOverlap(ArrayRef<CrossbarWeight> lhs, ArrayRef<CrossbarWeight> rhs) {
+unsigned countResidentWeightOverlap(ArrayRef<ResidentWeight> lhs, ArrayRef<ResidentWeight> rhs) {
   unsigned overlap = 0;
-  for (const CrossbarWeight& weight : rhs)
-    if (containsCrossbarWeight(lhs, weight))
+  for (const ResidentWeight& weight : rhs)
+    if (containsResidentWeight(lhs, weight))
       ++overlap;
   return overlap;
 }
 
-size_t getCrossbarUnionSize(ArrayRef<CrossbarWeight> lhs, ArrayRef<CrossbarWeight> rhs) {
+size_t getResidentWeightUnionSize(ArrayRef<ResidentWeight> lhs, ArrayRef<ResidentWeight> rhs) {
   size_t size = lhs.size();
-  for (const CrossbarWeight& weight : rhs)
-    if (!containsCrossbarWeight(lhs, weight))
+  for (const ResidentWeight& weight : rhs)
+    if (!containsResidentWeight(lhs, weight))
       ++size;
   return size;
 }
 
-void insertCrossbarWeights(CrossbarUsage& usage, ArrayRef<CrossbarWeight> weights) {
-  for (const CrossbarWeight& weight : weights)
-    addCrossbarWeight(usage, weight);
+void insertResidentWeights(ResidentWeightSet& usage, ArrayRef<ResidentWeight> weights) {
+  for (const ResidentWeight& weight : weights)
+    addResidentWeight(usage, weight);
 }
 
-CrossbarUsage getComputeInstanceCrossbarUsage(const ComputeInstance& instance) {
-  CrossbarUsage usage;
+ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance) {
+  ResidentWeightSet usage;
   if (isa<SpatCompute>(instance.op))
-    return collectDistinctCrossbarWeights(instance.op);
+    return collectDistinctResidentWeights(instance.op);
 
   for (uint32_t lane = instance.laneStart; lane < instance.laneStart + instance.laneCount; ++lane)
-    insertCrossbarWeights(usage, collectDistinctCrossbarWeights(instance.op, lane));
+    insertResidentWeights(usage, collectDistinctResidentWeights(instance.op, lane));
   return usage;
 }
 
-ComputeGraph buildComputeGraph(Operation* entryOp) {
+ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance,
+                                                     const SchedulingTarget& target) {
+  ResidentWeightSet tiled;
+  for (const ResidentWeight& weight : getComputeInstanceResidentWeights(instance)) {
+    if (weight.opaqueValue || weight.sizes.size() < 2
+        || target.matrixRows == 0 || target.matrixColumns == 0) {
+      addResidentWeight(tiled, weight);
+      continue;
+    }
+
+    const size_t rowDim = weight.sizes.size() - 2;
+    const size_t columnDim = weight.sizes.size() - 1;
+    for (int64_t row = 0; row < weight.sizes[rowDim];
+         row += static_cast<int64_t>(target.matrixRows)) {
+      for (int64_t column = 0; column < weight.sizes[columnDim];
+           column += static_cast<int64_t>(target.matrixColumns)) {
+        ResidentWeight tile = weight;
+        tile.offsets[rowDim] += row * tile.strides[rowDim];
+        tile.offsets[columnDim] += column * tile.strides[columnDim];
+        tile.sizes[rowDim] =
+          std::min<int64_t>(target.matrixRows, weight.sizes[rowDim] - row);
+        tile.sizes[columnDim] =
+          std::min<int64_t>(target.matrixColumns, weight.sizes[columnDim] - column);
+        addResidentWeight(tiled, std::move(tile));
+      }
+    }
+  }
+  return tiled;
+}
+
+ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& target) {
   ComputeGraph graph;
+  SchedulerCostModel costModel {target};
 
   for (Region& region : entryOp->getRegions()) {
     for (Block& block : region) {
@@ -838,20 +887,26 @@ ComputeGraph buildComputeGraph(Operation* entryOp) {
             continue;
           ComputeInstance instance {spatCompute.getOperation(), 0, 1};
           size_t index = graph.nodes.size();
-          graph.nodes.push_back(
-            {instance, getComputeInstanceCost(instance), getComputeInstanceCrossbarUsage(instance), index});
+          graph.nodes.push_back({instance,
+                                 getComputeInstanceCost(instance, target),
+                                 getComputeInstanceResidentWeights(instance, target),
+                                 index});
           graph.instanceToIndex[instance] = index;
           continue;
         }
         if (auto batch = dyn_cast<SpatComputeBatch>(&op)) {
           if (isUsedAsWeightOnly(batch.getOperation()))
             continue;
-          size_t chunkCount = getBatchChunkTargetCount(batch);
+          size_t chunkCount =
+            getBatchChunkTargetCount(batch, target.processorCount);
           for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) {
-            ComputeInstance instance = getBatchChunkForIndex(batch, chunkIndex);
+            ComputeInstance instance = getBatchChunkForIndex(
+              batch, chunkIndex, target.processorCount);
             size_t index = graph.nodes.size();
-            graph.nodes.push_back(
-              {instance, getComputeInstanceCost(instance), getComputeInstanceCrossbarUsage(instance), index});
+            graph.nodes.push_back({instance,
+                                   getComputeInstanceCost(instance, target),
+                                   getComputeInstanceResidentWeights(instance, target),
+                                   index});
             graph.instanceToIndex[instance] = index;
           }
         }
@@ -863,12 +918,15 @@ ComputeGraph buildComputeGraph(Operation* entryOp) {
   for (const auto& [targetIndex, node] : llvm::enumerate(graph.nodes)) {
     llvm::SmallVector<Value, 4> inputs = getComputeInstanceInputs(node.instance);
     for (Value input : inputs) {
-      for (const ProducerValueRef& producerRef : collectProducerValueRefs(input, node.instance)) {
+      for (const ProducerValueRef& producerRef :
+           collectProducerValueRefs(input, node.instance,
+                                    target.processorCount)) {
         auto producerIt = graph.instanceToIndex.find(producerRef.instance);
         if (producerIt == graph.instanceToIndex.end())
           continue;
-        rawEdges.push_back(
-          {producerIt->second, targetIndex, getProducerTransferCost(input, node.instance, producerRef)});
+        rawEdges.push_back({producerIt->second,
+                            targetIndex,
+                            getProducerTransferCost(input, node.instance, producerRef, costModel)});
       }
     }
   }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp

@@ -11,9 +11,10 @@
 #include <vector>
 
 #include "ComputeInstance.hpp"
+#include "SchedulingTarget.hpp"
 #include "Utils.hpp"
 
-struct CrossbarWeight {
+struct ResidentWeight {
   mlir::Value root;
   mlir::Attribute rootAttr;
   llvm::SmallVector<int64_t, 4> offsets;
@@ -22,14 +23,14 @@ struct CrossbarWeight {
   mlir::Value opaqueValue;
   uint32_t opaqueLane = 0;
 
-  bool operator==(const CrossbarWeight& other) const {
+  bool operator==(const ResidentWeight& other) const {
     bool sameRoot = rootAttr && other.rootAttr ? rootAttr == other.rootAttr : root == other.root;
     return sameRoot && offsets == other.offsets && sizes == other.sizes && strides == other.strides
         && opaqueValue == other.opaqueValue && opaqueLane == other.opaqueLane;
   }
 };
 
-using CrossbarUsage = llvm::SmallVector<CrossbarWeight, 6>;
+using ResidentWeightSet = llvm::SmallVector<ResidentWeight, 6>;
 
 namespace onnx_mlir {
 namespace spatial {
@@ -37,36 +38,47 @@ namespace spatial {
 struct ComputeGraphNode {
   ComputeInstance instance;
   Cost cost = 0;
-  CrossbarUsage crossbarUsage;
+  ResidentWeightSet residentWeights;
   size_t originalOrder = 0;
 };
 
+struct TransferCost {
+  Cost fixed = 0;
+  Cost networkFlits = 0;
+};
+
 struct ComputeGraphEdge {
   size_t source = 0;
   size_t target = 0;
-  Cost transferCost = 0;
+  TransferCost transferCost;
 };
 
 struct ComputeGraph {
   std::vector<ComputeGraphNode> nodes;
   std::vector<ComputeGraphEdge> edges;
-  std::vector<std::vector<std::pair<size_t, Cost>>> successors;
-  std::vector<std::vector<std::pair<size_t, Cost>>> predecessors;
+  std::vector<std::vector<std::pair<size_t, TransferCost>>> successors;
+  std::vector<std::vector<std::pair<size_t, TransferCost>>> predecessors;
   llvm::DenseMap<ComputeInstance, size_t> instanceToIndex;
 };
 
-ComputeGraph buildComputeGraph(mlir::Operation* entryOp);
+ComputeGraph buildComputeGraph(mlir::Operation* entryOp, const SchedulingTarget& target);
 bool verifyAcyclic(const ComputeGraph& graph);
 
 uint64_t countComputeBodyInstructions(mlir::Region& body);
 uint64_t countComputeBodyOperationInstances(mlir::Region& body);
-Cost getComputeInstanceCost(const ComputeInstance& instance);
-CrossbarUsage collectDistinctCrossbarWeights(mlir::Operation* owner, std::optional<uint32_t> lane = std::nullopt);
-CrossbarUsage getComputeInstanceCrossbarUsage(const ComputeInstance& instance);
-bool containsCrossbarWeight(llvm::ArrayRef<CrossbarWeight> usage, const CrossbarWeight& weight);
-unsigned countCrossbarOverlap(llvm::ArrayRef<CrossbarWeight> lhs, llvm::ArrayRef<CrossbarWeight> rhs);
-size_t getCrossbarUnionSize(llvm::ArrayRef<CrossbarWeight> lhs, llvm::ArrayRef<CrossbarWeight> rhs);
-void insertCrossbarWeights(CrossbarUsage& usage, llvm::ArrayRef<CrossbarWeight> weights);
+Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target);
+ResidentWeightSet collectDistinctResidentWeights(mlir::Operation* owner,
+                                                 std::optional<uint32_t> lane = std::nullopt);
+ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance);
+ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance,
+                                                     const SchedulingTarget& target);
+bool containsResidentWeight(llvm::ArrayRef<ResidentWeight> usage, const ResidentWeight& weight);
+unsigned countResidentWeightOverlap(llvm::ArrayRef<ResidentWeight> lhs,
+                                    llvm::ArrayRef<ResidentWeight> rhs);
+size_t getResidentWeightUnionSize(llvm::ArrayRef<ResidentWeight> lhs,
+                                  llvm::ArrayRef<ResidentWeight> rhs);
+void insertResidentWeights(ResidentWeightSet& usage,
+                           llvm::ArrayRef<ResidentWeight> weights);
 
 } // namespace spatial
 } // namespace onnx_mlir

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

@@ -2,23 +2,15 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 #include <algorithm>
-#include <limits>
 #include <optional>
 
 #include "ComputeInstanceUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 namespace spatial {
 
-size_t getSchedulingCpuBudget() {
-  if (coresCount.getValue() > 0)
-    return static_cast<size_t>(coresCount.getValue());
-  return std::numeric_limits<size_t>::max();
-}
-
 static BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkCount, size_t chunkIndex) {
   assert(laneCount > 0 && "laneCount must be positive");
   assert(chunkIndex < chunkCount && "chunkIndex out of range");
@@ -33,22 +25,27 @@ static BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkCount,
   return {static_cast<uint32_t>(start), static_cast<uint32_t>(count)};
 }
 
-size_t getBatchChunkTargetCount(SpatComputeBatch batch) {
+size_t getBatchChunkTargetCount(SpatComputeBatch batch, size_t processorCount) {
   int32_t laneCount = batch.getLaneCount();
   assert(laneCount > 0 && "laneCount must be positive");
-  return std::min(static_cast<size_t>(laneCount), getSchedulingCpuBudget());
+  assert(processorCount > 0 && "processorCount must be positive");
+  return std::min(static_cast<size_t>(laneCount), processorCount);
 }
 
-BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex) {
-  return getBatchChunkRange(batch.getLaneCount(), getBatchChunkTargetCount(batch), chunkIndex);
+BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex,
+                                  size_t processorCount) {
+  return getBatchChunkRange(batch.getLaneCount(),
+                            getBatchChunkTargetCount(batch, processorCount),
+                            chunkIndex);
 }
 
-size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane) {
+size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane,
+                                 size_t processorCount) {
   int32_t laneCount = batch.getLaneCount();
   assert(laneCount > 0 && "laneCount must be positive");
   assert(lane < static_cast<uint32_t>(laneCount) && "lane out of range");
 
-  size_t chunkCount = getBatchChunkTargetCount(batch);
+  size_t chunkCount = getBatchChunkTargetCount(batch, processorCount);
   size_t laneCountSize = static_cast<size_t>(laneCount);
   size_t baseChunkSize = laneCountSize / chunkCount;
   size_t remainder = laneCountSize % chunkCount;
@@ -61,17 +58,22 @@ size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane) {
   return remainder + ((laneIndex - largerChunkLanes) / baseChunkSize);
 }
 
-ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex) {
-  BatchChunkRange chunk = getBatchChunkRange(batch, chunkIndex);
+ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex,
+                                      size_t processorCount) {
+  BatchChunkRange chunk = getBatchChunkRange(batch, chunkIndex, processorCount);
   return {batch.getOperation(), chunk.laneStart, chunk.laneCount};
 }
 
-ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) {
-  return getBatchChunkForIndex(batch, getBatchChunkIndexForLane(batch, lane));
+ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane,
+                                     size_t processorCount) {
+  return getBatchChunkForIndex(
+    batch, getBatchChunkIndexForLane(batch, lane, processorCount),
+    processorCount);
 }
 
 llvm::SmallVector<ComputeInstance, 4>
-getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount) {
+getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart,
+                       uint32_t laneCount, size_t processorCount) {
   llvm::SmallVector<ComputeInstance, 4> chunks;
   if (laneCount == 0)
     return chunks;
@@ -80,11 +82,13 @@ getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t lane
   assert(laneEnd >= laneStart && "lane range overflow");
   assert(laneEnd <= static_cast<uint32_t>(batch.getLaneCount()) && "lane range out of bounds");
 
-  size_t firstChunk = getBatchChunkIndexForLane(batch, laneStart);
-  size_t lastChunk = getBatchChunkIndexForLane(batch, laneEnd - 1);
+  size_t firstChunk =
+    getBatchChunkIndexForLane(batch, laneStart, processorCount);
+  size_t lastChunk =
+    getBatchChunkIndexForLane(batch, laneEnd - 1, processorCount);
   chunks.reserve(lastChunk - firstChunk + 1);
   for (size_t chunkIndex = firstChunk; chunkIndex <= lastChunk; ++chunkIndex)
-    chunks.push_back(getBatchChunkForIndex(batch, chunkIndex));
+    chunks.push_back(getBatchChunkForIndex(batch, chunkIndex, processorCount));
   return chunks;
 }
 
@@ -150,7 +154,9 @@ static std::optional<ProducerValueRef> getResultfulBatchProducerValueRef(SpatCom
   };
 }
 
-std::optional<ProducerValueRef> getProducerValueRef(Value value, const ComputeInstance* consumerInstance) {
+std::optional<ProducerValueRef> getProducerValueRef(
+    Value value, const ComputeInstance* consumerInstance,
+    size_t processorCount) {
   Operation* op = value.getDefiningOp();
   if (!op)
     return std::nullopt;
@@ -187,7 +193,8 @@ std::optional<ProducerValueRef> getProducerValueRef(Value value, const ComputeIn
     if (batch.getNumResults() != 0)
       return getResultfulBatchProducerValueRef(batch, value, consumerInstance);
     uint32_t lane = cast<OpResult>(value).getResultNumber();
-    ComputeInstance instance = getBatchChunkForLane(batch, lane);
+    ComputeInstance instance =
+      getBatchChunkForLane(batch, lane, processorCount);
     size_t resultIndex = lane - instance.laneStart;
     return ProducerValueRef {instance, resultIndex};
   }
@@ -195,8 +202,11 @@ std::optional<ProducerValueRef> getProducerValueRef(Value value, const ComputeIn
   return std::nullopt;
 }
 
-std::optional<ComputeInstance> getComputeProducerInstance(Value value, const ComputeInstance* consumerInstance) {
-  if (std::optional<ProducerValueRef> producer = getProducerValueRef(value, consumerInstance))
+std::optional<ComputeInstance> getComputeProducerInstance(
+    Value value, const ComputeInstance* consumerInstance,
+    size_t processorCount) {
+  if (std::optional<ProducerValueRef> producer =
+        getProducerValueRef(value, consumerInstance, processorCount))
     return producer->instance;
   return std::nullopt;
 }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp

@@ -26,19 +26,25 @@ struct BatchChunkRange {
   uint32_t laneCount = 0;
 };
 
-size_t getSchedulingCpuBudget();
-size_t getBatchChunkTargetCount(SpatComputeBatch batch);
-BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex);
-size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane);
-ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex);
-ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane);
+size_t getBatchChunkTargetCount(SpatComputeBatch batch, size_t processorCount);
+BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex,
+                                  size_t processorCount);
+size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane,
+                                 size_t processorCount);
+ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex,
+                                      size_t processorCount);
+ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane,
+                                     size_t processorCount);
 llvm::SmallVector<ComputeInstance, 4>
-getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount);
+getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart,
+                       uint32_t laneCount, size_t processorCount);
 
 std::optional<ProducerValueRef> getProducerValueRef(mlir::Value value,
-                                                    const ComputeInstance* consumerInstance = nullptr);
+                                                    const ComputeInstance* consumerInstance,
+                                                    size_t processorCount);
 std::optional<ComputeInstance> getComputeProducerInstance(mlir::Value value,
-                                                          const ComputeInstance* consumerInstance = nullptr);
+                                                          const ComputeInstance* consumerInstance,
+                                                          size_t processorCount);
 
 llvm::SmallVector<mlir::Value, 4> getComputeInstanceInputs(const ComputeInstance& instance);
 llvm::SmallVector<mlir::Value, 4> getComputeInstanceWeights(const ComputeInstance& instance);

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp

@@ -13,6 +13,7 @@ namespace onnx_mlir {
 namespace spatial {
 
 struct MergeScheduleResult {
+  size_t processorCount = 0;
   std::vector<ComputeInstance> dominanceOrderCompute;
   llvm::DenseMap<ComputeInstance, size_t> computeToCpuMap;
   llvm::DenseMap<ComputeInstance, size_t> computeToCpuSlotMap;

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

@@ -9,7 +9,6 @@
 #include "ComputeGraph.hpp"
 #include "MergeSchedulingAnalysis.hpp"
 #include "PeftScheduler.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
@@ -18,8 +17,7 @@ namespace {
 
 void verifySchedule(const ComputeGraph& graph,
                     const MergeScheduleResult& result,
-                    unsigned long crossbarCapacity,
-                    size_t processorCount) {
+                    const SchedulingTarget& target) {
   llvm::DenseMap<size_t, std::vector<std::pair<size_t, size_t>>> tasksByCpu;
   tasksByCpu.reserve(result.cpuToLastComputeMap.size());
 
@@ -44,13 +42,13 @@ void verifySchedule(const ComputeGraph& graph,
       return lhs.second < rhs.second;
     });
 
-    CrossbarUsage usedCrossbars;
+    ResidentWeightSet residentWeights;
     for (size_t slot = 0; slot < scheduledTasks.size(); ++slot) {
       if (scheduledTasks[slot].first != slot)
         llvm::report_fatal_error("merge scheduling: CPU slots are not contiguous");
-      insertCrossbarWeights(usedCrossbars, graph.nodes[scheduledTasks[slot].second].crossbarUsage);
-      if (usedCrossbars.size() > crossbarCapacity)
-        llvm::report_fatal_error("merge scheduling: CPU crossbar capacity exceeded");
+      insertResidentWeights(residentWeights, graph.nodes[scheduledTasks[slot].second].residentWeights);
+      if (residentWeights.size() > target.residentWeightCapacity)
+        llvm::report_fatal_error("merge scheduling: processor resident-weight capacity exceeded");
     }
 
     const ComputeInstance expectedLast = graph.nodes[scheduledTasks.back().second].instance;
@@ -63,20 +61,21 @@ void verifySchedule(const ComputeGraph& graph,
 
   for (const ComputeGraphEdge& edge : graph.edges) {
     const ComputeInstance source = graph.nodes[edge.source].instance;
-    const ComputeInstance target = graph.nodes[edge.target].instance;
+    const ComputeInstance destination = graph.nodes[edge.target].instance;
     const size_t sourceCpu = result.computeToCpuMap.lookup(source);
-    const size_t targetCpu = result.computeToCpuMap.lookup(target);
+    const size_t targetCpu = result.computeToCpuMap.lookup(destination);
     const size_t sourceSlot = result.computeToCpuSlotMap.lookup(source);
-    const size_t targetSlot = result.computeToCpuSlotMap.lookup(target);
+    const size_t targetSlot = result.computeToCpuSlotMap.lookup(destination);
     const Time sourceStart = static_cast<Time>(result.computeToAestMap.lookup(source));
-    const Time targetStart = static_cast<Time>(result.computeToAestMap.lookup(target));
+    const Time targetStart =
+      static_cast<Time>(result.computeToAestMap.lookup(destination));
     if (sourceCpu == targetCpu && sourceSlot >= targetSlot)
       llvm::report_fatal_error("merge scheduling: same-CPU dependency order is invalid");
 
     Time earliestTargetStart = addOrMax(sourceStart, graph.nodes[edge.source].cost);
     if (sourceCpu != targetCpu)
       earliestTargetStart = addOrMax(
-        earliestTargetStart, getPeftTransferTime(edge.transferCost, sourceCpu, targetCpu, processorCount));
+        earliestTargetStart, getPeftTransferTime(edge.transferCost, sourceCpu, targetCpu, target));
     if (targetStart < earliestTargetStart) {
       std::string message = llvm::formatv("merge scheduling: dependency legality failed between tasks {0} and {1}",
                                           graph.nodes[edge.source].originalOrder,
@@ -89,30 +88,22 @@ void verifySchedule(const ComputeGraph& graph,
 
 } // namespace
 
-MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op)
-: entryOp(op) {
+MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op,
+                                                 const SchedulingTarget& schedulingTarget)
+: entryOp(op), target(schedulingTarget) {
   result = run();
 }
 
 MergeScheduleResult MergeSchedulingAnalysis::run() {
-  verifyExplicitPimCoreCount();
-  ComputeGraph graph = buildComputeGraph(entryOp);
+  ComputeGraph graph = buildComputeGraph(entryOp, target);
   if (!verifyAcyclic(graph))
     llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
 
-  size_t processorCount = 0;
-  if (coresCount.getValue() > 0)
-    processorCount = static_cast<size_t>(coresCount.getValue());
-
   MergeScheduleResult schedule = runPeftScheduler(
     graph, PeftScheduleOptions {
-      processorCount,
-      static_cast<unsigned long>(crossbarCountInCore.getValue()),
+      target,
       entryOp->getContext()});
-  verifySchedule(graph,
-                 schedule,
-                 static_cast<unsigned long>(crossbarCountInCore.getValue()),
-                 processorCount);
+  verifySchedule(graph, schedule, target);
   return schedule;
 }
 

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

@@ -3,17 +3,19 @@
 #include "mlir/IR/Operation.h"
 
 #include "MergeSchedule.hpp"
+#include "SchedulingTarget.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
 
 class MergeSchedulingAnalysis {
 public:
-  explicit MergeSchedulingAnalysis(mlir::Operation* op);
+  MergeSchedulingAnalysis(mlir::Operation* op, const SchedulingTarget& target);
   MergeScheduleResult& getResult() { return result; }
 
 private:
   mlir::Operation* entryOp = nullptr;
+  const SchedulingTarget& target;
   MergeScheduleResult result;
 
   MergeScheduleResult run();

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

@@ -4,8 +4,8 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 
-#include <cmath>
 #include <limits>
+#include <numeric>
 #include <optional>
 #include <queue>
 #include <tuple>
@@ -22,9 +22,7 @@ using namespace mlir;
 namespace {
 
 // Pressure means distinct weights exceed half the fleet's one-copy capacity.
-// The reserved headroom conservatively avoids greedy capacity fragmentation;
-// makespan, communication, instruction-path, and local-memory proxies remain stable.
-constexpr size_t kHighCrossbarPressureCapacityDivisor = 2;
+constexpr size_t kHighResidentWeightPressureCapacityDivisor = 2;
 
 struct ScheduledTask {
   size_t processor = std::numeric_limits<size_t>::max();
@@ -32,63 +30,22 @@ struct ScheduledTask {
   Time endTime = 0;
 };
 
-struct MeshModel {
-  size_t rows = 1;
-  size_t cols = 1;
-  long double averageDistance = 0.0L;
-
-  static MeshModel infer(size_t processorCount) {
-    MeshModel model;
-    if (processorCount == 0)
-      return model;
-
-    model.rows = static_cast<size_t>(std::sqrt(static_cast<long double>(processorCount)));
-    if (model.rows == 0)
-      model.rows = 1;
-    while (model.rows > 1 && processorCount % model.rows != 0)
-      --model.rows;
-    model.cols = (processorCount + model.rows - 1) / model.rows;
-
-    auto averageAxisDistance = [](size_t size) -> long double {
-      if (size <= 1)
-        return 0.0L;
-      return static_cast<long double>(size * size - 1) / (3.0L * static_cast<long double>(size));
-    };
-    model.averageDistance = averageAxisDistance(model.rows) + averageAxisDistance(model.cols);
-    return model;
-  }
-
-  std::pair<size_t, size_t> getCoord(size_t processor) const {
-    return {processor / cols, processor % cols};
-  }
-
-  size_t getDistance(size_t lhs, size_t rhs) const {
-    auto [lhsRow, lhsCol] = getCoord(lhs);
-    auto [rhsRow, rhsCol] = getCoord(rhs);
-    size_t rowDistance = lhsRow > rhsRow ? lhsRow - rhsRow : rhsRow - lhsRow;
-    size_t colDistance = lhsCol > rhsCol ? lhsCol - rhsCol : rhsCol - lhsCol;
-    return rowDistance + colDistance;
-  }
-
-  Time scaleTransferCost(Time transferCost, size_t sourceProcessor, size_t targetProcessor) const {
-    if (sourceProcessor == targetProcessor || transferCost == 0)
-      return 0;
-    long double distance = static_cast<long double>(getDistance(sourceProcessor, targetProcessor));
-    long double scale = averageDistance > 0.0L ? distance / averageDistance : 1.0L;
-    scale = std::max(0.25L, scale);
-    return static_cast<Time>(std::ceil(static_cast<long double>(transferCost) * scale));
-  }
+struct TopologyModel {
+  const SchedulingTarget& target;
 
   size_t getCenterDistance(size_t processor) const {
-    auto [row, col] = getCoord(processor);
-    size_t centerRow = rows / 2;
-    size_t centerCol = cols / 2;
-    size_t rowDistance = row > centerRow ? row - centerRow : centerRow - row;
-    size_t colDistance = col > centerCol ? col - centerCol : centerCol - col;
-    return rowDistance + colDistance;
+    Cost total = 0;
+    for (size_t other = 0; other < target.processorCount; ++other)
+      total = checkedAdd(total, target.getInterProcessorLatencyNs(processor, other));
+    return static_cast<size_t>(total);
   }
 };
 
+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;
@@ -143,14 +100,14 @@ void verifyOctTableSize(size_t nodeCount, size_t processorCount) {
   }
 }
 
-bool hasHighCrossbarPressure(const ComputeGraph& graph, size_t processorCount, size_t crossbarCapacity) {
-  if (crossbarCapacity > std::numeric_limits<size_t>::max() / processorCount)
+bool hasHighResidentWeightPressure(const ComputeGraph& graph, size_t processorCount, size_t residentWeightCapacity) {
+  if (residentWeightCapacity > std::numeric_limits<size_t>::max() / processorCount)
     return false;
-  const size_t threshold = processorCount * crossbarCapacity / kHighCrossbarPressureCapacityDivisor;
-  CrossbarUsage distinctWeights;
+  const size_t threshold = processorCount * residentWeightCapacity / kHighResidentWeightPressureCapacityDivisor;
+  ResidentWeightSet distinctWeights;
   for (const ComputeGraphNode& node : graph.nodes) {
-    for (const CrossbarWeight& weight : node.crossbarUsage) {
-      if (!containsCrossbarWeight(distinctWeights, weight))
+    for (const ResidentWeight& weight : node.residentWeights) {
+      if (!containsResidentWeight(distinctWeights, weight))
         distinctWeights.push_back(weight);
       if (distinctWeights.size() > threshold)
         return true;
@@ -159,36 +116,36 @@ bool hasHighCrossbarPressure(const ComputeGraph& graph, size_t processorCount, s
   return false;
 }
 
-std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
+std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph& graph,
                                                     size_t processorCount,
-                                                    size_t crossbarCapacity,
-                                                    const MeshModel& mesh,
-                                                    bool preferCrossbarReuse) {
+                                                              size_t residentWeightCapacity,
+                                                              const TopologyModel& topology,
+                                                              bool preferWeightReuse) {
   std::vector<size_t> weightedTasks;
   for (size_t task = 0; task < graph.nodes.size(); ++task)
-    if (!graph.nodes[task].crossbarUsage.empty())
+    if (!graph.nodes[task].residentWeights.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();
+    if (graph.nodes[lhs].residentWeights.size() != graph.nodes[rhs].residentWeights.size())
+      return graph.nodes[lhs].residentWeights.size() > graph.nodes[rhs].residentWeights.size();
     return graph.nodes[lhs].originalOrder < graph.nodes[rhs].originalOrder;
   });
 
-  std::vector<CrossbarUsage> reservations(processorCount);
+  std::vector<ResidentWeightSet> 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)
+      size_t residentWeightUnion =
+        getResidentWeightUnionSize(reservations[processor], graph.nodes[task].residentWeights);
+      if (residentWeightUnion > residentWeightCapacity)
         continue;
-      size_t addedCrossbars = crossbarUnion - reservations[processor].size();
-      ReservationScore score {preferCrossbarReuse ? addedCrossbars : reservedLoad[processor],
-                              preferCrossbarReuse ? reservedLoad[processor] : addedCrossbars,
-                              mesh.getCenterDistance(processor),
+      size_t addedWeights = residentWeightUnion - reservations[processor].size();
+      ReservationScore score {preferWeightReuse ? addedWeights : reservedLoad[processor],
+                              preferWeightReuse ? reservedLoad[processor] : addedWeights,
+                              topology.getCenterDistance(processor),
                               processor};
       if (!bestScore || score < *bestScore) {
         bestProcessor = processor;
@@ -200,13 +157,13 @@ std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
         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(),
+                      graph.nodes[task].residentWeights.size(),
                       processorCount,
-                      crossbarCapacity)
+                      residentWeightCapacity)
           .str();
       llvm::report_fatal_error(llvm::StringRef(message));
     }
-    insertCrossbarWeights(reservations[bestProcessor], graph.nodes[task].crossbarUsage);
+    insertResidentWeights(reservations[bestProcessor], graph.nodes[task].residentWeights);
     reservedLoad[bestProcessor] = addOrMax(reservedLoad[bestProcessor], graph.nodes[task].cost);
   }
   return reservations;
@@ -214,8 +171,8 @@ std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
 
 using LanePublicationSignatures = llvm::SmallVector<llvm::SmallVector<int64_t, 8>, 8>;
 
-FailureOr<LanePublicationSignatures>
-buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCache& publicationCache) {
+FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeBatch batch,
+                                                                    GraphBatchPublicationCache& publicationCache) {
   LanePublicationSignatures signatures(batch.getLaneCount());
   for (auto [resultIndex, result] : llvm::enumerate(batch.getResults())) {
     auto publicationMap = getGraphBatchPublicationMap(batch, resultIndex, publicationCache);
@@ -231,8 +188,8 @@ buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCach
       auto sourceOffsets = blueprint.getFragmentSourceOffsets();
       auto fragmentStrides = blueprint.getFragmentStrides();
       auto outputType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
-      if (!operandIndices || !sourceSlots || !sourceOffsets || !fragmentStrides
-          || !outputType || !outputType.hasStaticShape())
+      if (!operandIndices || !sourceSlots || !sourceOffsets || !fragmentStrides || !outputType
+          || !outputType.hasStaticShape())
         return blueprint.emitOpError("PEFT publication compatibility requires complete static fragment metadata"),
                failure();
 
@@ -241,11 +198,9 @@ buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCach
       int64_t rank = outputType.getRank();
       if (rank <= 0 || fragmentOffsets.size() != fragmentSizes.size()
           || fragmentOffsets.size() != fragmentStrides->size()
-          || fragmentOffsets.size() != operandIndices->size() * rank
-          || sourceSlots->size() != operandIndices->size()
+          || fragmentOffsets.size() != operandIndices->size() * rank || sourceSlots->size() != operandIndices->size()
           || sourceOffsets->size() != operandIndices->size())
-        return blueprint.emitOpError("PEFT publication compatibility found inconsistent fragment metadata"),
-               failure();
+        return blueprint.emitOpError("PEFT publication compatibility found inconsistent fragment metadata"), failure();
 
       llvm::SmallVector<llvm::SmallVector<size_t, 2>, 8> fragmentsByLane(batch.getLaneCount());
       for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
@@ -287,19 +242,83 @@ buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCach
 
 } // namespace
 
-Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount) {
-  return MeshModel::infer(processorCount).scaleTransferCost(transferCost, sourceProcessor, targetProcessor);
+std::vector<size_t> mapLogicalProcessorsToPhysicalCores(ArrayRef<Cost> logicalTrafficFlits,
+                                                        const SchedulingTarget& target) {
+  const size_t processorCount = target.processorCount;
+  assert(logicalTrafficFlits.size() == processorCount * processorCount
+         && "logical traffic matrix must cover every processor pair");
+
+  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;
+      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.processorCount;
+  const size_t processorCount = options.target.processorCount;
   if (processorCount == 0)
     llvm::report_fatal_error("PEFT scheduler: processor count must be positive");
-  MeshModel mesh = MeshModel::infer(processorCount);
-  const bool preferCrossbarReuse = hasHighCrossbarPressure(graph, processorCount, options.crossbarCapacity);
-  std::vector<CrossbarUsage> capacityReservations =
-    planCrossbarReservations(graph, processorCount, options.crossbarCapacity, mesh, preferCrossbarReuse);
+  TopologyModel topology {options.target};
+  const bool preferWeightReuse =
+    hasHighResidentWeightPressure(graph, processorCount, options.target.residentWeightCapacity);
+  std::vector<ResidentWeightSet> capacityReservations = planResidentWeightReservations(
+    graph, processorCount, options.target.residentWeightCapacity, topology, preferWeightReuse);
 
   verifyOctTableSize(nodeCount, processorCount);
   std::vector<std::vector<size_t>> reverseLevels = buildReverseLevels(graph);
@@ -317,7 +336,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
       std::vector<Time> maxVals(processorCount, 0);
 
       for (const auto& [succ, comm] : graph.successors[task]) {
-        Time valDifferentCpu = addOrMax(minOctPlusComp[succ], comm);
+        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);
@@ -378,7 +397,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
   }
 
   std::vector<char> scheduled(nodeCount, false);
-  std::vector<CrossbarUsage> processorCrossbars(processorCount);
+  std::vector<ResidentWeightSet> processorResidentWeights(processorCount);
   std::vector<ScheduledTask> schedules(nodeCount);
   std::vector<std::vector<size_t>> tasksByProcessor(processorCount);
 
@@ -394,23 +413,25 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     Time bestEft = 0;
     Time bestOeft = std::numeric_limits<Time>::max();
     unsigned int bestOverlapCount = 0;
+    size_t bestTaskCount = std::numeric_limits<size_t>::max();
     size_t bestCenterDistance = std::numeric_limits<size_t>::max();
-    size_t smallestCrossbarUnion = std::numeric_limits<size_t>::max();
-    bool crossbarRejected = false;
+    size_t smallestResidentWeightUnion = std::numeric_limits<size_t>::max();
+    bool residentWeightRejected = false;
 
     for (size_t processor = 0; processor < processorCount; ++processor) {
-      unsigned int overlapCount = countCrossbarOverlap(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;
+      unsigned int overlapCount =
+        countResidentWeightOverlap(processorResidentWeights[processor], graph.nodes[task].residentWeights);
+      size_t residentWeightUnion =
+        getResidentWeightUnionSize(capacityReservations[processor], graph.nodes[task].residentWeights);
+      smallestResidentWeightUnion = std::min(smallestResidentWeightUnion, residentWeightUnion);
+      if (!graph.nodes[task].residentWeights.empty() && residentWeightUnion > options.target.residentWeightCapacity) {
+        residentWeightRejected = true;
         continue;
       }
       Time dataReady = 0;
       for (const auto& [pred, comm] : graph.predecessors[task]) {
         const ScheduledTask& predSchedule = schedules[pred];
-        Time commPenalty = getPeftTransferTime(comm, predSchedule.processor, processor, processorCount);
+        Time commPenalty = getPeftTransferTime(comm, predSchedule.processor, processor, options.target);
         dataReady = std::max(dataReady, addOrMax(predSchedule.endTime, commPenalty));
       }
 
@@ -437,9 +458,10 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
 
       Time eft = addOrMax(est, computeCost);
       Time oeft = addOrMax(eft, oct[task * processorCount + processor]);
-      size_t centerDistance = mesh.getCenterDistance(processor);
-      bool betterCrossbarChoice =
-        preferCrossbarReuse ? overlapCount > bestOverlapCount : overlapCount < bestOverlapCount;
+      size_t centerDistance = topology.getCenterDistance(processor);
+      size_t taskCount = tasksByProcessor[processor].size();
+      bool betterResidentWeightChoice =
+        preferWeightReuse ? overlapCount > bestOverlapCount : overlapCount < bestOverlapCount;
 
       if (oeft < bestOeft || (oeft == bestOeft && eft < bestEft)
           || (oeft == bestOeft && eft == bestEft && est < bestEst)) {
@@ -448,40 +470,52 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
         bestEft = eft;
         bestOeft = oeft;
         bestOverlapCount = overlapCount;
+        bestTaskCount = taskCount;
         bestCenterDistance = centerDistance;
       }
-      else if (oeft == bestOeft && eft == bestEft && est == bestEst
+      else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount < bestTaskCount) {
+        bestProcessor = processor;
+        bestEst = est;
+        bestEft = eft;
+        bestOeft = oeft;
+        bestOverlapCount = overlapCount;
+        bestTaskCount = taskCount;
+        bestCenterDistance = centerDistance;
+      }
+      else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount == bestTaskCount
                && centerDistance < bestCenterDistance) {
         bestProcessor = processor;
         bestEst = est;
         bestEft = eft;
         bestOeft = oeft;
         bestOverlapCount = overlapCount;
+        bestTaskCount = taskCount;
         bestCenterDistance = centerDistance;
       }
-      else if (oeft == bestOeft && eft == bestEft && est == bestEst
-               && centerDistance == bestCenterDistance && betterCrossbarChoice) {
+      else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount == bestTaskCount
+               && centerDistance == bestCenterDistance && betterResidentWeightChoice) {
         bestProcessor = processor;
         bestEst = est;
         bestEft = eft;
         bestOeft = oeft;
         bestOverlapCount = overlapCount;
+        bestTaskCount = taskCount;
         bestCenterDistance = centerDistance;
       }
     }
 
     if (bestProcessor == std::numeric_limits<size_t>::max()) {
-      if (crossbarRejected) {
+      if (residentWeightRejected) {
         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 crossbar capacity {5}",
+                        "smallest processor union is {4}, exceeding resident-weight capacity {5}",
                         graph.nodes[task].originalOrder,
                         instance.laneStart,
                         instance.laneStart + instance.laneCount,
-                        graph.nodes[task].crossbarUsage.size(),
-                        smallestCrossbarUnion,
-                        options.crossbarCapacity)
+                        graph.nodes[task].residentWeights.size(),
+                        smallestResidentWeightUnion,
+                        options.target.residentWeightCapacity)
             .str();
         llvm::report_fatal_error(llvm::StringRef(message));
       }
@@ -495,8 +529,8 @@ 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);
+    insertResidentWeights(capacityReservations[bestProcessor], graph.nodes[task].residentWeights);
+    insertResidentWeights(processorResidentWeights[bestProcessor], graph.nodes[task].residentWeights);
 
     // 3. CRITICAL FIX: Topological Append
     // Because the readyQueue pops in strict topological order, simply pushing to the
@@ -576,6 +610,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
 
   // 6. Populate Final Result
   MergeScheduleResult result;
+  result.processorCount = processorCount;
   result.dominanceOrderCompute.reserve(nodeCount);
 
   for (size_t task : scheduledOrder)

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

@@ -4,18 +4,32 @@
 
 #include "ComputeGraph.hpp"
 #include "MergeSchedule.hpp"
+#include "SchedulingTarget.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
 
 struct PeftScheduleOptions {
-  size_t processorCount = 0;
-  unsigned long crossbarCapacity = 0;
+  SchedulingTarget target;
   mlir::MLIRContext* context = nullptr;
 };
 
-Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount);
+inline Time getPeftTransferTime(const TransferCost& transferCost,
+                                size_t sourceProcessor,
+                                size_t targetProcessor,
+                                const SchedulingTarget& target) {
+  if (sourceProcessor == targetProcessor)
+    return 0;
+  return checkedAdd(transferCost.fixed,
+                    checkedMultiply(transferCost.networkFlits, target.averageInterProcessorLatencyNs));
+}
+
+// PEFT assigns logical processors. Physical core IDs are chosen only after
+// materialization exposes the exact transfer traffic.
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options);
 
+std::vector<size_t> mapLogicalProcessorsToPhysicalCores(llvm::ArrayRef<Cost> logicalTrafficFlits,
+                                                        const SchedulingTarget& target);
+
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp

@@ -0,0 +1,54 @@
+#pragma once
+
+#include <cassert>
+#include <cstddef>
+#include <vector>
+
+#include "Utils.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+struct SchedulingTarget {
+  size_t processorCount = 0;
+  size_t residentWeightCapacity = 0;
+  std::vector<Cost> interProcessorLatencyNs;
+  Cost averageInterProcessorLatencyNs = 0;
+
+  Cost computeBitwidth = 8;
+  Cost processorPeriodNs = 1;
+  Cost localMemoryWidthBytes = 64;
+  Cost localMemoryReadLatencyCycles = 1;
+  Cost localMemoryWriteLatencyCycles = 1;
+  Cost transferWidthBytes = 8;
+  Cost vectorWidth = 16;
+  Cost vectorLatencyCycles = 4;
+
+  Cost matrixRows = 128;
+  Cost matrixColumns = 128;
+  Cost matrixPeriodNs = 1;
+  Cost matrixInputResolutionBits = 1;
+  Cost matrixInputLatencyCycles = 1;
+  Cost matrixInputParallelism = 128;
+  Cost matrixReadLatencyNs = 30;
+  Cost matrixSampleLatencyCycles = 1;
+  Cost matrixOutputLatencyCycles = 10;
+  Cost matrixOutputParallelism = 2;
+  Cost matrixShiftLatencyCycles = 1;
+  Cost matrixBufferLatencyCycles = 1;
+  Cost matrixInputBufferLatencyCycles = 0;
+  bool matrixPipeline = true;
+
+  Cost getInterProcessorLatencyNs(size_t source,
+                                  size_t destination) const {
+    assert(source < processorCount && destination < processorCount
+           && "processor index out of range");
+    assert(interProcessorLatencyNs.size()
+             == processorCount * processorCount
+           && "incomplete inter-processor latency matrix");
+    return interProcessorLatencyNs[source * processorCount + destination];
+  }
+};
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/TrivialGraphComputeMergePass.cpp

@@ -36,10 +36,14 @@ static bool hasOnlyStructuralAttrs(ComputeOp op) {
   });
 }
 
-static bool hasCapacityFor(Operation* producer, Operation* consumer) {
-  CrossbarUsage producerWeights = collectDistinctCrossbarWeights(producer);
-  CrossbarUsage consumerWeights = collectDistinctCrossbarWeights(consumer);
-  return getCrossbarUnionSize(producerWeights, consumerWeights) <= static_cast<size_t>(crossbarCountInCore.getValue());
+static bool hasCapacityFor(Operation* producer, Operation* consumer,
+                           size_t residentWeightCapacity) {
+  ResidentWeightSet producerWeights =
+    collectDistinctResidentWeights(producer);
+  ResidentWeightSet consumerWeights =
+    collectDistinctResidentWeights(consumer);
+  return getResidentWeightUnionSize(producerWeights, consumerWeights)
+      <= residentWeightCapacity;
 }
 
 template <typename ConsumerOp>
@@ -156,8 +160,11 @@ static void mapExternalArguments(OldOp oldOp, NewOp newOp, IRMapping& mapper, bo
 }
 
 struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
-  MergeTrivialScalarComputes(MLIRContext *context, TrivialGraphMergeStats *stats)
-      : OpRewritePattern(context), stats(stats) {}
+  MergeTrivialScalarComputes(MLIRContext *context,
+                             TrivialGraphMergeStats *stats,
+                             size_t residentWeightCapacity)
+      : OpRewritePattern(context), stats(stats),
+        residentWeightCapacity(residentWeightCapacity) {}
 
   LogicalResult matchAndRewrite(SpatGraphCompute consumer, PatternRewriter& rewriter) const override {
     SpatGraphCompute producer;
@@ -166,7 +173,8 @@ struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
       if (candidate && candidate->getBlock() == consumer->getBlock() && hasOnlyStructuralAttrs(candidate)
           && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
           && isExclusivelyConsumedBy(candidate, consumer)
-          && hasCapacityFor(candidate, consumer) && hasNoNestedArgumentCaptures(candidate)
+          && hasCapacityFor(candidate, consumer, residentWeightCapacity)
+          && hasNoNestedArgumentCaptures(candidate)
           && hasNoNestedArgumentCaptures(consumer)) {
         producer = candidate;
         break;
@@ -201,6 +209,7 @@ struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
 
 private:
   TrivialGraphMergeStats *stats;
+  size_t residentWeightCapacity;
 };
 
 static bool isLaneIndex(Value value, Value lane, int64_t laneCount) {
@@ -397,8 +406,11 @@ static bool hasDirectLaneConsumers(SpatGraphComputeBatch producer, SpatGraphComp
 }
 
 struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
-  MergeTrivialBatchComputes(MLIRContext *context, TrivialGraphMergeStats *stats)
-      : OpRewritePattern(context), stats(stats) {}
+  MergeTrivialBatchComputes(MLIRContext *context,
+                            TrivialGraphMergeStats *stats,
+                            size_t residentWeightCapacity)
+      : OpRewritePattern(context), stats(stats),
+        residentWeightCapacity(residentWeightCapacity) {}
 
   LogicalResult matchAndRewrite(SpatGraphComputeBatch consumer, PatternRewriter& rewriter) const override {
     SpatGraphComputeBatch producer;
@@ -409,7 +421,8 @@ struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
           && candidate.getLaneCount() == consumer.getLaneCount() && hasOnlyStructuralAttrs(candidate)
           && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
           && isExclusivelyConsumedBy(candidate, consumer)
-          && hasCapacityFor(candidate, consumer) && hasDirectLaneConsumers(candidate, consumer)
+          && hasCapacityFor(candidate, consumer, residentWeightCapacity)
+          && hasDirectLaneConsumers(candidate, consumer)
           && succeeded(fragments = collectPublishedFragments(candidate))) {
         producer = candidate;
         break;
@@ -469,11 +482,16 @@ struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
 
 private:
   TrivialGraphMergeStats *stats;
+  size_t residentWeightCapacity;
 };
 
 struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMergePass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TrivialGraphComputeMergePass)
 
+  TrivialGraphComputeMergePass() = default;
+  explicit TrivialGraphComputeMergePass(size_t residentWeightCapacity)
+      : residentWeightCapacity(residentWeightCapacity) {}
+
   StringRef getArgument() const override { return "pim-trivial-graph-compute-merge"; }
   StringRef getDescription() const override {
     return "Inline linear exclusive graph compute chains while preserving fan-in boundaries.";
@@ -481,10 +499,17 @@ struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMerge
 
   void runOnOperation() override {
     ModuleOp module = getOperation();
+    if (residentWeightCapacity == 0) {
+      module.emitError(
+        "TrivialGraphComputeMerge requires an explicit valid resident-weight capacity");
+      signalPassFailure();
+      return;
+    }
     TrivialGraphMergeStats stats;
     std::tie(stats.scalarBefore, stats.batchBefore) = countGraphComputes(module);
     RewritePatternSet patterns(&getContext());
-    patterns.add<MergeTrivialScalarComputes, MergeTrivialBatchComputes>(&getContext(), &stats);
+    patterns.add<MergeTrivialScalarComputes, MergeTrivialBatchComputes>(
+      &getContext(), &stats, residentWeightCapacity);
     if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
       signalPassFailure();
       return;
@@ -510,6 +535,9 @@ struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMerge
         signalPassFailure();
     }
   }
+
+private:
+  size_t residentWeightCapacity = 0;
 };
 
 } // namespace
@@ -519,4 +547,10 @@ std::unique_ptr<Pass> createTrivialGraphComputeMergePass() {
   return std::make_unique<spatial::TrivialGraphComputeMergePass>();
 }
 
+std::unique_ptr<Pass> createTrivialGraphComputeMergePass(
+    size_t residentWeightCapacity) {
+  return std::make_unique<spatial::TrivialGraphComputeMergePass>(
+    residentWeightCapacity);
+}
+
 } // namespace onnx_mlir