Meta register all foreach ops #112281

isuruf · 2023-10-27T20:00:58Z

Stack from ghstack (oldest at bottom):

cc @mruberry @ZainRizvi @ezyang @eellison @bdhirsh

[ghstack-poisoned]

pytorch-bot · 2023-10-27T20:01:01Z

🔗 Helpful Links

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

📄 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 7d4f4aa with merge base d64bc8f ():

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

pull / linux-focal-py3.8-clang10-onnx / test (default, 1, 2, linux.2xlarge) (gh)

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

lezcano

Approving as it's a strict improvement, but see the comment

torch/_meta_registrations.py

lezcano

It looks like you missed a couple ops here and there tho

[ghstack-poisoned]

lezcano · 2023-11-03T10:45:10Z

torch/_meta_registrations.py

+            scalar_op = getattr(aten, scalar_op_name).default
+            scalar_fn = decomposition_table.get(scalar_op, None)
+            if not scalar_fn:
+                scalar_fn = meta_table[scalar_op]


Isn't it enough to use the scalar_op here?

Yes, fixed now.

lezcano · 2023-11-03T10:48:09Z

torch/_meta_registrations.py

+    )
+
+
+@register_meta_foreach(
    [
        aten._foreach_abs_.default,


Now we can merge this list and the out-of-place one, right?

lezcano · 2023-11-03T10:48:25Z

Looks good overall, just two comments

[ghstack-poisoned]

isuruf · 2023-11-03T17:18:43Z

@lezcano, I changed this PR to a generic implementation for all foreach operations. Let me know what you think.

[ghstack-poisoned]

torch/_meta_registrations.py

lezcano · 2023-11-06T09:45:20Z

torch/_meta_registrations.py

-def meta__foreach_binop__list(self, other, alpha=1):
-    _check_foreach_binop_tensor_lists(self, other)
-
+def _check_meta_foreach(*args, _nlists=1, **kwargs):


kwarg is not used. Also, why using a _ in a parameter?

lezcano · 2023-11-06T09:51:20Z

torch/_meta_registrations.py

+        if i == 0:
+            nelem = len(arg)
+        else:


This check + the nelem > 0 should be outside the loop. ATM if _nlists == 1 you are not executing the nelem > 0 check.
Also, try to match the errors from the eager API whenever possible.

lezcano · 2023-11-06T09:54:19Z

torch/_meta_registrations.py

+def register_meta_foreach(ops):
+    for op_name, nlists in ops:
+        for overload in ["default", "Scalar", "Tensor", "List", "ScalarList"]:
+            op = getattr(getattr(aten, op_name), overload, None)


OpOverloadPacket already has a method that returns the list of overloads if I'm not mistaken.

lezcano · 2023-11-06T09:56:55Z

torch/_meta_registrations.py

+        ("_foreach_addcdiv", 2),
+        ("_foreach_addcmul", 2),
+        ("_foreach_lerp", 2),


Is it not possible to see that these should have a 2 rather than a 1 from the overloads directly?

[ghstack-poisoned]

isuruf · 2023-11-13T15:24:27Z

This is ready for another review.

lezcano · 2023-11-13T15:36:17Z

torch/_meta_registrations.py

-        len(self) == len(tensor1) and len(self) == len(tensor2),
-        lambda: "All input tensor lists must have the same length",
-    )
+    nlists = 0


I assume we must have nlists > 0?

And if we have a few of them, we may need to perform some sort of "all lists have the same length" check?

Yes, we need nlists > 0. Since the first argument is always a list and the dispatcher will dispatch only if the first argument is a list, I didn't add a check.

I added a check to see if all lists have the same length.

Add an assert nlists > 0 to be on the safe side.

Also, have a look at the checks that foreach functions execute in eager mode and replicate them here please.

Add an assert nlists > 0 to be on the safe side.

Sure.

Also, have a look at the checks that foreach functions execute in eager mode and replicate them here please.

Sorry, I couldn't find one. Do you have an example check?

find the definition of the _foreach_* functions within ATen and see which checks they implement

torch/_meta_registrations.py

[ghstack-poisoned]

lezcano

Ok, this LGTM! Thank you for the big push in this non-trivial issue :)

