Initial version of Dynamo capture for HigherOrderOperator #99988
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. [ghstack-poisoned]
Closed
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/99988
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ 2 Unrelated FailuresAs of commit 12c0d14: BROKEN TRUNK - The following jobs failed but were present on the merge base 0093df7:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
zou3519
added a commit
that referenced
this pull request
Apr 25, 2023
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. ghstack-source-id: 04eea8b7cfb6b5079c7615fea91ff0f1cc027b8d Pull Request resolved: #99988
zou3519
changed the title
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. [ghstack-poisoned]
zou3519
changed the title
zou3519
added a commit
that referenced
this pull request
Apr 26, 2023
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. ghstack-source-id: 272297cfe16f878c609cba59b805cfadf8cf4301 Pull Request resolved: #99988
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. [ghstack-poisoned]
zou3519
added a commit
that referenced
this pull request
Apr 26, 2023
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. ghstack-source-id: 8a3bae02fe0464c6fc8a2a0599f1b81412063830 Pull Request resolved: #99988
Comment on lines
+468
to
+471
| if len(self.tracers) > 1: | ||
| unimplemented( | ||
| "accessing attribute of nn.Module inside HigherOrderOperator" | ||
| ) |
There was a problem hiding this comment.
- Why here?
- You may have to let constant source through, or add a comment about constant source.
There was a problem hiding this comment.
Why here?
I'm not sure where else to do it. I've considered doing it in SubgraphTracer.create_proxy when it sees a get_attr call, but we can't actually tell what the get_attr call is for (e.g. if it is for a Tensor, or for a body function (which is represented as a GraphModule)) without adding more infra.
You may have to let constant source through, or add a comment about constant source.
Good point, let me play around with it.
There was a problem hiding this comment.
@voznesenskym do you have an example of what it means to have constant source? I am not sure how to construct a repro that gets here
zou3519
added a commit
that referenced
this pull request
Apr 28, 2023
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. ghstack-source-id: 8a3bae02fe0464c6fc8a2a0599f1b81412063830 Pull Request resolved: #99988
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. cc soumith voznesenskym penguinwu anijain2305 @EikanWang jgong5 @Guobing-Chen @XiaobingSuper zhuhaozhe blzheng @Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
|
Phew, resolved most of the comments. This should be ready for another pass-through if you folks are interested |
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. cc soumith voznesenskym penguinwu anijain2305 @EikanWang jgong5 @Guobing-Chen @XiaobingSuper zhuhaozhe blzheng @Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
bdhirsh
reviewed
May 1, 2023
|
|
||
| body_name = add_subgraph( | ||
| "body", torch.fx.GraphModule(tx.output.nn_modules, body_graph) | ||
| ) |
There was a problem hiding this comment.
(probably dumb) nit: it looks like add_subgraph hardcodes the name of the submodule to be cond_{name}_{i}. Is add_subgraph general infra meant for use by all higher order ops?
There was a problem hiding this comment.
Yes, this is a pre-existing problem. We'll fix that in a follow-up somewhere (tracked in #100278)
bdhirsh
reviewed
May 1, 2023
torch/_dynamo/output_graph.py
Outdated
| # input of the root SubgraphTracer. | ||
| # | ||
| # After that happens, since the variable is now being tracked, | ||
| # we are now back to case 1. |
There was a problem hiding this comment.
Confirming my understanding - is this note effectively saying that:
(1) we will always be in case 1 by the time we get to this code.
(2) Any "free variable" we encounter is guaranteed to already have been lifted into a graph input of the root graph. The additional work we need to do is only that we need to lift this free variable into inputs (recursively) of every nested higher-order-op subgraph that we're currently inside of, but not the root graph.
There was a problem hiding this comment.
Yes! I'll make that clearer
There was a problem hiding this comment.
There's a small caveat to 2:
- any "free variable" we encounter is either already a graph input of the root graph, or some bound local variable of some subgraph.
- The additional work we need to do is to lift the free variable into inputs (recursively) of every nested higher-order-op subgraph until we hit the subgraph where the free variable is bound
|
Really cool! |
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. cc soumith voznesenskym penguinwu anijain2305 @EikanWang jgong5 @Guobing-Chen @XiaobingSuper zhuhaozhe blzheng @Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
This PR introduces a `wrap(body_fn, *args)` higher order operator The semantics of `wrap(body_fn, *args)` is to just run `body_fn(*args)` Underneath Dynamo, this PR makes it so that we rewrite calls to `wrap(body_fn, *args)` with `wrap(new_fn, *new_args)` where `new_fn` has no free variables. This PR does not update cond/map to use the new mechanism yet (we do not support nn.Modues yet, will come in the future). The design we take is: - OutputGraph represents the graph being built by Dynamo that may be compiled and executed. - OutputGraph owns a root SubgraphTracer, where it builds the FX graph. - OutputGraph may own multiple nested SubgraphTracers. - When we need to trace the body function of a HigherOrderOperator, we construct a new SubgraphTracer to build the graph of the body function. Mechanically, when Dynamo sees a new `wrap` HigherOrderOperator with a body function, it: - Creates a new SubgraphTracer via OutputGraph.new_subtracer - Executes the body function This captures the body function into the graph on the new SubgraphTracer while modifying the state of the OutputGraph. For example, the OutputGraph may receive new GraphArgs, new guards, and new side effects. If capture of the body function fails, then Dynamo graph breaks on the HigherOrderOperator. Test Plan: - added test/dynamo/test_higher_order_ops.py Future: - We're not actually able to tell Dynamo to completely graph break on the HigherOrderOperator. Instead, when we do graph break, Dynamo begins introspecting `HigherOrderOperator.__call__`. It should probably not do this. - Ideally we would error out on new SideEffects. I don't know how to do this yet. - We don't support dealing with nn.Modules yet (e.g. calling nn.Modules or accessing attributes of tracked nn.Modules from a body_fn). There's an open question on what should actually happen here - Ideally we would rewrite map/cond to use the new mechanism but we need to fix the previous bullet point before we can get there. cc soumith voznesenskym penguinwu anijain2305 @EikanWang jgong5 @Guobing-Chen @XiaobingSuper zhuhaozhe blzheng @Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
|
@pytorchbot merge -f "infra failures" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
@pytorchbot revert -m "breaking internal builds" -c ghfirst |
|
@pytorchbot successfully started a revert job. Check the current status here. |
|
@zou3519 your PR has been successfully reverted. |
pytorchmergebot
added a commit
that referenced
this pull request
May 3, 2023
…9988)" This reverts commit 4c99f9c. Reverted #99988 on behalf of https://github.com/atalman due to breaking internal builds ([comment](#99988 (comment)))
zou3519
added a commit
that referenced
this pull request
May 3, 2023
Original PR #99988 The problem was that we added `wrap` to torch._ops which actually puts it on `torch.ops.wrap` which is a namespace that can be open-registered to. The fix is that we now shove `wrap` into a new file [ghstack-poisoned]
zou3519
added a commit
that referenced
this pull request
May 3, 2023
Original PR #99988 The problem was that we added `wrap` to torch._ops which actually puts it on `torch.ops.wrap` which is a namespace that can be open-registered to. The fix is that we now shove `wrap` into a new file ghstack-source-id: 82d9ba3725f27afb4dd5f2b113e2dcdb9137ef79 Pull Request resolved: #100544
zou3519
added a commit
that referenced
this pull request
May 3, 2023
…Operator" Original PR #99988 The problem was that we added `wrap` to torch._ops which actually puts it on `torch.ops.wrap` which is a namespace that can be open-registered to. The fix is that we now shove `wrap` into a new file cc soumith voznesenskym penguinwu anijain2305 @EikanWang jgong5 @Guobing-Chen @XiaobingSuper zhuhaozhe blzheng @Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
zou3519
added a commit
that referenced
this pull request
May 3, 2023
Original PR #99988 The problem was that we added `wrap` to torch._ops which actually puts it on `torch.ops.wrap` which is a namespace that can be open-registered to. The fix is that we now shove `wrap` into a new file ghstack-source-id: 7f2ed64dd059da034d4576660774371887757423 Pull Request resolved: #100544
pytorchmergebot
pushed a commit
that referenced
this pull request
May 3, 2023
…100544) Original PR #99988 The problem was that we added `wrap` to torch._ops which actually puts it on `torch.ops.wrap` which is a namespace that can be open-registered to. The fix is that we now shove `wrap` into a new file Pull Request resolved: #100544 Approved by: https://github.com/voznesenskym
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack:
This PR introduces a
wrap(body_fn, *args)higher order operatorThe semantics of
wrap(body_fn, *args)is to just runbody_fn(*args)Underneath Dynamo, this PR makes it so that we rewrite calls to
wrap(body_fn, *args)withwrap(new_fn, *new_args)wherenew_fnhasno free variables. This PR does not update cond/map to use the new
mechanism yet (we do not support nn.Modues yet, will come in the future).
The design we take is:
compiled and executed.
construct a new SubgraphTracer to build the graph of the body function.
Mechanically, when Dynamo sees a new
wrapHigherOrderOperator with abody function, it:
This captures the body function into the graph on the new
SubgraphTracer while modifying the state of the OutputGraph. For
example, the OutputGraph may receive new GraphArgs, new guards, and new
side effects.
If capture of the body function fails, then Dynamo graph breaks on the
HigherOrderOperator.
Test Plan:
Future:
HigherOrderOperator. Instead, when we do graph break, Dynamo begins
introspecting
HigherOrderOperator.__call__. It should probably not dothis.
this yet.
or accessing attributes of tracked nn.Modules from a body_fn). There's
an open question on what should actually happen here
to fix the previous bullet point before we can get there.
cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire