[dynamo] Represent all cells as `NewCellVariable` #140153

StrongerXi · 2024-11-08T17:23:02Z

Stack from ghstack (oldest at bottom):

In addition to NewCellVariable, Dynamo has 3 ways of modeling cell objects:

For cells captured and created by the root frame, represent them as
their contents in root_tx.symbolic_locals, which LOAD_DEREF and
STORE_DEREF update directly, without going through SideEffects.
ClosureVariable: this is created when cells from (1) are captured
by a newly created function Dynamo is about to inline. It's a handle
with a name that redirects LOAD_DEREF and STORE_DEREF back (1),
to make root_tx.symbolic_locals up-to-date.
For cells that are captured by both the root frame and some
pre-existing function Dynamo is about to inline, represent those
cells as contents, and do not allow writes to them.

Note that (2) and (3) are mainly to conform with (1) -- to make sure
Dynamo has a consistent modeling of cells for the same cell objects.

In this patch, we represent all of these cells as NewCellVariable. The
main new code paths introduced are:

using NewCellVariable to model cell objects created by the root
frame (the cells are passed in as input to InstructionTranslator),
this is what allows us to get rid of all 3 legacy paths above.
adding a new AutoDerefLocalSource to deal with the python-code
level (guards) and bytecode level (codegen) auto-dereferencing
behavior, when accessing pre-existing python cells. This also
involves a tiny update to guard manager generation.
plumbing some extra info into LocalSource and CellVariable so that
we can still emit LOAD_DEREF, STORE_DEREF, LOAD_CLOSURE (instead
of make_cell, cell_contents attribute access, and LOAD_FAST),
which is important for readability, performance, and some
assumptions bytecode_transformation.py makes.

As a result, this patch removes a lot of the now-dead code paths and
TODOs. Notably, it significantly simplified the prune_dead_locals
function, which was duplicating a lot of the logic from
prune_dead_object_new; this conveniently closes #137123.

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

[ghstack-poisoned]

pytorch-bot · 2024-11-08T17:23:06Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/140153

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

[DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit ec4e620 with merge base f98c601 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

StrongerXi · 2024-11-08T20:11:20Z

Tests are failing, here's what I've found so far:

Python 3.9 and 3.10 fail because the frame object we pass to Dynamo's eval frame callback doesn't have f_func: FunctionType (this patch uses it to retrieve the cells captured by the root frame).
This patch exposed a bug in Dynamo's handling of out= keyword for torch operators like sort -- the Python semantics is in-place mutation on the underlying tensor object, but in Dynamo, we create a new TensorVariable (through wrap_fx_proxy) and tries to replace instances of the old TensorVariable in symbolic_locals with this new one. The replacement never accounted for variables that are not in symbolic_locals, but reachable from it (e.g., within a TupleVariable).
A torch._dynamo.exc.Unsupported: reconstruct: NewCellVariable() failure ~~that's 3.11 specific. I have some idea and will look more~~. This patch exposed a small bug in OutputGraph codegen order -- in general codegen_save_tempvars should run first, to allocate and cache source for newly created objects.

StrongerXi · 2024-11-08T22:20:54Z

I'll try to create separate issues for (2) and (3) above, and fix them. (2) feels a little annoying.

[ghstack-poisoned]

StrongerXi · 2024-11-11T23:27:43Z

~~Plumb closure: Tuple[CellType] from different versions of CPython all the way to InstructionTranslator.~~
Moved to #140436 as a somewhat orthogonal change.

[ghstack-poisoned]

StrongerXi · 2024-11-13T00:07:57Z

Rebase.

[ghstack-poisoned]

StrongerXi · 2024-11-13T12:54:51Z

torch/_dynamo/source.py

+@dataclasses.dataclass(frozen=True)
+class AutoDerefLocalSource(ChainedSource):


This is mainly to make root-frame cells (the ones Python generates LOAD_DEREF and STORE_DEREF for) play well with guards. I think we can remove it if we make our f_locals (or whatever it becomes) contain cell objects without dereferencing them, which might be do-able as part of #140063 (comment)? @jansel @williamwen42

StrongerXi · 2024-11-13T12:57:15Z

torch/_dynamo/guards.py

+        elif istype(source, AutoDerefLocalSource):
+            # Guard checks run on f_locals, in which the python level
+            # auto-dereferenced cell objects are also dereferenced (e.g., rather
+            # than `f_locals` being `{ 'cell' : <cell object of int> }`, it'll
+            # be `{ 'cell' : <int> }`. So the guard manager is the same as the
+            # base guard manager.
+            assert isinstance(base_guard_manager, GuardManager)  # tame mypy
+            out = base_guard_manager


See https://github.com/pytorch/pytorch/pull/140153/files#r1840220615

StrongerXi · 2024-11-13T12:58:26Z

test/dynamo/test_higher_order_ops.py

+    def forward(self, s0: "Sym(s0)", s1: "Sym(s1)", L_y_: "f32[s0, s1]", s2: "Sym(s2)", L_x_: "f32[s2, s0]"):
        l_y_ = L_y_
+        l_x_ = L_x_

        wrap_body_1 = self.wrap_body_1
-        wrap = torch.ops.higher_order.wrap(wrap_body_1, s0, s1, l_x_, s2, l_y_);  wrap_body_1 = s0 = s1 = l_x_ = s2 = l_y_ = None
-        getitem: "f32[s0, s2]" = wrap[0];  wrap = None
+        wrap = torch.ops.higher_order.wrap(wrap_body_1, s2, s0, l_x_, s1, l_y_);  wrap_body_1 = s2 = s0 = l_x_ = s1 = l_y_ = None
+        getitem: "f32[s2, s1]" = wrap[0];  wrap = None
        return (getitem,)

    class wrap_body_1(torch.nn.Module):
-        def forward(self, s0: "Sym(s0)", s1: "Sym(s1)", l_x_: "f32[s0, s1]", s2: "Sym(s2)", l_y_: "f32[s1, s2]"):
+        def forward(self, s2: "Sym(s2)", s0: "Sym(s0)", l_x_: "f32[s2, s0]", s1: "Sym(s1)", l_y_: "f32[s0, s1]"):
            wrap_body_0 = self.wrap_body_0
-            wrap = torch.ops.higher_order.wrap(wrap_body_0, s0, s1, l_x_, s2, l_y_);  wrap_body_0 = s0 = s1 = l_x_ = s2 = l_y_ = None
-            getitem: "f32[s0, s2]" = wrap[0];  wrap = None
+            wrap = torch.ops.higher_order.wrap(wrap_body_0, s2, s0, l_x_, s1, l_y_);  wrap_body_0 = s2 = s0 = l_x_ = s1 = l_y_ = None
+            getitem: "f32[s2, s1]" = wrap[0];  wrap = None
            return (getitem,)

        class wrap_body_0(torch.nn.Module):
-            def forward(self, s0: "Sym(s0)", s1: "Sym(s1)", l_x_: "f32[s0, s1]", s2: "Sym(s2)", l_y_: "f32[s1, s2]"):
-                matmul: "f32[s0, s2]" = l_x_ @ l_y_;  l_x_ = l_y_ = None
+            def forward(self, s2: "Sym(s2)", s0: "Sym(s0)", l_x_: "f32[s2, s0]", s1: "Sym(s1)", l_y_: "f32[s0, s1]"):
+                matmul: "f32[s2, s1]" = l_x_ @ l_y_;  l_x_ = l_y_ = None


@ydwu4 does this change matter?

The order doesn't matter for this hop as long as it's deterministic.

[ghstack-poisoned]

…140154) Now that all cells are modeled as `NewCellVariable` in Dynamo, we no longer need to put cell variables into this special `closure_cells`, rather we just merge `closure_cells` with `symbolic_locals`. This allows us to merge and remove some code paths, notably make `LOAD_CLOSURE` the same as `LOAD_FAST`, and `LOAD_DEREF` & `STORE_DEREF` the same for inlining or regular `InstructionTranslator`. Pull Request resolved: pytorch#140154 Approved by: https://github.com/jansel ghstack dependencies: pytorch#140330, pytorch#140152, pytorch#140436, pytorch#140435, pytorch#140153

