[mlir][SCCP] Add support for propagating constants across inter-region control flow

[rriddle]

Summary

This is possible by adding two new ControlFlowInterface additions:

• A new interface, RegionBranchOpInterface

This interface allows for region holding operations to describe how control flows between regions. This interface initially contains two methods:

• getSuccessorEntryOperands

Returns the operands of this operation used as the entry arguments when entering the region at index, which was specified as a successor by getSuccessorRegions. when entering. These operands should correspond 1-1 with the successor inputs specified in getSuccessorRegions, and may be a subset of the entry arguments for that region.

• getSuccessorRegions

Returns the viable successors of a region, or the possible successor when branching from the parent op. This allows for describing which regions may be executed when entering an operation, and which regions are executed after having executed another region of the parent op. For example, a structured loop operation may always enter into the loop body region. The loop body region may branch back to itself, or exit to the operation.

• A trait, ReturnLike

This trait signals that a terminator exits a region and forwards all of its operands as "exiting" values.

These additions allow for performing more general dataflow analysis in the presence of region holding operations.

Depends On D78397

[bondhugula]

It's fantastic to have the same conditional constant propagation work for higher-order control flow ops and across nested regions! I think the revision title should be updated to better reflect that. "Region-based control flow" doesn't immediately imply that this works across nested regions / in and out of regions and not for just intra-region.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

94–99

ReturnLike trait ops should also produce zero results I believe - but its results are probably harmless / have no meaning since blocks with ReturnLike terminators won't properly dominate anything and so those results would be unused anyway. We aren't doing a zero result trait check on IsTerminator as well, but they are I guess allowed to have results (a way to send values to all successors for say branch like operations).

[rriddle]

Resolve comments

[bondhugula]

may be executed when entering an operation, and which regions are executed by "branching" from other regions

This could be confusing since one can't directly branch from another region (as in using a 'br'). Change "by branching from other regions" to "after having executed another region of the parent op". Also, this commit summary has more information than the ODS description of getSuccessorRegions in some ways. Consider making the ODS doc desc. more comprehensive.

The operands may correspond to a subset of the entry arguments to the region.

I think you meant the other way? The entry arguments to the region may correspond to a subset of the operands.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

93–94

Is this the right file to put this trait? Ops that don't implement the op interface above should have access to this trait? (Plus, see the comment on ReturnOp below.)

mlir/lib/Transforms/SCCP.cpp

526–528

You missed marking ReturnOp ReturnLike. As such, you will return early here and miss propagating lattice states to successor regions. Although there is no registered op that will exercise this, this could happen let's say in a special func-like op that has two regions (sort of two functions) and the semantics of the op were to execute the two regions one after the other.

[rriddle]

Resolve comments

[rriddle]

In D78447#1991402, @bondhugula wrote:

may be executed when entering an operation, and which regions are executed by "branching" from other regions

This could be confusing since one can't directly branch from another region (as in using a 'br'). Change "by branching from other regions" to "after having executed another region of the parent op". Also, this commit summary has more information than the ODS description of getSuccessorRegions in some ways. Consider making the ODS doc desc. more comprehensive.

The operands may correspond to a subset of the entry arguments to the region.

I think you meant the other way? The entry arguments to the region may correspond to a subset of the operands.

No, I mean the operands returned by this method may correspond to a subset of the entry arguments. This method shouldn't return more operands than there are entry block arguments.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

93–94

Yes this is the right place. OpDefinition.h is not the right place to put every possible trait, and is something that I've been purposely moving away from. It is better to keep the interfaces and traits that pertain to a specific abstraction in one place, e.g., the traits *and* interfaces for side effect modeling are in SideEffects.h. Realistically the files in this directory should remove the redundant "Interfaces" from the end of the name.

mlir/lib/Transforms/SCCP.cpp

526–528

Such an op isn't possible ATM given the constraints of ReturnOp. I intend to add interprocedural support here, at which point having it on ReturnOp would be useful. Also as a policy I like to avoid trying to mark every possible op that could take advantage of something. IMO it is better to separate adding a feature, and using it in every possible place. It keeps this patch more focused.

[rriddle]

mlir/lib/Transforms/SCCP.cpp

526–528

To clarify, I do intend to add this trait/interface to the things that can use it in a followup.

[bondhugula]

In D78447#1991521, @rriddle wrote:

In D78447#1991402, @bondhugula wrote:

may be executed when entering an operation, and which regions are executed by "branching" from other regions

