[mlir] Optimize Operation field count memory usage

[rriddle]

Summary

Various components within an operation (e.g. regions, results,
successors) have a small number of elements in the overwhelmingly
common case. This commit refactors Operation such that in the
common case, the counts of these fields are stored within a u8,
and in the extreme case the counts are stored in a trailing
object that allows for a full u32 number of elements. This allows
for dropping the size of Operation by 8 bytes (in an x64 we are
now down to 56 bytes).

[rriddle]

@lattner I did some quick benchmarking and it didn't seem to have a negative impact on compile time, but can you try on your circt benchmarks?

[lattner]

awesome, I'll take a look as soon as we get your other work merged into the CIRCT llvm submodule

mlir/include/mlir/IR/Operation.h

725

It'd be really awesome to make this be a full 'unsigned char' for the bool value, since it would make the accessors a single load+compare, instead of load+and+compare. Similarly hasOperandStorage is also going to be hot. Do you think that is possible without bloating things?

[rriddle]

mlir/include/mlir/IR/Operation.h

725

I don't think we can as we are on a word boundary right now. We could make one of these a full char (given that the region/result/successor counts only use 3 bytes), but the other would need to steal a bit from somewhere else (orderIndex?). I'm not sure if it's worth it or not to go that route though.

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, detail::OperandStorage,
+                                    detail::LargeOperationFieldCounts,
                                     BlockOperand, Region, OpOperand> {
 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);
 
public:
   operand_type_iterator operand_type_begin() { return operand_begin(); }
   operand_type_iterator operand_type_end() { return operand_end(); }
   operand_type_range getOperandTypes() { return getOperands().getTypes(); }
 
   //===--------------------------------------------------------------------===//
   // 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)); }
 
   /// Support result iteration.
   using result_range = ResultRange;
   using result_iterator = result_range::iterator;
 
   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(); }
   OpResult getOpResult(unsigned idx) { return getResult(idx); }
 
   /// Support result type iteration.
   using result_type_iterator = result_range::type_iterator;
   using result_type_range = result_range::type_range;
for (auto attr : getAttrs())
         attrs.push_back(attr);
     setAttrs(attrs.getDictionary(getContext()));
   }
 
   //===--------------------------------------------------------------------===//
   // 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<Region> getRegions() {
-    auto *regions = getTrailingObjects<Region>();
-    return {regions, numRegions};
+    if (!hasRegions())
+      return MutableArrayRef<Region>();
+    return {getTrailingObjects<Region>(), 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];
   }
 
   //===--------------------------------------------------------------------===//
   // Successors
   //===--------------------------------------------------------------------===//
 
+  /// Returns the block operands held by this operation.
   MutableArrayRef<BlockOperand> getBlockOperands() {
-    return {getTrailingObjects<BlockOperand>(), numSuccs};
+    return {getTrailingObjects<BlockOperand>(), getNumSuccessors()};
   }
 
   // Successor iteration.
   using succ_iterator = SuccessorRange::iterator;
   succ_iterator successor_begin() { return getSuccessors().begin(); }
   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());
     return getBlockOperands()[index].get();
   }
   void setSuccessor(Block *block, unsigned index);
 
   //===--------------------------------------------------------------------===//
private:
   /// Update the order index of this operation of this operation if necessary,
   /// potentially recomputing the order of the parent block.
   void updateOrderIfNecessary();
 
   /// Returns true if this operation has a valid order.
   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.
   ~Operation();
 
   /// Returns the additional size necessary for allocating the given objects
   /// before an Operation in-memory.
   static size_t prefixAllocSize(unsigned numOutOfLineResults,
private:
   }
 
   /// Returns the operand storage object.
   detail::OperandStorage &getOperandStorage() {
     assert(hasOperandStorage && "expected operation to have operand storage");
     return *getTrailingObjects<detail::OperandStorage>();
   }
 
+  /// Returns the large field count storage.
+  detail::LargeOperationFieldCounts &getLargeFieldCountStorage() {
+    assert(hasLargeFieldCounts &&
+           "expected operation to have large field counts");
+    return *getTrailingObjects<detail::LargeOperationFieldCounts>();
+  }
+
   /// 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)
     // the inline results.
     return reinterpret_cast<detail::OutOfLineOpResult *>(getInlineOpResult(
                detail::OpResultImpl::getMaxInlineResults() - 1)) -
            ++resultNumber;
   }