…ytorch#140155) This is no longer needed now that we've replaced `ClosureVariable` with `NewCellVariable`, i.e., Dynamo now treats `LOAD_CLOSURE` the same as `LOAD_FAST`. Pull Request resolved: pytorch#140155 Approved by: https://github.com/jansel, https://github.com/williamwen42 ghstack dependencies: pytorch#140330, pytorch#140152, pytorch#140436, pytorch#140435, pytorch#140153, pytorch#140154

In addition to `NewCellVariable`, Dynamo has 3 ways of modeling cell objects: 1. For cells captured and created by the root frame, represent them as their contents in `root_tx.symbolic_locals`, which `LOAD_DEREF` and `STORE_DEREF` update directly, without going through `SideEffects`. 2. `ClosureVariable`: this is created when cells from (1) are captured by a newly created function Dynamo is about to inline. It's a handle with a name that redirects `LOAD_DEREF` and `STORE_DEREF` back (1), to make `root_tx.symbolic_locals` up-to-date. 3. For cells that are captured by both the root frame and some pre-existing function Dynamo is about to inline, represent those cells as contents, and do not allow writes to them. Note that (2) and (3) are mainly to conform with (1) -- to make sure Dynamo has a consistent modeling of cells for the same cell objects. In this patch, we represent all of these cells as `NewCellVariable`. The main new code paths introduced are: - using `NewCellVariable` to model cell objects created by the root frame (the cells are passed in as input to `InstructionTranslator`), this is what allows us to get rid of all 3 legacy paths above. - adding a new `AutoDerefLocalSource` to deal with the python-code level (guards) and bytecode level (codegen) auto-dereferencing behavior, when accessing pre-existing python cells. This also involves a tiny update to guard manager generation. - plumbing some extra info into `LocalSource` and `CellVariable` so that we can still emit `LOAD_DEREF`, `STORE_DEREF`, `LOAD_CLOSURE` (instead of `make_cell`, `cell_contents` attribute access, and `LOAD_FAST`), which is important for readability, performance, and some assumptions `bytecode_transformation.py` makes. As a result, this patch removes a lot of the now-dead code paths and TODOs. Notably, it significantly simplified the `prune_dead_locals` function, which was duplicating a lot of the logic from `prune_dead_object_new`; this conveniently closes pytorch#137123. Pull Request resolved: pytorch#140153 Approved by: https://github.com/jansel ghstack dependencies: pytorch#140330, pytorch#140152, pytorch#140436, pytorch#140435

…140154) Now that all cells are modeled as `NewCellVariable` in Dynamo, we no longer need to put cell variables into this special `closure_cells`, rather we just merge `closure_cells` with `symbolic_locals`. This allows us to merge and remove some code paths, notably make `LOAD_CLOSURE` the same as `LOAD_FAST`, and `LOAD_DEREF` & `STORE_DEREF` the same for inlining or regular `InstructionTranslator`. Pull Request resolved: pytorch#140154 Approved by: https://github.com/jansel ghstack dependencies: pytorch#140330, pytorch#140152, pytorch#140436, pytorch#140435, pytorch#140153

…ytorch#140155) This is no longer needed now that we've replaced `ClosureVariable` with `NewCellVariable`, i.e., Dynamo now treats `LOAD_CLOSURE` the same as `LOAD_FAST`. Pull Request resolved: pytorch#140155 Approved by: https://github.com/jansel, https://github.com/williamwen42 ghstack dependencies: pytorch#140330, pytorch#140152, pytorch#140436, pytorch#140435, pytorch#140153, pytorch#140154

Update

2a4009d

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo labels Nov 8, 2024

StrongerXi linked an issue Nov 8, 2024 that may be closed by this pull request

Investigate making prune_dead_locals more aggressive by tracing liveness from variables in SideEffects #137123

Closed

Update

34528b0

