WithFields for Tuples #9533
WithFields for Tuples #9533
Conversation
src/relay/ir/expr.cc
Outdated
| @@ -71,6 +71,29 @@ Tuple::Tuple(tvm::Array<relay::Expr> fields, Span span) { | |||
| data_ = std::move(n); | |||
| } | |||
|
|
|||
| Tuple Tuple::WithFields(Optional<Array<Expr>> opt_fields, Optional<Span> opt_span) { | |||
| Array<Expr> fields = opt_fields.value_or(get()->fields); | |||
There was a problem hiding this comment.
make it as a global function https://github.com/apache/tvm/blob/main/include/tvm/ir/attrs.h#L371, so that you do std::move and COW when there is a single copy.
There was a problem hiding this comment.
To elaborate, x.WithFields in the current case might mutate x if it is the only copy.
auto backup_fields = x.fields;
y = x.WithFields(..., ...)
// if x is the only copy, it will change x itself. Which may not be the intention of the API.
// this check will fail, and that could against the original intention of x.WithFields(that x should not change)
CHECK(x->fields.same_as(backup_fields));instead if we do
y = WithFields(std::move(x), .., ...)_and x is the only copy, we know for certainly that x will not be used in the future and it is safe to re-use the memory space
There was a problem hiding this comment.
The intention is to mutate if x is the only copy.
There was a problem hiding this comment.
I see, thanks for the explaination, if it is a mution semantics, then we might want to rename it as SetXYZ(or COWSet?), like Array.Set.
It might be worth thinking a bit about the tradeoff since the original semantics on the python and c++ side is WithAttrs do not mutate the original value that is passed in(unless it is moved in and ensures that there is no followup), the api name With is also deliberately chosen over SetXYZ to show that intention(that we do not modify the original value). So it would be good to make it align with the WithAttrs semantics.
It is indeed quite subtle difference but these are very interesting design choices here.
A common usage pattern of with_attrs may ends up look like this:
x = PrimFunc(args, body)
# y is a new copy that comes with a new attr
# x remains as the original case
y = x.with_attr("new_attr", value)
record([x, y])So in the c++ API
// y do not reuse x' sapce
y = WithAttrs(x, attr)
record({y, x})Instead if you do
// compiler could decide that y can reuse x, because x is no longer used, the call becomes a move
y = WithAttrs(x, attr)
// x is not used as a followup workAnd with a deliberate move where x is the only copy, we ensures that x's space get reused
y = WithAttrs(std::move(x), attr)
// x is not used as a followup workThere was a problem hiding this comment.
OK, thanks for the feedback. I just changed the signature to mirror WithAttrs and I removed the COW functionality. Now it will return the argument tuple if all the fields passed in are the same, otherwise it creates a copy and returns that.
src/relay/ir/expr.cc
Outdated
| bool all_fields_unchanged = true; | ||
| if (fields.size() == tuple->fields.size()) { | ||
| for (uint i = 0; i < fields.size(); i++) { | ||
| all_fields_unchanged &= fields[i] == (tuple->fields[i]); |
There was a problem hiding this comment.
consider use same_as to avoid future overload of operator ==
src/relay/ir/expr.cc
Outdated
| if (all_fields_unchanged) { | ||
| return std::move(tuple); | ||
| } else { | ||
| return Tuple(std::move(fields), std::move(span)); |
There was a problem hiding this comment.
it is useful to actually invoke COW here.
Line 374 in a6c948a
If the tuple is moved in and a single copy, it is actually useful to directly reuse that part of the memory. This will allow calls like below to only create one copy instead of two
x = Tuple(fields, span);
x = WithFields(std::move(x), new_fields, span)
This is also the reason why WithAttrs is a global function(so there won't be ambiguity in such case).
There was a problem hiding this comment.
I see, Mark and I were confused about what you were saying before. Changed it back to use COW
electriclilies
requested review from
anijain2305,
icemelon,
jroesch,
junrushao,
jwfromm,
MarisaKirisame,
mbrookhart,
slyubomirsky,
vinx13,
wweic,
yzhliu,
zhiics and
ZihengJiang
as code owners
|
I think this is ready for review now, pending CI. @mbs-octoml @mikepapadim PTAL! |
|
LGTM |
src/relay/ir/expr.cc
Outdated
| TupleNode* cow_tuple_node = tuple.CopyOnWrite(); | ||
| cow_tuple_node->fields = fields; | ||
| cow_tuple_node->span = span; | ||
| return GetRef<Tuple>(cow_tuple_node); |
There was a problem hiding this comment.
nit, you can directly return tuple here
There was a problem hiding this comment.
Thanks for all your help with the subtleties!
|
Hmm, looking into the MSVC error, not sure if it is related to include order or something |
|
@tqchen Thanks for looking! I'm going to resubmit at EOD and see if it goes away, it looked like it might be flaky to me but I wanted to see if anything else broke further along in CI since it's still running. |
|
no worries, we are not in a hurry. this seems was a compiler error |
This is a draft PR implementing WithFields on tuples.
WithFields is parallel to WithAttrs in behavior. It returns this* if all the fields passed in are the same as the ones already in this. Otherwise, it creates a copy of the tuple with the new fields.
The motivation for WithFields is to ensure that spans, checked types, and other fields on nodes are propagated correctly during passes so we can rely on that information. If we decide to go ahead with this change, I will also implement WithFields for all other relay exprs.