[mlir] Give each OpResult its own use list

[rriddle]

Summary

This revision removes the multi use-list to ensure that each result gets its own. This decision was made by doing some extensive benchmarking of programs that actually use multiple results. This results in a size increase of 1-word per result >1, but the common case of 1-result remains unaffected. A side benefit is that 0-result operations now shrink by 1-word.

[mehdi_amini]

My pass execution goes from 123.4811 seconds to 0.4536 seconds... This is just amazing :)

[jpienaar]

In D78701#1999550, @mehdi_amini wrote:

My pass execution goes from 123.4811 seconds to 0.4536 seconds... This is just amazing :)

270x speedup nothing to sneeze at ;-)

[mehdi_amini]

mlir/lib/IR/OperationSupport.cpp

141

"Otherwise" seems misplaced here

[rriddle]

Resolve comments

mlir/include/mlir/IR/Operation.h

 #include "llvm/ADT/Twine.h"
 
 namespace mlir {
 /// 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 IRMultiObjectWithUseList<OpOperand>,
-      public llvm::ilist_node_with_parent<Operation, Block>,
-      private llvm::TrailingObjects<Operation, detail::TrailingOpResult,
-                                    BlockOperand, Region,
-                                    detail::OperandStorage> {
+    : public llvm::ilist_node_with_parent<Operation, Block>,
+      private llvm::TrailingObjects<Operation, detail::InLineOpResult,
+                                    detail::TrailingOpResult, BlockOperand,
+                                    Region, detail::OperandStorage> {
 public:
   /// Create a new Operation with the specific fields.
   static Operation *create(Location location, OperationName name,
                            ArrayRef<Type> resultTypes, ArrayRef<Value> operands,
                            ArrayRef<NamedAttribute> attributes,
                            ArrayRef<Block *> successors, unsigned numRegions,
                            bool resizableOperandList);
 
public:
   ///         return WalkResult::advance();
   ///       });
   template <typename FnT, typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
     return detail::walkOperations(this, std::forward<FnT>(callback));
   }
 
   //===--------------------------------------------------------------------===//
+  // Uses
+  //===--------------------------------------------------------------------===//
+
+  /// Drop all uses of results of this operation.
+  void dropAllUses() {
+    for (OpResult result : getOpResults())
+      result.dropAllUses();
+  }
+
+  /// This class implements a use iterator for the Operation. This iterates over
+  /// all uses of all results.
+  class UseIterator final
+      : public llvm::iterator_facade_base<
+            UseIterator, std::forward_iterator_tag, OpOperand> {
+  public:
+    /// Initialize UseIterator for op, specify end to return iterator to last
+    /// use.
+    explicit UseIterator(Operation *op, bool end = false);
+
+    UseIterator &operator++();
+    OpOperand *operator->() const { return use.getOperand(); }
+    OpOperand &operator*() const { return *use.getOperand(); }
+
+    bool operator==(const UseIterator &rhs) const { return use == rhs.use; }
+    bool operator!=(const UseIterator &rhs) const { return !(*this == rhs); }
+
+  private:
+    void skipOverResultsWithNoUsers();
+
+    /// The operation whose uses are being iterated over.
+    Operation *op;
+    /// The result of op who's uses are being iterated over.
+    Operation::result_iterator res;
+    /// The use of the result.
+    Value::use_iterator use;
+  };
+  using use_iterator = UseIterator;
+  using use_range = iterator_range<use_iterator>;
+
+  use_iterator use_begin() { return use_iterator(this); }
+  use_iterator use_end() { return use_iterator(this, /*end=*/true); }
+
+  /// Returns a range of all uses, which is useful for iterating over all uses.
+  use_range getUses() { return {use_begin(), use_end()}; }
+
+  /// Returns true if this operation has exactly one use.
+  bool hasOneUse() { return llvm::hasSingleElement(getUses()); }
+
+  /// Returns true if this operation has no uses.
+  bool use_empty() {
+    return llvm::all_of(getOpResults(),
+                        [](OpResult result) { return result.use_empty(); });
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Users
+  //===--------------------------------------------------------------------===//
+
+  using user_iterator = ValueUserIterator<use_iterator, OpOperand>;
+  using user_range = iterator_range<user_iterator>;
+
+  user_iterator user_begin() { return user_iterator(use_begin()); }
+  user_iterator user_end() { return user_iterator(use_end()); }
+
+  /// Returns a range of all users.
+  user_range getUsers() { return {user_begin(), user_end()}; }
+
+  //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
 
   /// Emit an error with the op name prefixed, like "'dim' op " which is
   /// convenient for verifiers.
   InFlightDiagnostic emitOpError(const Twine &message = {});
 
   /// Emit an error about fatal conditions with this operation, reporting up to
private:
   // allocated with malloc.
   ~Operation();
 
   /// Returns the operand storage object.
   detail::OperandStorage &getOperandStorage() {
     return *getTrailingObjects<detail::OperandStorage>();
   }
 
-  /// Returns a raw pointer to the storage for the given trailing result. The
-  /// given result number should be 0-based relative to the trailing results,
-  /// and not all of the results of the operation. This method should generally
-  /// only be used by the 'Value' classes.
-  detail::TrailingOpResult *getTrailingResult(unsigned trailingResultNumber) {
-    return getTrailingObjects<detail::TrailingOpResult>() +
-           trailingResultNumber;
+  /// Returns a pointer to the use list for the given trailing result.
+  detail::TrailingOpResult *getTrailingResult(unsigned resultNumber) {
+    return getTrailingObjects<detail::TrailingOpResult>() + resultNumber;
+  }
+
+  /// Returns a pointer to the use list for the given inline result.
+  detail::InLineOpResult *getInlineResult(unsigned resultNumber) {
+    return getTrailingObjects<detail::InLineOpResult>() + resultNumber;
   }
 
   /// Provide a 'getParent' method for ilist_node_with_parent methods.
   /// We mark it as a const function because ilist_node_with_parent specifically
   /// requires a 'getParent() const' method. Once ilist_node removes this
   /// constraint, we should drop the const to fit the rest of the MLIR const
   /// model.
   Block *getParent() const { return block; }
private:
   NamedAttributeList attrs;
 
   // allow ilist_traits access to 'block' field.
   friend struct llvm::ilist_traits<Operation>;
 
   // allow block to access the 'orderIndex' field.
   friend class Block;
 
-  // allow value to access the 'getTrailingResult' method.
+  // allow value to access the 'ResultStorage' methods.
   friend class Value;
 
   // allow ilist_node_with_parent to access the 'getParent' method.
   friend class llvm::ilist_node_with_parent<Operation, Block>;
 
   // This stuff is used by the TrailingObjects template.
-  friend llvm::TrailingObjects<Operation, detail::TrailingOpResult,
-                               BlockOperand, Region, detail::OperandStorage>;
+  friend llvm::TrailingObjects<Operation, detail::InLineOpResult,
+                               detail::TrailingOpResult, BlockOperand, Region,
+                               detail::OperandStorage>;
+  size_t numTrailingObjects(OverloadToken<detail::InLineOpResult>) const {
+    return OpResult::getNumInline(
+        const_cast<Operation *>(this)->getNumResults());
+  }
   size_t numTrailingObjects(OverloadToken<detail::TrailingOpResult>) const {
     return OpResult::getNumTrailing(
         const_cast<Operation *>(this)->getNumResults());
   }
   size_t numTrailingObjects(OverloadToken<BlockOperand>) const {
     return numSuccs;
   }
   size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }

mlir/include/mlir/IR/OperationSupport.h

private:
   friend llvm::TrailingObjects<OperandStorage, ResizableStorage, OpOperand>;
   size_t numTrailingObjects(OverloadToken<ResizableStorage>) const {
     return resizable ? 1 : 0;
   }
 };
 } // end namespace detail
 
 //===----------------------------------------------------------------------===//
-// TrailingOpResult
+// ResultStorage
 //===----------------------------------------------------------------------===//
 
 namespace detail {
-/// This class provides the implementation for a trailing operation result.
-struct TrailingOpResult {
-  /// The only element is the trailing result number, or the offset from the
-  /// beginning of the trailing array.
+/// This class provides the implementation for an in-line operation result. This
+/// is an operation result whose number can be stored inline inside of the bits
+/// of an Operation*.
+struct InLineOpResult : public IRObjectWithUseList<OpOperand> {};
+/// This class provides the implementation for an out-of-line operation result.
+/// This is an operation result whose number cannot be stored inline inside of
+/// the bits of an Operation*.
+struct TrailingOpResult : public IRObjectWithUseList<OpOperand> {
+  TrailingOpResult(uint64_t trailingResultNumber)
+      : trailingResultNumber(trailingResultNumber) {}
+
+  /// Returns the parent operation of this trailing result.
+  Operation *getOwner();
+
+  /// Return the proper result number of this op result.
+  unsigned getResultNumber() {
+    return trailingResultNumber + OpResult::getMaxInlineResults();
+  }
+
+  /// The trailing result number, or the offset from the beginning of the
+  /// trailing array.
   uint64_t trailingResultNumber;
 };
 } // end namespace detail
 
 //===----------------------------------------------------------------------===//
 // OpPrintingFlags
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/IR/UseDefLists.h

protected:
   OperandType *getFirstUse() const { return firstUse; }
 
 private:
   template <typename DerivedT, typename IRValueTy> friend class IROperand;
   OperandType *firstUse = nullptr;
 };
 
 //===----------------------------------------------------------------------===//
-// IRMultiObjectWithUseList
-//===----------------------------------------------------------------------===//
-
-/// This class represents multiple IR objects with a single use list. This class
-/// provides wrapper functionality for manipulating the uses of a single object.
-template <typename OperandType>
-class IRMultiObjectWithUseList : public IRObjectWithUseList<OperandType> {
-public:
-  using BaseType = IRObjectWithUseList<OperandType>;
-  using ValueType = typename OperandType::ValueType;
-
-  /// Drop all uses of `value` from their respective owners.
-  void dropAllUses(ValueType value) {
-    assert(this == OperandType::getUseList(value) &&
-           "value not attached to this use list");
-    for (OperandType &use : llvm::make_early_inc_range(getUses(value)))
-      use.drop();
-  }
-  using BaseType::dropAllUses;
-
-  /// Replace all uses of `oldValue` with the new value, updating anything in
-  /// the IR that uses 'this' to use the other value instead.  When this returns
-  /// there are zero uses of 'this'.
-  void replaceAllUsesWith(ValueType oldValue, ValueType newValue) {
-    assert(this == OperandType::getUseList(oldValue) &&
-           "value not attached to this use list");
-    assert((!newValue || this != OperandType::getUseList(newValue)) &&
-           "cannot RAUW a value with itself");
-    for (OperandType &use : llvm::make_early_inc_range(getUses(oldValue)))
-      use.set(newValue);
-  }
-  using BaseType::replaceAllUsesWith;
-
-  //===--------------------------------------------------------------------===//
-  // Uses
-  //===--------------------------------------------------------------------===//
-
-  using filtered_use_iterator = FilteredValueUseIterator<OperandType>;
-  using filtered_use_range = iterator_range<filtered_use_iterator>;
-
-  filtered_use_iterator use_begin(ValueType value) const {
-    return filtered_use_iterator(this->getFirstUse(), value);
-  }
-  filtered_use_iterator use_end(ValueType) const { return use_end(); }
-  filtered_use_range getUses(ValueType value) const {
-    return {use_begin(value), use_end(value)};
-  }
-  bool hasOneUse(ValueType value) const {
-    return llvm::hasSingleElement(getUses(value));
-  }
-  bool use_empty(ValueType value) const {
-    return use_begin(value) == use_end(value);
-  }
-  using BaseType::getUses;
-  using BaseType::hasOneUse;
-  using BaseType::use_begin;
-  using BaseType::use_empty;
-  using BaseType::use_end;
-
-  //===--------------------------------------------------------------------===//
-  // Users
-  //===--------------------------------------------------------------------===//
-
-  using filtered_user_iterator =
-      ValueUserIterator<filtered_use_iterator, OperandType>;
-  using filtered_user_range = iterator_range<filtered_user_iterator>;
-
-  filtered_user_iterator user_begin(ValueType value) const {
-    return {use_begin(value)};
-  }
-  filtered_user_iterator user_end(ValueType value) const { return {use_end()}; }
-  filtered_user_range getUsers(ValueType value) const {
-    return {user_begin(value), user_end(value)};
-  }
-  using BaseType::getUsers;
-  using BaseType::user_begin;
-  using BaseType::user_end;
-};
-
-//===----------------------------------------------------------------------===//
 // IROperand
 //===----------------------------------------------------------------------===//
 
 /// A reference to a value, suitable for use as an operand of an operation.
 /// IRValueTy is the root type to use for values this tracks. Derived operand
 /// types are expected to provide the following:
 ///  * static IRObjectWithUseList *getUseList(IRValueTy value);
 ///    - Provide the use list that is attached to the given value.
public:
   void set(ValueType newValue) {
     // It isn't worth optimizing for the case of switching operands on a single
     // value.
     removeFromCurrent();
     value = newValue;
     insertIntoCurrent();
   }
 
+  /// Returns true if this operand contains the given value.
+  bool is(ValueType other) const { return value == other; }
+
   /// Return the owner of this operand.
   Operation *getOwner() { return owner; }
   Operation *getOwner() const { return owner; }
 
   /// \brief Remove this use of the operand.
   void drop() {
     removeFromCurrent();
     value = nullptr;

 /// This class provides an opaque type erased wrapper around a `Value`.
 class OpaqueValue {
 public:
   /// Implicit conversion from 'Value'.
   OpaqueValue(Value value);
   OpaqueValue(std::nullptr_t = nullptr) : impl(nullptr) {}
   OpaqueValue(const OpaqueValue &) = default;
   OpaqueValue &operator=(const OpaqueValue &) = default;
+  bool operator==(const OpaqueValue &other) const { return impl == other.impl; }
   operator bool() const { return impl; }
 
   /// Implicit conversion back to 'Value'.
   operator Value() const;
 
 private:
   void *impl;
 };
 } // namespace detail
 
-/// A reference to a value, suitable for use as an operand of an operation.
+/// This class represents an operand of an operation. Instances of this class
+/// contain a reference to a specific `Value`.
 class OpOperand : public IROperand<OpOperand, detail::OpaqueValue> {
 public:
   using IROperand<OpOperand, detail::OpaqueValue>::IROperand;
 
   /// Provide the use list that is attached to the given value.
   static IRObjectWithUseList<OpOperand> *getUseList(Value value);
 
   /// Return the current value being used by this operand.
   Value get() const;
 
   /// Set the operand to the given value.
   void set(Value value);
 
   /// Return which operand this is in the operand list of the User.
   unsigned getOperandNumber();
 };
 
 //===----------------------------------------------------------------------===//
 // ValueUseIterator
 //===----------------------------------------------------------------------===//
 
-namespace detail {
-/// A base iterator class that allows for iterating over the uses of a value.
-/// This is templated to allow for derived iterators to override specific
-/// iterator methods.
-template <typename DerivedT, typename OperandType>
-class ValueUseIteratorImpl
-    : public llvm::iterator_facade_base<DerivedT, std::forward_iterator_tag,
+/// An iterator class that allows for iterating over the uses of an IR operand
+/// type.
+template <typename OperandType>
+class ValueUseIterator
+    : public llvm::iterator_facade_base<ValueUseIterator<OperandType>,
+                                        std::forward_iterator_tag,
                                         OperandType> {
 public:
-  template <typename T>
-  ValueUseIteratorImpl(const ValueUseIteratorImpl<T, OperandType> &other)
-      : current(other.getOperand()) {}
-  ValueUseIteratorImpl(OperandType *current = nullptr) : current(current) {}
+  ValueUseIterator(OperandType *current = nullptr) : current(current) {}
 
+  /// Returns the user that owns this use.
   Operation *getUser() const { return current->getOwner(); }
-  OperandType *getOperand() const { return current; }
 
+  /// Returns the current operands.
+  OperandType *getOperand() const { return current; }
   OperandType &operator*() const { return *current; }
 
-  using llvm::iterator_facade_base<DerivedT, std::forward_iterator_tag,
+  using llvm::iterator_facade_base<ValueUseIterator<OperandType>,
+                                   std::forward_iterator_tag,
                                    OperandType>::operator++;
-  ValueUseIteratorImpl &operator++() {
+  ValueUseIterator &operator++() {
     assert(current && "incrementing past end()!");
     current = (OperandType *)current->getNextOperandUsingThisValue();
     return *this;
   }
 
-  bool operator==(const ValueUseIteratorImpl &rhs) const {
+  bool operator==(const ValueUseIterator &rhs) const {
     return current == rhs.current;
   }
 
 protected:
   OperandType *current;
 };
 
-} // end namespace detail
-
-/// An iterator over all of the uses of an IR object.
-template <typename OperandType>
-class ValueUseIterator
-    : public detail::ValueUseIteratorImpl<ValueUseIterator<OperandType>,
-                                          OperandType> {
-public:
-  using detail::ValueUseIteratorImpl<ValueUseIterator<OperandType>,
-                                     OperandType>::ValueUseIteratorImpl;
-};
-
-/// This class represents an iterator of the uses of a IR object that optionally
-/// filters on a specific sub-value. This allows for filtering the uses of an
-/// IRMultiObjectWithUseList.
-template <typename OperandType>
-class FilteredValueUseIterator
-    : public detail::ValueUseIteratorImpl<FilteredValueUseIterator<OperandType>,
-                                          OperandType> {
-public:
-  using BaseT =
-      detail::ValueUseIteratorImpl<FilteredValueUseIterator<OperandType>,
-                                   OperandType>;
-
-  FilteredValueUseIterator() = default;
-  FilteredValueUseIterator(const ValueUseIterator<OperandType> &it)
-      : BaseT(it), filterVal(nullptr) {}
-  FilteredValueUseIterator(OperandType *current,
-                           typename OperandType::ValueType filterVal)
-      : BaseT(current), filterVal(filterVal) {
-    findNextValid();
-  }
-
-  using BaseT::operator++;
-  FilteredValueUseIterator<OperandType> &operator++() {
-    BaseT::operator++();
-    findNextValid();
-    return *this;
-  }
-
-private:
-  void findNextValid() {
-    if (!filterVal)
-      return;
-    while (this->current && ((OperandType *)this->current)->get() != filterVal)
-      BaseT::operator++();
-  }
-
-  /// An optional value to use to filter specific uses.
-  typename OperandType::ValueType filterVal;
-};
-
 //===----------------------------------------------------------------------===//
 // ValueUserIterator
 //===----------------------------------------------------------------------===//
 
-/// An iterator over all users of a ValueBase.
+/// An iterator over the users of an IRObject. This is a wrapper iterator around
+/// a specific use iterator.
 template <typename UseIteratorT, typename OperandType>
 class ValueUserIterator final
     : public llvm::mapped_iterator<UseIteratorT,
                                    Operation *(*)(OperandType &)> {
   static Operation *unwrap(OperandType &value) { return value.getOwner(); }
 
 public:
   using pointer = Operation *;
   using reference = Operation *;
 
-  /// Initializes the result type iterator to the specified result iterator.
+  /// Initializes the user iterator to the specified use iterator.
   ValueUserIterator(UseIteratorT it)
       : llvm::mapped_iterator<UseIteratorT, Operation *(*)(OperandType &)>(
             it, &unwrap) {}
   Operation *operator->() { return **this; }
 };
 
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/Value.h

public:
   void
   replaceAllUsesExcept(Value newValue,
                        const SmallPtrSetImpl<Operation *> &exceptions) const;
 
   //===--------------------------------------------------------------------===//
   // Uses
 
   /// This class implements an iterator over the uses of a value.
-  using use_iterator = FilteredValueUseIterator<OpOperand>;
+  using use_iterator = ValueUseIterator<OpOperand>;
   using use_range = iterator_range<use_iterator>;
 
   use_iterator use_begin() const;
   use_iterator use_end() const { return use_iterator(); }
 
   /// Returns a range of all uses, which is useful for iterating over all uses.
   use_range getUses() const { return {use_begin(), use_end()}; }
 
public:
   }
 
   /// Returns the operation that owns this result.
   Operation *getOwner() const;
 
   /// Returns the number of this result.
   unsigned getResultNumber() const;
 
+  /// Returns the maximum number of results that can be stored inline.
+  static unsigned getMaxInlineResults() {
+    return static_cast<unsigned>(Kind::TrailingOpResult);
+  }
+
 private:
   /// Given a number of operation results, returns the number that need to be
+  /// stored inline.
+  static unsigned getNumInline(unsigned numResults);
+
+  /// Given a number of operation results, returns the number that need to be
   /// stored as trailing.
   static unsigned getNumTrailing(unsigned numResults);
 
-  /// Allow access to `create` and `destroy`.
+  /// Allow access to constructor.
   friend Operation;
 };
 
 /// Make Value hashable.
 inline ::llvm::hash_code hash_value(Value arg) {
   return ::llvm::hash_value(arg.ownerAndKind.getOpaqueValue());
 }
 

mlir/lib/IR/Operation.cpp

 Operation *Operation::create(Location location, OperationName name,
                              ArrayRef<Type> resultTypes,
                              ArrayRef<Value> operands,
                              NamedAttributeList attributes,
                              ArrayRef<Block *> successors, unsigned numRegions,
                              bool resizableOperandList) {
   // We only need to allocate additional memory for a subset of results.
   unsigned numTrailingResults = OpResult::getNumTrailing(resultTypes.size());
+  unsigned numInlineResults = OpResult::getNumInline(resultTypes.size());
   unsigned numSuccessors = successors.size();
   unsigned numOperands = operands.size();
 
   // Compute the byte size for the operation and the operand storage.
-  auto byteSize = totalSizeToAlloc<detail::TrailingOpResult, BlockOperand,
-                                   Region, detail::OperandStorage>(
-      numTrailingResults, numSuccessors, numRegions,
+  auto byteSize =
+      totalSizeToAlloc<detail::InLineOpResult, detail::TrailingOpResult,
+                       BlockOperand, Region, detail::OperandStorage>(
+          numInlineResults, numTrailingResults, numSuccessors, numRegions,
           /*detail::OperandStorage*/ 1);
   byteSize += llvm::alignTo(detail::OperandStorage::additionalAllocSize(
                                 numOperands, resizableOperandList),
                             alignof(Operation));
   void *rawMem = malloc(byteSize);
 
   // Create the new Operation.
   auto op = ::new (rawMem) Operation(location, name, resultTypes, numSuccessors,
                                      numRegions, attributes);
 
   assert((numSuccessors == 0 || !op->isKnownNonTerminator()) &&
          "unexpected successors in a non-terminator operation");
 
-  // Initialize the trailing results.
-  if (LLVM_UNLIKELY(numTrailingResults > 0)) {
-    // We initialize the trailing results with their result number. This makes
-    // 'getResultNumber' checks much more efficient. The main purpose for these
-    // results is to give an anchor to the main operation anyways, so this is
-    // purely an optimization.
-    auto *trailingResultIt = op->getTrailingObjects<detail::TrailingOpResult>();
-    for (unsigned i = 0; i != numTrailingResults; ++i, ++trailingResultIt)
-      trailingResultIt->trailingResultNumber = i;
-  }
+  // Initialize the results.
+  for (unsigned i = 0; i < numInlineResults; ++i)
+    new (op->getInlineResult(i)) detail::InLineOpResult();
+  for (unsigned i = 0; i < numTrailingResults; ++i)
+    new (op->getTrailingResult(i)) detail::TrailingOpResult(i);
 
   // Initialize the regions.
   for (unsigned i = 0; i != numRegions; ++i)
     new (&op->getRegion(i)) Region(op);
 
   // Initialize the operands.
   new (&op->getOperandStorage())
       detail::OperandStorage(numOperands, resizableOperandList);
if (region.empty())
     region.push_back(new Block);
 
   Block &block = region.back();
   if (!block.empty() && block.back().isKnownTerminator())
     return;
 
   block.push_back(buildTerminatorOp());
 }
+
+//===----------------------------------------------------------------------===//
+// UseIterator
+//===----------------------------------------------------------------------===//
+
+Operation::UseIterator::UseIterator(Operation *op, bool end)
+    : op(op), res(end ? op->result_end() : op->result_begin()) {
+  // Only initialize current use if there are results/can be uses.
+  if (op->getNumResults())
+    skipOverResultsWithNoUsers();
+}
+
+Operation::UseIterator &Operation::UseIterator::operator++() {
+  // We increment over uses, if we reach the last use then move to next
+  // result.
+  if (use != (*res).use_end())
+    ++use;
+  if (use == (*res).use_end()) {
+    ++res;
+    skipOverResultsWithNoUsers();
+  }
+  return *this;
+}
+
+void Operation::UseIterator::skipOverResultsWithNoUsers() {
+  while (res != op->result_end() && (*res).use_empty())
+    ++res;
+
+  // If we are at the last result, then set use to first use of
+  // first result (sentinel value used for end).
+  if (res == op->result_end())
+    use = {};
+  else
+    use = (*res).use_begin();
+}

mlir/lib/IR/OperationSupport.cpp

void detail::OperandStorage::grow(ResizableStorage &resizeUtil,
 
   // Destroy the original operands and update the resizable storage pointer.
   for (auto &operand : operands)
     operand.~OpOperand();
   resizeUtil.setDynamicStorage(newStorage);
 }
 
 //===----------------------------------------------------------------------===//
+// ResultStorage
+//===----------------------------------------------------------------------===//
+
+/// Returns the parent operation of this trailing result.
+Operation *detail::TrailingOpResult::getOwner() {
+  // We need to do some arithmetic to get the operation pointer. Move the
+  // trailing owner to the start of the array.
+  TrailingOpResult *trailingIt = this - trailingResultNumber;
+
+  // Move the owner past the inline op results to get to the operation.
+  auto *inlineResultIt = reinterpret_cast<InLineOpResult *>(trailingIt) -
+                         OpResult::getMaxInlineResults();
+  return reinterpret_cast<Operation *>(inlineResultIt) - 1;
+}
+
+//===----------------------------------------------------------------------===//
 // Operation Value-Iterators
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // TypeRange
 
 TypeRange::TypeRange(ArrayRef<Type> types)
     : TypeRange(types.data(), types.size()) {}

mlir/lib/IR/Value.cpp

 //===- Value.cpp - MLIR Value Classes -------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/ADT/SmallPtrSet.h"
+
 using namespace mlir;
+using namespace mlir::detail;
 
 /// Construct a value.
-Value::Value(detail::BlockArgumentImpl *impl)
+Value::Value(BlockArgumentImpl *impl)
     : ownerAndKind(impl, Kind::BlockArgument) {}
 Value::Value(Operation *op, unsigned resultNo) {
   assert(op->getNumResults() > resultNo && "invalid result number");
   if (LLVM_LIKELY(canPackResultInline(resultNo))) {
     ownerAndKind = {op, static_cast<Kind>(resultNo)};
     return;
   }
 
-  // If we can't pack the result directly, we need to represent this as a
-  // trailing result.
-  unsigned trailingResultNo =
-      resultNo - static_cast<unsigned>(Kind::TrailingOpResult);
-  ownerAndKind = {op->getTrailingResult(trailingResultNo),
-                  Kind::TrailingOpResult};
+  // If we can't pack the result directly, grab the use list from the parent op.
+  unsigned trailingNo = resultNo - OpResult::getMaxInlineResults();
+  ownerAndKind = {op->getTrailingResult(trailingNo), Kind::TrailingOpResult};
 }
 
 /// Return the type of this value.
 Type Value::getType() const {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getType();
 
   // If this is an operation result, query the parent operation.

 //===----------------------------------------------------------------------===//
 // Value::UseLists
 //===----------------------------------------------------------------------===//
 
 /// Provide the use list that is attached to this value.
 IRObjectWithUseList<OpOperand> *Value::getUseList() const {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getImpl();
-  return cast<OpResult>().getOwner();
+  if (getKind() != Kind::TrailingOpResult) {
+    OpResult result = cast<OpResult>();
+    return result.getOwner()->getInlineResult(result.getResultNumber());
   }
 
-/// Drop all uses of this object from their respective owners.
-void Value::dropAllUses() const {
-  if (BlockArgument arg = dyn_cast<BlockArgument>())
-    return arg.getImpl()->dropAllUses();
-  Operation *owner = cast<OpResult>().getOwner();
-  if (owner->hasSingleResult)
-    return owner->dropAllUses();
-  return owner->dropAllUses(*this);
+  // Otherwise this is a trailing operation result, which contains a use list.
+  return reinterpret_cast<TrailingOpResult *>(ownerAndKind.getPointer());
 }
 
+/// Drop all uses of this object from their respective owners.
+void Value::dropAllUses() const { return getUseList()->dropAllUses(); }
+
 /// Replace all uses of 'this' value with the new value, updating anything in
 /// the IR that uses 'this' to use the other value instead.  When this returns
 /// there are zero uses of 'this'.
 void Value::replaceAllUsesWith(Value newValue) const {
-  if (BlockArgument arg = dyn_cast<BlockArgument>())
-    return arg.getImpl()->replaceAllUsesWith(newValue);
-  Operation *owner = cast<OpResult>().getOwner();
-  IRMultiObjectWithUseList<OpOperand> *useList = owner;
-  if (owner->hasSingleResult)
-    return useList->replaceAllUsesWith(newValue);
-  useList->replaceAllUsesWith(*this, newValue);
+  return getUseList()->replaceAllUsesWith(newValue);
 }
 
 /// Replace all uses of 'this' value with the new value, updating anything in
 /// the IR that uses 'this' to use the other value instead except if the user is
 /// listed in 'exceptions' .
 void Value::replaceAllUsesExcept(
     Value newValue, const SmallPtrSetImpl<Operation *> &exceptions) const {
   for (auto &use : llvm::make_early_inc_range(getUses())) {
     if (exceptions.count(use.getOwner()) == 0)
       use.set(newValue);
   }
 }
 
 //===--------------------------------------------------------------------===//
 // Uses
 
 auto Value::use_begin() const -> use_iterator {
-  if (BlockArgument arg = dyn_cast<BlockArgument>())
-    return arg.getImpl()->use_begin();
-  Operation *owner = cast<OpResult>().getOwner();
-  return owner->hasSingleResult ? use_iterator(owner->use_begin())
-                                : owner->use_begin(*this);
+  return getUseList()->use_begin();
 }
 
 /// Returns true if this value has exactly one use.
-bool Value::hasOneUse() const {
-  if (BlockArgument arg = dyn_cast<BlockArgument>())
-    return arg.getImpl()->hasOneUse();
-  Operation *owner = cast<OpResult>().getOwner();
-  return owner->hasSingleResult ? owner->hasOneUse() : owner->hasOneUse(*this);
-}
+bool Value::hasOneUse() const { return getUseList()->hasOneUse(); }
 
 /// Returns true if this value has no uses.
-bool Value::use_empty() const {
-  if (BlockArgument arg = dyn_cast<BlockArgument>())
-    return arg.getImpl()->use_empty();
-  Operation *owner = cast<OpResult>().getOwner();
-  return owner->hasSingleResult ? owner->use_empty() : owner->use_empty(*this);
-}
+bool Value::use_empty() const { return getUseList()->use_empty(); }
 
 //===----------------------------------------------------------------------===//
 // OpResult
 //===----------------------------------------------------------------------===//
 
 /// Returns the operation that owns this result.
 Operation *OpResult::getOwner() const {
   // If the result is in-place, the `owner` is the operation.
+  void *owner = ownerAndKind.getPointer();
   if (LLVM_LIKELY(getKind() != Kind::TrailingOpResult))
-    return reinterpret_cast<Operation *>(ownerAndKind.getPointer());
+    return static_cast<Operation *>(owner);
 
-  // Otherwise, we need to do some arithmetic to get the operation pointer.
-  // Move the trailing owner to the start of the array.
-  auto *trailingIt =
-      static_cast<detail::TrailingOpResult *>(ownerAndKind.getPointer());
-  trailingIt -= trailingIt->trailingResultNumber;
-
-  // This point is the first trailing object after the operation. So all we need
-  // to do here is adjust for the operation size.
-  return reinterpret_cast<Operation *>(trailingIt) - 1;
+  // Otherwise, query the trailing result for the owner.
+  return static_cast<TrailingOpResult *>(owner)->getOwner();
 }
 
 /// Return the result number of this result.
 unsigned OpResult::getResultNumber() const {
   // If the result is in-place, we can use the kind directly.
   if (LLVM_LIKELY(getKind() != Kind::TrailingOpResult))
     return static_cast<unsigned>(ownerAndKind.getInt());
-  // Otherwise, we add the number of inline results to the trailing owner.
-  auto *trailingIt =
-      static_cast<detail::TrailingOpResult *>(ownerAndKind.getPointer());
-  unsigned trailingNumber = trailingIt->trailingResultNumber;
-  return trailingNumber + static_cast<unsigned>(Kind::TrailingOpResult);
+  // Otherwise, query the trailing result.
+  auto *result = static_cast<TrailingOpResult *>(ownerAndKind.getPointer());
+  return result->getResultNumber();
+}
+
+/// Given a number of operation results, returns the number that need to be
+/// stored inline.
+unsigned OpResult::getNumInline(unsigned numResults) {
+  return std::min(numResults, getMaxInlineResults());
 }
 
 /// Given a number of operation results, returns the number that need to be
 /// stored as trailing.
 unsigned OpResult::getNumTrailing(unsigned numResults) {
   // If we can pack all of the results, there is no need for additional storage.
-  if (numResults <= static_cast<unsigned>(Kind::TrailingOpResult))
-    return 0;
-  return numResults - static_cast<unsigned>(Kind::TrailingOpResult);
+  unsigned maxInline = getMaxInlineResults();
+  return numResults <= maxInline ? 0 : numResults - maxInline;
 }
 
 //===----------------------------------------------------------------------===//
 // BlockOperand
 //===----------------------------------------------------------------------===//
 
 /// Provide the use list that is attached to the given block.
 IRObjectWithUseList<BlockOperand> *BlockOperand::getUseList(Block *value) {

 
 /// Provide the use list that is attached to the given value.
 IRObjectWithUseList<OpOperand> *OpOperand::getUseList(Value value) {
   return value.getUseList();
 }
 
 /// Return the current value being used by this operand.
 Value OpOperand::get() const {
-  return IROperand<OpOperand, detail::OpaqueValue>::get();
+  return IROperand<OpOperand, OpaqueValue>::get();
 }
 
 /// Set the operand to the given value.
 void OpOperand::set(Value value) {
-  IROperand<OpOperand, detail::OpaqueValue>::set(value);
+  IROperand<OpOperand, OpaqueValue>::set(value);
 }
 
 /// Return which operand this is in the operand list.
 unsigned OpOperand::getOperandNumber() {
   return this - &getOwner()->getOpOperands()[0];
 }
 
 //===----------------------------------------------------------------------===//
-// detail::OpaqueValue
+// OpaqueValue
 //===----------------------------------------------------------------------===//
 
 /// Implicit conversion from 'Value'.
-detail::OpaqueValue::OpaqueValue(Value value)
-    : impl(value.getAsOpaquePointer()) {}
+OpaqueValue::OpaqueValue(Value value) : impl(value.getAsOpaquePointer()) {}
 
 /// Implicit conversion back to 'Value'.
-detail::OpaqueValue::operator Value() const {
+OpaqueValue::operator Value() const {
   return Value::getFromOpaquePointer(impl);
 }