[mlir][IR] Refactor the internal implementation of Value

The current implementation of Value involves a pointer int pair with several different kinds of owners, i.e. BlockArgumentImpl*, Operation *, TrailingOpResult*. This design arose from the desire to save memory overhead for operations that have a very small number of results (generally 0-2). There are, unfortunately, many problematic aspects of the current implementation that make Values difficult to work with or just inefficient.

Operation result types are stored as a separate array on the Operation. This is very inefficient for many reasons: we use TupleType for multiple results, which can lead to huge amounts of memory usage if multi-result operations change types frequently(they do). It also means that simple methods like Value::getType/Value::setType now require complex logic to get to the desired type.

Value only has one pointer bit free, severely limiting the ability to use it in things like PointerUnion/PointerIntPair. Given that we store the kind of a Value along with the "owner" pointer, we only leave one bit free for users of Value. This creates situations where we end up nesting PointerUnions to be able to use Value in one.

As noted above, most of the methods in Value need to branch on at least 3 different cases which is both inefficient, possibly error prone, and verbose. The current storage of results also creates problems for utilities like ValueRange/TypeRange, which want to efficiently store base pointers to ranges (of which Operation* isn't really useful as one).

This revision greatly simplifies the implementation of Value by the introduction of a new ValueImpl class. This class contains all of the state shared between all of the various derived value classes; i.e. the use list, the type, and the kind. This shared implementation class provides several large benefits:

* Most of the methods on value are now branchless, and often one-liners.

* The "kind" of the value is now stored in ValueImpl instead of Value
This frees up all of Value's pointer bits, allowing for users to take full advantage of PointerUnion/PointerIntPair/etc. It also allows for storing more operation results as "inline", 6 now instead of 2, freeing up 1 word per new inline result.

* Operation result types are now stored in the result, instead of a side array
This drops the size of zero-result operations by 1 word. It also removes the memory crushing use of TupleType for operations results (which could lead up to hundreds of megabytes of "dead" TupleTypes in the context). This also allowed restructured ValueRange, making it simpler and one word smaller.

This revision does come with two conceptual downsides:
* Operation::getResultTypes no longer returns an ArrayRef<Type>
This conceptually makes some usages slower, as the iterator increment is slightly more complex.
* OpResult::getOwner is slightly more expensive, as it now requires a little bit of arithmetic

From profiling, neither of the conceptual downsides have resulted in any perceivable hit to performance. Given the advantages of the new design, most compiles are slightly faster.

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

Author: River707

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

@@ -77,7 +77,7 @@ class LLVMTypeConverter : public TypeConverter {
   /// supported LLVM IR type.  In particular, if more than one value is
   /// returned, create an LLVM IR structure type with elements that correspond
   /// to each of the MLIR types converted with `convertType`.
-  Type packFunctionResults(ArrayRef<Type> types);
+  Type packFunctionResults(TypeRange types);
 
   /// Convert a type in the context of the default or bare pointer calling
   /// convention. Calling convention sensitive types, such as MemRefType and

mlir/include/mlir/IR/Operation.h

@@ -272,18 +272,21 @@ class alignas(8) Operation final
   //===--------------------------------------------------------------------===//
 
   /// Return the number of results held by this operation.
-  unsigned getNumResults();
+  unsigned getNumResults() { return numResults; }
 
   /// Get the 'idx'th result of this operation.
-  OpResult getResult(unsigned idx) { return OpResult(this, idx); }
+  OpResult getResult(unsigned idx) { return OpResult(getOpResultImpl(idx)); }
 
   /// Support result iteration.
   using result_range = ResultRange;
   using result_iterator = result_range::iterator;
 
   result_iterator result_begin() { return getResults().begin(); }
   result_iterator result_end() { return getResults().end(); }
-  result_range getResults() { return result_range(this); }
+  result_range getResults() {
+    return numResults == 0 ? result_range(nullptr, 0)
+                           : result_range(getInlineOpResult(0), numResults);
+  }
 
   result_range getOpResults() { return getResults(); }
   OpResult getOpResult(unsigned idx) { return getResult(idx); }
@@ -293,7 +296,7 @@ class alignas(8) Operation final
   using result_type_range = result_range::type_range;
   result_type_iterator result_type_begin() { return getResultTypes().begin(); }
   result_type_iterator result_type_end() { return getResultTypes().end(); }
-  result_type_range getResultTypes();
+  result_type_range getResultTypes() { return getResults().getTypes(); }
 
   //===--------------------------------------------------------------------===//
   // Attributes
@@ -620,7 +623,7 @@ class alignas(8) Operation final
   bool hasValidOrder() { return orderIndex != kInvalidOrderIdx; }
 
 private:
-  Operation(Location location, OperationName name, TypeRange resultTypes,
+  Operation(Location location, OperationName name, unsigned numResults,
             unsigned numSuccessors, unsigned numRegions,
             DictionaryAttr attributes, bool hasOperandStorage);
 
@@ -630,17 +633,17 @@ class alignas(8) Operation final
 
   /// Returns the additional size necessary for allocating the given objects
   /// before an Operation in-memory.
-  static size_t prefixAllocSize(unsigned numTrailingResults,
+  static size_t prefixAllocSize(unsigned numOutOfLineResults,
                                 unsigned numInlineResults) {
-    return sizeof(detail::TrailingOpResult) * numTrailingResults +
-           sizeof(detail::InLineOpResult) * numInlineResults;
+    return sizeof(detail::OutOfLineOpResult) * numOutOfLineResults +
+           sizeof(detail::InlineOpResult) * numInlineResults;
   }
   /// Returns the additional size allocated before this Operation in-memory.
   size_t prefixAllocSize() {
     unsigned numResults = getNumResults();
-    unsigned numTrailingResults = OpResult::getNumTrailing(numResults);
+    unsigned numOutOfLineResults = OpResult::getNumTrailing(numResults);
     unsigned numInlineResults = OpResult::getNumInline(numResults);
-    return prefixAllocSize(numTrailingResults, numInlineResults);
+    return prefixAllocSize(numOutOfLineResults, numInlineResults);
   }
 
   /// Returns the operand storage object.
@@ -649,20 +652,29 @@ class alignas(8) Operation final
     return *getTrailingObjects<detail::OperandStorage>();
   }
 
-  /// Returns a pointer to the use list for the given trailing result.
-  detail::TrailingOpResult *getTrailingResult(unsigned resultNumber) {
-    // Trailing results are stored in reverse order after(before in memory) the
-    // inline results.
-    return reinterpret_cast<detail::TrailingOpResult *>(
-               getInlineResult(OpResult::getMaxInlineResults() - 1)) -
+  /// Returns a pointer to the use list for the given out-of-line result.
+  detail::OutOfLineOpResult *getOutOfLineOpResult(unsigned resultNumber) {
+    // Out-of-line results are stored in reverse order after (before in memory)
+    // the inline results.
+    return reinterpret_cast<detail::OutOfLineOpResult *>(getInlineOpResult(
+               detail::OpResultImpl::getMaxInlineResults() - 1)) -
            ++resultNumber;
   }
 
   /// Returns a pointer to the use list for the given inline result.
-  detail::InLineOpResult *getInlineResult(unsigned resultNumber) {
+  detail::InlineOpResult *getInlineOpResult(unsigned resultNumber) {
     // Inline results are stored in reverse order before the operation in
     // memory.
-    return reinterpret_cast<detail::InLineOpResult *>(this) - ++resultNumber;
+    return reinterpret_cast<detail::InlineOpResult *>(this) - ++resultNumber;
+  }
+
+  /// Returns a pointer to the use list for the given result, which may be
+  /// either inline or out-of-line.
+  detail::OpResultImpl *getOpResultImpl(unsigned resultNumber) {
+    unsigned maxInlineResults = detail::OpResultImpl::getMaxInlineResults();
+    if (resultNumber < maxInlineResults)
+      return getInlineOpResult(resultNumber);
+    return getOutOfLineOpResult(resultNumber - maxInlineResults);
   }
 
   /// Provide a 'getParent' method for ilist_node_with_parent methods.
@@ -683,24 +695,15 @@ class alignas(8) Operation final
   /// O(1) local dominance checks between operations.
   mutable unsigned orderIndex = 0;
 
+  const unsigned numResults;
   const unsigned numSuccs;
-  const unsigned numRegions : 30;
+  const unsigned numRegions : 31;
 
   /// This bit signals whether this operation has an operand storage or not. The
   /// operand storage may be elided for operations that are known to never have
   /// operands.
   bool hasOperandStorage : 1;
 
-  /// This holds the result types of the operation. There are three different
-  /// states recorded here:
-  /// - 0 results : The type below is null.
-  /// - 1 result  : The single result type is held here.
-  /// - N results : The type here is a tuple holding the result types.
-  /// Note: We steal a bit for 'hasSingleResult' from somewhere else so that we
-  /// can use 'resultType` in an ArrayRef<Type>.
-  bool hasSingleResult : 1;
-  Type resultType;
-
   /// This holds the name of the operation.
   OperationName name;
 

mlir/include/mlir/IR/OperationSupport.h

@@ -555,36 +555,6 @@ class OperandStorage final
 };
 } // end namespace detail
 
-//===----------------------------------------------------------------------===//
-// ResultStorage
-//===----------------------------------------------------------------------===//
-
-namespace detail {
-/// 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 alignas(8) 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 alignas(8) 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
 //===----------------------------------------------------------------------===//
@@ -757,63 +727,52 @@ class MutableOperandRange {
 
 /// This class implements the result iterators for the Operation class.
 class ResultRange final
-    : public llvm::indexed_accessor_range<ResultRange, Operation *, OpResult,
-                                          OpResult, OpResult> {
+    : public llvm::detail::indexed_accessor_range_base<
+          ResultRange, detail::OpResultImpl *, OpResult, OpResult, OpResult> {
 public:
-  using indexed_accessor_range<ResultRange, Operation *, OpResult, OpResult,
-                               OpResult>::indexed_accessor_range;
-  ResultRange(Operation *op);
+  using RangeBaseT::RangeBaseT;
 
   /// Returns the types of the values within this range.
-  using type_iterator = ArrayRef<Type>::iterator;
-  using type_range = ArrayRef<Type>;
-  type_range getTypes() const;
+  using type_iterator = ValueTypeIterator<iterator>;
+  using type_range = ValueTypeRange<ResultRange>;
+  type_range getTypes() const { return {begin(), end()}; }
   auto getType() const { return getTypes(); }
 
 private:
-  /// See `llvm::indexed_accessor_range` for details.
-  static OpResult dereference(Operation *op, ptrdiff_t index);
+  /// See `llvm::detail::indexed_accessor_range_base` for details.
+  static detail::OpResultImpl *offset_base(detail::OpResultImpl *object,
+                                           ptrdiff_t index) {
+    return object->getNextResultAtOffset(index);
+  }
+  /// See `llvm::detail::indexed_accessor_range_base` for details.
+  static OpResult dereference_iterator(detail::OpResultImpl *object,
+                                       ptrdiff_t index) {
+    return offset_base(object, index);
+  }
 
-  /// Allow access to `dereference_iterator`.
-  friend llvm::indexed_accessor_range<ResultRange, Operation *, OpResult,
-                                      OpResult, OpResult>;
+  /// Allow access to `offset_base` and `dereference_iterator`.
+  friend RangeBaseT;
 };
 
 //===----------------------------------------------------------------------===//
 // ValueRange
 
-namespace detail {
-/// The type representing the owner of a ValueRange. This is either a list of
-/// values, operands, or an Operation+start index for results.
-struct ValueRangeOwner {
-  ValueRangeOwner(const Value *owner) : ptr(owner), startIndex(0) {}
-  ValueRangeOwner(OpOperand *owner) : ptr(owner), startIndex(0) {}
-  ValueRangeOwner(Operation *owner, unsigned startIndex)
-      : ptr(owner), startIndex(startIndex) {}
-  bool operator==(const ValueRangeOwner &rhs) const { return ptr == rhs.ptr; }
-
-  /// The owner pointer of the range. The owner has represents three distinct
-  /// states:
-  /// const Value *: The owner is the base to a contiguous array of Value.
-  /// OpOperand *  : The owner is the base to a contiguous array of operands.
-  /// void*        : This owner is an Operation*. It is marked as void* here
-  ///                because the definition of Operation is not visible here.
-  PointerUnion<const Value *, OpOperand *, void *> ptr;
-
-  /// Ths start index into the range. This is only used for Operation* owners.
-  unsigned startIndex;
-};
-} // end namespace detail
-
 /// This class provides an abstraction over the different types of ranges over
 /// Values. In many cases, this prevents the need to explicitly materialize a
 /// SmallVector/std::vector. This class should be used in places that are not
 /// suitable for a more derived type (e.g. ArrayRef) or a template range
 /// parameter.
 class ValueRange final
     : public llvm::detail::indexed_accessor_range_base<
-          ValueRange, detail::ValueRangeOwner, Value, Value, Value> {
+          ValueRange,
+          PointerUnion<const Value *, OpOperand *, detail::OpResultImpl *>,
+          Value, Value, Value> {
 public:
+  /// The type representing the owner of a ValueRange. This is either a list of
+  /// values, operands, or results.
+  using OwnerT =
+      PointerUnion<const Value *, OpOperand *, detail::OpResultImpl *>;
+
   using RangeBaseT::RangeBaseT;
 
   template <typename Arg,
@@ -841,8 +800,6 @@ class ValueRange final
   auto getType() const { return getTypes(); }
 
 private:
-  using OwnerT = detail::ValueRangeOwner;
-
   /// See `llvm::detail::indexed_accessor_range_base` for details.
   static OwnerT offset_base(const OwnerT &owner, ptrdiff_t index);
   /// See `llvm::detail::indexed_accessor_range_base` for details.

mlir/include/mlir/IR/TypeRange.h

@@ -34,11 +34,11 @@ class ValueTypeRange;
 /// a SmallVector/std::vector. This class should be used in places that are not
 /// suitable for a more derived type (e.g. ArrayRef) or a template range
 /// parameter.
-class TypeRange
-    : public llvm::detail::indexed_accessor_range_base<
-          TypeRange,
-          llvm::PointerUnion<const Value *, const Type *, OpOperand *>, Type,
-          Type, Type> {
+class TypeRange : public llvm::detail::indexed_accessor_range_base<
+                      TypeRange,
+                      llvm::PointerUnion<const Value *, const Type *,
+                                         OpOperand *, detail::OpResultImpl *>,
+                      Type, Type, Type> {
 public:
   using RangeBaseT::RangeBaseT;
   TypeRange(ArrayRef<Type> types = llvm::None);
@@ -64,7 +64,9 @@ class TypeRange
   /// * A pointer to the first element of an array of values.
   /// * A pointer to the first element of an array of types.
   /// * A pointer to the first element of an array of operands.
-  using OwnerT = llvm::PointerUnion<const Value *, const Type *, OpOperand *>;
+  /// * A pointer to the first element of an array of results.
+  using OwnerT = llvm::PointerUnion<const Value *, const Type *, OpOperand *,
+                                    detail::OpResultImpl *>;
 
   /// See `llvm::detail::indexed_accessor_range_base` for details.
   static OwnerT offset_base(OwnerT object, ptrdiff_t index);
@@ -111,6 +113,18 @@ class ValueTypeRange final
   template <typename Container>
   ValueTypeRange(Container &&c) : ValueTypeRange(c.begin(), c.end()) {}
 
+  /// Return the type at the given index.
+  Type operator[](size_t index) const {
+    assert(index < size() && "invalid index into type range");
+    return *(this->begin() + index);
+  }
+
+  /// Return the size of this range.
+  size_t size() const { return llvm::size(*this); }
+
+  /// Return first type in the range.
+  Type front() { return (*this)[0]; }
+
   /// Compare this range with another.
   template <typename OtherT>
   bool operator==(const OtherT &other) const {

mlir/include/mlir/IR/TypeUtilities.h

@@ -57,8 +57,7 @@ LogicalResult verifyCompatibleShape(Type type1, Type type2);
 
 /// Returns success if the given two arrays have the same number of elements and
 /// each pair wise entries have compatible shape.
-LogicalResult verifyCompatibleShapes(ArrayRef<Type> types1,
-                                     ArrayRef<Type> types2);
+LogicalResult verifyCompatibleShapes(TypeRange types1, TypeRange types2);
 
 //===----------------------------------------------------------------------===//
 // Utility Iterators