diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h --- a/mlir/include/mlir/IR/Operation.h +++ b/mlir/include/mlir/IR/Operation.h @@ -21,15 +21,14 @@ #include "llvm/ADT/Twine.h" namespace mlir { -/// Operation is a basic unit of execution within a function. Operations can -/// be nested within other operations effectively forming a tree. Child -/// operations are organized into operation blocks represented by a 'Block' -/// class. -class Operation final +/// 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, - private llvm::TrailingObjects { + private llvm::TrailingObjects { public: /// Create a new Operation with the specific fields. static Operation *create(Location location, OperationName name, @@ -654,6 +653,21 @@ // allocated with malloc. ~Operation(); + /// Returns the additional size necessary for allocating the given objects + /// before an Operation in-memory. + static size_t prefixAllocSize(unsigned numTrailingResults, + unsigned numInlineResults) { + return sizeof(detail::TrailingOpResult) * numTrailingResults + + sizeof(detail::InLineOpResult) * numInlineResults; + } + /// Returns the additional size allocated before this Operation in-memory. + size_t prefixAllocSize() { + unsigned numResults = getNumResults(); + unsigned numTrailingResults = OpResult::getNumTrailing(numResults); + unsigned numInlineResults = OpResult::getNumInline(numResults); + return prefixAllocSize(numTrailingResults, numInlineResults); + } + /// Returns the operand storage object. detail::OperandStorage &getOperandStorage() { assert(hasOperandStorage && "expected operation to have operand storage"); @@ -662,12 +676,18 @@ /// Returns a pointer to the use list for the given trailing result. detail::TrailingOpResult *getTrailingResult(unsigned resultNumber) { - return getTrailingObjects() + resultNumber; + // Trailing results are stored in reverse order after(before in memory) the + // inline results. + return reinterpret_cast( + getInlineResult(OpResult::getMaxInlineResults() - 1)) - + ++resultNumber; } /// Returns a pointer to the use list for the given inline result. detail::InLineOpResult *getInlineResult(unsigned resultNumber) { - return getTrailingObjects() + resultNumber; + // Inline results are stored in reverse order before the operation in + // memory. + return reinterpret_cast(this) - ++resultNumber; } /// Provide a 'getParent' method for ilist_node_with_parent methods. @@ -725,17 +745,8 @@ friend class llvm::ilist_node_with_parent; // This stuff is used by the TrailingObjects template. - friend llvm::TrailingObjects; - size_t numTrailingObjects(OverloadToken) const { - return OpResult::getNumInline( - const_cast(this)->getNumResults()); - } - size_t numTrailingObjects(OverloadToken) const { - return OpResult::getNumTrailing( - const_cast(this)->getNumResults()); - } size_t numTrailingObjects(OverloadToken) const { return numSuccs; } 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 @@ -570,11 +570,11 @@ /// This class provides the implementation for an in-line operation result. This /// is an operation result whose number can be stored inline inside of the bits /// of an Operation*. -struct InLineOpResult : public IRObjectWithUseList {}; +struct alignas(8) InLineOpResult : public IRObjectWithUseList {}; /// This class provides the implementation for an out-of-line operation result. /// This is an operation result whose number cannot be stored inline inside of /// the bits of an Operation*. -struct TrailingOpResult : public IRObjectWithUseList { +struct alignas(8) TrailingOpResult : public IRObjectWithUseList { TrailingOpResult(uint64_t trailingResultNumber) : trailingResultNumber(trailingResultNumber) {} diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp --- a/mlir/lib/IR/Operation.cpp +++ b/mlir/lib/IR/Operation.cpp @@ -119,16 +119,18 @@ needsOperandStorage = !abstractOp->hasTrait(); } - // Compute the byte size for the operation and the operand storage. - auto byteSize = - totalSizeToAlloc( - numInlineResults, numTrailingResults, numSuccessors, numRegions, - needsOperandStorage ? 1 : 0); - byteSize += - llvm::alignTo(detail::OperandStorage::additionalAllocSize(numOperands), - alignof(Operation)); - void *rawMem = malloc(byteSize); + // 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( + numSuccessors, numRegions, needsOperandStorage ? 1 : 0) + + detail::OperandStorage::additionalAllocSize(numOperands); + size_t prefixByteSize = llvm::alignTo( + Operation::prefixAllocSize(numTrailingResults, numInlineResults), + alignof(Operation)); + char *mallocMem = reinterpret_cast(malloc(byteSize + prefixByteSize)); + void *rawMem = mallocMem + prefixByteSize; // Create the new Operation. Operation *op = @@ -200,8 +202,12 @@ /// Destroy this operation or one of its subclasses. void Operation::destroy() { + // Operations may have additional prefixed allocation, which needs to be + // accounted for here when computing the address to free. + char *rawMem = reinterpret_cast(this) - + llvm::alignTo(prefixAllocSize(), alignof(Operation)); this->~Operation(); - free(this); + free(rawMem); } /// Return the context this operation is associated with. 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 @@ -382,14 +382,15 @@ /// Returns the parent operation of this trailing result. Operation *detail::TrailingOpResult::getOwner() { - // We need to do some arithmetic to get the operation pointer. Move the - // trailing owner to the start of the array. - TrailingOpResult *trailingIt = this - trailingResultNumber; + // We need to do some arithmetic to get the operation pointer. Trailing + // results are stored in reverse order before the inline results of the + // operation, so move the trailing owner up to the start of the array. + TrailingOpResult *trailingIt = this + (trailingResultNumber + 1); // Move the owner past the inline op results to get to the operation. - auto *inlineResultIt = reinterpret_cast(trailingIt) - + auto *inlineResultIt = reinterpret_cast(trailingIt) + OpResult::getMaxInlineResults(); - return reinterpret_cast(inlineResultIt) - 1; + return reinterpret_cast(inlineResultIt); } //===----------------------------------------------------------------------===//