[llvm-branch-commits] [mlir] 47364f9 - [mlir][IR] Move the storage for results to before the Operation instead of after.

River Riddle via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Fri Dec 4 21:12:42 PST 2020


Author: River Riddle
Date: 2020-12-04T21:01:42-08:00
New Revision: 47364f95e810f96fd300ceaa095719f76683e6fa

URL: https://github.com/llvm/llvm-project/commit/47364f95e810f96fd300ceaa095719f76683e6fa
DIFF: https://github.com/llvm/llvm-project/commit/47364f95e810f96fd300ceaa095719f76683e6fa.diff

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

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.

Differential Revision: https://reviews.llvm.org/D92687

Added: 
    

Modified: 
    mlir/include/mlir/IR/Operation.h
    mlir/include/mlir/IR/OperationSupport.h
    mlir/lib/IR/Operation.cpp
    mlir/lib/IR/OperationSupport.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 5ed5cd9930fa..0c295027923d 100644
--- 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<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,
@@ -654,6 +653,21 @@ class Operation final
   // 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 @@ class Operation final
 
   /// 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.
@@ -725,17 +745,8 @@ class Operation final
   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;
   }

diff  --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 88857f174783..343637e457b5 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -570,11 +570,11 @@ 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) {}
 

diff  --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 60b0f87932ee..1795dd947fa0 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -119,16 +119,18 @@ Operation *Operation::create(Location location, OperationName name,
       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),
+      alignof(Operation));
+  char *mallocMem = reinterpret_cast<char *>(malloc(byteSize + prefixByteSize));
+  void *rawMem = mallocMem + prefixByteSize;
 
   // Create the new Operation.
   Operation *op =
@@ -200,8 +202,12 @@ Operation::~Operation() {
 
 /// 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.

diff  --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp
index 32d64ac32272..4a06c5c5d600 100644
--- a/mlir/lib/IR/OperationSupport.cpp
+++ b/mlir/lib/IR/OperationSupport.cpp
@@ -382,14 +382,15 @@ MutableArrayRef<OpOperand> detail::OperandStorage::resize(Operation *owner,
 
 /// 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);
 }
 
 //===----------------------------------------------------------------------===//


        


More information about the llvm-branch-commits mailing list