[mlir] Fix handling of some region branch terminator successors

[Mogball]

Summary

When RegionBranchOpInterface::getSuccessorRegions is called for anything other than the parent op, it expects the operands of the terminator of the source region to be passed, not the operands of the parent op. This was not always respected.

This fixes a bug in integer range inference and ForwardDataFlowSolver and changes scf.while to allow narrowing of successors using constant inputs.

Fixes #55873

[mehdi_amini]

LG but please wait for one another review here, I only skim through it.

[krzysz00]

I think the semantics of DataFlowAnalysis::getSuccessorsForOperands() could be clarified ... but I'm not sure it matters, since there're plans to scrap that entire framework.

mlir/lib/Analysis/IntRangeAnalysis.cpp

224

Looking here and below, is there call for a helper to iterate over the block terminators directly?

[Mogball]

Yeah someone actually reported issues internally to me yesterday just as I was wrapping up this fix. The old framework will go away soon^TM.

mlir/lib/Analysis/IntRangeAnalysis.cpp

224

Don't think so..

[Mogball]

This comment was removed by Mogball.

[krzysz00]

So yeah, from the perspective of the integer range stuff, feel free to land
this.

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

let extraClassDeclaration = [{
     ConditionOp getConditionOp();
     YieldOp getYieldOp();
     Block::BlockArgListType getBeforeArguments();
     Block::BlockArgListType getAfterArguments();
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
-  let hasRegionVerifier = 1;
+  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 def YieldOp : SCF_Op<"yield", [NoSideEffect, ReturnLike, Terminator,
                                ParentOneOf<["ExecuteRegionOp, ForOp",

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

static_assert(!ConcreteOp::template hasTrait<OpTrait::ZeroRegions>(),
                   "expected operation to have non-zero regions");
     return detail::verifyTypesAlongControlFlowEdges($_op);
   }];
   let verifyWithRegions = 1;
 
   let extraClassDeclaration = [{
     /// Convenience helper in case none of the operands is known.
     void getSuccessorRegions(Optional<unsigned> index,
-                             SmallVectorImpl<RegionSuccessor> &regions) {
-       SmallVector<Attribute, 2> nullAttrs(getOperation()->getNumOperands());
-       getSuccessorRegions(index, nullAttrs, regions);
-    }
+                             SmallVectorImpl<RegionSuccessor> &regions);
 
     /// Return `true` if control flow originating from the given region may
     /// eventually branch back to the same region. (Maybe after passing through
     /// other regions.)
     bool isRepetitiveRegion(unsigned index);
   }];
 }
 

mlir/lib/Analysis/DataFlowAnalysis.cpp

