[Mlir-commits] [mlir] SparseAnalysis: support ReturnLike terminators (PR #140797)

[Jeremy Kun]

================
@@ -299,6 +330,28 @@ class SparseForwardDataFlowAnalysis
     setAllToEntryStates(resultLattices);
   }
 
+  /// Visit a region terminator. This is intended for non-control-flow
+  /// region-bearing ops whose terminators determine the lattice values of the
+  /// parent op's results. By default the terminator's operand lattices are
+  /// forwarded to the parent result lattices, if there is a 1-1
+  /// correspondence.
+  virtual LogicalResult visitNonControlFlowTerminator(
+      Operation *terminatorOp,
+      ArrayRef<const StateT *> terminatorOperandLattices,
+      ArrayRef<StateT *> parentResultLattices) {
+    // ReturnLike terminators forward their lattice values to the results of the
+    // parent op.
+    if (terminatorOp->hasTrait<OpTrait::ReturnLike>() &&
+        terminatorOperandLattices.size() == parentResultLattices.size()) {
+      for (const auto &[operandLattice, resultLattice] :
+           llvm::zip(terminatorOperandLattices, parentResultLattices)) {
+        propagateIfChanged(resultLattice, resultLattice->join(*operandLattice));
+      }
+    }
+
+    return success();
----------------
j2kun wrote:

I would add to this the documentation on `RegionSuccessor`:

```
/// This interface assumes that the values from the current region that are used
/// to populate the successor inputs are the operands of the return-like
/// terminator operations in the blocks within this region.
```

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