[mlir] Optimize the allocation of resizable operand lists

[rriddle]

Summary

This revision optimizes resizable operand lists by allocating them in the same location as the trailing operands. This has numerous benefits:

• If the operation has at least one operand at construction time, there is 0 additional memory overhead to the resizable storage.
• There is less pointer arithmetic necessary as the resizable storage is now only used when the operands are dynamically allocated.

[mehdi_amini]

Why do even keep the resizable=false possibility instead of nuking this and making it always resizable? Is there a cost to it after this change?

[mehdi_amini]

I have some questions about possible follow-up improvements, but LGTM as is!

mlir/include/mlir/IR/OperationSupport.h

 //===----------------------------------------------------------------------===//
 // OperandStorage
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// A utility class holding the information necessary to dynamically resize
 /// operands.
 struct ResizableStorage {
-  ResizableStorage(OpOperand *opBegin, unsigned numOperands)
-      : firstOpAndIsDynamic(opBegin, false), capacity(numOperands) {}
-
+  ResizableStorage() : firstOp(nullptr), capacity(0) {}
+  ResizableStorage(ResizableStorage &&other)
+      : firstOp(other.firstOp), capacity(other.capacity) {
+    other.firstOp = nullptr;
+  }
   ~ResizableStorage() { cleanupStorage(); }
 
   /// Cleanup any allocated storage.
-  void cleanupStorage() {
-    // If the storage is dynamic, then we need to free the storage.
-    if (isStorageDynamic())
-      free(firstOpAndIsDynamic.getPointer());
-  }
+  void cleanupStorage() { free(firstOp); }
 
   /// Sets the storage pointer to a new dynamically allocated block.
   void setDynamicStorage(OpOperand *opBegin) {
     /// Cleanup the old storage if necessary.
     cleanupStorage();
-    firstOpAndIsDynamic.setPointerAndInt(opBegin, true);
+    firstOp = opBegin;
   }
 
-  /// Returns the current storage pointer.
-  OpOperand *getPointer() { return firstOpAndIsDynamic.getPointer(); }
-
-  /// Returns if the current storage of operands is in the trailing objects is
-  /// in a dynamically allocated memory block.
-  bool isStorageDynamic() const { return firstOpAndIsDynamic.getInt(); }
-
-  /// A pointer to the first operand element. This is either to the trailing
-  /// objects storage, or a dynamically allocated block of memory.
-  llvm::PointerIntPair<OpOperand *, 1, bool> firstOpAndIsDynamic;
+  /// A pointer to the first operand element in a dynamically allocated block.
+  OpOperand *firstOp;
 
-  // The maximum number of operands that can be currently held by the storage.
+  /// The maximum number of operands that can be currently held by the storage.
   unsigned capacity;
 };
 
 /// This class handles the management of operation operands. Operands are
 /// stored similarly to the elements of a SmallVector except for two key
 /// differences. The first is the inline storage, which is a trailing objects
 /// array. The second is that being able to dynamically resize the operand list
 /// is optional.
 class OperandStorage final
