diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h --- a/mlir/include/mlir/IR/OperationSupport.h +++ b/mlir/include/mlir/IR/OperationSupport.h @@ -347,37 +347,27 @@ /// 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 in a dynamically allocated block. + OpOperand *firstOp; - /// A pointer to the first operand element. This is either to the trailing - /// objects storage, or a dynamically allocated block of memory. - llvm::PointerIntPair firstOpAndIsDynamic; - - // 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; }; @@ -387,25 +377,20 @@ /// array. The second is that being able to dynamically resize the operand list /// is optional. class OperandStorage final - : private llvm::TrailingObjects { + : private llvm::TrailingObjects { public: OperandStorage(unsigned numOperands, bool resizable) - : numOperands(numOperands), resizable(resizable) { - // Initialize the resizable storage. - if (resizable) { - new (&getResizableStorage()) - ResizableStorage(getTrailingObjects(), 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(); } @@ -426,8 +411,10 @@ /// Returns the additional size necessary for allocating this object. static size_t additionalAllocSize(unsigned numOperands, bool resizable) { - return additionalSizeToAlloc(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(numOperands), + sizeof(ResizableStorage)); } /// Returns if this storage is resizable. @@ -439,30 +426,34 @@ /// Returns the current pointer for the raw operands array. OpOperand *getRawOperands() { - return resizable ? getResizableStorage().getPointer() - : getTrailingObjects(); + return storageIsDynamic ? getResizableStorage().firstOp + : getTrailingObjects(); } /// Returns the resizable operand utility class. ResizableStorage &getResizableStorage() { - assert(resizable); - return *getTrailingObjects(); + // We represent the resizable storage in the same location as the first + // operand. + assert(storageIsDynamic); + return *reinterpret_cast( + getTrailingObjects()); } /// 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; - size_t numTrailingObjects(OverloadToken) 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; }; } // end namespace detail diff --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp --- a/mlir/lib/IR/OperationSupport.cpp +++ b/mlir/lib/IR/OperationSupport.cpp @@ -86,13 +86,25 @@ // Otherwise, we need to be resizable. assert(resizable && "Only resizable operations may add operands"); - // Grow the capacity if necessary. - auto &resizeUtil = getResizableStorage(); - if (resizeUtil.capacity < operands.size()) - grow(resizeUtil, operands.size()); + // 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)