[TensorExpr] Fuser: only require input shapes to be known (output shapes can be inferred). #43171
Conversation
…pes can be inferred). [ghstack-poisoned]
facebook-github-bot
added
the
oncall: jit
💊 CI failures summary and remediationsAs of commit d317823 (more details on the Dr. CI page): 💚 💚 Looks good so far! There are no failures yet. 💚 💚 This comment was automatically generated by Dr. CI (expand for details).Follow this link to opt-out of these comments for your Pull Requests.Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group. This comment has been revised 14 times. |
Because the logic is that we are merging new nodes into an existing fusion group - all the nodes in the fusion group have to be supported, that's an invariant (we start from a fusible node and only add other fusible nodes there). This check should not be needed at all, but since we've had it before I replaced it with an assert. |
|
|
||
| bool allShapesAreKnown(Node* node) { | ||
| // TODO: Relax the checks to support dynamic shapes | ||
| for (Value* output : node->outputs()) { |
There was a problem hiding this comment.
Just because you know the input shapes doesnt mean you know the output shapes:
def index(i: int):
x = torch.tensor(1, 2, 3, 4)
out = x[i]
There was a problem hiding this comment.
It does for the subset of ops NNC supports though, at least currently. When we add more complicated ops, we would need to adjust the fuser logic as well, but as it stands now, I think we always can infer output shapes when we know input shapes.
| // Symbolic checks | ||
| REQ(canHandle(producer)); | ||
| REQ((canHandle(consumer) || consumer->kind() == getTensorExprSymbol())); | ||
| AT_ASSERT(canHandle(consumer) || consumer->kind() == getTensorExprSymbol()); |
There was a problem hiding this comment.
nit: AT_ASSERT -> TORCH_INTERNAL_ASSERT
…(output shapes can be inferred)." Differential Revision: [D23178228](https://our.internmc.facebook.com/intern/diff/D23178228) [ghstack-poisoned]
…(output shapes can be inferred)." Differential Revision: [D23178228](https://our.internmc.facebook.com/intern/diff/D23178228) [ghstack-poisoned]
…(output shapes can be inferred)." Differential Revision: [D23178228](https://our.internmc.facebook.com/intern/diff/D23178228) [ghstack-poisoned]
|
@ZolotukhinM merged this pull request in 8dc4b41. |
…pes can be inferred). ghstack-source-id: ea617d8 Pull Request resolved: pytorch/pytorch#43171
4 participants
Stack from ghstack:
Differential Revision: D23178228