Revert "Remove use of tuple for multiresult type storage"

This reverts commit 08f0764.

Author: jpienaar

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, Type,
+      private llvm::TrailingObjects<Operation, BlockOperand, Region,
                                     detail::OperandStorage> {
 public:
   /// Create a new Operation with the specific fields.
@@ -651,7 +651,7 @@ class alignas(8) Operation final
 
   /// 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,18 +695,11 @@ class alignas(8) Operation final
   /// 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 empty and instead the result types are held
-  ///               in trailing storage.
+  /// - 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;
-
-  /// 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;
+  Type resultType;
 
   /// This holds the name of the operation.
   OperationName name;
@@ -727,15 +720,12 @@ class alignas(8) Operation final
   friend class llvm::ilist_node_with_parent<Operation, Block>;
 
   // This stuff is used by the TrailingObjects template.
-  friend llvm::TrailingObjects<Operation, BlockOperand, Region, Type,
+  friend llvm::TrailingObjects<Operation, BlockOperand, Region,
                                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) {

mlir/lib/IR/Operation.cpp

@@ -123,11 +123,8 @@ 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, 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) +
+      totalSizeToAlloc<BlockOperand, Region, detail::OperandStorage>(
+          numSuccessors, numRegions, needsOperandStorage ? 1 : 0) +
       detail::OperandStorage::additionalAllocSize(numOperands);
   size_t prefixByteSize = llvm::alignTo(
       Operation::prefixAllocSize(numTrailingResults, numInlineResults),
@@ -175,18 +172,13 @@ 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()) {
-    resultTypeOrSize.size = 0;
-  } else {
-    // If there is a single result it is stored in-place, otherwise use trailing
-    // type storage.
+  if (!resultTypes.empty()) {
+    // If there is a single result it is stored in-place, otherwise use a tuple.
     hasSingleResult = resultTypes.size() == 1;
-    if (hasSingleResult) {
-      resultTypeOrSize.type = resultTypes.front();
-    } else {
-      resultTypeOrSize.size = resultTypes.size();
-      llvm::copy(resultTypes, getTrailingObjects<Type>());
-    }
+    if (hasSingleResult)
+      resultType = resultTypes.front();
+    else
+      resultType = TupleType::get(location->getContext(), resultTypes);
   }
 }
 
@@ -553,15 +545,17 @@ void Operation::dropAllDefinedValueUses() {
 
 /// Return the number of results held by this operation.
 unsigned Operation::getNumResults() {
-  if (hasSingleResult)
-    return 1;
-  return resultTypeOrSize.size;
+  if (!resultType)
+    return 0;
+  return hasSingleResult ? 1 : resultType.cast<TupleType>().size();
 }
 
 auto Operation::getResultTypes() -> result_type_range {
+  if (!resultType)
+    return llvm::None;
   if (hasSingleResult)
-    return resultTypeOrSize.type;
-  return ArrayRef<Type>(getTrailingObjects<Type>(), resultTypeOrSize.size);
+    return resultType;
+  return resultType.cast<TupleType>().getTypes();
 }
 
 void Operation::setSuccessor(Block *block, unsigned index) {

mlir/lib/IR/Value.cpp

@@ -39,21 +39,31 @@ Type Value::getType() const {
   OpResult result = cast<OpResult>();
   Operation *owner = result.getOwner();
   if (owner->hasSingleResult)
-    return owner->resultTypeOrSize.type;
-  return owner->getResultTypes()[result.getResultNumber()];
+    return owner->resultType;
+  return owner->resultType.cast<TupleType>().getType(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.
   Operation *owner = result.getOwner();
-  if (owner->hasSingleResult)
-    owner->resultTypeOrSize.type = newType;
-  else
-    owner->getTrailingObjects<Type>()[result.getResultNumber()] = newType;
+  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 this value is the result of an Operation, return the operation that