[Util] Filter non-returning region terminators (#24795)

`Explorer::walkReturnOps` treated every block terminator in a region
branch operation as a return to the parent. For scf.while, this included
the after-region yield that branches back to the before region, causing
loop-carried values to contribute to result affinities.

Skip terminators that do not implement the region branch terminator
interface, as well as region branch terminators that cannot branch to
the parent operation.

Fixes #24787

Signed-off-by: zhangpan <98025960+zhangpan2001@users.noreply.github.com>

compiler/src/iree/compiler/Dialect/Stream/Transforms/test/annotate_affinities.mlir

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/annotate_affinities.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/annotate_affinities.mlir
index 40ba703..58f1023 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/annotate_affinities.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/annotate_affinities.mlir
@@ -1657,18 +1657,13 @@
     scf.yield %cst_a, %t : tensor<1xi32>, tensor<1xi32>
   // CHECK: } attributes {
   // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_c>]]
-  // The while result's producer affinity is currently over-approximated as the
-  // union of both loop-carried affinities ([dev_a, dev_c]) instead of just
-  // [dev_c] (the single result only forwards %arg1). Re-enable this check once
-  // that producer-side over-approximation is fixed (https://github.com/iree-org/iree/issues/24787).
-  // COM: CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  // The single result only forwards %arg1, so its producer affinity is @dev_c.
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
   }
   // CHECK: flow.tensor.transfer
+  // The transfer consumes the while result, so its operand affinity is @dev_c.
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
   // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
-  // The transfer consumes the while result, so its operand affinity inherits
-  // the same over-approximated union. Re-enable this check once the
-  // producer-side over-approximation is fixed (https://github.com/iree-org/iree/issues/24787).
-  // COM: CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
   %while_c = flow.tensor.transfer %while : tensor<1xi32> to #hal.device.promise<@dev_c>
   // CHECK: util.return
   // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]

compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp

diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp b/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp
index 9b921c2..8c327aa 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp
@@ -569,8 +569,22 @@
     auto enumerateTerminatorOps = [&](Region &region) {
       for (auto &block : region) {
         if (auto *terminatorOp = block.getTerminator()) {
-          // TODO(benvanik): ensure this terminator can return to parent? this
-          // region op interface confuses me.
+          auto branchTerminatorOp =
+              dyn_cast<RegionBranchTerminatorOpInterface>(terminatorOp);
+          if (!branchTerminatorOp) {
+            continue;
+          }
+          // Region branch terminators may either return to the parent
+          // operation or branch to another region. Only the former can
+          // provide operands for the parent operation's results.
+          SmallVector<RegionSuccessor, 2> successors;
+          regionOp.getSuccessorRegions(RegionBranchPoint(branchTerminatorOp),
+                                       successors);
+          if (llvm::none_of(successors, [&](RegionSuccessor successor) {
+                return successor == parentOp;
+              })) {
+            continue;
+          }
           LLVM_DEBUG({
             llvm::dbgs() << "  == emitting region branch terminator op ";
             terminatorOp->print(llvm::dbgs(), asmState);