[dynamo] implement framelocals mapping as c++ object #140063

williamwen42 · 2024-11-07T23:16:04Z

Stack from ghstack (oldest at bottom):

Implements #93753 - move frame local guard accessors to C++.

Before, we used dict accessors on a Python dict representing the frame's fastlocals that we manually build. We move this accessor to C++ and additionally use the fastlocal index whenever possible.

Some implementation notes:

FrameLocalsMapping is now initialized as a C++ vector of PyObjects. We do not just use the frame's localsplus/fastlocals buffer because we also unbox cells.
FrameLocalsMapping can still be converted into a Python dict representing the frame's fastlocals, but it is done lazily.
We update LeafGuard, GuardAccessor, and GuardManager's check_nopybind methods to accept FrameLocalsMapping. By default, we convert the FrameLocalsMapping to a Python dict and run the original check_nopybind on it, but in some cases, conversion is not needed.
We add a new guard accessor FrameLocalsGuardAccessor, which is similar to DictGetItemGuardAccessor but has special handling for FrameLocalsMapping. We create a separate class to emphasize different use cases, but we could probably combine these two (can do in a follow up)

dynamo_guard_eval.py microbenchmark update:

713.2us -> 630.0us (3.10)
598.8us -> 530.7us (3.12)

Other followups:

Add FrameLocalsMapping version for check_verbose_nopybind in order to match behavior between check_nopybind and check_verbose_nopybind. This can prevent difficult debugging situations where guards fail (check_nopybind returns false) but no guard error message is generated (check_verbose_nopybind succeeds).
Rewrite the SHAPE_ENV guard into C++ - it is a fairly common guard that results in FrameLocalsMapping needing to convert to a dict

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

[ghstack-poisoned]

pytorch-bot · 2024-11-07T23:16:08Z

🔗 Helpful Links

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

📄 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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit c01e3fd with merge base 7ab3177 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

pull / linux-jammy-py3.10-clang15-asan / test (default, 1, 6, lf.linux.4xlarge) (gh) (detected as infra flaky with no log or failing log classifier)

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

Part of implementing #93753. Implement `FrameLocalsMapping` as a C++ `unordered_map` over a `py::dict`. This PR ensures we can pass around the C++ object around C Dynamo properly. It is still converted back into a `py::dict` for guards and convert_frame, but the next step will involve removing the `py::dict` conversion for guards. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ghstack-source-id: be903d9 Pull Request resolved: #140063

Part of implementing #93753. Implement `FrameLocalsMapping` as a C++ `unordered_map` over a `py::dict`. This PR ensures we can pass around the C++ object around C Dynamo properly. It is still converted back into a `py::dict` for guards and convert_frame, but the next step will involve removing the `py::dict` conversion for guards. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

jansel · 2024-11-09T18:10:39Z

torch/csrc/dynamo/framelocals_mapping.h

 #endif

+void framelocals_mapping_free(FrameLocalsMapping* map);
+
+// Borrowed reference


Here you say borrowed, but above you use py::reinterpret_steal which seems wrong if this is borrowed.

jansel · 2024-11-09T18:20:27Z

torch/csrc/dynamo/framelocals_mapping.h

