Revert " [MLIR] Revamp RegionBranchOpInterface " (#165356)

Reverts #161575

Broke Windows on ARM buildbot build, needs investigations.

Author: joker-eph

flang/lib/Optimizer/Dialect/FIROps.cpp

@@ -4484,7 +4484,7 @@ void fir::IfOp::getSuccessorRegions(
     llvm::SmallVectorImpl<mlir::RegionSuccessor> &regions) {
   // The `then` and the `else` region branch back to the parent operation.
   if (!point.isParent()) {
-    regions.push_back(mlir::RegionSuccessor(getOperation(), getResults()));
+    regions.push_back(mlir::RegionSuccessor(getResults()));
     return;
   }
 
@@ -4494,8 +4494,7 @@ void fir::IfOp::getSuccessorRegions(
   // Don't consider the else region if it is empty.
   mlir::Region *elseRegion = &this->getElseRegion();
   if (elseRegion->empty())
-    regions.push_back(
-        mlir::RegionSuccessor(getOperation(), getOperation()->getResults()));
+    regions.push_back(mlir::RegionSuccessor());
   else
     regions.push_back(mlir::RegionSuccessor(elseRegion));
 }
@@ -4514,7 +4513,7 @@ void fir::IfOp::getEntrySuccessorRegions(
     if (!getElseRegion().empty())
       regions.emplace_back(&getElseRegion());
     else
-      regions.emplace_back(getOperation(), getOperation()->getResults());
+      regions.emplace_back(getResults());
   }
 }
 

mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h

@@ -397,7 +397,7 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// itself.
   virtual void visitRegionBranchControlFlowTransfer(
       RegionBranchOpInterface branch, RegionBranchPoint regionFrom,
-      RegionSuccessor regionTo, const AbstractDenseLattice &after,
+      RegionBranchPoint regionTo, const AbstractDenseLattice &after,
       AbstractDenseLattice *before) {
     meet(before, after);
   }
@@ -526,7 +526,7 @@ class DenseBackwardDataFlowAnalysis
   /// and "to" regions.
   virtual void visitRegionBranchControlFlowTransfer(
       RegionBranchOpInterface branch, RegionBranchPoint regionFrom,
-      RegionSuccessor regionTo, const LatticeT &after, LatticeT *before) {
+      RegionBranchPoint regionTo, const LatticeT &after, LatticeT *before) {
     AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchControlFlowTransfer(
         branch, regionFrom, regionTo, after, before);
   }
@@ -571,7 +571,7 @@ class DenseBackwardDataFlowAnalysis
   }
   void visitRegionBranchControlFlowTransfer(
       RegionBranchOpInterface branch, RegionBranchPoint regionForm,
-      RegionSuccessor regionTo, const AbstractDenseLattice &after,
+      RegionBranchPoint regionTo, const AbstractDenseLattice &after,
       AbstractDenseLattice *before) final {
     visitRegionBranchControlFlowTransfer(branch, regionForm, regionTo,
                                          static_cast<const LatticeT &>(after),

mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h

@@ -286,7 +286,7 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
   /// and propagating therefrom.
   virtual void
   visitRegionSuccessors(ProgramPoint *point, RegionBranchOpInterface branch,
-                        RegionSuccessor successor,
+                        RegionBranchPoint successor,
                         ArrayRef<AbstractSparseLattice *> lattices);
 };
 

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

@@ -644,13 +644,6 @@ def ForallOp : SCF_Op<"forall", [
 
     /// Returns true if the mapping specified for this forall op is linear.
     bool usesLinearMapping();
-
-    /// RegionBranchOpInterface
-
-    OperandRange getEntrySuccessorOperands(RegionSuccessor successor) {
-      return getInits();
-    }
-
   }];
 }
 

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

@@ -25,7 +25,7 @@ include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
 
 def AlternativesOp : TransformDialectOp<"alternatives",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
