[JIT] Normalize op aliases#38735
Conversation
[ghstack-poisoned]
💊 CI failures summary and remediationsAs of commit 1d1b87f (more details on the Dr. CI page):
ci.pytorch.org: 1 failedThis 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 15 times. |
|
I previously thought this numpy-compatible project is for logical compatible, so instead of adding 'absolute', we can reuse 'abs' and fix the non-numpy-compatible stuff inside if any. Now I see that we also need naming compatible, so adding normalization layer and support alias is reasonable to me now. 👍 Thanks. |
|
|
||
| graph = torch.jit.script(test_multiple).graph | ||
| FileCheck().check_count("prim::If", 3, exactly=True).run(graph) | ||
| print(torch.jit.script(test_multiple).code) |
There was a problem hiding this comment.
Thank you for cleaning this up!
| """ | ||
| FileCheck().run(graph_str, parse_ir(graph_str)) | ||
|
|
||
| def test_norm_aliases(self): |
There was a problem hiding this comment.
This test is cool but either in this PR or a follow-up we should make it extensible just like the method of registering aliases in the JIT is. That way when people add a new alias they can also easily extend this test.
|
This is super cool, thanks @eellison! I left a few comments requesting clarification on how to add and test new aliases. |
Follow up to my comment #36597 This adds a pass to convert op aliases into a normalized form. Having two ops generated in our IR that do the same thing makes the IR harder for downstream consumers of the IR, such as TorchScript passes but also ONNX, glow, etc. Another solution would have been to fix our code generation to only emit `aten::abs` from the start. This seems trickier, and doesn't really buy us much if we still have to expose `aten::absolute` in C++, as @glaringlee of the C++ API thinks we should. Bike shedding: maybe this should be `CanonicalizeOps` instead [ghstack-poisoned]
| namespace torch { | ||
| namespace jit { | ||
|
|
||
| static const std::unordered_map<Symbol, Symbol> alias_map = { |
There was a problem hiding this comment.
Sorry, my comment here didn't get saved, I guess. What about a comment here explaining how people should add (and test) new aliases?
There was a problem hiding this comment.
I added a comment. I would prefer not to put testing details in source code because it makes our code more verbose and the details would likely grow stale. If someone wants to look at how something is tested the typical thing is to look at git blame or ask around.
Follow up to my comment #36597 This adds a pass to convert op aliases into a normalized form. Having two ops generated in our IR that do the same thing makes the IR harder for downstream consumers of the IR, such as TorchScript passes but also ONNX, glow, etc. Another solution would have been to fix our code generation to only emit `aten::abs` from the start. This seems trickier, and doesn't really buy us much if we still have to expose `aten::absolute` in C++, as @glaringlee of the C++ API thinks we should. Bike shedding: maybe this should be `CanonicalizeOps` instead [ghstack-poisoned]
jamesr66a
left a comment
There was a problem hiding this comment.
Cool!
Do we have a plan for extending this to other use cases? Main one is e.g. aten::_convolution v.s. aten::conv*d
| namespace torch { | ||
| namespace jit { | ||
|
|
||
| // This passes converts aten ops to a normalized form. It is |
|
|
||
| auto cu = get_python_cu(); | ||
| auto name = c10::QualifiedName(qualname); | ||
| NormalizeOps(graph); |
There was a problem hiding this comment.
Is this also needed in _create_method_from_trace? Can we just embed this call in trace() (called by createGraphByTracing)? Maybe here:
https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/frontend/tracer.cpp#L454
| {aten::absolute_, aten::abs_}, | ||
| }; | ||
|
|
||
| void replaceNodeWithNewSymbol(Node* node, Symbol new_symbol) { |
There was a problem hiding this comment.
Would this be a useful method to expose on Node or Graph? I can definitely think of places I've wanted a function like this elsewhere
There was a problem hiding this comment.
Do you have any examples of other call sites ?
| // having multiple ops in our IR that do the same thing makes the IR more | ||
| // difficult to consumer for downstream user of the IR, such as our own | ||
| // optimization passes here, we convert op aliases into a standard form | ||
| bool normalizeOpAliases(graph_node_list_iterator& iter) { |
There was a problem hiding this comment.
static, or just wrap all the implementation code in an anonymous namespace
|
@jamesr66a no concrete plan to move over |
Follow up to my comment #36597 This adds a pass to convert op aliases into a normalized form. Having two ops generated in our IR that do the same thing makes the IR harder for downstream consumers of the IR, such as TorchScript passes but also ONNX, glow, etc. Another solution would have been to fix our code generation to only emit `aten::abs` from the start. This seems trickier, and doesn't really buy us much if we still have to expose `aten::absolute` in C++, as @glaringlee of the C++ API thinks we should. Bike shedding: maybe this should be `CanonicalizeOps` instead [ghstack-poisoned]
Summary: Pull Request resolved: pytorch#38735 Follow up to my comment pytorch#36597 This adds a pass to convert op aliases into a normalized form. Having two ops generated in our IR that do the same thing makes the IR harder for downstream consumers of the IR, such as TorchScript passes but also ONNX, glow, etc. Another solution would have been to fix our code generation to only emit `aten::abs` from the start. This seems trickier, and doesn't really buy us much if we still have to expose `aten::absolute` in C++, as glaringlee of the C++ API thinks we should. Bike shedding: maybe this should be `CanonicalizeOps` instead Test Plan: Imported from OSS Differential Revision: D21673108 Pulled By: eellison fbshipit-source-id: c328618907de1af22e07f57fd27fa619978c2817
Stack from ghstack:
Follow up to my comment #36597
This adds a pass to convert op aliases into a normalized form. Having two ops generated in our IR that do the same thing makes the IR harder for downstream consumers of the IR, such as TorchScript passes but also ONNX, glow, etc.
Another solution would have been to fix our code generation to only emit
aten::absfrom the start. This seems trickier, and doesn't really buy us much if we still have to exposeaten::absolutein C++, as @glaringlee of the C++ API thinks we should.Bike shedding: maybe this should be
CanonicalizeOpsinsteadDifferential Revision: D21673108