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 @@ -27,7 +27,7 @@ /// 'Block' class. class alignas(8) Operation final : public llvm::ilist_node_with_parent, - private llvm::TrailingObjects { public: /// Create a new Operation with the specific fields. @@ -651,7 +651,7 @@ /// 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( getInlineResult(OpResult::getMaxInlineResults() - 1)) - @@ -695,11 +695,18 @@ /// 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. + /// - 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 @@ friend class llvm::ilist_node_with_parent; // This stuff is used by the TrailingObjects template. - friend llvm::TrailingObjects; size_t numTrailingObjects(OverloadToken) const { return numSuccs; } size_t numTrailingObjects(OverloadToken) const { return numRegions; } + size_t numTrailingObjects(OverloadToken) 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 --- a/mlir/lib/IR/Operation.cpp +++ b/mlir/lib/IR/Operation.cpp @@ -123,8 +123,11 @@ // into account the size of the operation, its trailing objects, and its // prefixed objects. size_t byteSize = - totalSizeToAlloc( - numSuccessors, numRegions, needsOperandStorage ? 1 : 0) + + totalSizeToAlloc( + 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 @@ 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()); + } } } @@ -545,17 +553,15 @@ /// Return the number of results held by this operation. unsigned Operation::getNumResults() { - if (!resultType) - return 0; - return hasSingleResult ? 1 : resultType.cast().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().getTypes(); + return resultTypeOrSize.type; + return ArrayRef(getTrailingObjects(), resultTypeOrSize.size); } void Operation::setSuccessor(Block *block, unsigned index) { diff --git a/mlir/lib/IR/Value.cpp b/mlir/lib/IR/Value.cpp --- a/mlir/lib/IR/Value.cpp +++ b/mlir/lib/IR/Value.cpp @@ -39,31 +39,21 @@ OpResult result = cast(); Operation *owner = result.getOwner(); if (owner->hasSingleResult) - return owner->resultType; - return owner->resultType.cast().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()) return arg.setType(newType); - OpResult result = cast(); - // If the owner has a single result, simply update it directly. + OpResult result = cast(); Operation *owner = result.getOwner(); - if (owner->hasSingleResult) { - owner->resultType = newType; - return; - } - unsigned resultNo = result.getResultNumber(); - - // Otherwise, rebuild the tuple if the new type is different from the current. - auto curTypes = owner->resultType.cast().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()[result.getResultNumber()] = newType; } /// If this value is the result of an Operation, return the operation that