[fx] Move map_aggregate to C++ #148243

jansel · 2025-03-01T05:20:00Z

Stack from ghstack (oldest at bottom):

Microbenchmarking fx.symbolic_trace(lambda x: functools.reduce(operator.add, [x, *range(100000)])), before:

30603618 function calls (29403419 primitive calls) in 13.744 seconds

after:

25203549 function calls (24403352 primitive calls) in 12.090 seconds

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

[ghstack-poisoned]

pytorch-bot · 2025-03-01T05:20:03Z

🔗 Helpful Links

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

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

✅ No Failures

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

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

ghstack-source-id: 651ad20 Pull Request resolved: #148243

[ghstack-poisoned]

XuehaiPan

I wonder if there is any plan to reimplement map_aggregate to reuse the pytree infra.

XuehaiPan · 2025-03-02T08:53:17Z

torch/fx/node.py

-    elif isinstance(a, list):
-        result = immutable_list([map_aggregate(elem, fn) for elem in a])
-    elif isinstance(a, dict):
-        result = immutable_dict([(k, map_aggregate(v, fn)) for k, v in a.items()])


One blocker to use the pytree tree_map is that pytree uses type(a) is list / type(a) is dict rather than isinstance(a, list) / isinstance(a, dict). How do we want to support user subclasses, such as collections.UserList / collections.UserDict? Or we can just support list (immutable_list), dict (immutable_dict), defaultdict, OrderedDict only.

jansel · 2025-03-02T17:20:07Z

@XuehaiPan map_aggregate is a backwards compatibility surface for FX and the public API and behavior can't be changed. It implicitly converts all list subclasses to immutable_list.

pytorchmergebot · 2025-03-02T18:54:07Z

Starting merge as part of PR stack under #148261

Microbenchmarking `fx.symbolic_trace(lambda x: functools.reduce(operator.add, [x, *range(100000)]))`, before: ``` 25203549 function calls (24403352 primitive calls) in 12.090 seconds ``` after: ``` 24303536 function calls (23503339 primitive calls) in 10.726 seconds ``` Pull Request resolved: #148260 Approved by: https://github.com/oulgen ghstack dependencies: #148243

Microbenchmarking `fx.symbolic_trace(lambda x: functools.reduce(operator.add, [x, *range(100000)]))`, before: ``` 24303536 function calls (23503339 primitive calls) in 10.726 seconds ``` after: ``` 20003454 function calls (19203257 primitive calls) in 8.936 seconds ``` Pull Request resolved: #148261 Approved by: https://github.com/oulgen ghstack dependencies: #148243, #148260

This reverts commit edaff88. Reverted #148243 on behalf of https://github.com/jovianjaison due to breaking internal builds [T216910920] ([comment](#148243 (comment)))

pytorchmergebot · 2025-03-04T19:40:35Z

@jansel your PR has been successfully reverted.

…ytorch#148303) Pull Request resolved: pytorch#148303 Approved by: https://github.com/Skylion007 ghstack dependencies: pytorch#148243, pytorch#148260, pytorch#148261