+
+typedef struct VISIBILITY_HIDDEN FrameLocalsMapping {
+ private:
+  std::unordered_map<std::string, PyObject*> _map;


Is a std::unordered_map actually much faster than a Python dict?

The thing I was suggesting in #93753 was replacing:

f_locals["foo"]

with

fastlocals[3]

where the conversion to figure out that the variable "foo" is stored at index 3 can be done ahead of time. Thus making the lookup just a single load instruction (rather than a hashtable lookup).

Since both PyDict and unordered_map are hashtables, I'd expect similar performance.

You should compute the fastlocals offset ahead of time when the guard object is constructed, and store it on the guard. So when we evaluate the guard you can replace the hashtable lookup with an array lookup.

nit: you could use pybind11 to do this casting for you?

I don't believe there is a guarantee on the ordering of the fastlocals names (but in practice, this may be true). The safe thing to do would be to add a check that fastlocals name order is the same.

I'd also think that a C++ unordered has less overhead than using Python dict, although I'd have to measure to be sure. The point of this PR is not to introduce the complete optimization, but rather, is a step to moving framelocals/fastlocals lookup to C++. I plan on incrementally introducing optimizations in followup PRs. For this PR, we just need to make sure there are no regressions.

I don't believe there is a guarantee on the ordering of the fastlocals names (but in practice, this may be true). The safe thing to do would be to add a check that fastlocals name order is the same.

Isn't the order defined on the code object which is a constant? How could the order change?

Ah, I hadn't considered that.

I think directly using fastlocals will actually be easier to implement than a unordered_map. This also matches what CPython does, I believe:
LOAD_FAST 3 is just push(fastlocals[3])

jansel

See comments above.

Part of implementing #93753. Implement `FrameLocalsMapping` as a C++ `unordered_map` over a `py::dict`. This PR ensures we can pass around the C++ object around C Dynamo properly. It is still converted back into a `py::dict` for guards and convert_frame, but the next step will involve removing the `py::dict` conversion for guards. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

TODO: update PR description Benchmark update: - 713.2us -> 630.0us (3.10) - 598.8us -> 530.7us (3.12) cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ghstack-source-id: 67b5fac Pull Request resolved: #140063

Move frame local guard accessors to C++. Before, we used dict accessors on a Python dict representing the frame's fastlocals that we manually build. We move this accessor to C++ and additionally use the fastlocal index whenever possible. Some implementation notes: - `FrameLocalsMapping` is now initialized as a C++ vector of `PyObject`s. We do not just use the frame's localsplus/fastlocals buffer because we also unbox cells. - `FrameLocalsMapping` can still be converted into a Python dict representing the frame's fastlocals, but it is done lazily. - We update `LeafGuard`, `GuardAccessor`, and `GuardManager`'s `check_nopybind` methods to accept `FrameLocalsMapping`. By default, we convert the `FrameLocalsMapping` to a Python dict and run the original `check_nopybind` on it, but in some cases, conversion is not needed. - We add a new guard accessor `FrameLocalsGuardAccessor`, which is similar to `DictGetItemGuardAccessor` but has special handling for `FrameLocalsMapping`. dynamo_guard_eval.py microbenchmark update: - 713.2us -> 630.0us (3.10) - 598.8us -> 530.7us (3.12) cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ghstack-source-id: da32c91 Pull Request resolved: #140063

jansel · 2024-12-14T03:16:01Z

torch/_dynamo/guards.py

@@ -616,9 +617,15 @@ class GuardManagerType(enum.Enum):
    DICT_SUBCLASS_GUARD_MANAGER = 3


+@functools.lru_cache(None)


Do we need a limit here?

jansel · 2024-12-14T03:17:27Z

torch/_dynamo/guards.py

+                framelocals_names_reversed = code_framelocals_names_reversed_cached(
+                    self.f_code
+                )
+                framelocals_idx = (
+                    len(framelocals_names_reversed)
+                    - framelocals_names_reversed.index(source.local_name)
+                    - 1
+                )


Should the first one return a dict so you can compute framelocals_idx in O(1) time?

ghstack-source-id: f1f02aa Pull Request resolved: pytorch/pytorch#140063

williamwen42 · 2024-12-16T22:29:17Z

@pytorchbot rebase

pytorchmergebot · 2024-12-16T22:30:44Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[ghstack-poisoned]

ghstack-source-id: 9b5e82d Pull Request resolved: #140063

pytorchmergebot · 2024-12-16T22:31:04Z

Successfully rebased gh/williamwen42/160/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/140063)

williamwen42 · 2024-12-17T18:45:27Z

@pytorchbot merge

pytorchmergebot · 2024-12-17T18:48:20Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Summary: Implements pytorch/pytorch#93753 - move frame local guard accessors to C++. Before, we used dict accessors on a Python dict representing the frame's fastlocals that we manually build. We move this accessor to C++ and additionally use the fastlocal index whenever possible. Some implementation notes: - `FrameLocalsMapping` is now initialized as a C++ vector of `PyObject`s. We do not just use the frame's localsplus/fastlocals buffer because we also unbox cells. - `FrameLocalsMapping` can still be converted into a Python dict representing the frame's fastlocals, but it is done lazily. - We update `LeafGuard`, `GuardAccessor`, and `GuardManager`'s `check_nopybind` methods to accept `FrameLocalsMapping`. By default, we convert the `FrameLocalsMapping` to a Python dict and run the original `check_nopybind` on it, but in some cases, conversion is not needed. - We add a new guard accessor `FrameLocalsGuardAccessor`, which is similar to `DictGetItemGuardAccessor` but has special handling for `FrameLocalsMapping`. We create a separate class to emphasize different use cases, but we could probably combine these two (can do in a follow up) dynamo_guard_eval.py microbenchmark update: - 713.2us -> 630.0us (3.10) - 598.8us -> 530.7us (3.12) Other followups: - Add `FrameLocalsMapping` version for `check_verbose_nopybind` in order to match behavior between `check_nopybind` and `check_verbose_nopybind`. This can prevent difficult debugging situations where guards fail (`check_nopybind` returns false) but no guard error message is generated (`check_verbose_nopybind` succeeds). - Rewrite the `SHAPE_ENV` guard into C++ - it is a fairly common guard that results in `FrameLocalsMapping` needing to convert to a dict X-link: pytorch/pytorch#140063 Approved by: https://github.com/jansel ghstack dependencies: #142117, #142430 Reviewed By: huydhn Differential Revision: D67355697 Pulled By: williamwen42 fbshipit-source-id: 930855f49db629b2bc2d7b7ff14c065bd044e894

