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 @@ -28,6 +28,7 @@ class alignas(8) Operation final : public llvm::ilist_node_with_parent, private llvm::TrailingObjects { public: /// Create a new Operation with the specific fields. @@ -269,8 +270,18 @@ // Results //===--------------------------------------------------------------------===// - /// Return the number of results held by this operation. - unsigned getNumResults() { return numResults; } + /// Returns true if this operation has results. + bool hasResults() { + // Note: We can always rely on comparing small field counts to zero. + return smallNumResults != 0; + } + + /// Returns the number of results held by this operation. + unsigned getNumResults() { + if (LLVM_UNLIKELY(hasLargeFieldCounts)) + return getLargeFieldCountStorage().numResults; + return smallNumResults; + } /// Get the 'idx'th result of this operation. OpResult getResult(unsigned idx) { return OpResult(getOpResultImpl(idx)); } @@ -282,8 +293,8 @@ result_iterator result_begin() { return getResults().begin(); } result_iterator result_end() { return getResults().end(); } result_range getResults() { - return numResults == 0 ? result_range(nullptr, 0) - : result_range(getInlineOpResult(0), numResults); + return hasResults() ? result_range(getInlineOpResult(0), getNumResults()) + : result_range(nullptr, 0); } result_range getOpResults() { return getResults(); } @@ -414,18 +425,28 @@ // Blocks //===--------------------------------------------------------------------===// + /// Returns true if this operation has any regions. + bool hasRegions() { + // Note: We can always rely on comparing small field counts to zero. + return smallNumRegions != 0; + } /// Returns the number of regions held by this operation. - unsigned getNumRegions() { return numRegions; } + unsigned getNumRegions() { + if (LLVM_UNLIKELY(hasLargeFieldCounts)) + return getLargeFieldCountStorage().numRegions; + return smallNumRegions; + } /// Returns the regions held by this operation. MutableArrayRef getRegions() { - auto *regions = getTrailingObjects(); - return {regions, numRegions}; + if (!hasRegions()) + return MutableArrayRef(); + return {getTrailingObjects(), getNumRegions()}; } /// Returns the region held by this operation at position 'index'. Region &getRegion(unsigned index) { - assert(index < numRegions && "invalid region index"); + assert(index < getNumRegions() && "invalid region index"); return getRegions()[index]; } @@ -433,8 +454,9 @@ // Successors //===--------------------------------------------------------------------===// + /// Returns the block operands held by this operation. MutableArrayRef getBlockOperands() { - return {getTrailingObjects(), numSuccs}; + return {getTrailingObjects(), getNumSuccessors()}; } // Successor iteration. @@ -443,8 +465,17 @@ succ_iterator successor_end() { return getSuccessors().end(); } SuccessorRange getSuccessors() { return SuccessorRange(this); } - bool hasSuccessors() { return numSuccs != 0; } - unsigned getNumSuccessors() { return numSuccs; } + /// Returns true if this operation has any successors. + bool hasSuccessors() { + // Note: We can always rely on comparing small field counts to zero. + return smallNumSuccs != 0; + } + /// Returns the number of successors held by this operation. + unsigned getNumSuccessors() { + if (LLVM_UNLIKELY(hasLargeFieldCounts)) + return getLargeFieldCountStorage().numSuccs; + return smallNumSuccs; + } Block *getSuccessor(unsigned index) { assert(index < getNumSuccessors()); @@ -601,9 +632,10 @@ bool hasValidOrder() { return orderIndex != kInvalidOrderIdx; } private: - Operation(Location location, OperationName name, unsigned numResults, - unsigned numSuccessors, unsigned numRegions, - DictionaryAttr attributes, bool hasOperandStorage); + Operation(Location location, OperationName name, unsigned smallNumResults, + unsigned smallNumSuccs, unsigned smallNumRegions, + DictionaryAttr attributes, bool hasLargeFieldCounts, + bool hasOperandStorage); // Operations are deleted through the destroy() member because they are // allocated with malloc. @@ -630,6 +662,13 @@ return *getTrailingObjects(); } + /// Returns the large field count storage. + detail::LargeOperationFieldCounts &getLargeFieldCountStorage() { + assert(hasLargeFieldCounts && + "expected operation to have large field counts"); + return *getTrailingObjects(); + } + /// Returns a pointer to the use list for the given out-of-line result. detail::OutOfLineOpResult *getOutOfLineOpResult(unsigned resultNumber) { // Out-of-line results are stored in reverse order after (before in memory) @@ -673,9 +712,17 @@ /// O(1) local dominance checks between operations. mutable unsigned orderIndex = 0; - const unsigned numResults; - const unsigned numSuccs; - const unsigned numRegions : 31; + /// Counts for various fields when the `hasLargeFieldCounts` is false. When + /// true, the counts are stored in the trailing `LargeOperationFieldCounts` + /// storage; and all of these fields are assigned to a non-zero value (this + /// allows for efficient empty checks). + uint8_t smallNumRegions; + uint8_t smallNumResults; + uint8_t smallNumSuccs; + + /// A flag indicating if the operation has a large number of elements for one + /// of its fields, i.e. regions, results, or successors. + bool hasLargeFieldCounts : 1; /// This bit signals whether this operation has an operand storage or not. The /// operand storage may be elided for operations that are known to never have @@ -701,15 +748,22 @@ friend class llvm::ilist_node_with_parent; // This stuff is used by the TrailingObjects template. - friend llvm::TrailingObjects; size_t numTrailingObjects(OverloadToken) const { return hasOperandStorage ? 1 : 0; } + size_t + numTrailingObjects(OverloadToken) const { + return hasLargeFieldCounts ? 1 : 0; + } size_t numTrailingObjects(OverloadToken) const { - return numSuccs; + return const_cast(this)->getNumSuccessors(); + } + size_t numTrailingObjects(OverloadToken) const { + return const_cast(this)->getNumRegions(); } - size_t numTrailingObjects(OverloadToken) const { return numRegions; } }; inline raw_ostream &operator<<(raw_ostream &os, const Operation &op) { 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 @@ -568,6 +568,29 @@ }; } // end namespace detail +//===----------------------------------------------------------------------===// +// OperandStorage +//===----------------------------------------------------------------------===// + +namespace detail { +/// This class is used to provide storage for the counts of various Operation +/// fields when those fields contain an abnormally large number of elements. +/// This allows for keeping the size of Operation small in the overwhelmingly +/// common case, but retaining support for large counts. +struct LargeOperationFieldCounts { + LargeOperationFieldCounts(unsigned numRegions, unsigned numResults, + unsigned numSuccs) + : numRegions(numRegions), numResults(numResults), numSuccs(numSuccs) {} + + /// The number of regions held by the operation. + unsigned numRegions; + /// The number of results held by the operation. + unsigned numResults; + /// The number of successors held by the operation. + unsigned numSuccs; +}; +} // end namespace detail + //===----------------------------------------------------------------------===// // OpPrintingFlags //===----------------------------------------------------------------------===// 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 @@ -113,6 +113,18 @@ unsigned numOperands = operands.size(); unsigned numResults = resultTypes.size(); + // Compute the "small" counts for the relevant fields. If any of the counts + // are too big, we need to use the large field storage. In the case of a large + // storage, avoid using 0 for the small counts so that we can retain efficient + // "empty" checks. + unsigned maxNumSmallFields = std::numeric_limits::max(); + bool hasLargeFieldCount = numRegions > maxNumSmallFields || + numResults > maxNumSmallFields || + numSuccessors > maxNumSmallFields; + uint8_t smallNumRegions = hasLargeFieldCount ? 0xFF : numRegions; + uint8_t smallNumResults = hasLargeFieldCount ? 0xFF : numResults; + uint8_t smallNumSuccs = hasLargeFieldCount ? 0xFF : numSuccessors; + // If the operation is known to have no operands, don't allocate an operand // storage. bool needsOperandStorage = true; @@ -124,9 +136,11 @@ // 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( - needsOperandStorage ? 1 : 0, numSuccessors, numRegions, numOperands); + size_t byteSize = totalSizeToAlloc( + needsOperandStorage ? 1 : 0, hasLargeFieldCount ? 1 : 0, numSuccessors, + numRegions, numOperands); size_t prefixByteSize = llvm::alignTo( Operation::prefixAllocSize(numTrailingResults, numInlineResults), alignof(Operation)); @@ -134,13 +148,19 @@ void *rawMem = mallocMem + prefixByteSize; // Create the new Operation. - Operation *op = - ::new (rawMem) Operation(location, name, numResults, numSuccessors, - numRegions, attributes, needsOperandStorage); + Operation *op = ::new (rawMem) + Operation(location, name, smallNumResults, smallNumSuccs, smallNumRegions, + attributes, hasLargeFieldCount, needsOperandStorage); assert((numSuccessors == 0 || op->mightHaveTrait()) && "unexpected successors in a non-terminator operation"); + // Initialize the field counts. + if (hasLargeFieldCount) { + new (&op->getLargeFieldCountStorage()) detail::LargeOperationFieldCounts( + numRegions, numResults, numSuccessors); + } + // Initialize the results. auto resultTypeIt = resultTypes.begin(); for (unsigned i = 0; i < numInlineResults; ++i, ++resultTypeIt) @@ -168,12 +188,14 @@ return op; } -Operation::Operation(Location location, OperationName name, unsigned numResults, - unsigned numSuccessors, unsigned numRegions, - DictionaryAttr attributes, bool hasOperandStorage) - : location(location), numResults(numResults), numSuccs(numSuccessors), - numRegions(numRegions), hasOperandStorage(hasOperandStorage), name(name), - attrs(attributes) { +Operation::Operation(Location location, OperationName name, + unsigned smallNumResults, unsigned smallNumSuccs, + unsigned smallNumRegions, DictionaryAttr attributes, + bool hasLargeFieldCounts, bool hasOperandStorage) + : location(location), smallNumRegions(smallNumRegions), + smallNumResults(smallNumResults), smallNumSuccs(smallNumSuccs), + hasLargeFieldCounts(hasLargeFieldCounts), + hasOperandStorage(hasOperandStorage), name(name), attrs(attributes) { assert(attributes && "unexpected null attribute dictionary"); #ifndef NDEBUG if (!getDialect() && !getContext()->allowsUnregisteredDialects()) @@ -611,7 +633,7 @@ auto *newOp = cloneWithoutRegions(mapper); // Clone the regions. - for (unsigned i = 0; i != numRegions; ++i) + for (unsigned i = 0, e = getNumRegions(); i != e; ++i) getRegion(i).cloneInto(&newOp->getRegion(i), mapper); return newOp;