[Integrate] Cherry-pick llvm/llvm-project@41f6566 (#22470)

Still carrying two reverts as #22366.

New cherry-picks:
- [MLIR] Revamp RegionBranchOpInterface
(llvm/llvm-project@41f6566)
- Corresponding torch-mlir fix
(llvm/torch-mlir#4358)

Also, IREE fixes on `RegionBranchPoint` and `RegionSuccessor`.

ci-extra: test_torch, windows_x64_msvc

Signed-off-by: Yu-Zhewen <[email protected]>

Author: Yu-Zhewen

compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp

@@ -3085,8 +3085,8 @@ std::pair<unsigned, unsigned> AsyncExecuteOp::getTiedResultsIndexAndLength() {
 }
 
 OperandRange
-AsyncExecuteOp::getEntrySuccessorOperands(RegionBranchPoint point) {
-  assert(point.getRegionOrNull() == &getBody() && "invalid region index");
+AsyncExecuteOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+  assert(successor.getSuccessor() == &getBody() && "invalid region index");
   return getResourceOperands();
 }
 
@@ -3096,7 +3096,7 @@ void AsyncExecuteOp::getSuccessorRegions(
   // return the correct RegionSuccessor purely based on the index being None or
   // 0.
   if (!point.isParent()) {
-    regions.push_back(RegionSuccessor(getResults()));
+    regions.push_back(RegionSuccessor(getOperation(), getResults()));
   } else {
     regions.push_back(RegionSuccessor(&getBody(), getBody().getArguments()));
   }
@@ -3240,8 +3240,8 @@ LogicalResult AsyncConcurrentOp::verify() {
 }
 
 OperandRange
-AsyncConcurrentOp::getEntrySuccessorOperands(RegionBranchPoint point) {
-  assert(point == &getBody() && "invalid region index");
+AsyncConcurrentOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+  assert(successor.getSuccessor() == &getBody() && "invalid region index");
   return getResourceOperands();
 }
 
@@ -3251,7 +3251,7 @@ void AsyncConcurrentOp::getSuccessorRegions(
   // return the correct RegionSuccessor purely based on the index being None or
   // 0.
   if (!point.isParent()) {
-    regions.push_back(RegionSuccessor(getResults()));
+    regions.push_back(RegionSuccessor(getOperation(), getResults()));
   } else {
     regions.push_back(RegionSuccessor(&getBody(), getBody().getArguments()));
   }
@@ -4049,8 +4049,9 @@ LogicalResult CmdExecuteOp::verify() {
   return success();
 }
 
-OperandRange CmdExecuteOp::getEntrySuccessorOperands(RegionBranchPoint point) {
-  assert(point == &getBody() && "invalid region index");
+OperandRange
+CmdExecuteOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+  assert(successor.getSuccessor() == &getBody() && "invalid region index");
   return getResourceOperands();
 }
 
@@ -4060,7 +4061,8 @@ void CmdExecuteOp::getSuccessorRegions(
   // return the correct RegionSuccessor purely based on the index being None or
   // 0.
   if (!point.isParent()) {
-    regions.push_back(RegionSuccessor({}));
+    regions.push_back(
+        RegionSuccessor(getOperation(), Operation::result_range(nullptr, 0)));
   } else {
     regions.push_back(RegionSuccessor(&getBody(), getBody().getArguments()));
   }
@@ -4130,7 +4132,8 @@ void CmdSerialOp::getSuccessorRegions(
   // return the correct RegionSuccessor purely based on the index being None or
   // 0.
   if (!point.isParent()) {
-    regions.push_back(RegionSuccessor({}));
+    regions.push_back(
+        RegionSuccessor(getOperation(), Operation::result_range(nullptr, 0)));
   } else {
     regions.push_back(RegionSuccessor(&getBody(), {}));
   }
@@ -4155,7 +4158,8 @@ void CmdConcurrentOp::getSuccessorRegions(
   // return the correct RegionSuccessor purely based on the index being None or
   // 0.
   if (!point.isParent()) {
-    regions.push_back(RegionSuccessor({}));
+    regions.push_back(
+        RegionSuccessor(getOperation(), Operation::result_range(nullptr, 0)));
   } else {
     regions.push_back(RegionSuccessor(&getBody(), {}));
   }
@@ -4426,7 +4430,7 @@ LogicalResult DispatchWorkgroupSizeOp::verify() {
 //===----------------------------------------------------------------------===//
 
 MutableOperandRange
-YieldOp::getMutableSuccessorOperands(RegionBranchPoint point) {
+YieldOp::getMutableSuccessorOperands(RegionSuccessor successor) {
   return getResourceOperandsMutable();
 }
 

compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp

@@ -1161,7 +1161,7 @@ static IREE::Stream::AffinityAttr findLocalValueAffinity(Value value) {
       auto terminatorOp =
           cast<RegionBranchTerminatorOpInterface>(block.getTerminator());
       value = terminatorOp.getSuccessorOperands(
-          RegionBranchPoint::parent())[resultIndex];
+          RegionSuccessor(definingOp, definingOp->getResults()))[resultIndex];
     } else if (auto tiedOp =
                    dyn_cast<IREE::Util::TiedOpInterface>(definingOp)) {
       // If the producer is tied then try to get the operand.

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

@@ -599,7 +599,8 @@ TraversalResult Explorer::walkReturnOperands(Operation *parentOp,
   return walkReturnOps(parentOp, [&](Operation *returnOp) {
     if (auto terminatorOp =
             dyn_cast<RegionBranchTerminatorOpInterface>(returnOp)) {
-      return fn(terminatorOp.getSuccessorOperands(RegionBranchPoint::parent()));
+      return fn(terminatorOp.getSuccessorOperands(
+          RegionSuccessor(parentOp, parentOp->getResults())));
     } else {
       return fn(returnOp->getOperands());
     }
@@ -992,8 +993,9 @@ TraversalResult Explorer::walkTransitiveUses(Value value, UseWalkFn fn,
   // Move within/out-of a region.
   auto traverseRegionBranchOp = [&](RegionBranchTerminatorOpInterface branchOp,
                                     unsigned operandIdx) {
-    auto successorOperands =
-        branchOp.getSuccessorOperands(RegionBranchPoint::parent());
+    Operation *parentOp = branchOp.getOperation()->getParentOp();
+    auto successorOperands = branchOp.getSuccessorOperands(
+        RegionSuccessor(parentOp, parentOp->getResults()));
     unsigned beginIdx = successorOperands.getBeginOperandIndex();
     if (operandIdx < beginIdx ||
         operandIdx >= beginIdx + successorOperands.size()) {
@@ -1003,8 +1005,7 @@ TraversalResult Explorer::walkTransitiveUses(Value value, UseWalkFn fn,
                  << operandIdx << "\n");
       return TraversalResult::COMPLETE;
     }
-    auto result = branchOp.getSuccessorOperands(
-        RegionBranchPoint::parent())[operandIdx - beginIdx];
+    auto result = successorOperands[operandIdx - beginIdx];
     LLVM_DEBUG({
       llvm::dbgs() << "   + queuing region result ";
       result.printAsOperand(llvm::dbgs(), asmState);