better constant propagation through tuples #21378
Conversation
pytorchbot
added
oncall: jit
|
What is stopping us from representing tuples as constants? The code would be a lot cleaner if we could do that, and just apply the same rules we have for everything else to tuples. |
|
I think it would lead to a more complicated IR and a simplified constant propagation, which isn't necessarily a good trade off. It could also lead to worse memory if you're not careful with something like, Previously all uses of a would get pooled into a single constant, but with constants that can contain other contexts the pooling logic is much more difficult. It would also complicate the IR Parser. |
| if (inp) { | ||
| inputs.push_back(*inp); | ||
| } else { | ||
| inserted_tuple = false; |
There was a problem hiding this comment.
Couldn't you return nullptr here and get rid of this boolean? Why do you need to keep trying to constify-elements if any one of them fails?
| (*new_output)->setType(n->outputs()[i]->type()); | ||
| } | ||
| n->outputs()[i]->replaceAllUsesWith(*new_output); | ||
| } else if (outputs[i].isTuple() && n->outputs().size() == 1) { |
There was a problem hiding this comment.
Why does the tuple have to be the only output?
There was a problem hiding this comment.
In general it's a bug for an op to not have one output. It should be either None, T, or a Tuple. So this should never be hit. But i included it because it would be weird if one of the outputs of an op was constantified, and we still kept the op in the graph. and ran it again
| if (tuple_ops.count(n->kind())) { | ||
| return ( | ||
| std::all_of(n->inputs().begin(), n->inputs().end(), [&](Value* v) { | ||
| return v->node()->kind() == prim::Constant || tuples.count(v); |
There was a problem hiding this comment.
Won't tuples.count(v) return true for any tuple since contains all of them regardless of whether they are constant or not
There was a problem hiding this comment.
No, it will only return true for tuples which we've computed the ivalue for.
gh-metadata: pytorch pytorch 21500 gh/eellison/5/head
gh-metadata: pytorch pytorch 21501 gh/eellison/6/head
eellison
force-pushed
the
gh/eellison/2/head
branch
from
7b1806c to
3ce9a9f
Compare
|
Closing in favor of #21501 since i did something to mess up ghstack here |
Stack from ghstack:
This PR lets us constant prop through ops that forward tuples, & through builtins that return tuples. Previously, tuples inhibited constant prop and combined with lower tuples pass caused jitter for running constant propagation on compilation.