Merge branch 'TestRottoConDeadLock' of chef.heaplab.deib.polimi.it:nnicolosi/Raptor into TestRottoConDeadLock

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

@@ -77,9 +77,16 @@ static LogicalResult eraseOldGraph(func::FuncOp funcOp,
         rewriter.eraseOp(blueprint);
       continue;
     }
-    if (!op->use_empty())
-      return op->emitOpError(
-        "phase 2 cannot erase an old graph compute with live results");
+    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();
+        }
+      }
+    }
     rewriter.eraseOp(op);
   }
   return success();
@@ -100,6 +107,25 @@ static LogicalResult eraseDeferredSourceSelectors(
   return success();
 }
 
+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();
+    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())
+      unused.push_back(deferred);
+  });
+  for (SpatDeferredCommunicationOp deferred : llvm::reverse(unused))
+    rewriter.eraseOp(deferred);
+}
+
 static LogicalResult verifyDominance(func::FuncOp funcOp) {
   DominanceInfo dominance(funcOp);
   WalkResult result = funcOp.walk([&](Operation *op) {
@@ -119,6 +145,9 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) {
 LogicalResult realizeDeferredCommunication(
     func::FuncOp funcOp,
     const ScheduledComputeMaterializationResult &materialization) {
+  IRRewriter rewriter(funcOp.getContext());
+  eraseUnusedIdentityDeferredCommunications(funcOp, rewriter);
+
   auto transfers = buildDeferredTransferPlan(funcOp, materialization);
   if (failed(transfers))
     return funcOp.emitOpError(
@@ -134,7 +163,6 @@ LogicalResult realizeDeferredCommunication(
     return funcOp.emitOpError(
       "phase 2 failed to build sparse boundary programs");
 
-  IRRewriter rewriter(funcOp.getContext());
   if (failed(retargetDeferredPublications(funcOp, *transfers))
       || failed(replaceFinalGraphPublications(funcOp, *transfers)))
     return failure();

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

@@ -42,6 +42,17 @@ static bool hasCapacityFor(Operation* producer, Operation* consumer) {
   return getCrossbarUnionSize(producerWeights, consumerWeights) <= static_cast<size_t>(crossbarCountInCore.getValue());
 }
 
+template <typename ConsumerOp>
+static bool isUniqueGraphComputePredecessor(Operation *candidate, ConsumerOp consumer) {
+  llvm::SmallSetVector<Operation *, 4> predecessors;
+  for (Value input : consumer.getInputs()) {
+    Operation *producer = input.getDefiningOp();
+    if (producer && isGraphComputeLike(producer))
+      predecessors.insert(producer);
+  }
+  return predecessors.size() == 1 && predecessors.front() == candidate;
+}
+
 struct TrivialGraphMergeStats {
   size_t scalarBefore = 0;
   size_t batchBefore = 0;
@@ -153,7 +164,8 @@ struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
     for (Value input : consumer.getInputs()) {
       auto candidate = input.getDefiningOp<SpatGraphCompute>();
       if (candidate && candidate->getBlock() == consumer->getBlock() && hasOnlyStructuralAttrs(candidate)
-          && hasOnlyStructuralAttrs(consumer) && isExclusivelyConsumedBy(candidate, consumer)
+          && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
+          && isExclusivelyConsumedBy(candidate, consumer)
           && hasCapacityFor(candidate, consumer) && hasNoNestedArgumentCaptures(candidate)
           && hasNoNestedArgumentCaptures(consumer)) {
         producer = candidate;
@@ -319,7 +331,8 @@ struct FoldBatchLeadingUnitNormalization : OpRewritePattern<SpatGraphCompute> {
       ? SpatGraphComputeBatch()
       : consumer.getInputs().front().getDefiningOp<SpatGraphComputeBatch>();
     if (!producer || producer->getBlock() != consumer->getBlock() || !hasOnlyStructuralAttrs(producer)
-        || !hasOnlyStructuralAttrs(consumer) || !isExclusivelyConsumedBy(producer, consumer))
+        || !hasOnlyStructuralAttrs(consumer) || !isUniqueGraphComputePredecessor(producer, consumer)
+        || !isExclusivelyConsumedBy(producer, consumer))
       return failure();
     auto fragments = collectPublishedFragments(producer);
     if (!matchLeadingUnitNormalization(producer, consumer) || failed(fragments))
@@ -394,7 +407,8 @@ struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
       auto candidate = input.getDefiningOp<SpatGraphComputeBatch>();
       if (candidate && candidate->getBlock() == consumer->getBlock()
           && candidate.getLaneCount() == consumer.getLaneCount() && hasOnlyStructuralAttrs(candidate)
-          && hasOnlyStructuralAttrs(consumer) && isExclusivelyConsumedBy(candidate, consumer)
+          && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
+          && isExclusivelyConsumedBy(candidate, consumer)
           && hasCapacityFor(candidate, consumer) && hasDirectLaneConsumers(candidate, consumer)
           && succeeded(fragments = collectPublishedFragments(candidate))) {
         producer = candidate;
@@ -462,7 +476,7 @@ struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMerge
 
   StringRef getArgument() const override { return "pim-trivial-graph-compute-merge"; }
   StringRef getDescription() const override {
-    return "Inline graph computes with exclusive direct dependencies before PEFT scheduling.";
+    return "Inline linear exclusive graph compute chains while preserving fan-in boundaries.";
   }
 
   void runOnOperation() override {