Before: ![image](https://github.com/user-attachments/assets/3a9ed22b-ae33-41ec-a0db-01f4f3ca2ffe) After: ![image](https://github.com/user-attachments/assets/44c6e578-c63e-4a43-b3e0-d11d4bdbb6db) Pull Request resolved: pytorch#148288 Approved by: https://github.com/oulgen ghstack dependencies: pytorch#148243, pytorch#148260, pytorch#148261, pytorch#148303

…ch#148292) Before: 19502951 function calls (18702776 primitive calls) in 8.533 seconds After: 16402551 function calls (15602452 primitive calls) in 7.701 seconds Pull Request resolved: pytorch#148292 Approved by: https://github.com/oulgen ghstack dependencies: pytorch#148243, pytorch#148260, pytorch#148261, pytorch#148303, pytorch#148288

[ghstack-poisoned]

pytorchmergebot · 2025-03-10T16:05:09Z

Starting merge as part of PR stack under #148292

Microbenchmarking `fx.symbolic_trace(lambda x: functools.reduce(operator.add, [x, *range(100000)]))`, before: ``` 25203549 function calls (24403352 primitive calls) in 12.090 seconds ``` after: ``` 24303536 function calls (23503339 primitive calls) in 10.726 seconds ``` Pull Request resolved: #148260 Approved by: https://github.com/oulgen ghstack dependencies: #148243

Microbenchmarking `fx.symbolic_trace(lambda x: functools.reduce(operator.add, [x, *range(100000)]))`, before: ``` 24303536 function calls (23503339 primitive calls) in 10.726 seconds ``` after: ``` 20003454 function calls (19203257 primitive calls) in 8.936 seconds ``` Pull Request resolved: #148261 Approved by: https://github.com/oulgen ghstack dependencies: #148243, #148260

Before: ![image](https://github.com/user-attachments/assets/3a9ed22b-ae33-41ec-a0db-01f4f3ca2ffe) After: ![image](https://github.com/user-attachments/assets/44c6e578-c63e-4a43-b3e0-d11d4bdbb6db) Pull Request resolved: #148288 Approved by: https://github.com/oulgen ghstack dependencies: #148243, #148260, #148261

Before: 19502951 function calls (18702776 primitive calls) in 8.533 seconds After: 16402551 function calls (15602452 primitive calls) in 7.701 seconds Pull Request resolved: #148292 Approved by: https://github.com/oulgen ghstack dependencies: #148243, #148260, #148261, #148288

This reverts commit bec7bda. ghstack-source-id: 3604b06 Pull Request resolved: #150500

Update

23bae5a

[ghstack-poisoned]

pytorch-bot bot added the release notes: fx release notes category label Mar 1, 2025

jansel added a commit that referenced this pull request Mar 1, 2025

[fx] Move map_aggregate to C++

0b7b46a

ghstack-source-id: 651ad20 Pull Request resolved: #148243

facebook-github-bot added the fx label Mar 1, 2025

Update

5027a7b

[ghstack-poisoned]

This was referenced Mar 1, 2025

[fx] Move Node._update_args_kwargs to C++ #148260

Closed

[fx] Move Node._prepend/Node._remove_from_list to C++ #148261

Closed

Update

71326de

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo labels Mar 1, 2025

jansel added 3 commits March 1, 2025 17:37

Update

fbbc8f9

[ghstack-poisoned]

Update

7961de3

[ghstack-poisoned]

Update

9ccaff5

[ghstack-poisoned]

XuehaiPan reviewed Mar 2, 2025

View reviewed changes

jansel requested a review from oulgen March 2, 2025 17:26

jansel marked this pull request as ready for review March 2, 2025 17:31

jansel requested a review from zou3519 as a code owner March 2, 2025 17:31

oulgen approved these changes Mar 2, 2025

View reviewed changes

pytorchmergebot added the Merged label Mar 2, 2025

pytorchmergebot closed this in edaff88 Mar 2, 2025

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Mar 4, 2025

pytorchmergebot reopened this Mar 4, 2025

jansel added 2 commits March 8, 2025 19:28

Update

8f86de1

[ghstack-poisoned]

Update

d945153

[ghstack-poisoned]

Update

53f4232

[ghstack-poisoned]

jansel added the ciflow/trunk Trigger trunk jobs on your pull request label Mar 9, 2025

jansel added 3 commits March 9, 2025 18:03

Update

800c8b1

[ghstack-poisoned]

Update

cc0d3c0

[ghstack-poisoned]

Update

8348536

[ghstack-poisoned]

pytorchmergebot closed this in bec7bda Mar 10, 2025

clee2000 added a commit that referenced this pull request Apr 2, 2025

Revert "[fx] Move map_aggregate to C++ (#148243)"

5e34b41

This reverts commit bec7bda. ghstack-source-id: 3604b06 Pull Request resolved: #150500

github-actions bot deleted the gh/jansel/507/head branch April 12, 2025 02:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[fx] Move map_aggregate to C++ #148243

[fx] Move map_aggregate to C++ #148243

Uh oh!

jansel commented Mar 1, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Mar 1, 2025 •

edited

Loading

Uh oh!

XuehaiPan left a comment

Uh oh!

XuehaiPan Mar 2, 2025

Uh oh!

jansel commented Mar 2, 2025 •

edited

Loading

Uh oh!

pytorchmergebot commented Mar 2, 2025

Uh oh!

pytorchmergebot commented Mar 4, 2025

Uh oh!

pytorchmergebot commented Mar 10, 2025

Uh oh!

Uh oh!

[fx] Move map_aggregate to C++ #148243

[fx] Move map_aggregate to C++ #148243

Uh oh!

Conversation

jansel commented Mar 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Mar 1, 2025 • 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/148243

✅ No Failures

Uh oh!

XuehaiPan left a comment

Choose a reason for hiding this comment

Uh oh!

XuehaiPan Mar 2, 2025

Choose a reason for hiding this comment

Uh oh!

jansel commented Mar 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorchmergebot commented Mar 2, 2025

Uh oh!

pytorchmergebot commented Mar 4, 2025

Uh oh!

pytorchmergebot commented Mar 10, 2025

Uh oh!

Uh oh!

jansel commented Mar 1, 2025 •

edited

Loading

pytorch-bot bot commented Mar 1, 2025 •

edited

Loading

jansel commented Mar 2, 2025 •

edited

Loading