[dynamo] implement framelocals mapping as c++ object

b8ced06

[ghstack-poisoned]

This was referenced Nov 7, 2024

[dynamo] switch to get_framelocals_mapping for 3.11 #139950

Closed

[dynamo] switch to get_framelocals_mapping for 3.10 and below #140037

Closed

pytorch-bot bot added the module: dynamo label Nov 7, 2024

williamwen42 added topic: not user facing topic category ciflow/trunk Trigger trunk jobs on your pull request labels Nov 7, 2024

williamwen42 added a commit that referenced this pull request Nov 7, 2024

[dynamo] implement framelocals mapping as c++ object

5ed8ea1

ghstack-source-id: be903d9 Pull Request resolved: #140063

williamwen42 mentioned this pull request Nov 8, 2024

[dynamo] use raw framelocals mapping c++ object in guards.cpp #140184

Closed

williamwen42 requested review from jansel and anijain2305 November 8, 2024 22:10

jansel requested changes Nov 9, 2024

View reviewed changes

StrongerXi mentioned this pull request Nov 13, 2024

[dynamo] Represent all cells as NewCellVariable #140153

Closed

pytorch-bot bot added the ciflow/inductor label Nov 22, 2024

williamwen42 added 4 commits November 27, 2024 13:42

williamwen42 mentioned this pull request Dec 5, 2024

[dynamo] format eval_frame.c #142117

Closed

williamwen42 mentioned this pull request Dec 9, 2024

[dynamo] add benchmark for guard eval #142430

Closed

williamwen42 added 2 commits December 10, 2024 23:25

williamwen42 added a commit that referenced this pull request Dec 11, 2024

[dynamo] implement framelocals mapping as c++ object

7fb47c1

ghstack-source-id: 67b5fac Pull Request resolved: #140063

williamwen42 added a commit that referenced this pull request Dec 13, 2024

[dynamo] implement framelocals mapping as c++ object

2c146c5

ghstack-source-id: da32c91 Pull Request resolved: #140063

williamwen42 requested review from Skylion007 and jansel December 13, 2024 22:54

jansel approved these changes Dec 14, 2024

View reviewed changes

Esquains pushed a commit to Esquains/study1 that referenced this pull request Dec 15, 2024

[dynamo] implement framelocals mapping as c++ object

073c2be

ghstack-source-id: f1f02aa Pull Request resolved: pytorch/pytorch#140063

This was referenced Dec 16, 2024

[dynamo, guards] Implement FrameLocalsMapping version of check_verbose_nopybind #143308

Open

[dynamo, guards] Move SHAPE_ENV guard to C++ #143309

Open

Update

c01e3fd

[ghstack-poisoned]

pytorchmergebot pushed a commit that referenced this pull request Dec 16, 2024

[dynamo] implement framelocals mapping as c++ object

48d525c

ghstack-source-id: 9b5e82d Pull Request resolved: #140063

pytorchmergebot added the merging label Dec 17, 2024

pytorchmergebot added the Merged label Dec 17, 2024

pytorchmergebot closed this in 18261e9 Dec 17, 2024

pytorchmergebot removed the merging label Dec 17, 2024

williamwen42 mentioned this pull request Dec 17, 2024

Don't populate f_locals to check guards #93753

Closed

github-actions bot deleted the gh/williamwen42/160/head branch January 18, 2025 02:03

		@@ -616,9 +617,15 @@ class GuardManagerType(enum.Enum):
		DICT_SUBCLASS_GUARD_MANAGER = 3


		@functools.lru_cache(None)

[dynamo] implement framelocals mapping as c++ object #140063

[dynamo] implement framelocals mapping as c++ object #140063

Uh oh!

Conversation

williamwen42 commented Nov 7, 2024 • edited by pytorchmergebot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 7, 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/140063

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jansel left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

williamwen42 commented Dec 16, 2024

Uh oh!

pytorchmergebot commented Dec 16, 2024

Uh oh!

pytorchmergebot commented Dec 16, 2024

Uh oh!

williamwen42 commented Dec 17, 2024

Uh oh!

pytorchmergebot commented Dec 17, 2024

Merge started

Uh oh!

Uh oh!

williamwen42 commented Nov 7, 2024 •

edited by pytorchmergebot

Loading

pytorch-bot bot commented Nov 7, 2024 •

edited

Loading

jansel Nov 9, 2024 •

edited

Loading

jansel left a comment •

edited

Loading