[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

mlir/include/mlir/IR/BlockSupport.h

@@ -21,8 +21,6 @@
 namespace mlir {
 class Block;
 
-using BlockOperand = IROperandImpl<Block>;
-
 //===----------------------------------------------------------------------===//
 // Predecessors
 //===----------------------------------------------------------------------===//

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) {

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,6 +177,36 @@ private:
   }
 };
 
+//===----------------------------------------------------------------------===//
+// BlockOperand
+//===----------------------------------------------------------------------===//
+
+/// Terminator operations can have Block operands to represent successors.
+class BlockOperand : public IROperand {
+public:
+  using IROperand::IROperand;
+
+  /// Return the current block being used by this operand.
+  Block *get();
+
+  /// Set the current value being used by this operand.
+  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;
+};
+
+//===----------------------------------------------------------------------===//
+// OpOperand
+//===----------------------------------------------------------------------===//
+
 /// A reference to a value, suitable for use as an operand of an operation.
 class OpOperand : public IROperand {
 public:
@@ -184,23 +223,9 @@ public:
   unsigned getOperandNumber();
 };
 
-/// 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 {
-public:
-  IROperandImpl(Operation *owner) : IROperand(owner) {}
-  IROperandImpl(Operation *owner, IRValueTy *value) : IROperand(owner, value) {}
-
-  /// Return the current value being used by this operand.
-  IRValueTy *get() { return (IRValueTy *)IROperand::get(); }
-
-  /// Set the current value being used by this operand.
-  void set(IRValueTy *newValue) { IROperand::set(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>
@@ -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

mlir/lib/IR/Operation.cpp

@@ -74,36 +74,6 @@ unsigned OpResult::getResultNumber() const {
   return std::distance(resList.begin(), llvm::find(resList, *this));
 }
 
-//===----------------------------------------------------------------------===//
-// OpOperand
-//===----------------------------------------------------------------------===//
-
-OpOperand::OpOperand(Operation *owner, Value value)
-    : IROperand(owner, value.impl) {}
-
-/// Return the current value being used by this operand.
-Value OpOperand::get() { return (detail::ValueImpl *)IROperand::get(); }
-
-/// Set the current value being used by this operand.
-void OpOperand::set(Value newValue) { IROperand::set(newValue.impl); }
-
-/// Return which operand this is in the operand list.
-unsigned OpOperand::getOperandNumber() {
-  return this - &getOwner()->getOpOperands()[0];
-}
-
-//===----------------------------------------------------------------------===//
-// BlockOperand
-//===----------------------------------------------------------------------===//
-
-// TODO: This namespace is only required because of a bug in GCC<7.0.
-namespace mlir {
-/// Return which operand this is in the operand list.
-template <> unsigned BlockOperand::getOperandNumber() {
-  return this - &getOwner()->getBlockOperands()[0];
-}
-} // end namespace mlir
-
 //===----------------------------------------------------------------------===//
 // Operation
 //===----------------------------------------------------------------------===//
@@ -159,10 +129,10 @@ Operation *Operation::create(Location location, OperationName name,
   unsigned numOperands = operands.size() - numSuccessors;
 
   // Compute the byte size for the operation and the operand storage.
-  auto byteSize = totalSizeToAlloc<OpResult, BlockOperand, unsigned, Region,
-                                   detail::OperandStorage>(
-      resultTypes.size(), numSuccessors, numSuccessors, numRegions,
-      /*detail::OperandStorage*/ 1);
+  auto byteSize =
+      totalSizeToAlloc<OpResult, BlockOperand, Region, detail::OperandStorage>(
+          resultTypes.size(), numSuccessors, numRegions,
+          /*detail::OperandStorage*/ 1);
   byteSize += llvm::alignTo(detail::OperandStorage::additionalAllocSize(
                                 numOperands, resizableOperandList),
                             alignof(Operation));
@@ -212,9 +182,7 @@ Operation *Operation::create(Location location, OperationName name,
   assert(!op->isKnownNonTerminator() &&
          "Unexpected nullptr in operand list when creating non-terminator.");
   auto instBlockOperands = op->getBlockOperands();
-  unsigned *succOperandCountIt = op->getTrailingObjects<unsigned>();
-  unsigned *succOperandCountE = succOperandCountIt + numSuccessors;
-  (void)succOperandCountE;
+  unsigned *succOperandCount = nullptr;
 
   for (; operandIt != operandE; ++operandIt) {
     // If we encounter a sentinel branch to the next operand update the count
@@ -222,22 +190,15 @@ Operation *Operation::create(Location location, OperationName name,
     if (!operands[operandIt]) {
       assert(currentSuccNum < numSuccessors);
 
-      // After the first iteration update the successor operand count
-      // variable.
-      if (currentSuccNum != 0) {
-        ++succOperandCountIt;
-        assert(succOperandCountIt != succOperandCountE &&
-               "More sentinel operands than successors.");
-      }
-
       new (&instBlockOperands[currentSuccNum])
           BlockOperand(op, successors[currentSuccNum]);
-      *succOperandCountIt = 0;
+      succOperandCount =
+          &instBlockOperands[currentSuccNum].numSuccessorOperands;
       ++currentSuccNum;
       continue;
     }
     new (&opOperands[nextOperand++]) OpOperand(op, operands[operandIt]);
-    ++(*succOperandCountIt);
+    ++(*succOperandCount);
   }
 
   // Verify that the amount of sentinel operands is equivalent to the number of
@@ -607,10 +568,12 @@ unsigned Operation::getSuccessorOperandIndex(unsigned index) {
   // Count the number of operands for each of the successors after, and
   // including, the one at 'index'. This is based upon the assumption that all
   // non successor operands are placed at the beginning of the operand list.
-  auto *successorOpCountBegin = getTrailingObjects<unsigned>();
+  auto blockOperands = getBlockOperands().drop_front(index);
   unsigned postSuccessorOpCount =
-      std::accumulate(successorOpCountBegin + index,
-                      successorOpCountBegin + getNumSuccessors(), 0u);
+      std::accumulate(blockOperands.begin(), blockOperands.end(), 0u,
+                      [](unsigned cur, const BlockOperand &operand) {
+                        return cur + operand.numSuccessorOperands;
+                      });
   return getNumOperands() - postSuccessorOpCount;
 }
 
@@ -619,10 +582,10 @@ Operation::decomposeSuccessorOperandIndex(unsigned operandIndex) {
   assert(!isKnownNonTerminator() && "only terminators may have successors");
   assert(operandIndex < getNumOperands());
   unsigned currentOperandIndex = getNumOperands();
-  auto *successorOperandCounts = getTrailingObjects<unsigned>();
+  auto blockOperands = getBlockOperands();
   for (unsigned i = 0, e = getNumSuccessors(); i < e; i++) {
     unsigned successorIndex = e - i - 1;
-    currentOperandIndex -= successorOperandCounts[successorIndex];
+    currentOperandIndex -= blockOperands[successorIndex].numSuccessorOperands;
     if (currentOperandIndex <= operandIndex)
       return std::make_pair(successorIndex, operandIndex - currentOperandIndex);
   }

mlir/lib/IR/Value.cpp

@@ -32,6 +32,39 @@ Region *Value::getParentRegion() {
   return cast<BlockArgument>()->getOwner()->getParent();
 }
 
+//===----------------------------------------------------------------------===//
+// BlockOperand
+//===----------------------------------------------------------------------===//
+
+/// Return the current block being used by this operand.
+Block *BlockOperand::get() { return static_cast<Block *>(IROperand::get()); }
+
+/// Set the current value being used by this operand.
+void BlockOperand::set(Block *block) { IROperand::set(block); }
+
+/// Return which operand this is in the operand list.
+unsigned BlockOperand::getOperandNumber() {
+  return this - &getOwner()->getBlockOperands()[0];
+}
+
+//===----------------------------------------------------------------------===//
+// OpOperand
+//===----------------------------------------------------------------------===//
+
+OpOperand::OpOperand(Operation *owner, Value value)
+    : IROperand(owner, value.impl) {}
+
+/// Return the current value being used by this operand.
+Value OpOperand::get() { return (detail::ValueImpl *)IROperand::get(); }
+
+/// Set the current value being used by this operand.
+void OpOperand::set(Value newValue) { IROperand::set(newValue.impl); }
+
+/// Return which operand this is in the operand list.
+unsigned OpOperand::getOperandNumber() {
+  return this - &getOwner()->getOpOperands()[0];
+}
+
 //===----------------------------------------------------------------------===//
 // IRObjectWithUseList implementation.
 //===----------------------------------------------------------------------===//