remove dead logic

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

@@ -9,7 +9,6 @@
 
 #include <cstddef>
 #include <cstdint>
-#include <functional>
 #include <optional>
 #include <utility>
 
@@ -54,7 +53,6 @@ public:
     replaceExternalUses();
     if (failed(eraseOldScheduledOps()))
       return failure();
-    moveExternalUsersBeforeReturn();
     return success();
   }
 
@@ -97,18 +95,6 @@ private:
          | static_cast<uint32_t>(channelInfo.targetCoreId);
   }
 
-  void collectExternalUsers(Operation* op) {
-    if (!externalUsersToMove.insert(op).second)
-      return;
-    for (Value result : op->getResults()) {
-      for (Operation* user : result.getUsers()) {
-        if (oldComputeOps.contains(user) || isa<func::ReturnOp>(user))
-          continue;
-        collectExternalUsers(user);
-      }
-    }
-  }
-
   void collectScheduledTasks() {
     for (ComputeInstance scheduledInstance : schedule->dominanceOrderCompute) {
       oldComputeOps.insert(scheduledInstance.op);
@@ -151,25 +137,22 @@ private:
         auto& remoteInputs = remoteInputsByTask[task.computeInstance];
         remoteInputs.resize(taskInputs.size());
         for (auto [inputIndex, input] : llvm::enumerate(taskInputs)) {
-          auto producerRef = getProducerValueRef(input);
-          if (producerRef) {
+          if (auto producerRef = getProducerValueRef(input)) {
             auto producerIt = taskByComputeInstance.find(producerRef->instance);
-            if (producerIt != taskByComputeInstance.end()) {
-              if (producerIt->second.cpu != cpu) {
-                ChannelInfo info {
-                  (*nextChannelId)++,
-                  static_cast<int32_t>(producerIt->second.cpu),
-                  static_cast<int32_t>(cpu),
-                };
-                remoteInputs[inputIndex] = info;
-                auto& perResultChannels = remoteSendsByTask[producerRef->instance];
-                if (perResultChannels.empty())
-                  perResultChannels.resize(getComputeInstanceOutputTypes(producerIt->second.computeInstance).size());
-                perResultChannels[producerRef->resultIndex].push_back(
-                  {info, task.computeInstance, inputIndex, task.orderWithinCpu, 0});
-              }
-              continue;
+            if (producerIt->second.cpu != cpu) {
+              ChannelInfo info {
+                (*nextChannelId)++,
+                static_cast<int32_t>(producerIt->second.cpu),
+                static_cast<int32_t>(cpu),
+              };
+              remoteInputs[inputIndex] = info;
+              auto& perResultChannels = remoteSendsByTask[producerRef->instance];
+              if (perResultChannels.empty())
+                perResultChannels.resize(getComputeInstanceOutputTypes(producerIt->second.computeInstance).size());
+              perResultChannels[producerRef->resultIndex].push_back(
+                {info, task.computeInstance, inputIndex, task.orderWithinCpu, 0});
             }
+            continue;
           }
           if (seenExternalInputsByCpu[cpu].insert(input).second)
             cpuExternalInputs[cpu].push_back(input);
@@ -183,8 +166,6 @@ private:
             if (oldComputeOps.contains(useOwner))
               continue;
             hasExternalUser = true;
-            if (!isa<func::ReturnOp>(useOwner))
-              collectExternalUsers(useOwner);
           }
           if (hasExternalUser)
             cpuExternalOutputs[cpu].push_back({task.computeInstance, resultIndex});
@@ -407,7 +388,8 @@ private:
               if (producerIt->second.cpu == cpu) {
                 auto producedIt = producedValuesByTask.find(producerRef->instance);
                 if (producedIt == producedValuesByTask.end() || producedIt->second.size() <= producerRef->resultIndex) {
-                  task.computeInstance.op->emitOpError("missing local producer value during per-cpu merge materialization")
+                  task.computeInstance.op->emitOpError(
+                    "missing local producer value during per-cpu merge materialization")
                     << " consumerCpu=" << cpu << " producerCpu=" << producerIt->second.cpu
                     << " producerLaneStart=" << producerRef->instance.laneStart
                     << " producerLaneCount=" << producerRef->instance.laneCount;
@@ -586,18 +568,6 @@ private:
     return success();
   }
 
-  void moveExternalUsersBeforeReturn() {
-    SmallVector<Operation*> orderedUsersToMove;
-    for (Operation& op : func.getBody().front()) {
-      if (&op == returnOp.getOperation())
-        break;
-      if (externalUsersToMove.contains(&op))
-        orderedUsersToMove.push_back(&op);
-    }
-    for (Operation* op : orderedUsersToMove)
-      op->moveBefore(returnOp);
-  }
-
   func::FuncOp func;
   const MergeScheduleResult* schedule = nullptr;
   int64_t* nextChannelId = nullptr;
@@ -610,7 +580,6 @@ private:
   DenseMap<size_t, SmallVector<ScheduledTask>> tasksByCpu;
   SmallVector<size_t> orderedCpus;
   DenseSet<size_t> seenCpus;
-  DenseSet<Operation*> externalUsersToMove;
   DenseMap<ComputeInstance, SmallVector<SmallVector<RemoteSendInfo>>> remoteSendsByTask;
   DenseMap<ComputeInstance, SmallVector<std::optional<ChannelInfo>>> remoteInputsByTask;
   DenseMap<size_t, SmallVector<Value>> cpuExternalInputs;

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

@@ -13,7 +13,6 @@
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -28,9 +27,7 @@
 #include <cstdint>
 #include <cstdlib>
 #include <fstream>
-#include <functional>
 #include <iterator>
-#include <limits>
 #include <memory>
 #include <optional>
 #include <tuple>
@@ -39,13 +36,11 @@
 
 #include "MaterializeMergeSchedule.hpp"
 #include "PostMergeCompaction.hpp"
-#include "RegularOpCompaction.hpp"
 #include "Scheduling/ComputeInstanceUtils.hpp"
 #include "Scheduling/MergeSchedulingAnalysis.hpp"
 #include "src/Accelerators/PIM/Common/IR/CompactAsmUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -53,10 +48,8 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 using namespace onnx_mlir::compact_asm;
-using ProducerValueRef = spatial::ProducerValueRef;
 using SpatCompute = spatial::SpatCompute;
 using SpatComputeBatch = spatial::SpatComputeBatch;
-using spatial::getOriginalSpatCompute;
 using spatial::getProducerValueRef;
 
 bool isMergeProfilingEnabled() { return std::getenv("RAPTOR_PROFILE_MERGE") != nullptr; }
@@ -303,7 +296,7 @@ void emitMotifProfile(func::FuncOp funcOp) {
     }
 
     for (Value input : compute.getInputs()) {
-      auto parent = getOriginalSpatCompute(input.getDefiningOp());
+      auto parent = dyn_cast<SpatCompute>(input.getDefiningOp());
       if (!parent || parent == compute)
         continue;
       auto parentIt = computeToIndex.find(parent);

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

@@ -22,7 +22,7 @@ size_t getBatchChunkTargetCount(int32_t laneCount) {
 }
 
 ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex) {
-  size_t totalLanes = static_cast<size_t>(batch.getLaneCount());
+  size_t totalLanes = batch.getLaneCount();
   size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
   size_t baseChunkSize = totalLanes / chunkCount;
   size_t largeChunkCount = totalLanes % chunkCount;
@@ -33,7 +33,7 @@ ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex)
 }
 
 ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) {
-  size_t totalLanes = static_cast<size_t>(batch.getLaneCount());
+  size_t totalLanes = batch.getLaneCount();
   size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
   size_t baseChunkSize = totalLanes / chunkCount;
   size_t largeChunkCount = totalLanes % chunkCount;
@@ -47,32 +47,11 @@ ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) {
   return getBatchChunkForIndex(batch, chunkIndex);
 }
 
-SpatCompute getOriginalSpatCompute(Operation *op) {
-  if (!op)
-    return {};
-
-  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
-    op = extract.getSource().getDefiningOp();
-    if (!op)
-      return {};
-  }
-
-  return dyn_cast<SpatCompute>(op);
-}
-
 std::optional<ProducerValueRef> getProducerValueRef(Value value) {
   Operation *op = value.getDefiningOp();
   if (!op)
     return std::nullopt;
 
-  //TODO Extract Slice is not the only global non compute operation. There are other legal op
-  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
-    value = extract.getSource();
-    op = value.getDefiningOp();
-    if (!op)
-      return std::nullopt;
-  }
-
   if (auto compute = dyn_cast<SpatCompute>(op)) {
     return ProducerValueRef {
       ComputeInstance {compute.getOperation(), 0, 1},
@@ -81,9 +60,9 @@ std::optional<ProducerValueRef> getProducerValueRef(Value value) {
   }
 
   if (auto batch = dyn_cast<SpatComputeBatch>(op)) {
-    uint32_t lane = static_cast<uint32_t>(cast<OpResult>(value).getResultNumber());
+    uint32_t lane = cast<OpResult>(value).getResultNumber();
     ComputeInstance instance = getBatchChunkForLane(batch, lane);
-    size_t resultIndex = static_cast<size_t>(lane - instance.laneStart);
+    size_t resultIndex = lane - instance.laneStart;
     return ProducerValueRef {instance, resultIndex};
   }
 

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

@@ -26,7 +26,6 @@ size_t getBatchChunkTargetCount(int32_t laneCount);
 ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex);
 ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane);
 
-SpatCompute getOriginalSpatCompute(mlir::Operation *op);
 std::optional<ProducerValueRef> getProducerValueRef(mlir::Value value);
 std::optional<ComputeInstance> getComputeProducerInstance(mlir::Value value);
 

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

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

validation/validate.py

@@ -67,7 +67,7 @@ def main():
                     help="Core count to pass to Raptor. Required for PIM validation.")
     ap.add_argument("--pim-merge-scheduler", choices=("peft", "dcp"), default="peft",
                     help="Scheduler used by the Spatial merge-compute-nodes pass.")
-    ap.add_argument("--command-timeout-seconds", type=float, default=6000000000000000.0,
+    ap.add_argument("--command-timeout-seconds", type=float, default=60.0,
                     help="Per-subprocess timeout in seconds for compiler, runner, and simulator commands.")
     ap.add_argument("--clean", action="store_true",
                     help="Remove generated validation artifacts under each model workspace and exit.")