fix vgg16 diamonds

updat ops validations

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();