[mlir] Merge the successor operand count into BlockOperand.

Summary: The successor operand counts are directly tied to block operands anyways, and this simplifies the trailing objects of Operation(i.e. one less computation to perform).

Reviewed By: mehdi_amini

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

3 files changed, 50 insertions, 27 deletions

diff --git a/mlir/include/mlir/IR/BlockSupport.h b/mlir/include/mlir/IR/BlockSupport.h
index bc6a8245c45..6025dd9bb1e 100644
--- a/mlir/include/mlir/IR/BlockSupport.h
+++ b/mlir/include/mlir/IR/BlockSupport.h
@@ -21,8 +21,6 @@
 namespace mlir {
 class Block;
 
-using BlockOperand = IROperandImpl<Block>;
-
 //===----------------------------------------------------------------------===//
 // Predecessors
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 9ef1636d3d0..7586fc88337 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -20,17 +20,14 @@
 #include "llvm/ADT/Twine.h"
 
 namespace mlir {
-/// Terminator operations can have Block operands to represent successors.
-using BlockOperand = IROperandImpl<Block>;
-
 /// Operation is a basic unit of execution within a function. Operations can
 /// be nested within other operations effectively forming a tree. Child
 /// operations are organized into operation blocks represented by a 'Block'
 /// class.
 class Operation final
     : public llvm::ilist_node_with_parent<Operation, Block>,
-      private llvm::TrailingObjects<Operation, OpResult, BlockOperand, unsigned,
-                                    Region, detail::OperandStorage> {
+      private llvm::TrailingObjects<Operation, OpResult, BlockOperand, Region,
+                                    detail::OperandStorage> {
 public:
   /// Create a new Operation with the specific fields.
   static Operation *create(Location location, OperationName name,
@@ -406,7 +403,7 @@ public:
   unsigned getNumSuccessorOperands(unsigned index) {
     assert(!isKnownNonTerminator() && "only terminators may have successors");
     assert(index < getNumSuccessors());
-    return getTrailingObjects<unsigned>()[index];
+    return getBlockOperands()[index].numSuccessorOperands;
   }
 
   Block *getSuccessor(unsigned index) {
@@ -422,7 +419,7 @@ public:
     assert(opIndex < getNumSuccessorOperands(succIndex));
     getOperandStorage().eraseOperand(getSuccessorOperandIndex(succIndex) +
                                      opIndex);
-    --getTrailingObjects<unsigned>()[succIndex];
+    --getBlockOperands()[succIndex].numSuccessorOperands;
   }
 
   /// Get the index of the first operand of the successor at the provided
@@ -626,8 +623,8 @@ private:
   friend class llvm::ilist_node_with_parent<Operation, Block>;
 
   // This stuff is used by the TrailingObjects template.
-  friend llvm::TrailingObjects<Operation, OpResult, BlockOperand, unsigned,
-                               Region, detail::OperandStorage>;
+  friend llvm::TrailingObjects<Operation, OpResult, BlockOperand, Region,
+                               detail::OperandStorage>;
   size_t numTrailingObjects(OverloadToken<OpResult>) const {
     return numResults;
   }
@@ -635,7 +632,6 @@ private:
     return numSuccs;
   }
   size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }
-  size_t numTrailingObjects(OverloadToken<unsigned>) const { return numSuccs; }
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Operation &op) {
diff --git a/mlir/include/mlir/IR/UseDefLists.h b/mlir/include/mlir/IR/UseDefLists.h
index 05720ed39af..222b50f8960 100644
--- a/mlir/include/mlir/IR/UseDefLists.h
+++ b/mlir/include/mlir/IR/UseDefLists.h
@@ -19,12 +19,17 @@
 
 namespace mlir {
 
+class Block;
 class IROperand;
 class Operation;
 class Value;
 template <typename OperandType> class ValueUseIterator;
 template <typename OperandType> class ValueUserIterator;
 
+//===----------------------------------------------------------------------===//
+// IRObjectWithUseList
+//===----------------------------------------------------------------------===//
+
 class IRObjectWithUseList {
 public:
   ~IRObjectWithUseList() {
@@ -76,6 +81,10 @@ private:
   IROperand *firstUse = nullptr;
 };
 
+//===----------------------------------------------------------------------===//
+// IROperand
+//===----------------------------------------------------------------------===//
+
 /// A reference to a value, suitable for use as an operand of an operation.
 class IROperand {
 public:
@@ -168,40 +177,56 @@ private:
   }
 };
 
-/// A reference to a value, suitable for use as an operand of an operation.
-class OpOperand : public IROperand {
+//===----------------------------------------------------------------------===//
+// BlockOperand
+//===----------------------------------------------------------------------===//
+
+/// Terminator operations can have Block operands to represent successors.
+class BlockOperand : public IROperand {
 public:
-  OpOperand(Operation *owner) : IROperand(owner) {}
-  OpOperand(Operation *owner, Value value);
+  using IROperand::IROperand;
 
-  /// Return the current value being used by this operand.
-  Value get();
+  /// Return the current block being used by this operand.
+  Block *get();
 
   /// Set the current value being used by this operand.
-  void set(Value newValue);
+  void set(Block *block);
 
   /// Return which operand this is in the operand list of the User.
   unsigned getOperandNumber();
+
+private:
+  /// The number of OpOperands that correspond with this block operand.
+  unsigned numSuccessorOperands = 0;
+
+  /// Allow access to 'numSuccessorOperands'.
+  friend Operation;
 };
 
-/// A reference to a value, suitable for use as an operand of an operation,
-/// operation, etc.  IRValueTy is the root type to use for values this tracks,
-/// and SSAUserTy is the type that will contain operands.
-template <typename IRValueTy> class IROperandImpl : public IROperand {
+//===----------------------------------------------------------------------===//
+// OpOperand
+//===----------------------------------------------------------------------===//
+
+/// A reference to a value, suitable for use as an operand of an operation.
+class OpOperand : public IROperand {
 public:
-  IROperandImpl(Operation *owner) : IROperand(owner) {}
-  IROperandImpl(Operation *owner, IRValueTy *value) : IROperand(owner, value) {}
+  OpOperand(Operation *owner) : IROperand(owner) {}
+  OpOperand(Operation *owner, Value value);
 
   /// Return the current value being used by this operand.
-  IRValueTy *get() { return (IRValueTy *)IROperand::get(); }
+  Value get();
 
   /// Set the current value being used by this operand.
-  void set(IRValueTy *newValue) { IROperand::set(newValue); }
+  void set(Value newValue);
 
   /// Return which operand this is in the operand list of the User.
   unsigned getOperandNumber();
 };
 
+//===----------------------------------------------------------------------===//
+// ValueUseIterator
+//===----------------------------------------------------------------------===//
+
 /// An iterator over all uses of a ValueBase.
 template <typename OperandType>
 class ValueUseIterator
@@ -255,6 +280,10 @@ inline bool IRObjectWithUseList::hasOneUse() const {
   return firstUse && firstUse->getNextOperandUsingThisValue() == nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// ValueUserIterator
+//===----------------------------------------------------------------------===//
+
 /// An iterator over all users of a ValueBase.
 template <typename OperandType>
 class ValueUserIterator final