Skip to content

Commit e88d1c4

Browse files
swolchokfacebook-github-bot
authored andcommitted
[PyTorch] Add tuple inline storage (pytorch#64066)
Summary: Pull Request resolved: pytorch#64066 I noticed a bunch of time being spent heap-allocating Tuples in the unpickler. 1-, 2-, and 3-element Tuples are apparently common enough that they get their own bytecode instructions, so I decided to try also giving them their own representation. We store up to 3 IValues inline in `Tuple` rather than doing a second heap allocation for a `std::vector<IValue>`. ghstack-source-id: 140695395 Test Plan: Added automated tests for TupleElements. Pixel 3 before: https://www.internalfb.com/intern/aibench/details/761596366576284 Pixel 3 after: https://www.internalfb.com/intern/aibench/details/591414145082422 We went from 347 ms to 302 ms. Reviewed By: dhruvbird Differential Revision: D30592622 fbshipit-source-id: 93625c54c9dca5f765ef6d5c191944179cb281a8
1 parent f8f9a47 commit e88d1c4

30 files changed

+623
-118
lines changed

aten/src/ATen/core/ivalue.cpp

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -51,19 +51,19 @@ TORCH_API c10::intrusive_ptr<ConstantString> ConstantString::create(
5151
}
5252

5353
bool operator==(const ivalue::Tuple& lhs, const ivalue::Tuple& rhs) {
54-
return lhs.elements_.size() == rhs.elements_.size() &&
54+
return lhs.size() == rhs.size() &&
5555
// see [container equality]
5656
std::equal(
57-
lhs.elements_.cbegin(),
58-
lhs.elements_.cend(),
59-
rhs.elements_.cbegin(),
57+
lhs.elements().cbegin(),
58+
lhs.elements().cend(),
59+
rhs.elements().cbegin(),
6060
_fastEqualsForContainer);
6161
}
6262

6363
TupleTypePtr Tuple::type() const {
6464
if (!type_) {
6565
type_ = TupleType::create(
66-
fmap(elements_, [&](const IValue& v) { return v.type(); }));
66+
fmap(elements(), [&](const IValue& v) { return v.type(); }));
6767
}
6868
return type_;
6969
}

0 commit comments

Comments
 (0)