private:
   /// This holds information about the source location the operation was defined
   /// or derived from.
   Location location;
 
   /// Relative order of this operation in its parent block. Used for
   /// 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
   /// operands.
   bool hasOperandStorage : 1;
 
   /// This holds the name of the operation.
   OperationName name;
private:
 
   // 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, detail::OperandStorage, BlockOperand,
+  friend llvm::TrailingObjects<Operation, detail::OperandStorage,
+                               detail::LargeOperationFieldCounts, BlockOperand,
                                Region, OpOperand>;
   size_t numTrailingObjects(OverloadToken<detail::OperandStorage>) const {
     return hasOperandStorage ? 1 : 0;
   }
+  size_t
+  numTrailingObjects(OverloadToken<detail::LargeOperationFieldCounts>) const {
+    return hasLargeFieldCounts ? 1 : 0;
+  }
   size_t numTrailingObjects(OverloadToken<BlockOperand>) const {
-    return numSuccs;
+    return const_cast<Operation *>(this)->getNumSuccessors();
+  }
+  size_t numTrailingObjects(OverloadToken<Region>) const {
+    return const_cast<Operation *>(this)->getNumRegions();
   }
-  size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, const Operation &op) {
   const_cast<Operation &>(op).print(os, OpPrintingFlags().useLocalScope());
   return os;
 }
 
 } // end namespace mlir

mlir/include/mlir/IR/OperationSupport.h

private:
   /// The number of operands within the storage.
   unsigned numOperands;
   /// A pointer to the operand storage.
   OpOperand *operandStorage;
 };
 } // 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
 //===----------------------------------------------------------------------===//
 
 /// Set of flags used to control the behavior of the various IR print methods
 /// (e.g. Operation::Print).
 class OpPrintingFlags {
 public:
   OpPrintingFlags();

mlir/lib/IR/Operation.cpp

Operation *Operation::create(Location location, OperationName name,
 
   // We only need to allocate additional memory for a subset of results.
   unsigned numTrailingResults = OpResult::getNumTrailing(resultTypes.size());
   unsigned numInlineResults = OpResult::getNumInline(resultTypes.size());
   unsigned numSuccessors = successors.size();
   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<uint8_t>::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;
   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<detail::OperandStorage, BlockOperand, Region, OpOperand>(
-          needsOperandStorage ? 1 : 0, numSuccessors, numRegions, numOperands);
+  size_t byteSize = totalSizeToAlloc<detail::OperandStorage,
+                                     detail::LargeOperationFieldCounts,
+                                     BlockOperand, Region, OpOperand>(
+      needsOperandStorage ? 1 : 0, hasLargeFieldCount ? 1 : 0, numSuccessors,
+      numRegions, 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 *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<OpTrait::IsTerminator>()) &&
          "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)
     new (op->getInlineOpResult(i)) detail::InlineOpResult(*resultTypeIt, i);
   for (unsigned i = 0; i < numTrailingResults; ++i, ++resultTypeIt) {
     new (op->getOutOfLineOpResult(i))
         detail::OutOfLineOpResult(*resultTypeIt, i);
   }
Operation *Operation::create(Location location, OperationName name,
   // Initialize the successors.
   auto blockOperands = op->getBlockOperands();
   for (unsigned i = 0; i != numSuccessors; ++i)
     new (&blockOperands[i]) BlockOperand(op, successors[i]);
 
   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())
     llvm::report_fatal_error(
         name.getStringRef() +
         " created with unregistered dialect. If this is intended, please call "
         "allowUnregisteredDialects() on the MLIRContext, or use "
         "-allow-unregistered-dialect with the MLIR tool used.");

 /// values outside of the operation using the map that is provided (leaving
 /// them alone if no entry is present).  Replaces references to cloned
 /// sub-operations to the corresponding operation that is copied, and adds
 /// those mappings to the map.
 Operation *Operation::clone(BlockAndValueMapping &mapper) {
   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;
 }
 
 Operation *Operation::clone() {
   BlockAndValueMapping mapper;
   return clone(mapper);