lezcano · 2023-11-15T10:03:14Z

One of the errors is real though

2023-11-15T04:28:44.5602104Z NotImplementedError: aten::foreach_maximum.List: attempted to run this operator with Meta tensors, but there was no abstract impl or Meta kernel registered.

[ghstack-poisoned]

lezcano · 2023-11-15T15:09:29Z

torch/_meta_registrations.py

-        aten._foreach_add_.List,
-        aten._foreach_sub_.List,
-        aten._foreach_mul_.List,
-        aten._foreach_div_.List,
        aten._foreach_maximum_.List,


what's the issue with this one, and the others that cannot be registered in a generic way?

foreach_maximum: There is no aten.maximum_ method or any torch.maximum_ method.

foreach_pow: Only foreach_pow has a ScalarAndTensor and needed special handling.

foreach_addc{div, mul}: The foreach_addcmul and addcmul have different signatures. The alpha parameter is a keyword argument that is a list in foreach_addcmul and a scalar in addcmul. There are also different names there. scalars in foreach_addcmul and value in addcmul.

Fair enough, mind leaving comments?

[ghstack-poisoned]

isuruf · 2023-11-20T21:49:15Z

This is ready now.

lezcano · 2023-11-21T09:22:04Z

@pytorchbot merge

pytorchmergebot · 2023-11-21T09:24:42Z

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team

Raised by workflow job

isuruf · 2023-11-21T14:18:34Z

@pytorchbot label "module: tests" "module: meta tensors"

isuruf · 2023-11-21T14:19:34Z

@pytorchbot label "release notes: Meta API"

isuruf · 2023-11-21T14:20:15Z

@pytorchbot merge

pytorchmergebot · 2023-11-21T14:22:22Z

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

Meta register all unary foreach ops

d722f67

[ghstack-poisoned]

isuruf requested review from mruberry and ngimel as code owners October 27, 2023 20:00

This was referenced Oct 27, 2023

[pt2] enable meta tests for foreach ops #112280

Closed

Fix foreach_abs dtypes for outplace #112282

Closed

Fix meta registrations for foreach_{div, mul} #112283

Closed

pytorchbot added the open source label Oct 27, 2023

lezcano approved these changes Oct 30, 2023

View reviewed changes

torch/_meta_registrations.py Outdated Show resolved Hide resolved

lezcano reviewed Oct 30, 2023

View reviewed changes

isuruf added 5 commits October 30, 2023 17:59

Update on "Meta register all unary foreach ops"

124029f

[ghstack-poisoned]

Update on "Meta register all unary foreach ops"

204e4f9

[ghstack-poisoned]

Update on "Meta register all unary foreach ops"

a49ffa8

[ghstack-poisoned]

Update on "Meta register all unary foreach ops"

f7cd7a5

[ghstack-poisoned]

Update on "Meta register all unary foreach ops"

d50d314

[ghstack-poisoned]

lezcano reviewed Nov 3, 2023

View reviewed changes

lezcano self-requested a review November 3, 2023 10:48

Update on "Meta register all unary foreach ops"

cd84609

[ghstack-poisoned]

Update on "Meta register all unary foreach ops"

f89291c

[ghstack-poisoned]

lezcano changed the title ~~Meta register all unary foreach ops~~ Meta register all foreach ops Nov 6, 2023

lezcano reviewed Nov 6, 2023

View reviewed changes

Update on "Meta register all foreach ops"

0f43d5d

[ghstack-poisoned]

lezcano reviewed Nov 13, 2023

View reviewed changes

isuruf added 3 commits November 14, 2023 23:30

Update on "Meta register all foreach ops"

6a87c5f

[ghstack-poisoned]

