[mlir] Support getSuccessorInputs from parent op

[Mogball]

Summary

Ops that implement RegionBranchOpInterface are allowed to indicate that they can branch back to themselves in getSuccessorRegions, but there is no API that allows them to specify the forwarded operands. This patch enables that by changing getSuccessorEntryOperands to accept None.

Fixes #54928

mlir/include/mlir/Dialect/SCF/SCFOps.td

OpOperand &getOpOperandForResult(OpResult opResult) {
       return getOperation()->getOpOperand(
         getNumControlOperands() + opResult.getResultNumber());
     }
 
     /// Return operands used when entering the region at 'index'. These operands
     /// correspond to the loop iterator operands, i.e., those exclusing the
     /// induction variable. LoopOp only has one region, so 0 is the only valid
     /// value for `index`.
-    OperandRange getSuccessorEntryOperands(unsigned index);
+    OperandRange getSuccessorEntryOperands(Optional<unsigned> index);
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
   let hasRegionVerifier = 1;
 }
 
let description = [{
     ```
   }];
 
   let arguments = (ins Variadic<AnyType>:$inits);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$before, SizedRegion<1>:$after);
 
   let extraClassDeclaration = [{
-    OperandRange getSuccessorEntryOperands(unsigned index);
+    OperandRange getSuccessorEntryOperands(Optional<unsigned> index);
     ConditionOp getConditionOp();
     YieldOp getYieldOp();
     Block::BlockArgListType getBeforeArguments();
     Block::BlockArgListType getAfterArguments();
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let description = [{
     edges. By default, type equality is used.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Returns the operands of this operation used as the entry arguments when
         entering the region at `index`, which was specified as a successor of
-        this operation by `getSuccessorRegions`. These operands should
-        correspond 1-1 with the successor inputs specified in
+        this operation by `getSuccessorRegions`, or the operands forwarded to
+        the operation's results when it branches back to itself. These operands
+        should correspond 1-1 with the successor inputs specified in
         `getSuccessorRegions`.
       }],
       "::mlir::OperandRange", "getSuccessorEntryOperands",
-      (ins "unsigned":$index), [{}], /*defaultImplementation=*/[{
+      (ins "::llvm::Optional<unsigned>":$index), [{}], 
+      /*defaultImplementation=*/[{
         auto operandEnd = this->getOperation()->operand_end();
         return ::mlir::OperandRange(operandEnd, operandEnd);
       }]
     >,
     InterfaceMethod<[{
         Returns the viable successors of a region at `index`, or the possible
         successors when branching from the parent op if `index` is None. These
         are the regions that may be selected during the flow of control. If

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

auto getOperandIndexIfPred =
     return llvm::None;
   };
 
   // Check branches from the parent operation.
   Optional<unsigned> regionIndex;
   if (region) {
     // Determine the actual region number from the passed region.
     regionIndex = region->getRegionNumber();
-    if (Optional<unsigned> operandIndex =
-            getOperandIndexIfPred(/*predIndex=*/llvm::None)) {
-      collectUnderlyingAddressValues(
-          branch.getSuccessorEntryOperands(*regionIndex)[*operandIndex],
-          maxDepth, visited, output);
-    }
+  }
+        branch.getSuccessorEntryOperands(regionIndex)[*operandIndex], maxDepth,
+        visited, output);
   }
   // Check branches from each child region.
   Operation *op = branch.getOperation();
   for (int i = 0, e = op->getNumRegions(); i != e; ++i) {
     if (Optional<unsigned> operandIndex = getOperandIndexIfPred(i)) {
       for (Block &block : op->getRegion(i)) {
         Operation *term = block.getTerminator();
         // Try to determine possible region-branch successor operands for the

mlir/lib/Analysis/DataFlowAnalysis.cpp

void ForwardDataFlowSolver::visitRegionBranchOperation(
   // If the interface identified that no region will be executed. Mark
   // any results of this operation as overdefined, as we can't reason about
   // them.
   // TODO: If we had an interface to detect pass through operands, we could
   // resolve some results based on the lattice state of the operands. We could
   // also allow for the parent operation to have itself as a region successor.
   if (successors.empty())
     return markAllPessimisticFixpoint(branch, branch->getResults());
-  return visitRegionSuccessors(
-      branch, successors, operandLattices, [&](Optional<unsigned> index) {
-        assert(index && "expected valid region index");
-        return branch.getSuccessorEntryOperands(*index);
+  return visitRegionSuccessors(branch, successors, operandLattices,
+                               [&](Optional<unsigned> index) {
+                                 return branch.getSuccessorEntryOperands(index);
                                });
 }
 
 void ForwardDataFlowSolver::visitRegionSuccessors(
     Operation *parentOp, ArrayRef<RegionSuccessor> regionSuccessors,
     ArrayRef<AbstractLatticeElement *> operandLattices,
     function_ref<OperandRange(Optional<unsigned>)> getInputsForRegion) {
   for (const RegionSuccessor &it : regionSuccessors) {
     Region *region = it.getSuccessor();

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   results.add<AffineForEmptyLoopFolder>(context);
 }
 
 /// Return operands used when entering the region at 'index'. These operands
 /// correspond to the loop iterator operands, i.e., those excluding the
 /// induction variable. AffineForOp only has one region, so zero is the only
 /// valid value for `index`.
-OperandRange AffineForOp::getSuccessorEntryOperands(unsigned index) {
-  assert(index == 0 && "invalid region index");
+OperandRange AffineForOp::getSuccessorEntryOperands(Optional<unsigned> index) {
+  assert(!index || *index == 0 && "invalid region index");
 
   // The initial operands map to the loop arguments after the induction
-  // variable.
+  // variable or are forwarded to the results when the trip count is zero.
   return getIterOperands();
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.

mlir/lib/Dialect/Async/IR/Async.cpp

 }
 
 //===----------------------------------------------------------------------===//
 /// ExecuteOp
 //===----------------------------------------------------------------------===//
 
 constexpr char kOperandSegmentSizesAttr[] = "operand_segment_sizes";
 
-OperandRange ExecuteOp::getSuccessorEntryOperands(unsigned index) {
-  assert(index == 0 && "invalid region index");
+OperandRange ExecuteOp::getSuccessorEntryOperands(Optional<unsigned> index) {
+  assert(index && *index == 0 && "invalid region index");
   return operands();
 }
 
 bool ExecuteOp::areTypesCompatible(Type lhs, Type rhs) {
   const auto getValueOrTokenType = [](Type type) {
     if (auto value = type.dyn_cast<ValueType>())
       return value.getValueType();
     return type;

mlir/lib/Dialect/SCF/SCF.cpp

ForOp mlir::scf::getForInductionVarOwner(Value val) {
   auto *containingOp = ivArg.getOwner()->getParentOp();
   return dyn_cast_or_null<ForOp>(containingOp);
 }
 
 /// Return operands used when entering the region at 'index'. These operands
 /// correspond to the loop iterator operands, i.e., those excluding the
 /// induction variable. LoopOp only has one region, so 0 is the only valid value
 /// for `index`.
-OperandRange ForOp::getSuccessorEntryOperands(unsigned index) {
-  assert(index == 0 && "invalid region index");
+OperandRange ForOp::getSuccessorEntryOperands(Optional<unsigned> index) {
+  assert(index && *index == 0 && "invalid region index");
 
   // The initial operands map to the loop arguments after the induction
   // variable.
   return getInitArgs();
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
return emitOpError() << "needs to have type " << reduceType
                          << " (the type of the enclosing ReduceOp)";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // WhileOp
 //===----------------------------------------------------------------------===//
 
-OperandRange WhileOp::getSuccessorEntryOperands(unsigned index) {
-  assert(index == 0 &&
+OperandRange WhileOp::getSuccessorEntryOperands(Optional<unsigned> index) {
+  assert(index && *index == 0 &&
          "WhileOp is expected to branch only to the first region");
 
   return getInits();
 }
 
 ConditionOp WhileOp::getConditionOp() {
   return cast<ConditionOp>(getBefore().front().getTerminator());
 }

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

for (Operation &op : *getBodyBlock()) {
 
     SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
     iface.getEffectsOnValue(getBodyBlock()->getArgument(0), nestedEffects);
     for (const auto &effect : nestedEffects)
       effects.emplace_back(effect.getEffect(), getRoot(), effect.getResource());
   }
 }
 
-OperandRange transform::SequenceOp::getSuccessorEntryOperands(unsigned index) {
-  assert(index == 0 && "unexpected region index");
+OperandRange
+transform::SequenceOp::getSuccessorEntryOperands(Optional<unsigned> index) {
+  assert(index && *index == 0 && "unexpected region index");
   if (getOperation()->getNumOperands() == 1)
     return getOperation()->getOperands();
   return OperandRange(getOperation()->operand_end(),
                       getOperation()->operand_end());
 }
 
 void transform::SequenceOp::getSuccessorRegions(
     Optional<unsigned> index, ArrayRef<Attribute> operands,

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

 //===- ControlFlowInterfaces.cpp - ControlFlow Interfaces -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include <utility>
 
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ControlFlowInterfaces
 //===----------------------------------------------------------------------===//
verifyTypesAlongAllEdges(Operation *op, Optional<unsigned> sourceNo,
   return success();
 }
 
 /// Verify that types match along control flow edges described the given op.
 LogicalResult detail::verifyTypesAlongControlFlowEdges(Operation *op) {
   auto regionInterface = cast<RegionBranchOpInterface>(op);
 
   auto inputTypesFromParent = [&](Optional<unsigned> regionNo) -> TypeRange {
-    if (regionNo.hasValue()) {
-      return regionInterface.getSuccessorEntryOperands(regionNo.getValue())
-          .getTypes();
-    }
-
-    // If the successor of a parent op is the parent itself
-    // RegionBranchOpInterface does not have an API to query what the entry
-    // operands will be in that case. Vend out the result types of the op in
-    // that case so that type checking succeeds for this case.
-    return op->getResultTypes();
+    return regionInterface.getSuccessorEntryOperands(regionNo).getTypes();
   };
 
   // Verify types along control flow edges originating from the parent.
   if (failed(verifyTypesAlongAllEdges(op, llvm::None, inputTypesFromParent)))
     return failure();
 
   // RegionBranchOpInterface should not be implemented by Ops that do not have
   // attached regions.

mlir/test/Analysis/test-alias-analysis.mlir

func.func @region_loop_control_flow(%arg: memref<2xf32>, %loopI0 : index,
   %result = scf.for %i0 = %loopI0 to %loopI1 step %loopI2 iter_args(%si = %0) -> (memref<8x64xf32>) {
     scf.yield %si : memref<8x64xf32>
   } {test.ptr = "for_alloca"}
   return
 }
 
 // -----
 
+// CHECK-LABEL: Testing : "region_loop_zero_trip_count"
+// CHECK-DAG: alloca_1#0 <-> alloca_2#0: NoAlias
+// CHECK-DAG: alloca_1#0 <-> for_alloca#0: MustAlias
+// CHECK-DAG: alloca_1#0 <-> for_alloca.region0#0: MayAlias
+// CHECK-DAG: alloca_1#0 <-> for_alloca.region0#1: MayAlias
+
+// CHECK-DAG: alloca_2#0 <-> for_alloca#0: NoAlias
+// CHECK-DAG: alloca_2#0 <-> for_alloca.region0#0: MayAlias
+// CHECK-DAG: alloca_2#0 <-> for_alloca.region0#1: MayAlias
+
+// CHECK-DAG: for_alloca#0 <-> for_alloca.region0#0: MayAlias
+// CHECK-DAG: for_alloca#0 <-> for_alloca.region0#1: MayAlias
+
+// CHECK-DAG: for_alloca.region0#0 <-> for_alloca.region0#1: MayAlias
+func.func @region_loop_zero_trip_count() attributes {test.ptr = "func"} {
+  %0 = memref.alloca() {test.ptr = "alloca_1"} : memref<i32>
+  %1 = memref.alloca() {test.ptr = "alloca_2"} : memref<i32>
+  %result = affine.for %i = 0 to 0 iter_args(%si = %0) -> (memref<i32>) {
+    affine.yield %si : memref<i32>
+  } {test.ptr = "for_alloca"}
+  return
+}
+
+// -----
+
 // CHECK-LABEL: Testing : "view_like"
 // CHECK-DAG: alloc_1#0 <-> view#0: NoAlias
 
 // CHECK-DAG: alloca_1#0 <-> view#0: MustAlias
 
 // CHECK-DAG: view#0 <-> func.region0#0: NoAlias
 // CHECK-DAG: view#0 <-> func.region0#1: NoAlias
 func.func @view_like(%arg: memref<2xf32>, %size: index) attributes {test.ptr = "func"} {

mlir/test/Transforms/sccp-structured.mlir

func.func @loop_region_branch_terminator_op(%arg1 : i32) {
   }
   return
 }
 
 /// Check that propgation happens for affine.for -- tests its region branch op
 /// interface as well.
 
 // CHECK-LABEL: func @affine_loop_one_iter(
-func.func @affine_loop_one_iter(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
+func.func @affine_loop_one_iter() -> i32 {
   // CHECK: %[[C1:.*]] = arith.constant 1 : i32
   %s0 = arith.constant 0 : i32
   %s1 = arith.constant 1 : i32
   %result = affine.for %i = 0 to 1 iter_args(%si = %s0) -> (i32) {
     %sn = arith.addi %si, %s1 : i32
     affine.yield %sn : i32
   }
   // CHECK: return %[[C1]] : i32
   return %result : i32
 }
 
 // CHECK-LABEL: func @affine_loop_zero_iter(
-func.func @affine_loop_zero_iter(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
-  // This exposes a crash in sccp/forward data flow analysis: https://github.com/llvm/llvm-project/issues/54928
+func.func @affine_loop_zero_iter() -> i32 {
+  // CHECK: %[[C1:.*]] = arith.constant 1 : i32
+  %s1 = arith.constant 1 : i32
+  %result = affine.for %i = 0 to 0 iter_args(%si = %s1) -> (i32) {
+   %sn = arith.addi %si, %si : i32
+   affine.yield %sn : i32
+  }
+  // CHECK: return %[[C1]] : i32
+  return %result : i32
+}
+
+// CHECK-LABEL: func @affine_loop_unknown_trip_count(
+func.func @affine_loop_unknown_trip_count(%ub: index) -> i32 {
   // CHECK: %[[C0:.*]] = arith.constant 0 : i32
   %s0 = arith.constant 0 : i32
-  // %result = affine.for %i = 0 to 0 iter_args(%si = %s0) -> (i32) {
-  //  %sn = arith.addi %si, %si : i32
-  //  affine.yield %sn : i32
-  // }
-  // return %result : i32
+  %result = affine.for %i = 0 to %ub iter_args(%si = %s0) -> (i32) {
+   %sn = arith.addi %si, %si : i32
+   affine.yield %sn : i32
+  }
   // CHECK: return %[[C0]] : i32
-  return %s0 : i32
+  return %result : i32
 }
 
 // CHECK-LABEL: func @while_loop_different_arg_count
 func.func @while_loop_different_arg_count() -> index {
   // CHECK-DAG: %[[TRUE:.*]] = arith.constant true
   // CHECK-DAG: %[[C0:.*]] = arith.constant 0
   // CHECK-DAG: %[[C1:.*]] = arith.constant 1
   %c0 = arith.constant 0 : index

mlir/test/lib/Dialect/Test/TestDialect.cpp

if (parser.parseKeyword("then") || parser.parseRegion(*thenRegion, {}, {}) ||
       parser.parseKeyword("else") || parser.parseRegion(*elseRegion, {}, {}) ||
       parser.parseKeyword("join") || parser.parseRegion(*joinRegion, {}, {}))
     return failure();
 
   return parser.resolveOperands(operandInfos, operandTypes,
                                 parser.getCurrentLocation(), result.operands);
 }
 
-OperandRange RegionIfOp::getSuccessorEntryOperands(unsigned index) {
-  assert(index < 2 && "invalid region index");
+OperandRange RegionIfOp::getSuccessorEntryOperands(Optional<unsigned> index) {
+  assert(index && *index < 2 && "invalid region index");
   return getOperands();
 }
 
 void RegionIfOp::getSuccessorRegions(
     Optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {
   // We always branch to the join region.
   if (index.hasValue()) {

 // AnyCondOp
 //===----------------------------------------------------------------------===//
 
 void AnyCondOp::getSuccessorRegions(Optional<unsigned> index,
                                     ArrayRef<Attribute> operands,
                                     SmallVectorImpl<RegionSuccessor> &regions) {
   // The parent op branches into the only region, and the region branches back
   // to the parent op.
-  if (index)
+  if (!index)
     regions.emplace_back(&getRegion());
   else
     regions.emplace_back(getResults());
 }
 
 void AnyCondOp::getRegionInvocationBounds(
     ArrayRef<Attribute> operands,
     SmallVectorImpl<InvocationBounds> &invocationBounds) {

mlir/test/lib/Dialect/Test/TestOps.td

::mlir::Block::BlockArgListType getThenArgs() {
       return getBody(0)->getArguments();
     }
     ::mlir::Block::BlockArgListType getElseArgs() {
       return getBody(1)->getArguments();
     }
     ::mlir::Block::BlockArgListType getJoinArgs() {
       return getBody(2)->getArguments();
     }
-    ::mlir::OperandRange getSuccessorEntryOperands(unsigned index);
+    ::mlir::OperandRange getSuccessorEntryOperands(
+        ::llvm::Optional<unsigned> index);
   }];
   let hasCustomAssemblyFormat = 1;
 }
 
 def AnyCondOp : TEST_Op<"any_cond",
       [DeclareOpInterfaceMethods<RegionBranchOpInterface,
                                  ["getRegionInvocationBounds"]>,
        RecursiveSideEffects]> {