-    : private llvm::TrailingObjects<OperandStorage, ResizableStorage,
-                                    OpOperand> {
+    : private llvm::TrailingObjects<OperandStorage, OpOperand> {
 public:
   OperandStorage(unsigned numOperands, bool resizable)
-      : numOperands(numOperands), resizable(resizable) {
-    // Initialize the resizable storage.
-    if (resizable) {
-      new (&getResizableStorage())
-          ResizableStorage(getTrailingObjects<OpOperand>(), numOperands);
-    }
-  }
+      : numOperands(numOperands), resizable(resizable),
+        storageIsDynamic(false) {}
 
   ~OperandStorage() {
     // Manually destruct the operands.
     for (auto &operand : getOperands())
       operand.~OpOperand();
 
-    // If the storage is resizable then destruct the utility.
-    if (resizable)
+    // If the storage is currently in a dynamic allocation, then destruct the
+    // resizable storage.
+    if (storageIsDynamic)
       getResizableStorage().~ResizableStorage();
   }
 
   /// Replace the operands contained in the storage with the ones provided in
   /// 'operands'.
   void setOperands(Operation *owner, ValueRange operands);
 
   /// Erase an operand held by the storage.
   void eraseOperand(unsigned index);
 
   /// Get the operation operands held by the storage.
   MutableArrayRef<OpOperand> getOperands() {
     return {getRawOperands(), size()};
   }
 
   /// Return the number of operands held in the storage.
   unsigned size() const { return numOperands; }
 
   /// Returns the additional size necessary for allocating this object.
   static size_t additionalAllocSize(unsigned numOperands, bool resizable) {
-    return additionalSizeToAlloc<ResizableStorage, OpOperand>(resizable ? 1 : 0,
-                                                              numOperands);
+    // The trailing storage must be able to hold enough space for the number of
+    // operands, or at least the resizable storage if necessary.
+    return std::max(additionalSizeToAlloc<OpOperand>(numOperands),
+                    sizeof(ResizableStorage));
   }
 
   /// Returns if this storage is resizable.
   bool isResizable() const { return resizable; }
 
 private:
   /// Clear the storage and destroy the current operands held by the storage.
   void clear() { numOperands = 0; }
 
   /// Returns the current pointer for the raw operands array.
   OpOperand *getRawOperands() {
-    return resizable ? getResizableStorage().getPointer()
+    return storageIsDynamic ? getResizableStorage().firstOp
                             : getTrailingObjects<OpOperand>();
   }
 
   /// Returns the resizable operand utility class.
   ResizableStorage &getResizableStorage() {
-    assert(resizable);
-    return *getTrailingObjects<ResizableStorage>();
+    // We represent the resizable storage in the same location as the first
+    // operand.
+    assert(storageIsDynamic);
+    return *reinterpret_cast<ResizableStorage *>(
+        getTrailingObjects<OpOperand>());
   }
 
   /// Grow the internal resizable operand storage.
   void grow(ResizableStorage &resizeUtil, size_t minSize);
 
   /// The current number of operands, and the current max operand capacity.
-  unsigned numOperands : 31;
+  unsigned numOperands : 30;
 
   /// Whether this storage is resizable or not.
   bool resizable : 1;
 
-  // This stuff is used by the TrailingObjects template.
-  friend llvm::TrailingObjects<OperandStorage, ResizableStorage, OpOperand>;
-  size_t numTrailingObjects(OverloadToken<ResizableStorage>) const {
-    return resizable ? 1 : 0;
-  }
+  /// If the storage is resizable, this indicates if the operands array is
+  /// currently in the resizable storage or the trailing array.
+  bool storageIsDynamic : 1;
+
+  /// This stuff is used by the TrailingObjects template.
+  friend llvm::TrailingObjects<OperandStorage, OpOperand>;
 };
 } // end namespace detail
 
 //===----------------------------------------------------------------------===//
 // ResultStorage
 //===----------------------------------------------------------------------===//
 
 namespace detail {

mlir/lib/IR/OperationSupport.cpp

if (operands.size() <= numOperands) {
     for (unsigned i = 0; i != numOperands; ++i)
       opOperands[i].set(operands[i]);
     return;
   }
 
   // Otherwise, we need to be resizable.
   assert(resizable && "Only resizable operations may add operands");
 
-  // Grow the capacity if necessary.
+  // If the storage isn't already dynamic, we need to allocate a new buffer for
+  // it.
+  OpOperand *opBegin = nullptr;
+  if (!storageIsDynamic) {
+    // Grow a new storage first to avoid overwriting the existing operands.
+    ResizableStorage newStorage;
+    grow(newStorage, operands.size());
+    opBegin = newStorage.firstOp;
+    storageIsDynamic = true;
+    new (&getResizableStorage()) ResizableStorage(std::move(newStorage));
+  } else {
+    // Otherwise, grow the existing storage if necessary.
     auto &resizeUtil = getResizableStorage();
     if (resizeUtil.capacity < operands.size())
       grow(resizeUtil, operands.size());
+    opBegin = resizeUtil.firstOp;
+  }
 
   // Set the operands.
-  OpOperand *opBegin = getRawOperands();
   for (unsigned i = 0; i != numOperands; ++i)
     opBegin[i].set(operands[i]);
   for (unsigned e = operands.size(); numOperands != e; ++numOperands)
     new (&opBegin[numOperands]) OpOperand(owner, operands[numOperands]);
 }
 
 /// Erase an operand held by the storage.
 void detail::OperandStorage::eraseOperand(unsigned index) {