This could be confusing since one can't directly branch from another region (as in using a 'br'). Change "by branching from other regions" to "after having executed another region of the parent op". Also, this commit summary has more information than the ODS description of getSuccessorRegions in some ways. Consider making the ODS doc desc. more comprehensive.

The operands may correspond to a subset of the entry arguments to the region.

I think you meant the other way? The entry arguments to the region may correspond to a subset of the operands.

No, I mean the operands returned by this method may correspond to a subset of the entry arguments. This method shouldn't return more operands than there are entry block arguments.

Oh, I was confused thinking about the parent op's operands. Consider changing:
"The operands may correspond to a subset of the entry arguments to the region." -> ".... entry arguments of the containing region"

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

179

I think this should be called getSuccessorRegionEntryOperands or getSuccRegionEntryOperands.

179

Since this op has exactly one region and so index can only be zero, update the doc comment to reflect that.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

74

Region *region or else it gives the impression that this is a successor to that successor region.

75

successorInputs -> inputs likewise. This is already a RegionSuccessor class.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

105

Spell out "used" here, i.e., "used as arguments" or "used as the entry block arguments".

107

typo: specified

121

Nit: that either correspond

mlir/lib/Dialect/LoopOps/LoopOps.cpp

343

index.hasValue() to avoid confusion. What if index has a value and is zero? I can't recall.

mlir/lib/Transforms/SCCP.cpp

526–528

It's just weird that the standard ReturnOp isn't marked ReturnLike. You could add the support to exploit it later, but no harm marking the more prevalent / core one? Yes, there isn't a need to add it in every other possible place now. OTOH, if ReturnLike doesn't work well for the ReturnOp and you need another trait, then this one should be renamed YieldLike and ReturnLike could be used for ReturnOp and for other declarative ops.

[rriddle]

Resolve comments

[rriddle]

Thanks for the review Uday! I appreciate the attention to detail.

mlir/lib/Transforms/SCCP.cpp

526–528

Agreed, went ahead and added it. Thanks.

[bondhugula]

mlir/lib/Dialect/LoopOps/LoopOps.cpp

357

What if there is no else block here? (empty else region)

362

Likewise, what if the else region has no blocks?

mlir/test/Transforms/sccp-structured.mlir

133

Every loop.if that you have in the test cases has an else. Add one or more without an else block?

[bondhugula]

mlir/lib/Transforms/SCCP.cpp

141

.. if they are in executable blocks.

473

I think this path is never tested.

mlir/test/Transforms/sccp-structured.mlir

133

Just realized that no values would be yielded without the else. And you are testing for empty regions at one place.

[rriddle]

mlir/lib/Dialect/LoopOps/LoopOps.cpp

357

I was originally undecided on how we should treat empty regions. For now, updated to disallow using empty regions as successors.

mlir/test/Transforms/sccp-structured.mlir

133

Do you have a suggestion on what you have in mind? An IfOp can only produce results if it has both a then and an else region. Given that IfOp is implicit capture, nothing would really be testing the structure of the IfOp that isn't tested elsewhere.

[rriddle]

mlir/test/Transforms/sccp-structured.mlir

133

Seems we crossed comments. Should no longer be necessary if we restrict that successors are non-empty.

[rriddle]

Resolve comments

[bondhugula]

mlir/test/Transforms/sccp-structured.mlir

133

Yes, nothing else to do on this for now.

[ftynse]

Looks like this broke gcc-5 build https://buildkite.com/mlir/mlir-core/builds/4336

