[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:

The documentation for `ReturnLike`.

```
/// This trait indicates that a terminator operation is "return-like". This
/// means that it exits its current region and forwards its operands as "exit"
/// values to the parent region. Operations with this trait are not permitted to
/// contain successors or produce results.
```

It is a bit ambiguous, but my read is that "forwarding its operands as exit values to the parent region" means the types do not change.

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