-        ["getEntrySuccessorOperands",
+        ["getEntrySuccessorOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,
      DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
@@ -624,7 +624,7 @@ def ForeachOp : TransformDialectOp<"foreach",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<RegionBranchOpInterface, [
-         "getEntrySuccessorOperands"]>,
+         "getSuccessorRegions", "getEntrySuccessorOperands"]>,
      SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">
     ]> {
   let summary = "Executes the body for each element of the payload";
@@ -1237,7 +1237,7 @@ def SelectOp : TransformDialectOp<"select",
 
 def SequenceOp : TransformDialectOp<"sequence",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
-        ["getEntrySuccessorOperands",
+        ["getEntrySuccessorOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,
      MatchOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>,

mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td

@@ -63,7 +63,7 @@ def KnobOp : Op<Transform_Dialect, "tune.knob", [
 
 def AlternativesOp : Op<Transform_Dialect, "tune.alternatives", [
   DeclareOpInterfaceMethods<RegionBranchOpInterface,
-        ["getEntrySuccessorOperands",
+        ["getEntrySuccessorOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,
   DeclareOpInterfaceMethods<TransformOpInterface>,
   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,

mlir/include/mlir/IR/Diagnostics.h

@@ -29,7 +29,6 @@ class MLIRContext;
 class Operation;
 class OperationName;
 class OpPrintingFlags;
-class OpWithFlags;
 class Type;
 class Value;
 
@@ -200,7 +199,6 @@ class Diagnostic {
 
   /// Stream in an Operation.
   Diagnostic &operator<<(Operation &op);
-  Diagnostic &operator<<(OpWithFlags op);
   Diagnostic &operator<<(Operation *op) { return *this << *op; }
   /// Append an operation with the given printing flags.
   Diagnostic &appendOp(Operation &op, const OpPrintingFlags &flags);

mlir/include/mlir/IR/Operation.h

@@ -1114,7 +1114,6 @@ class OpWithFlags {
       : op(op), theFlags(flags) {}
   OpPrintingFlags &flags() { return theFlags; }
   const OpPrintingFlags &flags() const { return theFlags; }
-  Operation *getOperation() const { return op; }
 
 private:
   Operation *op;

mlir/include/mlir/IR/Region.h

@@ -379,8 +379,6 @@ class RegionRange
   friend RangeBaseT;
 };
 
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, Region &region);
-
 } // namespace mlir
 
 #endif // MLIR_IR_REGION_H

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

@@ -15,16 +15,10 @@
 #define MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
 
 #include "mlir/IR/OpDefinition.h"
-#include "mlir/IR/Operation.h"
-#include "llvm/ADT/PointerUnion.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/Support/DebugLog.h"
-#include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 class BranchOpInterface;
 class RegionBranchOpInterface;
-class RegionBranchTerminatorOpInterface;
 
 /// This class models how operands are forwarded to block arguments in control
 /// flow. It consists of a number, denoting how many of the successors block
@@ -192,108 +186,92 @@ class RegionSuccessor {
 public:
   /// Initialize a successor that branches to another region of the parent
   /// operation.
-  /// TODO: the default value for the regionInputs is somehow broken.
-  /// A region successor should have its input correctly set.
   RegionSuccessor(Region *region, Block::BlockArgListType regionInputs = {})
-      : successor(region), inputs(regionInputs) {
-    assert(region && "Region must not be null");
-  }
+      : region(region), inputs(regionInputs) {}
   /// Initialize a successor that branches back to/out of the parent operation.
-  /// The target must be one of the recursive parent operations.
-  RegionSuccessor(Operation *successorOp, Operation::result_range results)
-      : successor(successorOp), inputs(ValueRange(results)) {
-    assert(successorOp && "Successor op must not be null");
-  }
+  RegionSuccessor(Operation::result_range results)
+      : inputs(ValueRange(results)) {}
+  /// Constructor with no arguments.
+  RegionSuccessor() : inputs(ValueRange()) {}
 
   /// Return the given region successor. Returns nullptr if the successor is the
   /// parent operation.
-  Region *getSuccessor() const { return dyn_cast<Region *>(successor); }
+  Region *getSuccessor() const { return region; }
 
   /// Return true if the successor is the parent operation.
-  bool isParent() const { return isa<Operation *>(successor); }
+  bool isParent() const { return region == nullptr; }
 
   /// Return the inputs to the successor that are remapped by the exit values of
   /// the current region.
   ValueRange getSuccessorInputs() const { return inputs; }
 
-  bool operator==(RegionSuccessor rhs) const {
-    return successor == rhs.successor && inputs == rhs.inputs;
-  }
-
-  friend bool operator!=(RegionSuccessor lhs, RegionSuccessor rhs) {
-    return !(lhs == rhs);
-  }
-
 private:
-  llvm::PointerUnion<Region *, Operation *> successor{nullptr};
+  Region *region{nullptr};
   ValueRange inputs;
 };
 
 /// This class represents a point being branched from in the methods of the
 /// `RegionBranchOpInterface`.
 /// One can branch from one of two kinds of places:
 /// * The parent operation (aka the `RegionBranchOpInterface` implementation)
-/// * A RegionBranchTerminatorOpInterface inside a region within the parent
-//    operation.
+/// * A region within the parent operation.
 class RegionBranchPoint {
 public:
   /// Returns an instance of `RegionBranchPoint` representing the parent
   /// operation.
   static constexpr RegionBranchPoint parent() { return RegionBranchPoint(); }
 
-  /// Creates a `RegionBranchPoint` that branches from the given terminator.
-  inline RegionBranchPoint(RegionBranchTerminatorOpInterface predecessor);
+  /// Creates a `RegionBranchPoint` that branches from the given region.
+  /// The pointer must not be null.
+  RegionBranchPoint(Region *region) : maybeRegion(region) {
+    assert(region && "Region must not be null");
+  }
+
+  RegionBranchPoint(Region &region) : RegionBranchPoint(&region) {}
 
   /// Explicitly stops users from constructing with `nullptr`.
   RegionBranchPoint(std::nullptr_t) = delete;
 
+  /// Constructs a `RegionBranchPoint` from the the target of a
+  /// `RegionSuccessor` instance.
+  RegionBranchPoint(RegionSuccessor successor) {
+    if (successor.isParent())
+      maybeRegion = nullptr;
+    else
+      maybeRegion = successor.getSuccessor();
+  }
+
+  /// Assigns a region being branched from.
+  RegionBranchPoint &operator=(Region &region) {
+    maybeRegion = &region;
+    return *this;
+  }
+
   /// Returns true if branching from the parent op.
-  bool isParent() const { return predecessor == nullptr; }
+  bool isParent() const { return maybeRegion == nullptr; }
 
-  /// Returns the terminator if branching from a region.
+  /// Returns the region if branching from a region.
   /// A null pointer otherwise.
-  Operation *getTerminatorPredecessorOrNull() const { return predecessor; }
+  Region *getRegionOrNull() const { return maybeRegion; }
 
   /// Returns true if the two branch points are equal.
   friend bool operator==(RegionBranchPoint lhs, RegionBranchPoint rhs) {
-    return lhs.predecessor == rhs.predecessor;
+    return lhs.maybeRegion == rhs.maybeRegion;
   }
 
 private:
   // Private constructor to encourage the use of `RegionBranchPoint::parent`.
-  constexpr RegionBranchPoint() = default;
+  constexpr RegionBranchPoint() : maybeRegion(nullptr) {}
 
   /// Internal encoding. Uses nullptr for representing branching from the parent
-  /// op and the region terminator being branched from otherwise.
-  Operation *predecessor = nullptr;
+  /// op and the region being branched from otherwise.
+  Region *maybeRegion;
 };
 
 inline bool operator!=(RegionBranchPoint lhs, RegionBranchPoint rhs) {
   return !(lhs == rhs);
 }
 
-inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
-                                     RegionBranchPoint point) {
-  if (point.isParent())
-    return os << "<from parent>";
-  return os << "<region #"
-            << point.getTerminatorPredecessorOrNull()
-                   ->getParentRegion()
-                   ->getRegionNumber()
-            << ", terminator "
-            << OpWithFlags(point.getTerminatorPredecessorOrNull(),
-                           OpPrintingFlags().skipRegions())
-            << ">";
-}
-
-inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
-                                     RegionSuccessor successor) {
-  if (successor.isParent())
-    return os << "<to parent>";
-  return os << "<to region #" << successor.getSuccessor()->getRegionNumber()
-            << " with " << successor.getSuccessorInputs().size() << " inputs>";
-}
-
 /// This class represents upper and lower bounds on the number of times a region
 /// of a `RegionBranchOpInterface` can be invoked. The lower bound is at least
 /// zero, but the upper bound may not be known.
@@ -370,10 +348,4 @@ struct ReturnLike : public TraitBase<ConcreteType, ReturnLike> {
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/ControlFlowInterfaces.h.inc"
 
-namespace mlir {
-inline RegionBranchPoint::RegionBranchPoint(
-    RegionBranchTerminatorOpInterface predecessor)
-    : predecessor(predecessor.getOperation()) {}
-} // namespace mlir
-
 #endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES_H