Remove use of tuple for multi-result operation result type storage

[jpienaar]

Summary

Move the result types to trailing type op storage instead. This
results in each operation having its own types recorded vs deduped tuple type,
but comes at benefit that every mutation doesn't incurs uniquing. The latter
ran into cases whereupdating result type of operation led to very large memory
usage.

[mehdi_amini]

mlir/lib/IR/Operation.cpp

560–561

I don't think this is legal: you can only access the element of a union through the member it has been set IIRC.

[jpienaar]

Only access union element that was set for case

[jpienaar]

I know River had a better/alternate approach in mind here too

mlir/lib/IR/Operation.cpp

560–561

Good, point, had kept it to close to the original logic and didn't think about that. Updated here and below.

[NathanielMcVicar]

It doesn't look like MSVC likes this use of type. Break in the buildbot here: https://lab.llvm.org/buildbot/#/builders/13/builds/5148. Thanks!

FAILED: tools/mlir/lib/IR/CMakeFiles/obj.MLIRIR.dir/Operation.cpp.obj 
...
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\IR\Operation.cpp(174): error C2280: 'mlir::Operation::<unnamed-type-resultTypeOrSize>::<unnamed-type-resultTypeOrSize>(void)': attempting to reference a deleted function
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/IR/Operation.h(709): note: compiler has generated 'mlir::Operation::<unnamed-type-resultTypeOrSize>::<unnamed-type-resultTypeOrSize>' here
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/IR/Operation.h(709): note: 'mlir::Operation::<unnamed-type-resultTypeOrSize>::<unnamed-type-resultTypeOrSize>(void)': function was implicitly deleted because 'mlir::Operation::<unnamed-type-resultTypeOrSize>' has a variant data member 'mlir::Operation::<unnamed-type-resultTypeOrSize>::type' with a non-trivial default constructor
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/IR/Operation.h(706): note: see declaration of 'mlir::Operation::<unnamed-type-resultTypeOrSize>::type'

[jpienaar]

Thanks, reverted,

• smime.p7s3 KBDownload

mlir/include/mlir/IR/Operation.h

 
 namespace mlir {
 /// Operation is a basic unit of execution within MLIR. Operations can
 /// be nested within `Region`s held by other operations effectively forming a
 /// tree. Child operations are organized into operation blocks represented by a
 /// '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.
   static Operation *create(Location location, OperationName name,
                            TypeRange resultTypes, ValueRange operands,
                            ArrayRef<NamedAttribute> attributes,
                            BlockRange successors, unsigned numRegions);
 
private:
   /// Returns the operand storage object.
   detail::OperandStorage &getOperandStorage() {
     assert(hasOperandStorage && "expected operation to have operand storage");
     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
+    // Trailing results are stored in reverse order after (before in memory) the
     // inline results.
     return reinterpret_cast<detail::TrailingOpResult *>(
                getInlineResult(OpResult::getMaxInlineResults() - 1)) -
            ++resultNumber;
   }
 
   /// Returns a pointer to the use list for the given inline result.
   detail::InLineOpResult *getInlineResult(unsigned resultNumber) {
private:
   /// 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>.
+  /// - 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;
 
   /// This holds general named attributes for the operation.
   DictionaryAttr 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 '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, 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) {
   op.print(os, OpPrintingFlags().useLocalScope());
   return os;
 }
 
 } // end namespace mlir

mlir/lib/IR/Operation.cpp

if (operands.empty()) {
     if (const AbstractOperation *abstractOp = name.getAbstractOperation())
       needsOperandStorage = !abstractOp->hasTrait<OpTrait::ZeroOperands>();
   }
 
   // Compute the byte size for the operation and the operand storage. This takes
   // 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),
       alignof(Operation));
   char *mallocMem = reinterpret_cast<char *>(malloc(byteSize + prefixByteSize));
   void *rawMem = mallocMem + prefixByteSize;
 
   // Create the new Operation.
Operation::Operation(Location location, OperationName name,
                      unsigned numRegions, DictionaryAttr attributes,
                      bool hasOperandStorage)
     : location(location), numSuccs(numSuccessors), numRegions(numRegions),
       hasOperandStorage(hasOperandStorage), hasSingleResult(false), name(name),
       attrs(attributes) {
   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>());
+    }
   }
 }
 
 // Operations are deleted through the destroy() member because they are
 // allocated via malloc.
 Operation::~Operation() {
   assert(block == nullptr && "operation destroyed but still in a block");
 
void Operation::dropAllDefinedValueUses() {
   dropAllUses();
 
   for (auto &region : getRegions())
     for (auto &block : region)
       block.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) {
   assert(index < getNumSuccessors());
   getBlockOperands()[index].set(block);
 }
 
 /// Attempt to fold this operation using the Op's registered foldHook.

mlir/lib/IR/Value.cpp

 Type Value::getType() const {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getType();
 
   // If this is an operation result, query the parent operation.
   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
 /// defines it.
 Operation *Value::getDefiningOp() const {
   if (auto result = dyn_cast<OpResult>())
     return result.getOwner();
   return nullptr;