[Mlir-commits] [mlir] [mlir][Transforms] Fix crash in `reconcile-unrealized-casts` (PR #158067)

[Mehdi Amini]

================
@@ -3337,39 +3335,106 @@ void mlir::reconcileUnrealizedCasts(
   // Process ops in the worklist bottom-to-top.
   while (!worklist.empty()) {
     UnrealizedConversionCastOp castOp = worklist.pop_back_val();
-    if (castOp->use_empty()) {
-      // DCE: If the op has no users, erase it. Add the operands to the
-      // worklist to find additional DCE opportunities.
-      enqueueOperands(castOp);
-      if (remainingCastOps)
-        erasedOps.insert(castOp.getOperation());
-      castOp->erase();
-      continue;
-    }
 
     // Traverse the chain of input cast ops to see if an op with the same
     // input types can be found.
     UnrealizedConversionCastOp nextCast = castOp;
     while (nextCast) {
       if (nextCast.getInputs().getTypes() == castOp.getResultTypes()) {
+        if (llvm::any_of(nextCast.getInputs(), [&](Value v) {
+              return v.getDefiningOp() == castOp;
+            })) {
+          // Ran into a cycle.
+          break;
+        }
+
         // Found a cast where the input types match the output types of the
-        // matched op. We can directly use those inputs and the matched op can
-        // be removed.
-        enqueueOperands(castOp);
+        // matched op. We can directly use those inputs.
         castOp.replaceAllUsesWith(nextCast.getInputs());
-        if (remainingCastOps)
-          erasedOps.insert(castOp.getOperation());
-        castOp->erase();
         break;
       }
       nextCast = getInputCast(nextCast);
     }
   }
 
-  if (remainingCastOps)
-    for (UnrealizedConversionCastOp op : castOps)
-      if (!erasedOps.contains(op.getOperation()))
+  // A set of all alive cast ops. I.e., ops whose results are (transitively)
+  // used by an op that is not a cast op.
+  DenseSet<Operation *> liveOps;
+
+  // Helper function that marks the given op and transitively reachable input
+  // cast ops as alive.
+  auto markOpLive = [&](Operation *op) {
+    SmallVector<Operation *> worklist;
+    worklist.push_back(op);
+    while (!worklist.empty()) {
+      Operation *op = worklist.pop_back_val();
+      if (liveOps.insert(op).second) {
+        // Successfully inserted: process reachable input cast ops.
+        for (Value v : op->getOperands())
+          if (auto castOp = v.getDefiningOp<UnrealizedConversionCastOp>())
+            if (isCastOpFn(castOp))
+              worklist.push_back(castOp);
+      }
+    }
+  };
+
+  // Find all alive cast ops.
+  for (UnrealizedConversionCastOp op : castOps) {
----------------
joker-eph wrote:

```suggestion
  for (UnrealizedConversionCastOp op : castOps) {
     // The op may have been marked live already as being an operand of another live cast op.
     if (liveOps.contains(op.getOperation()) continue;
```

https://github.com/llvm/llvm-project/pull/158067