mlir/include/mlir/Dialect/LoopOps/LoopOps.h

 
 #ifndef MLIR_LOOPOPS_OPS_H_
 #define MLIR_LOOPOPS_OPS_H_
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffects.h"
 
 namespace mlir {
 namespace loop {
 
 class TerminatorOp;
 

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

 //===- Ops.td - Loop operation definitions ---------------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Defines MLIR loop operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LOOP_OPS
 #define LOOP_OPS
 
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffects.td"
 
 def LoopOps_Dialect : Dialect {
   let name = "loop";
   let cppNamespace = "";
 }
 
 // Base class for Loop dialect ops.
 class Loop_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<LoopOps_Dialect, mnemonic, traits> {
   // For every standard op, there needs to be a:
   //   * void print(OpAsmPrinter &p, ${C++ class of Op} op)
   //   * LogicalResult verify(${C++ class of Op} op)
   //   * ParseResult parse${C++ class of Op}(OpAsmParser &parser,
   //                                         OperationState &result)
   // functions.
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
 def ForOp : Loop_Op<"for",
       [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+       DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";
   let description = [{
     The "loop.for" operation represents a loop taking 3 SSA value as operands
     that represent the lower bound, upper bound and step respectively.  The
     operation defines an SSA value for its induction variable. It has one
     region capturing the loop body. The induction variable is represented as an
let extraClassDeclaration = [{
     /// Does the operation hold operands for loop-carried values
     bool hasIterOperands() {
       return getOperation()->getNumOperands() > getNumControlOperands();
     }
     /// Get Number of loop-carried values
     unsigned getNumIterOperands() {
       return getOperation()->getNumOperands() - getNumControlOperands();
     }
+
+    /// 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);
   }];
 }
 
 def IfOp : Loop_Op<"if",
-      [SingleBlockImplicitTerminator<"YieldOp">, RecursiveSideEffects]> {
+      [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+       SingleBlockImplicitTerminator<"YieldOp">, RecursiveSideEffects]> {
   let summary = "if-then-else operation";
   let description = [{
     The `loop.if` operation represents an if-then-else construct for
     conditionally executing two regions of code. The operand to an if operation
     is a boolean value. For example:
 
     ```mlir
     loop.if %b  {
let description = [{
     loop.reduce.return %res : f32
     ```
   }];
 
   let arguments = (ins AnyType:$result);
   let assemblyFormat = "$result attr-dict `:` type($result)";
 }
 
-def YieldOp : Loop_Op<"yield", [NoSideEffect, Terminator]> {
+def YieldOp : Loop_Op<"yield", [NoSideEffect, ReturnLike, Terminator]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "loop.yield" yields an SSA value from a loop dialect op region and
     terminates the regions. The semantics of how the values are yielded is
     defined by the parent operation.
     If "loop.yield" has any operands, the operands must match the parent
     operation's results.
     If the parent operation defines no values, then the "loop.yield" may be

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 def RemFOp : FloatArithmeticOp<"remf"> {
   let summary = "floating point division remainder operation";
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
-def ReturnOp : Std_Op<"return", [NoSideEffect, HasParent<"FuncOp">,
+def ReturnOp : Std_Op<"return", [NoSideEffect, HasParent<"FuncOp">, ReturnLike,
                                  Terminator]> {
   let summary = "return operation";
   let description = [{
     The `return` operation represents a return operation within a function.
     The operation takes variable number of operands and produces no results.
     The operand number and types must match the signature of the function
     that contains the operation.
 

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

 #ifndef MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
 #define MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
 
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
 class BranchOpInterface;
 
+//===----------------------------------------------------------------------===//
+// BranchOpInterface
+//===----------------------------------------------------------------------===//
+
 namespace detail {
 /// Erase an operand from a branch operation that is used as a successor
 /// operand. `operandIndex` is the operand within `operands` to be erased.
 void eraseBranchSuccessorOperand(OperandRange operands, unsigned operandIndex,
                                  Operation *op);
 
 /// Return the `BlockArgument` corresponding to operand `operandIndex` in some
 /// successor if `operandIndex` is within the range of `operands`, or None if
 /// `operandIndex` isn't a successor operand index.
 Optional<BlockArgument>
 getBranchSuccessorArgument(Optional<OperandRange> operands,
                            unsigned operandIndex, Block *successor);
 
 /// Verify that the given operands match those of the given successor block.
 LogicalResult verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
                                             Optional<OperandRange> operands);
 } // namespace detail
 
+//===----------------------------------------------------------------------===//
+// RegionBranchOpInterface
+//===----------------------------------------------------------------------===//
+
+/// This class represents a successor of a region. A region successor can either
+/// be another region, or the parent operation. If the successor is a region,
+/// this class accepts the destination region, as well as a set of arguments
+/// from that region that will be populated by values from the current region.
+/// If the successor is the parent operation, this class accepts an optional set
+/// of results that will be populated by values from the current region.
+class RegionSuccessor {
+public:
+  /// Initialize a successor that branches to another region of the parent
+  /// operation.
+  RegionSuccessor(Region *region, Block::BlockArgListType regionInputs = {})
+      : region(region), inputs(regionInputs) {}
+  /// Initialize a successor that branches back to/out of the parent operation.
+  RegionSuccessor(Optional<Operation::result_range> results = {})
+      : region(nullptr), inputs(results ? ValueRange(*results) : ValueRange()) {
+  }
+
+  /// Return the given region successor. Returns nullptr if the successor is the
+  /// parent operation.
+  Region *getSuccessor() const { return region; }
+
+  /// Return the inputs to the successor that are remapped by the exit values of
+  /// the current region.
+  ValueRange getSuccessorInputs() const { return inputs; }
+
+private:
+  Region *region;
+  ValueRange inputs;
+};
+
+//===----------------------------------------------------------------------===//
+// ControlFlow Interfaces
+//===----------------------------------------------------------------------===//
+
 #include "mlir/Interfaces/ControlFlowInterfaces.h.inc"
+
+//===----------------------------------------------------------------------===//
+// ControlFlow Traits
+//===----------------------------------------------------------------------===//
+
+namespace OpTrait {
+/// 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.
+template <typename ConcreteType>
+struct ReturnLike : public TraitBase<ConcreteType, ReturnLike> {
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(ConcreteType::template hasTrait<IsTerminator>(),
+                  "expected operation to be a terminator");
+    static_assert(ConcreteType::template hasTrait<ZeroResult>(),
+                  "expected operation to have zero results");
+    static_assert(ConcreteType::template hasTrait<ZeroSuccessor>(),
+                  "expected operation to have zero successors");
+    return success();
+  }
+};
+} // namespace OpTrait
+
 } // end namespace mlir
 
 #endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES_H

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

for (unsigned i = 0, e = $_op->getNumSuccessors(); i != e; ++i) {
       Optional<OperandRange> operands = concreteOp.getSuccessorOperands(i);
       if (failed(detail::verifyBranchSuccessorOperands($_op, i, operands)))
         return failure();
     }
     return success();
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// RegionBranchOpInterface
+//===----------------------------------------------------------------------===//
+
+def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
+  let description = [{
+    This interface provides information for region operations that contain
+    branching behavior between held regions, i.e. this interface allows for
+    expressing control flow information for region holding operations.
+  }];
+  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 by
+        `getSuccessorRegions`. These operands should correspond 1-1 with the
+        successor inputs specified in `getSuccessorRegions`, and may corre
+      }],
+      "OperandRange", "getSuccessorEntryOperands",
+      (ins "unsigned":$index), [{}], /*defaultImplementation=*/[{
+        auto operandEnd = this->getOperation()->operand_end();
+        return 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
+        `index` is None, `operands` is a set of optional attributes that
+        either correspond to a constant value for each operand of this
+        operation, or null if that operand is not a constant. If `index` is
+        valid, `operands` corresponds to the exit values of the region at
+        `index`. Only a region, i.e. a valid `index`, may use the parent
+        operation as a successor. This method allows for describing which
+        regions may be executed when entering an operation, and which regions
+        are executed after having executed another region of the parent op. The
+        successor region must be non-empty.
+      }],
+      "void", "getSuccessorRegions",
+      (ins "Optional<unsigned>":$index, "ArrayRef<Attribute>":$operands,
+           "SmallVectorImpl<RegionSuccessor> &":$regions)
+    >
+  ];
+}
+
+//===----------------------------------------------------------------------===//
+// ControlFlow Traits
+//===----------------------------------------------------------------------===//
+
+// Op is "return-like".
+def ReturnLike : NativeOpTrait<"ReturnLike">;
+
 #endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES

mlir/lib/Dialect/LoopOps/LoopOps.cpp

ForOp mlir::loop::getForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   if (!ivArg)
     return ForOp();
   assert(ivArg.getOwner() && "unlinked block argument");
   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 exclusing 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");
+
+  // The initial operands map to the loop arguments after the induction
+  // variable.
+  return initArgs();
+}
+
+/// 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.
+void ForOp::getSuccessorRegions(Optional<unsigned> index,
+                                ArrayRef<Attribute> operands,
+                                SmallVectorImpl<RegionSuccessor> &regions) {
+  // If the predecessor is the ForOp, branch into the body using the iterator
+  // arguments.
+  if (!index.hasValue()) {
+    regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+    return;
+  }
+
+  // Otherwise, the loop may branch back to itself or the parent operation.
+  assert(index.getValue() == 0 && "expected loop region");
+  regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+  regions.push_back(RegionSuccessor(getResults()));
+}
+
 //===----------------------------------------------------------------------===//
 // IfOp
 //===----------------------------------------------------------------------===//
 
 void IfOp::build(Builder *builder, OperationState &result, Value cond,
                  bool withElseRegion) {
   build(builder, result, /*resultTypes=*/llvm::None, cond, withElseRegion);
 }
if (!elseRegion.empty()) {
     p.printRegion(elseRegion,
                   /*printEntryBlockArgs=*/false,
                   /*printBlockTerminators=*/printBlockTerminators);
   }
 
   p.printOptionalAttrDict(op.getAttrs());
 }
 
+/// 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.
+void IfOp::getSuccessorRegions(Optional<unsigned> index,
+                               ArrayRef<Attribute> operands,
+                               SmallVectorImpl<RegionSuccessor> &regions) {
+  // The `then` and the `else` region branch back to the parent operation.
+  if (index.hasValue()) {
+    regions.push_back(RegionSuccessor(getResults()));
+    return;
+  }
+
+  // Don't consider the else region if it is empty.
+  Region *elseRegion = &this->elseRegion();
+  if (elseRegion->empty())
+    elseRegion = nullptr;
+
+  // Otherwise, the successor is dependent on the condition.
+  bool condition;
+  if (auto condAttr = operands.front().dyn_cast_or_null<IntegerAttr>()) {
+    condition = condAttr.getValue().isOneValue();
+  } else if (auto condAttr = operands.front().dyn_cast_or_null<BoolAttr>()) {
+    condition = condAttr.getValue();
+  } else {
+    // If the condition isn't constant, both regions may be executed.
+    regions.push_back(RegionSuccessor(&thenRegion()));
+    regions.push_back(RegionSuccessor(elseRegion));
+    return;
+  }
+
+  // Add the successor regions using the condition.
+  regions.push_back(RegionSuccessor(condition ? &thenRegion() : elseRegion));
+}
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(Builder *builder, OperationState &result, ValueRange lbs,
                        ValueRange ubs, ValueRange steps, ValueRange initVals) {
   result.addOperands(lbs);
   result.addOperands(ubs);

mlir/lib/Transforms/SCCP.cpp

 
 private:
   /// Replace the given value with a constant if the corresponding lattice
   /// represents a constant. Returns success if the value was replaced, failure
   /// otherwise.
   LogicalResult replaceWithConstant(OpBuilder &builder, OperationFolder &folder,
                                     Value value);
 
+  /// Visit the users of the given IR that reside within executable blocks.
+  template <typename T>
+  void visitUsers(T &value) {
+    for (Operation *user : value.getUsers())
+      if (isBlockExecutable(user->getBlock()))
+        visitOperation(user);
+  }
+
   /// Visit the given operation and compute any necessary lattice state.
   void visitOperation(Operation *op);
 
   /// Visit the given operation, which defines regions, and compute any
   /// necessary lattice state. This also resolves the lattice state of both the
   /// operation results and any nested regions.
-  void visitRegionOperation(Operation *op);
+  void visitRegionOperation(Operation *op,
+                            ArrayRef<Attribute> constantOperands);
+
+  /// Visit the given set of region successors, computing any necessary lattice
+  /// state. The provided function returns the input operands to the region at
+  /// the given index. If the index is 'None', the input operands correspond to
+  /// the parent operation results.
+  void visitRegionSuccessors(
+      Operation *parentOp, ArrayRef<RegionSuccessor> regionSuccessors,
+      function_ref<OperandRange(Optional<unsigned>)> getInputsForRegion);
 
   /// Visit the given terminator operation and compute any necessary lattice
   /// state.
   void visitTerminatorOperation(Operation *op,
                                 ArrayRef<Attribute> constantOperands);
 
   /// Visit the given block and compute any necessary lattice state.
   void visitBlock(Block *block);
void markAllOverdefined(ValuesT values) {
     for (auto value : values)
       markOverdefined(value);
   }
   template <typename ValuesT>
   void markAllOverdefined(Operation *op, ValuesT values) {
     markAllOverdefined(values);
     opWorklist.push_back(op);
   }
+  template <typename ValuesT>
+  void markAllOverdefinedAndVisitUsers(ValuesT values) {
+    for (auto value : values) {
+      auto &lattice = latticeValues[value];
+      if (!lattice.isOverdefined()) {
+        lattice.markOverdefined();
+        visitUsers(value);
+      }
+    }
+  }
 
   /// Returns true if the given value was marked as overdefined.
   bool isOverdefined(Value value) const;
 
   /// Merge in the given lattice 'from' into the lattice 'to'. 'owner'
   /// corresponds to the parent operation of 'to'.
   void meet(Operation *owner, LatticeValue &to, const LatticeValue &from);
 
for (Region &region : regions) {
     // as overdefined.
     markAllOverdefined(entryBlock->getArguments());
   }
 }
 
 void SCCPSolver::solve() {
   while (!blockWorklist.empty() || !opWorklist.empty()) {
     // Process any operations in the op worklist.
-    while (!opWorklist.empty()) {
-      Operation *op = opWorklist.pop_back_val();
-
-      // Visit all of the live users to propagate changes to this operation.
-      for (Operation *user : op->getUsers()) {
-        if (isBlockExecutable(user->getBlock()))
-          visitOperation(user);
-      }
-    }
+    while (!opWorklist.empty())
+      visitUsers(*opWorklist.pop_back_val());
 
     // Process any blocks in the block worklist.
     while (!blockWorklist.empty())
       visitBlock(blockWorklist.pop_back_val());
   }
 }
 
 void SCCPSolver::rewrite(MLIRContext *context,
void SCCPSolver::visitOperation(Operation *op) {
 
   // If this is a terminator operation, process any control flow lattice state.
   if (op->isKnownTerminator())
     visitTerminatorOperation(op, operandConstants);
 
   // Process region holding operations. The region visitor processes result
   // values, so we can exit afterwards.
   if (op->getNumRegions())
-    return visitRegionOperation(op);
+    return visitRegionOperation(op, operandConstants);
 
   // If this op produces no results, it can't produce any constants.
   if (op->getNumResults() == 0)
     return;
 
   // If all of the results of this operation are already overdefined, bail out
   // early.
   auto isOverdefinedFn = [&](Value value) { return isOverdefined(value); };
for (unsigned i = 0, e = foldResults.size(); i != e; ++i) {
     OpFoldResult foldResult = foldResults[i];
     if (Attribute foldAttr = foldResult.dyn_cast<Attribute>())
       meet(op, resultLattice, LatticeValue(foldAttr, opDialect));
     else
       meet(op, resultLattice, latticeValues[foldResult.get<Value>()]);
   }
 }
 
-void SCCPSolver::visitRegionOperation(Operation *op) {
+void SCCPSolver::visitRegionOperation(Operation *op,
+                                      ArrayRef<Attribute> constantOperands) {
+  // Check to see if we can reason about the internal control flow of this
+  // region operation.
+  auto regionInterface = dyn_cast<RegionBranchOpInterface>(op);
+  if (!regionInterface) {
+    // If we can't, conservatively mark all regions as executable.
     for (Region &region : op->getRegions()) {
       if (region.empty())
         continue;
       Block *entryBlock = &region.front();
       markBlockExecutable(entryBlock);
       markAllOverdefined(entryBlock->getArguments());
     }
 
     // Don't try to simulate the results of a region operation as we can't
-  // guarantee that folding will be out-of-place. We don't allow in-place folds
-  // as the desire here is for simulated execution, and not general folding.
+    // guarantee that folding will be out-of-place. We don't allow in-place
+    // folds as the desire here is for simulated execution, and not general
+    // folding.
     return markAllOverdefined(op, op->getResults());
   }
 
+  // Check to see which regions are executable.
+  SmallVector<RegionSuccessor, 1> successors;
+  regionInterface.getSuccessorRegions(/*index=*/llvm::None, constantOperands,
+                                      successors);
+
+  // 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 markAllOverdefined(op, op->getResults());
+  return visitRegionSuccessors(op, successors, [&](Optional<unsigned> index) {
+    assert(index && "expected valid region index");
+    return regionInterface.getSuccessorEntryOperands(*index);
+  });
+}
+
+void SCCPSolver::visitRegionSuccessors(
+    Operation *parentOp, ArrayRef<RegionSuccessor> regionSuccessors,
+    function_ref<OperandRange(Optional<unsigned>)> getInputsForRegion) {
+  for (const RegionSuccessor &it : regionSuccessors) {
+    Region *region = it.getSuccessor();
+    ValueRange succArgs = it.getSuccessorInputs();
+
+    // Check to see if this is the parent operation.
+    if (!region) {
+      ResultRange results = parentOp->getResults();
+      if (llvm::all_of(results, [&](Value res) { return isOverdefined(res); }))
+        continue;
+
+      // Mark the results outside of the input range as overdefined.
+      if (succArgs.size() != results.size()) {
+        opWorklist.push_back(parentOp);
+        if (succArgs.empty())
+          return markAllOverdefined(results);
+
+        unsigned firstResIdx = succArgs[0].cast<OpResult>().getResultNumber();
+        markAllOverdefined(results.take_front(firstResIdx));
+        markAllOverdefined(results.drop_front(firstResIdx + succArgs.size()));
+      }
+
+      // Update the lattice for any operation results.
+      OperandRange operands = getInputsForRegion(/*index=*/llvm::None);
+      for (auto it : llvm::zip(succArgs, operands))
+        meet(parentOp, latticeValues[std::get<0>(it)],
+             latticeValues[std::get<1>(it)]);
+      return;
+    }
+    assert(!region->empty() && "expected region to be non-empty");
+    Block *entryBlock = &region->front();
+    markBlockExecutable(entryBlock);
+
+    // If all of the arguments are already overdefined, the arguments have
+    // already been fully resolved.
+    auto arguments = entryBlock->getArguments();
+    if (llvm::all_of(arguments, [&](Value arg) { return isOverdefined(arg); }))
+      continue;
+
+    // Mark any arguments that do not receive inputs as overdefined, we won't be
+    // able to discern if they are constant.
+    if (succArgs.size() != arguments.size()) {
+      if (succArgs.empty()) {
+        markAllOverdefined(arguments);
+        continue;
+      }
+
+      unsigned firstArgIdx = succArgs[0].cast<BlockArgument>().getArgNumber();
+      markAllOverdefinedAndVisitUsers(arguments.take_front(firstArgIdx));
+      markAllOverdefinedAndVisitUsers(
+          arguments.drop_front(firstArgIdx + succArgs.size()));
+    }
+
+    // Update the lattice for arguments that have inputs from the predecessor.
+    OperandRange succOperands = getInputsForRegion(region->getRegionNumber());
+    for (auto it : llvm::zip(succArgs, succOperands)) {
+      LatticeValue &argLattice = latticeValues[std::get<0>(it)];
+      if (argLattice.meet(latticeValues[std::get<1>(it)]))
+        visitUsers(std::get<0>(it));
+    }
+  }
+}
+
 void SCCPSolver::visitTerminatorOperation(
     Operation *op, ArrayRef<Attribute> constantOperands) {
-  if (op->getNumSuccessors() == 0)
+  // If this operation has no successors, we treat it as an exiting terminator.
+  if (op->getNumSuccessors() == 0) {
+    // Check to see if the parent tracks region control flow.
+    Region *parentRegion = op->getParentRegion();
+    Operation *parentOp = parentRegion->getParentOp();
+    auto regionInterface = dyn_cast<RegionBranchOpInterface>(parentOp);
+    if (!regionInterface || !isBlockExecutable(parentOp->getBlock()))
       return;
 
+    // Query the set of successors from the current region.
+    SmallVector<RegionSuccessor, 1> regionSuccessors;
+    regionInterface.getSuccessorRegions(parentRegion->getRegionNumber(),
+                                        constantOperands, regionSuccessors);
+    if (regionSuccessors.empty())
+      return;
+
+    // If this terminator is not "region-like", conservatively mark all of the
+    // successor values as overdefined.
+    if (!op->hasTrait<OpTrait::ReturnLike>()) {
+      for (auto &it : regionSuccessors)
+        markAllOverdefinedAndVisitUsers(it.getSuccessorInputs());
+      return;
+    }
+
+    // Otherwise, propagate the operand lattice states to each of the
+    // successors.
+    OperandRange operands = op->getOperands();
+    return visitRegionSuccessors(parentOp, regionSuccessors,
+                                 [&](Optional<unsigned>) { return operands; });
+  }
+
   // Try to resolve to a specific successor with the constant operands.
   if (auto branch = dyn_cast<BranchOpInterface>(op)) {
     if (Block *singleSucc = branch.getSuccessorForOperands(constantOperands)) {
       markEdgeExecutable(op->getBlock(), singleSucc);
       return;
     }
   }
 
for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
 
     // Otherwise, meet the two lattice values.
     updatedLattice |= argLattice.meet(operandLattice->second);
     if (argLattice.isOverdefined())
       break;
   }
 
   // If the lattice was updated, visit any executable users of the argument.
-  if (updatedLattice) {
-    for (Operation *user : arg.getUsers())
-      if (isBlockExecutable(user->getBlock()))
-        visitOperation(user);
-  }
+  if (updatedLattice)
+    visitUsers(arg);
 }
 
 bool SCCPSolver::markBlockExecutable(Block *block) {
   bool marked = executableBlocks.insert(block).second;
   if (marked)
     blockWorklist.push_back(block);
   return marked;
 }

mlir/test/Transforms/sccp-structured.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="func(sccp)" -split-input-file | FileCheck %s
+
+/// Check that a constant is properly propagated when only one edge is taken.
+
+// CHECK-LABEL: func @simple(
+func @simple(%arg0 : i32) -> i32 {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK-NOT: loop.if
+  // CHECK: return %[[CST]] : i32
+
+  %cond = constant true
+  %res = loop.if %cond -> (i32) {
+    %1 = constant 1 : i32
+    loop.yield %1 : i32
+  } else {
+    loop.yield %arg0 : i32
+  }
+  return %res : i32
+}
+
+/// Check that a constant is properly propagated when both edges produce the
+/// same value.
+
+// CHECK-LABEL: func @simple_both_same(
+func @simple_both_same(%cond : i1) -> i32 {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK-NOT: loop.if
+  // CHECK: return %[[CST]] : i32
+
+  %res = loop.if %cond -> (i32) {
+    %1 = constant 1 : i32
+    loop.yield %1 : i32
+  } else {
+    %2 = constant 1 : i32
+    loop.yield %2 : i32
+  }
+  return %res : i32
+}
+
+/// Check that the arguments go to overdefined if the branch cannot detect when
+/// a specific successor is taken.
+
+// CHECK-LABEL: func @overdefined_unknown_condition(
+func @overdefined_unknown_condition(%cond : i1, %arg0 : i32) -> i32 {
+  // CHECK: %[[RES:.*]] = loop.if
+  // CHECK: return %[[RES]] : i32
+
+  %res = loop.if %cond -> (i32) {
+    %1 = constant 1 : i32
+    loop.yield %1 : i32
+  } else {
+    loop.yield %arg0 : i32
+  }
+  return %res : i32
+}
+
+/// Check that the arguments go to overdefined if there are conflicting
+/// constants.
+
+// CHECK-LABEL: func @overdefined_different_constants(
+func @overdefined_different_constants(%cond : i1) -> i32 {
+  // CHECK: %[[RES:.*]] = loop.if
+  // CHECK: return %[[RES]] : i32
+
+  %res = loop.if %cond -> (i32) {
+    %1 = constant 1 : i32
+    loop.yield %1 : i32
+  } else {
+    %2 = constant 2 : i32
+    loop.yield %2 : i32
+  }
+  return %res : i32
+}
+
+/// Check that arguments are properly merged across loop-like control flow.
+
+// CHECK-LABEL: func @simple_loop(
+func @simple_loop(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
+  // CHECK: %[[CST:.*]] = constant 0 : i32
+  // CHECK-NOT: loop.for
+  // CHECK: return %[[CST]] : i32
+
+  %s0 = constant 0 : i32
+  %result = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0) -> (i32) {
+    %sn = addi %si, %si : i32
+    loop.yield %sn : i32
+  }
+  return %result : i32
+}
+
+/// Check that arguments go to overdefined when loop backedges produce a
+/// conflicting value.
+
+// CHECK-LABEL: func @loop_overdefined(
+func @loop_overdefined(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
+  // CHECK: %[[RES:.*]] = loop.for
+  // CHECK: return %[[RES]] : i32
+
+  %s0 = constant 1 : i32
+  %result = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0) -> (i32) {
+    %sn = addi %si, %si : i32
+    loop.yield %sn : i32
+  }
+  return %result : i32
+}
+
+/// Test that we can properly propagate within inner control, and in situations
+/// where the executable edges within the CFG are sensitive to the current state
+/// of the analysis.
+
+// CHECK-LABEL: func @loop_inner_control_flow(
+func @loop_inner_control_flow(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK-NOT: loop.for
+  // CHECK-NOT: loop.if
+  // CHECK: return %[[CST]] : i32
+
+  %cst_1 = constant 1 : i32
+  %result = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %cst_1) -> (i32) {
+    %cst_20 = constant 20 : i32
+    %cond = cmpi "ult", %si, %cst_20 : i32
+    %inner_res = loop.if %cond -> (i32) {
+      %1 = constant 1 : i32
+      loop.yield %1 : i32
+    } else {
+      %si_inc = addi %si, %cst_1 : i32
+      loop.yield %si_inc : i32
+    }
+    loop.yield %inner_res : i32
+  }
+  return %result : i32
+}