if (op->getNumSuccessors() == 0) {
     if (isa<CallableOpInterface>(parentOp))
       return visitCallableTerminatorOperation(parentOp, op, operandLattices);
 
     // Otherwise, check to see if the parent tracks region control flow.
     auto regionInterface = dyn_cast<RegionBranchOpInterface>(parentOp);
     if (!regionInterface || !isBlockExecutable(parentOp->getBlock()))
       return;
 
-    // If the branch is a RegionBranchTerminatorOpInterface,
-    // construct the set of operand lattices as the set of non control-flow
-    // arguments of the parent and the values this op returns. This allows
-    // for the correct lattices to be passed to getSuccessorsForOperands()
-    // in cases such as scf.while.
-    ArrayRef<AbstractLatticeElement *> branchOpLattices = operandLattices;
-    SmallVector<AbstractLatticeElement *, 0> parentLattices;
-    if (auto regionTerminator =
-            dyn_cast<RegionBranchTerminatorOpInterface>(op)) {
-      parentLattices.reserve(regionInterface->getNumOperands());
-      for (Value parentOperand : regionInterface->getOperands()) {
-        AbstractLatticeElement *operandLattice =
-            analysis.lookupLatticeElement(parentOperand);
-        if (!operandLattice || operandLattice->isUninitialized())
-          return;
-        parentLattices.push_back(operandLattice);
-      }
-      unsigned regionNumber = parentRegion->getRegionNumber();
-      OperandRange iterArgs =
-          regionInterface.getSuccessorEntryOperands(regionNumber);
-      OperandRange terminatorArgs =
-          regionTerminator.getSuccessorOperands(regionNumber);
-      assert(iterArgs.size() == terminatorArgs.size() &&
-             "Number of iteration arguments for region should equal number of "
-             "those arguments defined by terminator");
-      if (!iterArgs.empty()) {
-        unsigned iterStart = iterArgs.getBeginOperandIndex();
-        unsigned terminatorStart = terminatorArgs.getBeginOperandIndex();
-        for (unsigned i = 0, e = iterArgs.size(); i < e; ++i)
-          parentLattices[iterStart + i] = operandLattices[terminatorStart + i];
-      }
-      branchOpLattices = parentLattices;
-    }
     // Query the set of successors of the current region using the current
     // optimistic lattice state.
     SmallVector<RegionSuccessor, 1> regionSuccessors;
     analysis.getSuccessorsForOperands(regionInterface,
                                       parentRegion->getRegionNumber(),
-                                      branchOpLattices, regionSuccessors);
+                                      operandLattices, regionSuccessors);
     if (regionSuccessors.empty())
       return;
 
     // Try to get "region-like" successor operands if possible in order to
     // propagate the operand states to the successors.
     if (isRegionReturnLike(op)) {
       auto getOperands = [&](Optional<unsigned> regionIndex) {
-        // Determine the individual region  successor operands for the given
+        // Determine the individual region successor operands for the given
         // region index (if any).
         return *getRegionBranchSuccessorOperands(op, regionIndex);
       };
-      return visitRegionSuccessors(parentOp, regionSuccessors, branchOpLattices,
+      return visitRegionSuccessors(parentOp, regionSuccessors, operandLattices,
                                    getOperands);
     }
 
     // If this terminator is not "region-like", conservatively mark all of the
     // successor values as having reached the pessimistic fixpoint.
     for (auto &it : regionSuccessors) {
       // If the successor is a region, mark the entry block as executable so
       // that we visit operations defined within. If the successor is the

mlir/lib/Analysis/IntRangeAnalysis.cpp

LogicalResult detail::IntRangeAnalysisImpl::getSuccessorsForOperands(
   }
   return failure();
 }
 
 void detail::IntRangeAnalysisImpl::getSuccessorsForOperands(
     RegionBranchOpInterface branch, Optional<unsigned> sourceIndex,
     ArrayRef<LatticeElement<IntRangeLattice> *> operands,
     SmallVectorImpl<RegionSuccessor> &successors) {
-  auto toConstantAttr = [&branch](auto enumPair) -> Attribute {
-    Optional<APInt> maybeConstValue =
-        enumPair.value()->getValue().value.getConstantValue();
+  // Get a type with which to construct a constant.
+  auto getOperandType = [branch, sourceIndex](unsigned index) {
+    // The types of all return-like operations are the same.
+    if (!sourceIndex)
+      return branch->getOperand(index).getType();
+
+    for (Block &block : branch->getRegion(*sourceIndex)) {
+      Operation *terminator = block.getTerminator();
+      if (getRegionBranchSuccessorOperands(terminator, *sourceIndex))
+        return terminator->getOperand(index).getType();
+    }
+    return Type();
+  };
 
-    if (maybeConstValue) {
-      return IntegerAttr::get(branch->getOperand(enumPair.index()).getType(),
+  auto toConstantAttr = [&getOperandType](auto enumPair) -> Attribute {
+    if (Optional<APInt> maybeConstValue =
+            enumPair.value()->getValue().value.getConstantValue()) {
+      return IntegerAttr::get(getOperandType(enumPair.index()),
                               *maybeConstValue);
     }
     return {};
   };
   SmallVector<Attribute> inferredConsts(
       llvm::map_range(llvm::enumerate(operands), toConstantAttr));
   branch.getSuccessorRegions(sourceIndex, inferredConsts, successors);
 }

mlir/lib/Dialect/SCF/SCF.cpp

 
 Block::BlockArgListType WhileOp::getAfterArguments() {
   return getAfter().front().getArguments();
 }
 
 void WhileOp::getSuccessorRegions(Optional<unsigned> index,
                                   ArrayRef<Attribute> operands,
                                   SmallVectorImpl<RegionSuccessor> &regions) {
-  (void)operands;
-
+  // The parent op always branches to the condition region.
   if (!index.hasValue()) {
     regions.emplace_back(&getBefore(), getBefore().getArguments());
     return;
   }
 
   assert(*index < 2 && "there are only two regions in a WhileOp");
-  if (*index == 0) {
-    regions.emplace_back(&getAfter(), getAfter().getArguments());
-    regions.emplace_back(getResults());
+  // The body region always branches back to the condition region.
+  if (*index == 1) {
+    regions.emplace_back(&getBefore(), getBefore().getArguments());
     return;
   }
 
-  regions.emplace_back(&getBefore(), getBefore().getArguments());
+  // Try to narrow the successor to the condition region.
+  assert(!operands.empty() && "expected at least one operand");
+  auto cond = operands[0].dyn_cast_or_null<BoolAttr>();
+  if (!cond || !cond.getValue())
+    regions.emplace_back(getResults());
+  if (!cond || cond.getValue())
+    regions.emplace_back(&getAfter(), getAfter().getArguments());
 }
 
 /// Parses a `while` op.
 ///
 /// op ::= `scf.while` assignments `:` function-type region `do` region
 ///         `attributes` attribute-dict
 /// initializer ::= /* empty */ | `(` assignment-list `)`
 /// assignment-list ::= assignment | assignment `,` assignment-list
if (auto yield = dyn_cast_or_null<TerminatorTy>(terminatorOperation))
     return yield;
 
   auto diag = op.emitOpError(errorMessage);
   if (terminatorOperation)
     diag.attachNote(terminatorOperation->getLoc()) << "terminator here";
   return nullptr;
 }
 
-LogicalResult scf::WhileOp::verifyRegions() {
+LogicalResult scf::WhileOp::verify() {
   auto beforeTerminator = verifyAndGetTerminator<scf::ConditionOp>(
       *this, getBefore(),
       "expects the 'before' region to terminate with 'scf.condition'");
   if (!beforeTerminator)
     return failure();
 
   auto afterTerminator = verifyAndGetTerminator<scf::YieldOp>(
       *this, getAfter(),

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/Interfaces/ControlFlowInterfaces.h"
-#include "mlir/IR/BuiltinTypes.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ControlFlowInterfaces
 //===----------------------------------------------------------------------===//
 

 /// the match itself).
 static LogicalResult
 verifyTypesAlongAllEdges(Operation *op, Optional<unsigned> sourceNo,
                          function_ref<Optional<TypeRange>(Optional<unsigned>)>
                              getInputsTypesForRegion) {
   auto regionInterface = cast<RegionBranchOpInterface>(op);
 
   SmallVector<RegionSuccessor, 2> successors;
-  unsigned numInputs;
-  if (sourceNo) {
-    Region &srcRegion = op->getRegion(sourceNo.getValue());
-    numInputs = srcRegion.getNumArguments();
-  } else {
-    numInputs = op->getNumOperands();
-  }
-  SmallVector<Attribute, 2> operands(numInputs, nullptr);
-  regionInterface.getSuccessorRegions(sourceNo, operands, successors);
+  regionInterface.getSuccessorRegions(sourceNo, successors);
 
   for (RegionSuccessor &succ : successors) {
     Optional<unsigned> succRegionNo;
     if (!succ.isParent())
       succRegionNo = succ.getSuccessor()->getRegionNumber();
 
     auto printEdgeName = [&](InFlightDiagnostic &diag) -> InFlightDiagnostic & {
       diag << "from ";
bool mlir::insideMutuallyExclusiveRegions(Operation *a, Operation *b) {
   return false;
 }
 
 bool RegionBranchOpInterface::isRepetitiveRegion(unsigned index) {
   Region *region = &getOperation()->getRegion(index);
   return isRegionReachable(region, region);
 }
 
+void RegionBranchOpInterface::getSuccessorRegions(
+    Optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+  unsigned numInputs = 0;
+  if (index) {
+    // If the predecessor is a region, get the number of operands from an
+    // exiting terminator in the region.
+    for (Block &block : getOperation()->getRegion(*index)) {
+      Operation *terminator = block.getTerminator();
+      if (getRegionBranchSuccessorOperands(terminator, *index)) {
+        numInputs = terminator->getNumOperands();
+        break;
+      }
+    }
+  } else {
+    // Otherwise, use the number of parent operation operands.
+    numInputs = getOperation()->getNumOperands();
+  }
+  SmallVector<Attribute, 2> operands(numInputs, nullptr);
+  getSuccessorRegions(index, operands, regions);
+}
+
 Region *mlir::getEnclosingRepetitiveRegion(Operation *op) {
   while (Region *region = op->getParentRegion()) {
     op = region->getParentOp();
     if (auto branchOp = dyn_cast<RegionBranchOpInterface>(op))
       if (branchOp.isRepetitiveRegion(region->getRegionNumber()))
         return region;
   }
   return nullptr;

mlir/test/Dialect/SCF/invalid.mlir

func.func @while_cross_region_type_mismatch() {
   }
 }
 
 // -----
 
 func.func @while_cross_region_type_mismatch() {
   %true = arith.constant true
   // expected-error@+1 {{'scf.while' op  along control flow edge from Region #0 to Region #1: source type #0 'i1' should match input type #0 'i32'}}
-  scf.while : () -> () {
+  %0 = scf.while : () -> (i1) {
     scf.condition(%true) %true : i1
   } do {
   ^bb0(%arg0: i32):
     scf.yield
   }
 }
 
 // -----

mlir/test/Interfaces/InferIntRangeInterface/infer-int-range-test-ops.mlir

 }
 
 // CHECK-LABEL: func @func_args_unbound
 // CHECK: test.reflect_bounds {smax = 9223372036854775807 : index, smin = -9223372036854775808 : index, umax = -1 : index, umin = 0 : index}
 func.func @func_args_unbound(%arg0 : index) -> index {
   %0 = test.reflect_bounds %arg0
   func.return %0 : index
 }
+
+// CHECK-LABEL: func @propagate_across_while_loop()
+func.func @propagate_across_while_loop() -> index {
+  // CHECK-DAG: %[[C0:.*]] = "test.constant"() {value = 0
+  // CHECK-DAG: %[[C1:.*]] = "test.constant"() {value = 1
+  %0 = test.with_bounds { umin = 0 : index, umax = 0 : index,
+                          smin = 0 : index, smax = 0 : index }
+  %1 = scf.while : () -> index {
+    %true = arith.constant true
+    // CHECK: scf.condition(%{{.*}}) %[[C0]]
+    scf.condition(%true) %0 : index
+  } do {
+  ^bb0(%i1: index):
+    scf.yield
+  }
+  // CHECK: return %[[C1]]
+  %2 = test.increment %1
+  return %2 : index
+}

mlir/test/Transforms/sccp-structured.mlir

func.func @affine_loop_zero_iter(%arg0 : index, %arg1 : index, %arg2 : index) -> 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
   // CHECK: return %[[C0]] : i32
   return %s0 : 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
+  %c1 = arith.constant 1 : index
+  // CHECK: %[[WHILE:.*]] = scf.while
+  %0 = scf.while (%arg3 = %c0, %arg4 = %c1) : (index, index) -> index {
+    %1 = arith.cmpi slt, %arg3, %c1 : index
+    // CHECK: scf.condition(%[[TRUE]]) %[[C1]]
+    scf.condition(%1) %arg4 : index
+  } do {
+  ^bb0(%arg3: index):
+    %1 = arith.muli %arg3, %c1 : index
+    // CHECK: scf.yield %[[C0]], %[[C1]]
+    scf.yield %c0, %1 : index, index
+  }
+  // CHECK: return %[[WHILE]]
+  return %0 : index
+}
+
+// CHECK-LABEL: func @while_loop_false_condition
+func.func @while_loop_false_condition(%arg0 : index) -> index {
+  // CHECK: %[[C0:.*]] = arith.constant 0
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = arith.muli %arg0, %c0 : index
+  %1 = scf.while (%arg1 = %0) : (index) -> index {
+    %2 = arith.cmpi slt, %arg1, %c0 : index
+    scf.condition(%2) %arg1 : index
+  } do {
+  ^bb0(%arg2 : index):
+    %3 = arith.addi %arg2, %c1 : index
+    scf.yield %3 : index
+  }
+  // CHECK: return %[[C0]]
+  func.return %1 : index
+}