[ghstack-poisoned]

See large code comment for details

TODO:
- tests


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]

See large code comment for details

TODO:
- tests


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]

See large code comment for details

TODO:
- tests


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]

See large code comment for details

TODO:
- tests


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]

See large code comment for details

TODO:
- tests


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]

See large code comment for details

TODO:
- tests


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]

```
# Note [Hopifying Context Managers]
#
# If the context manager class has been added to a opt-in dynamo config,
# we will convert it into a generic context manager HOP. When the
# HOP is later called in AOTAutograd, it will run the captured
# graph under the ctx.
#
# How does this work?
# - On enter:
#   - Enter into a new subtracer
# - The subtracer traces the body of the context manager into a graph
# - On exit:
#   - Exit the subtracer
#   - Grab the inner fx graph from the subtracer and install it onto
#     the outer fx graph
#   - Create node on the outer fx graph that calls this subgraph
#
# Some notes:
#
# 1. Determining the inputs and outputs
#
# One trickiness here is that a HOP requires a function as input,
# but, unlike functions, context managers don't have explicit inputs and
# outputs. For inputs, we rely on lifted_freevars, which the subtracer
# already tracks for ordinary HOPs. For outputs, ideally we might only
# return the outputs that are referenceable later, e.g. not temporaries,
# but doing that is hard, so instead we just return all intermediates
# and rely on AOTAutograd to trace through the HOP and DCE any
# unnecessary ops.
#
# 2. Fixing VariableTracker proxies
#
# Inserting the call to the subgraph into the outer fx graph creates
# fresh variable trackers and proxies, but the instruction translator
# still refers the original VTs, e.g. in its symbolic locals and these
# VTs still hold the inner fx graph's proxies, which is wrong.
#
# We fix this by updating the VTs to hold the new outer fx graph
# proxies (e.g. corresponding to the getitems of the call to the
# subgraph). To know what those VTs are, we maintain a mapping from
# fx graph node name to the VTs which is updated everytime a new
# TensorVariable is created.
 ```

From discussion in HOP meeting:
- Why do we need to put the body of the ctx in a HOP? Why not just insert two nodes into the graph?
   - There are pregrad passes and we don't want to match across the boundary
- Could there be issues for caching?
   - Yes, we store the args to reconstruct a ctx onto the GraphModule, but those args aren't necessarily hashable. In this PR we limit to allowing only constants args, though we leave an exception for AC to pass in policy_fn (we'll deal with that later).
 

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]