Remove use of tuple for multiresult type storage

Move the results in line with the op instead. This results in each
operation having its own types recorded vs single tuple type, but comes
at benefit that every mutation doesn't incurs uniquing. Ran into cases
where updating result type of operation led to very large memory usage.

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

3 files changed, 45 insertions, 39 deletions

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 63a94d9..5cc1639 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -27,7 +27,7 @@ namespace mlir {
 /// 'Block' class.
 class alignas(8) Operation final
     : public llvm::ilist_node_with_parent<Operation, Block>,
-      private llvm::TrailingObjects<Operation, BlockOperand, Region,
+      private llvm::TrailingObjects<Operation, BlockOperand, Region, Type,
                                     detail::OperandStorage> {
 public:
   /// Create a new Operation with the specific fields.
@@ -651,7 +651,7 @@ private:
 
   /// 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
+    // Trailing results are stored in reverse order after (before in memory) the
     // inline results.
     return reinterpret_cast<detail::TrailingOpResult *>(
                getInlineResult(OpResult::getMaxInlineResults() - 1)) -
@@ -695,11 +695,18 @@ private:
   /// 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>.
+  /// - N results : The type here is empty and instead the result types are held
+  ///               in trailing storage.
   bool hasSingleResult : 1;
-  Type resultType;
+
+  /// Union representing either the Type of a single result operation (if
+  /// hasSingleResult) or the number of result types for multi-result.
+  union {
+    // Type, set if single result Operation.
+    Type type = nullptr;
+    // Size, set if not a single result Operation.
+    unsigned size;
+  } resultTypeOrSize;
 
   /// This holds the name of the operation.
   OperationName name;
@@ -720,12 +727,15 @@ private:
   friend class llvm::ilist_node_with_parent<Operation, Block>;
 
   // This stuff is used by the TrailingObjects template.
-  friend llvm::TrailingObjects<Operation, BlockOperand, Region,
+  friend llvm::TrailingObjects<Operation, BlockOperand, Region, Type,
                                detail::OperandStorage>;
   size_t numTrailingObjects(OverloadToken<BlockOperand>) const {
     return numSuccs;
   }
   size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }
+  size_t numTrailingObjects(OverloadToken<Type>) const {
+    return hasSingleResult ? 0 : resultTypeOrSize.size;
+  }
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Operation &op) {
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index b1f2ad6..b074a67 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -123,8 +123,11 @@ Operation *Operation::create(Location location, OperationName name,
   // into account the size of the operation, its trailing objects, and its
   // prefixed objects.
   size_t byteSize =
-      totalSizeToAlloc<BlockOperand, Region, detail::OperandStorage>(
-          numSuccessors, numRegions, needsOperandStorage ? 1 : 0) +
+      totalSizeToAlloc<BlockOperand, Region, Type, detail::OperandStorage>(
+          numSuccessors, numRegions,
+          // Result type storage only needed if there is not 0 or 1 results.
+          resultTypes.size() == 1 ? 0 : resultTypes.size(),
+          needsOperandStorage ? 1 : 0) +
       detail::OperandStorage::additionalAllocSize(numOperands);
   size_t prefixByteSize = llvm::alignTo(
       Operation::prefixAllocSize(numTrailingResults, numInlineResults),
@@ -172,13 +175,18 @@ Operation::Operation(Location location, OperationName name,
   assert(attributes && "unexpected null attribute dictionary");
   assert(llvm::all_of(resultTypes, [](Type t) { return t; }) &&
          "unexpected null result type");
-  if (!resultTypes.empty()) {
-    // If there is a single result it is stored in-place, otherwise use a tuple.
+  if (resultTypes.empty()) {
+    resultTypeOrSize.size = 0;
+  } else {
+    // If there is a single result it is stored in-place, otherwise use trailing
+    // type storage.
     hasSingleResult = resultTypes.size() == 1;
-    if (hasSingleResult)
-      resultType = resultTypes.front();
-    else
-      resultType = TupleType::get(location->getContext(), resultTypes);
+    if (hasSingleResult) {
+      resultTypeOrSize.type = resultTypes.front();
+    } else {
+      resultTypeOrSize.size = resultTypes.size();
+      llvm::copy(resultTypes, getTrailingObjects<Type>());
+    }
   }
 }
 
@@ -545,17 +553,15 @@ void Operation::dropAllDefinedValueUses() {
 
 /// Return the number of results held by this operation.
 unsigned Operation::getNumResults() {
-  if (!resultType)
-    return 0;
-  return hasSingleResult ? 1 : resultType.cast<TupleType>().size();
+  if (hasSingleResult)
+    return 1;
+  return resultTypeOrSize.size;
 }
 
 auto Operation::getResultTypes() -> result_type_range {
-  if (!resultType)
-    return llvm::None;
   if (hasSingleResult)
-    return resultType;
-  return resultType.cast<TupleType>().getTypes();
+    return resultTypeOrSize.type;
+  return ArrayRef<Type>(getTrailingObjects<Type>(), resultTypeOrSize.size);
 }
 
 void Operation::setSuccessor(Block *block, unsigned index) {
diff --git a/mlir/lib/IR/Value.cpp b/mlir/lib/IR/Value.cpp
index fd7e5b5..c6d3e58 100644
--- a/mlir/lib/IR/Value.cpp
+++ b/mlir/lib/IR/Value.cpp
@@ -39,31 +39,21 @@ Type Value::getType() const {
   OpResult result = cast<OpResult>();
   Operation *owner = result.getOwner();
   if (owner->hasSingleResult)
-    return owner->resultType;
-  return owner->resultType.cast<TupleType>().getType(result.getResultNumber());
+    return owner->resultTypeOrSize.type;
+  return owner->getResultTypes()[result.getResultNumber()];
 }
 
 /// Mutate the type of this Value to be of the specified type.
 void Value::setType(Type newType) {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.setType(newType);
-  OpResult result = cast<OpResult>();
 
-  // If the owner has a single result, simply update it directly.
+  OpResult result = cast<OpResult>();
   Operation *owner = result.getOwner();
-  if (owner->hasSingleResult) {
-    owner->resultType = newType;
-    return;
-  }
-  unsigned resultNo = result.getResultNumber();
-
-  // Otherwise, rebuild the tuple if the new type is different from the current.
-  auto curTypes = owner->resultType.cast<TupleType>().getTypes();
-  if (curTypes[resultNo] == newType)
-    return;
-  auto newTypes = llvm::to_vector<4>(curTypes);
-  newTypes[resultNo] = newType;
-  owner->resultType = TupleType::get(newType.getContext(), newTypes);
+  if (owner->hasSingleResult)
+    owner->resultTypeOrSize.type = newType;
+  else
+    owner->getTrailingObjects<Type>()[result.getResultNumber()] = newType;
 }
 
 /// If this value is the result of an Operation, return the operation that