[Mlir-commits] [mlir] 08f0764 - Remove use of tuple for multiresult type storage

Jacques Pienaar llvmlistbot at llvm.org
Mon Mar 1 09:30:38 PST 2021


Author: Jacques Pienaar
Date: 2021-03-01T09:30:24-08:00
New Revision: 08f0764ff551c5aa2486c40871453e1ff40fb679

URL: https://github.com/llvm/llvm-project/commit/08f0764ff551c5aa2486c40871453e1ff40fb679
DIFF: https://github.com/llvm/llvm-project/commit/08f0764ff551c5aa2486c40871453e1ff40fb679.diff

LOG: Remove use of tuple for multiresult type storage

Move the results in line with the op instead. This results in each
operation having its own types recorded vs single tuple type, but comes
at benefit that every mutation doesn't incurs uniquing. Ran into cases
where updating result type of operation led to very large memory usage.

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 63a94d971da7..5cc1639b240d 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -27,7 +27,7 @@ namespace mlir {
 /// 'Block' class.
 class alignas(8) Operation final
     : public llvm::ilist_node_with_parent<Operation, Block>,
-      private llvm::TrailingObjects<Operation, BlockOperand, Region,
+      private llvm::TrailingObjects<Operation, BlockOperand, Region, Type,
                                     detail::OperandStorage> {
 public:
   /// Create a new Operation with the specific fields.
@@ -651,7 +651,7 @@ class alignas(8) Operation final
 
   /// Returns a pointer to the use list for the given trailing result.
   detail::TrailingOpResult *getTrailingResult(unsigned resultNumber) {
-    // Trailing results are stored in reverse order after(before in memory) the
+    // Trailing results are stored in reverse order after (before in memory) the
     // inline results.
     return reinterpret_cast<detail::TrailingOpResult *>(
                getInlineResult(OpResult::getMaxInlineResults() - 1)) -
@@ -695,11 +695,18 @@ class alignas(8) Operation final
   /// states recorded here:
   /// - 0 results : The type below is null.
   /// - 1 result  : The single result type is held here.
-  /// - N results : The type here is a tuple holding the result types.
-  /// Note: We steal a bit for 'hasSingleResult' from somewhere else so that we
-  /// can use 'resultType` in an ArrayRef<Type>.
+  /// - N results : The type here is empty and instead the result types are held
+  ///               in trailing storage.
   bool hasSingleResult : 1;
-  Type resultType;
+
+  /// Union representing either the Type of a single result operation (if
+  /// hasSingleResult) or the number of result types for multi-result.
+  union {
+    // Type, set if single result Operation.
+    Type type = nullptr;
+    // Size, set if not a single result Operation.
+    unsigned size;
+  } resultTypeOrSize;
 
   /// This holds the name of the operation.
   OperationName name;
@@ -720,12 +727,15 @@ class alignas(8) Operation final
   friend class llvm::ilist_node_with_parent<Operation, Block>;
 
   // This stuff is used by the TrailingObjects template.
-  friend llvm::TrailingObjects<Operation, BlockOperand, Region,
+  friend llvm::TrailingObjects<Operation, BlockOperand, Region, Type,
                                detail::OperandStorage>;
   size_t numTrailingObjects(OverloadToken<BlockOperand>) const {
     return numSuccs;
   }
   size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }
+  size_t numTrailingObjects(OverloadToken<Type>) const {
+    return hasSingleResult ? 0 : resultTypeOrSize.size;
+  }
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Operation &op) {

diff  --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index b1f2ad66efdb..b074a675f2ff 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -123,8 +123,11 @@ Operation *Operation::create(Location location, OperationName name,
   // 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) +
+      totalSizeToAlloc<BlockOperand, Region, Type, detail::OperandStorage>(
+          numSuccessors, numRegions,
+          // Result type storage only needed if there is not 0 or 1 results.
+          resultTypes.size() == 1 ? 0 : resultTypes.size(),
+          needsOperandStorage ? 1 : 0) +
       detail::OperandStorage::additionalAllocSize(numOperands);
   size_t prefixByteSize = llvm::alignTo(
       Operation::prefixAllocSize(numTrailingResults, numInlineResults),
@@ -172,13 +175,18 @@ Operation::Operation(Location location, OperationName name,
   assert(attributes && "unexpected null attribute dictionary");
   assert(llvm::all_of(resultTypes, [](Type t) { return t; }) &&
          "unexpected null result type");
-  if (!resultTypes.empty()) {
-    // If there is a single result it is stored in-place, otherwise use a tuple.
+  if (resultTypes.empty()) {
+    resultTypeOrSize.size = 0;
+  } else {
+    // If there is a single result it is stored in-place, otherwise use trailing
+    // type storage.
     hasSingleResult = resultTypes.size() == 1;
-    if (hasSingleResult)
-      resultType = resultTypes.front();
-    else
-      resultType = TupleType::get(location->getContext(), resultTypes);
+    if (hasSingleResult) {
+      resultTypeOrSize.type = resultTypes.front();
+    } else {
+      resultTypeOrSize.size = resultTypes.size();
+      llvm::copy(resultTypes, getTrailingObjects<Type>());
+    }
   }
 }
 
@@ -545,17 +553,15 @@ void Operation::dropAllDefinedValueUses() {
 
 /// Return the number of results held by this operation.
 unsigned Operation::getNumResults() {
-  if (!resultType)
-    return 0;
-  return hasSingleResult ? 1 : resultType.cast<TupleType>().size();
+  if (hasSingleResult)
+    return 1;
+  return resultTypeOrSize.size;
 }
 
 auto Operation::getResultTypes() -> result_type_range {
-  if (!resultType)
-    return llvm::None;
   if (hasSingleResult)
-    return resultType;
-  return resultType.cast<TupleType>().getTypes();
+    return resultTypeOrSize.type;
+  return ArrayRef<Type>(getTrailingObjects<Type>(), resultTypeOrSize.size);
 }
 
 void Operation::setSuccessor(Block *block, unsigned index) {

diff  --git a/mlir/lib/IR/Value.cpp b/mlir/lib/IR/Value.cpp
index fd7e5b5d64e5..c6d3e5857343 100644
--- a/mlir/lib/IR/Value.cpp
+++ b/mlir/lib/IR/Value.cpp
@@ -39,31 +39,21 @@ Type Value::getType() const {
   OpResult result = cast<OpResult>();
   Operation *owner = result.getOwner();
   if (owner->hasSingleResult)
-    return owner->resultType;
-  return owner->resultType.cast<TupleType>().getType(result.getResultNumber());
+    return owner->resultTypeOrSize.type;
+  return owner->getResultTypes()[result.getResultNumber()];
 }
 
 /// Mutate the type of this Value to be of the specified type.
 void Value::setType(Type newType) {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.setType(newType);
-  OpResult result = cast<OpResult>();
 
-  // If the owner has a single result, simply update it directly.
+  OpResult result = cast<OpResult>();
   Operation *owner = result.getOwner();
-  if (owner->hasSingleResult) {
-    owner->resultType = newType;
-    return;
-  }
-  unsigned resultNo = result.getResultNumber();
-
-  // Otherwise, rebuild the tuple if the new type is 
diff erent from the current.
-  auto curTypes = owner->resultType.cast<TupleType>().getTypes();
-  if (curTypes[resultNo] == newType)
-    return;
-  auto newTypes = llvm::to_vector<4>(curTypes);
-  newTypes[resultNo] = newType;
-  owner->resultType = TupleType::get(newType.getContext(), newTypes);
+  if (owner->hasSingleResult)
+    owner->resultTypeOrSize.type = newType;
+  else
+    owner->getTrailingObjects<Type>()[result.getResultNumber()] = newType;
 }
 
 /// If this value is the result of an Operation, return the operation that


        


More information about the Mlir-commits mailing list