[ghstack-poisoned]

This was referenced Nov 9, 2024

[dynamo] Fix bugs in side-effect pruning and codegen #140201

Closed

[dynamo] Restrict support for out= variants of torch operators #140202

Closed

Update

fe11191

[ghstack-poisoned]

StrongerXi added the topic: not user facing topic category label Nov 11, 2024

Update

e5bbcf0

[ghstack-poisoned]

StrongerXi mentioned this pull request Nov 11, 2024

[dynamo] Add a DynamoFrameType type above Python frame object #140330

Closed

Update

2699169

[ghstack-poisoned]

This was referenced Nov 12, 2024

[dynamo] Track from registered tensor hooks in prune_dead_object_new #140435

Closed

[dynamo] Identify pre-existing captured cells by cell id rather than content id #140436

Closed

StrongerXi added 2 commits November 12, 2024 14:31

Update

3dae280

[ghstack-poisoned]

Update

ed43202

[ghstack-poisoned]

StrongerXi added 2 commits November 12, 2024 19:17

Update

cf18ca2

[ghstack-poisoned]

Update

882ccd6

[ghstack-poisoned]

StrongerXi changed the title ~~[dynamo] Represent all cells as CellVariable~~ [dynamo] Represent all cells as NewCellVariable Nov 13, 2024

StrongerXi commented Nov 13, 2024

View reviewed changes

StrongerXi requested review from anijain2305, jansel and williamwen42 November 13, 2024 12:57

StrongerXi commented Nov 13, 2024

View reviewed changes

jansel approved these changes Nov 13, 2024

View reviewed changes

StrongerXi added 2 commits November 13, 2024 10:50

Update

31db48c

[ghstack-poisoned]

Update

ec4e620

[ghstack-poisoned]

pytorchmergebot closed this in ea1d11c Nov 15, 2024

pytorchmergebot added the Merged label Nov 15, 2024

github-actions bot deleted the gh/StrongerXi/28/head branch December 19, 2024 02:09

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[dynamo] Represent all cells as `NewCellVariable` #140153

[dynamo] Represent all cells as `NewCellVariable` #140153

Uh oh!

StrongerXi commented Nov 8, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Nov 8, 2024 •

edited

Loading

Uh oh!

StrongerXi commented Nov 8, 2024 •

edited

Loading

Uh oh!

StrongerXi commented Nov 8, 2024

Uh oh!

StrongerXi commented Nov 11, 2024 •

edited

Loading

Uh oh!

StrongerXi commented Nov 13, 2024

Uh oh!

StrongerXi Nov 13, 2024

Uh oh!

StrongerXi Nov 13, 2024

Uh oh!

StrongerXi Nov 13, 2024

Uh oh!

ydwu4 Nov 13, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

		@dataclasses.dataclass(frozen=True)
		class AutoDerefLocalSource(ChainedSource):

[dynamo] Represent all cells as NewCellVariable #140153

[dynamo] Represent all cells as NewCellVariable #140153

Uh oh!

Conversation

StrongerXi commented Nov 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/140153

❗ 1 Active SEVs

✅ No Failures

Uh oh!

StrongerXi commented Nov 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

StrongerXi commented Nov 8, 2024

Uh oh!

StrongerXi commented Nov 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

StrongerXi commented Nov 13, 2024

Uh oh!

StrongerXi Nov 13, 2024

Choose a reason for hiding this comment

Uh oh!

StrongerXi Nov 13, 2024

Choose a reason for hiding this comment

Uh oh!

StrongerXi Nov 13, 2024

Choose a reason for hiding this comment

Uh oh!

ydwu4 Nov 13, 2024

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

[dynamo] Represent all cells as `NewCellVariable` #140153

[dynamo] Represent all cells as `NewCellVariable` #140153

StrongerXi commented Nov 8, 2024 •

edited

Loading

pytorch-bot bot commented Nov 8, 2024 •

edited

Loading

StrongerXi commented Nov 8, 2024 •

edited

Loading

StrongerXi commented Nov 11, 2024 •

edited

Loading