[mlir][IR] Move the storage for results to before the Operation instead of after.

[rriddle]

Summary

Trailing objects are really nice for storing additional data inline with the main class, and is something that we heavily take advantage of for Operation(and many other classes). To get the address of the inline data you need to compute the address by doing some pointer arithmetic taking into account any objects stored before the object you want to access. Most classes keep the count of the number of objects, so this is relatively cheap to compute. This is not the case for results though, which have two different types(inline and trailing) that are not necessarily as cheap to compute as the count for other objects. This revision moves the storage for results to before the operation and stores them in reverse order. This allows for getting results to still be very fast given that they are never iterated directly in order, and also greatly improves the speed when accessing the other trailing objects of an operation(operands/regions/blocks/etc).

This reduced compile time when compiling a decently sized mlir module by about ~400ms, or 2.17s -> 1.76s.

[mehdi_amini]

mlir/include/mlir/IR/Operation.h

678

is "Trailing" still accurate?

[rriddle]

mlir/include/mlir/IR/Operation.h

678

From the perspective of english, it still works. From confusion related to TrailingObject/PrefixObject less so. I can change if you want, but didn't see a strong need. I kind of look at whether it is after/before as more of an internal implementation detail unrelated to the naming of the class.

mlir/include/mlir/IR/Operation.h

 #include "mlir/IR/Block.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/Region.h"
 #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<Operation, Block>,
-      private llvm::TrailingObjects<Operation, detail::InLineOpResult,
-                                    detail::TrailingOpResult, BlockOperand,
-                                    Region, detail::OperandStorage> {
+      private llvm::TrailingObjects<Operation, BlockOperand, Region,
+                                    detail::OperandStorage> {
 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);
 
   /// Overload of create that takes an existing MutableDictionaryAttr to avoid
private:
   Operation(Location location, OperationName name, TypeRange resultTypes,
             unsigned numSuccessors, unsigned numRegions,
             const MutableDictionaryAttr &attributes, 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 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");
     return *getTrailingObjects<detail::OperandStorage>();
   }
 
   /// Returns a pointer to the use list for the given trailing result.
   detail::TrailingOpResult *getTrailingResult(unsigned resultNumber) {
-    return getTrailingObjects<detail::TrailingOpResult>() + resultNumber;
+    // Trailing results are stored in reverse order after(before in memory) the
+    // inline results.
+    return reinterpret_cast<detail::TrailingOpResult *>(
+               getInlineResult(OpResult::getMaxInlineResults() - 1)) -
+           ++resultNumber;
   }
 
   /// Returns a pointer to the use list for the given inline result.
   detail::InLineOpResult *getInlineResult(unsigned resultNumber) {
-    return getTrailingObjects<detail::InLineOpResult>() + resultNumber;
+    // Inline results are stored in reverse order before the operation in
+    // memory.
+    return reinterpret_cast<detail::InLineOpResult *>(this) - ++resultNumber;
   }
 
   /// Provide a 'getParent' method for ilist_node_with_parent methods.
   /// We mark it as a const function because ilist_node_with_parent specifically
   /// requires a 'getParent() const' method. Once ilist_node removes this
   /// constraint, we should drop the const to fit the rest of the MLIR const
   /// model.
   Block *getParent() const { return block; }
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::InLineOpResult,
-                               detail::TrailingOpResult, BlockOperand, Region,
+  friend llvm::TrailingObjects<Operation, BlockOperand, Region,
                                detail::OperandStorage>;
-  size_t numTrailingObjects(OverloadToken<detail::InLineOpResult>) const {
-    return OpResult::getNumInline(
-        const_cast<Operation *>(this)->getNumResults());
-  }
-  size_t numTrailingObjects(OverloadToken<detail::TrailingOpResult>) const {
-    return OpResult::getNumTrailing(
-        const_cast<Operation *>(this)->getNumResults());
-  }
   size_t numTrailingObjects(OverloadToken<BlockOperand>) const {
     return numSuccs;
   }
   size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Operation &op) {
   op.print(os, OpPrintingFlags().useLocalScope());

mlir/include/mlir/IR/OperationSupport.h

 //===----------------------------------------------------------------------===//
 // ResultStorage
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// 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<OpOperand> {};
+struct alignas(8) InLineOpResult : public IRObjectWithUseList<OpOperand> {};
 /// 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<OpOperand> {
+struct alignas(8) TrailingOpResult : public IRObjectWithUseList<OpOperand> {
   TrailingOpResult(uint64_t trailingResultNumber)
       : trailingResultNumber(trailingResultNumber) {}
 
   /// Returns the parent operation of this trailing result.
   Operation *getOwner();
 
   /// Return the proper result number of this op result.
   unsigned getResultNumber() {

mlir/lib/IR/Operation.cpp

Operation *Operation::create(Location location, OperationName name,
   // 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.
-  auto byteSize =
-      totalSizeToAlloc<detail::InLineOpResult, detail::TrailingOpResult,
-                       BlockOperand, Region, detail::OperandStorage>(
-          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<BlockOperand, Region, detail::OperandStorage>(
+          numSuccessors, numRegions, needsOperandStorage ? 1 : 0) +
+      detail::OperandStorage::additionalAllocSize(numOperands);
+  size_t prefixByteSize = llvm::alignTo(
+      Operation::prefixAllocSize(numTrailingResults, numInlineResults),
+  char *mallocMem = reinterpret_cast<char *>(malloc(byteSize + prefixByteSize));
+  void *rawMem = mallocMem + prefixByteSize;
 
   // Create the new Operation.
   Operation *op =
       ::new (rawMem) Operation(location, name, resultTypes, numSuccessors,
                                numRegions, attributes, needsOperandStorage);
 
   assert((numSuccessors == 0 || !op->isKnownNonTerminator()) &&
          "unexpected successors in a non-terminator operation");
Operation::~Operation() {
 
   // Explicitly destroy the regions.
   for (auto &region : getRegions())
     region.~Region();
 }
 
 /// 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<char *>(this) -
+                 llvm::alignTo(prefixAllocSize(), alignof(Operation));
   this->~Operation();
-  free(this);
+  free(rawMem);
 }
 
 /// Return the context this operation is associated with.
 MLIRContext *Operation::getContext() { return location->getContext(); }
 
 /// Return the dialect this operation is associated with, or nullptr if the
 /// associated dialect is not registered.
 Dialect *Operation::getDialect() {

mlir/lib/IR/OperationSupport.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // ResultStorage
 //===----------------------------------------------------------------------===//
 
 /// 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<InLineOpResult *>(trailingIt) -
+  auto *inlineResultIt = reinterpret_cast<InLineOpResult *>(trailingIt) +
                          OpResult::getMaxInlineResults();
-  return reinterpret_cast<Operation *>(inlineResultIt) - 1;
+  return reinterpret_cast<Operation *>(inlineResultIt);
 }
 
 //===----------------------------------------------------------------------===//
 // Operation Value-Iterators
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // OperandRange