Update on "Meta register all foreach ops"

6c5c917

[ghstack-poisoned]

Update on "Meta register all foreach ops"

f79f7d4

[ghstack-poisoned]

lezcano approved these changes Nov 15, 2023

View reviewed changes

Update on "Meta register all foreach ops"

703d9ac

[ghstack-poisoned]

lezcano reviewed Nov 15, 2023

View reviewed changes

isuruf added 3 commits November 15, 2023 15:24

Update on "Meta register all foreach ops"

12d943d

[ghstack-poisoned]

Update on "Meta register all foreach ops"

ce2b225

[ghstack-poisoned]

Update on "Meta register all foreach ops"

7d4f4aa

[ghstack-poisoned]

lezcano approved these changes Nov 21, 2023

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Nov 21, 2023

pytorchmergebot added the merging label Nov 21, 2023

pytorchmergebot removed the merging label Nov 21, 2023

pytorch-bot bot added module: meta tensors module: tests Issues related to tests (not the torch.testing module) labels Nov 21, 2023

pytorch-bot bot added the release notes: Meta API release notes category label Nov 21, 2023

pytorchmergebot added the merging label Nov 21, 2023

pytorchmergebot added Merged and removed merging labels Nov 21, 2023

pytorchmergebot closed this in 4b7f9fa Nov 21, 2023

facebook-github-bot deleted the gh/isuruf/2/head branch November 24, 2023 15:27

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Meta register all foreach ops #112281

Meta register all foreach ops #112281

isuruf commented Oct 27, 2023 •

edited by pytorch-bot bot

pytorch-bot bot commented Oct 27, 2023 •

edited

lezcano left a comment

lezcano left a comment

lezcano Nov 3, 2023

isuruf Nov 3, 2023

lezcano Nov 3, 2023

lezcano commented Nov 3, 2023

isuruf commented Nov 3, 2023

lezcano Nov 6, 2023

lezcano Nov 6, 2023

lezcano Nov 6, 2023

lezcano Nov 6, 2023

isuruf commented Nov 13, 2023

lezcano Nov 13, 2023

lezcano Nov 13, 2023

isuruf Nov 13, 2023

lezcano Nov 13, 2023

isuruf Nov 13, 2023

lezcano Nov 13, 2023 •

edited

lezcano left a comment

lezcano commented Nov 15, 2023

lezcano Nov 15, 2023

isuruf Nov 15, 2023

lezcano Nov 15, 2023

isuruf commented Nov 20, 2023

lezcano commented Nov 21, 2023

pytorchmergebot commented Nov 21, 2023

isuruf commented Nov 21, 2023

isuruf commented Nov 21, 2023

isuruf commented Nov 21, 2023

pytorchmergebot commented Nov 21, 2023

Meta register all foreach ops #112281

Meta register all foreach ops #112281

Conversation

isuruf commented Oct 27, 2023 • edited by pytorch-bot bot

pytorch-bot bot commented Oct 27, 2023 • edited

🔗 Helpful Links

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

✅ You can merge normally! (1 Unrelated Failure)

lezcano left a comment

Choose a reason for hiding this comment

lezcano left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lezcano commented Nov 3, 2023

isuruf commented Nov 3, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

isuruf commented Nov 13, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lezcano Nov 13, 2023 • edited

Choose a reason for hiding this comment

lezcano left a comment

Choose a reason for hiding this comment

lezcano commented Nov 15, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

isuruf commented Nov 20, 2023

lezcano commented Nov 21, 2023

pytorchmergebot commented Nov 21, 2023

Merge failed

isuruf commented Nov 21, 2023

isuruf commented Nov 21, 2023

isuruf commented Nov 21, 2023

pytorchmergebot commented Nov 21, 2023

Merge started

isuruf commented Oct 27, 2023 •

edited by pytorch-bot bot

pytorch-bot bot commented Oct 27, 2023 •

edited

lezcano Nov 13, 2023 •

edited