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

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.

Differential Revision: https://reviews.llvm.org/D78447

Author: River707

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

@@ -17,6 +17,7 @@
 #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"
 

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

@@ -13,6 +13,7 @@
 #ifndef LOOP_OPS
 #define LOOP_OPS
 
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffects.td"
 
@@ -37,6 +38,7 @@ class Loop_Op<string mnemonic, list<OpTrait> traits = []> :
 
 def ForOp : Loop_Op<"for",
       [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+       DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";
@@ -169,11 +171,18 @@ def ForOp : Loop_Op<"for",
     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
@@ -385,7 +394,7 @@ def ReduceReturnOp :
   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

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

@@ -1865,7 +1865,7 @@ def RemFOp : FloatArithmeticOp<"remf"> {
 // 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 = [{

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

@@ -19,6 +19,10 @@
 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.
@@ -37,7 +41,69 @@ 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

@@ -90,4 +90,55 @@ def BranchOpInterface : OpInterface<"BranchOpInterface"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// 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

@@ -196,6 +196,39 @@ ForOp mlir::loop::getForInductionVarOwner(Value val) {
   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
 //===----------------------------------------------------------------------===//
@@ -298,6 +331,42 @@ static void print(OpAsmPrinter &p, IfOp op) {
   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
 //===----------------------------------------------------------------------===//