Graph partition with broadcast logic#280
Conversation
failing test_half due to the new partition logic
0b3c08f to
316ae16
Compare
csarofeen
left a comment
There was a problem hiding this comment.
I would like to go though this logic a little more carefully with you in the future, but don't see any major issues except one comment. Approve for now.
| self.assertGraphContains(t_jit3.graph_for(x, y, 2.0), FUSION_GROUP) | ||
|
|
||
| # Testing partition logic that is capable to avoid creating unsupported | ||
| # broadcasting semantics in CudaFusionGroup |
There was a problem hiding this comment.
Seems this test is figuring out what it can generate, not that it avoids unsupported cases.
There was a problem hiding this comment.
That is true.
If the partition does create unsupported case, it would error out. So it sorta checks the right thing.
This is not ideal tho, I can put some more lines there to traverse the fused graph.
There was a problem hiding this comment.
Just FYI, getting the bailout subgraph is trickier than what I thought it would be. I'm going to split the test into two instead of using the BailOut thing 🤦
There was a problem hiding this comment.
I'm going to leave it as-is for now.
Realized that relying on bailOut test is a poor strategy. I'll do a follow up PR to overhaul some related tests as well.
| n->outputs().size() == 1, | ||
| "not expecting multiple outputs from a node, graph partitioning logic needs to be updated"); | ||
| // assumes that if output is not a tensor type, it's not broadcasting | ||
| if (auto out_type = n->output(0)->type()->cast<TensorType>()) { |
There was a problem hiding this comment.
I don't spot anything wrong with it.
We assert on single output the line above and we check to see if the output is TensorType before we start checking broadcasting.
Are you suggesting to assert on TensorType instead? I didn't do that in the hope of making the code future proof-ish, in the case we might have nodes that extract information of the tensor like aten::size()
| // exists in input `shape` | ||
| for (const auto& opt_size : shape) { | ||
| // TODO: not sure if we need to check for output size != 1, since we | ||
| // are currently marking all size-1 dimension as broadcast in codegen. |
There was a problem hiding this comment.
Yes you should. We mark it as broadcast, but something is only "truly" a broadcast when it gets broadcasted. Broadcast op is almost like a "maybe broadcast".
There was a problem hiding this comment.
Doesn't have to be in this PR, can leave this TODO for later.
There was a problem hiding this comment.
Thanks, I'd like to keep it TODO for now.
I'm only worried that it would still trigger propagation failure even if it's not a real broadcast.
|
Merging this one as-is, please comment if I miss-read any review, I can accommodate them in later PR. |
Update partition logic resolves #190