Add support for nonzero, some improvements to reduce guards by ezyang · Pull Request #95387 · pytorch/pytorch

ezyang · 2023-02-23T18:15:13Z

Stack from ghstack (oldest at bottom):

This takes the strategy described in https://docs.google.com/document/d/1lFRYAJo5nrfxRhwIzGnfi2pbLpU6T4ytSRSuLJ5qebI/edit#

It is essentially #95222 but squashed and with changes that are unnecessary given that we assume nonzero returns > 1.

What's in the PR:

nonzero now supports meta propagation. When capture_dynamic_output_shape_ops, it will return a tensor with an unbacked SymInt representing the size in question.
The unbacked SymInt is UNSOUNDLY assumed to be not equal to 0/1. We will still error if you guard otherwise.
PrimTorch pointwise operators are updated to use empty_permuted, to avoid guarding on unbacked SymInt from empty_strided (tested in test_dynamic_pointwise_scalar)
Convolution is updated to skip backend selection if batch is unbacked, to avoid guarding on unbacked SymInt (tested in test_unbacked_batch_resnet)
I kept the helper utilities like definitely_true for working with possibly unbacked SymInts. They're not used right now but maybe someone will find them useful.
Added constrain_unify to let you specify two unbacked SymInts must have the same value

Signed-off-by: Edward Z. Yang ezyang@meta.com

cc @soumith @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]

pytorch-bot · 2023-02-23T18:20:43Z

🔗 Helpful Links

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

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

As of commit 8a25bc6:

NEW FAILURES - The following jobs have failed:

cuda11.7-py3.10-gcc7-sm80 / test (inductor_torchbench_smoketest_perf, 1, 1, linux.gcp.a100) (gh)

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

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

This takes the strategy described in https://docs.google.com/document/d/1lFRYAJo5nrfxRhwIzGnfi2pbLpU6T4ytSRSuLJ5qebI/edit# It is essentially #95222 but squashed and with changes that are unnecessary given that we assume nonzero returns > 1. What's in the PR: * nonzero now supports meta propagation. When `capture_dynamic_output_shape_ops`, it will return a tensor with an unbacked SymInt representing the size in question. * The unbacked SymInt is UNSOUNDLY assumed to be not equal to 0/1. We will still error if you guard otherwise. * PrimTorch pointwise operators are updated to use empty_permuted, to avoid guarding on unbacked SymInt from empty_strided (tested in `test_dynamic_pointwise_scalar`) * Convolution is updated to skip backend selection if batch is unbacked, to avoid guarding on unbacked SymInt (tested in `test_unbacked_batch_resnet`) * I kept the helper utilities like `definitely_true` for working with possibly unbacked SymInts. They're not used right now but maybe someone will find them useful. * Added `constrain_unify` to let you specify two unbacked SymInts must have the same value Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

voznesenskym · 2023-02-23T18:50:39Z

test/functorch/test_aotdispatch.py

+            # because that will reject known to be good tests see
+            # https://github.com/pytorch/pytorch/issues/94705
+            if op.name == "__getitem__":
+                self.skipTest("Dynamic output shape operation in trace")


Minor nit - which tests is this? Should we limit it to a list?

It's one of the sample inputs for getitem. The sample inputs are not named so there is no way to selectively select them.

voznesenskym · 2023-02-23T18:54:56Z

torch/_prims/__init__.py


-        return TensorMeta(device=device, shape=shape, strides=strides, dtype=dtype)
+        assert shape is not None
+        return torch.empty_permuted(shape, l2p_perm, device=device, dtype=dtype)  # type: ignore[return-value]


can you explain why this changed?

The headerdoc for this function also mentions strides, do we need to update it?

Previously we call empty_strided, which triggers a bunch of guards from looking at the strides and working out if they're contiguous or not. This empty_permuted bypasses that as it's guaranteed to be contiguous. I don't think this needs to change.

Chillee

Overall impression is that it looks good to me. Still need to look closer at some of the examples.

Chillee · 2023-02-23T19:13:58Z

torch/fx/experimental/symbolic_shapes.py

+# for backed SymInts, avoiding guards doesn't really matter in practice,
+# so I chose not to do it.
+
+def parallel_or(*args):


What is this stuff used for now? I don't think I see it used in this PR

It's not. I can delete it but it could be useful in the future so I didn't remove it.

Please keep, this is high value, even if unused, imo

voznesenskym

Someone smarter than I should look at the stride logic

This takes the strategy described in https://docs.google.com/document/d/1lFRYAJo5nrfxRhwIzGnfi2pbLpU6T4ytSRSuLJ5qebI/edit# It is essentially #95222 but squashed and with changes that are unnecessary given that we assume nonzero returns > 1. What's in the PR: * nonzero now supports meta propagation. When `capture_dynamic_output_shape_ops`, it will return a tensor with an unbacked SymInt representing the size in question. * The unbacked SymInt is UNSOUNDLY assumed to be not equal to 0/1. We will still error if you guard otherwise. * PrimTorch pointwise operators are updated to use empty_permuted, to avoid guarding on unbacked SymInt from empty_strided (tested in `test_dynamic_pointwise_scalar`) * Convolution is updated to skip backend selection if batch is unbacked, to avoid guarding on unbacked SymInt (tested in `test_unbacked_batch_resnet`) * I kept the helper utilities like `definitely_true` for working with possibly unbacked SymInts. They're not used right now but maybe someone will find them useful. * Added `constrain_unify` to let you specify two unbacked SymInts must have the same value Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 1fd1492 Pull Request resolved: #95387

voznesenskym · 2023-02-23T22:26:56Z

torch/fx/experimental/symbolic_shapes.py

+            lower = 2 if self.specialize_zero_one else 0
+            self.var_to_range[sympy_expr] = ValueRanges(lower, sympy.oo)


This takes the strategy described in https://docs.google.com/document/d/1lFRYAJo5nrfxRhwIzGnfi2pbLpU6T4ytSRSuLJ5qebI/edit# It is essentially pytorch/pytorch#95222 but squashed and with changes that are unnecessary given that we assume nonzero returns > 1. What's in the PR: * nonzero now supports meta propagation. When `capture_dynamic_output_shape_ops`, it will return a tensor with an unbacked SymInt representing the size in question. * The unbacked SymInt is UNSOUNDLY assumed to be not equal to 0/1. We will still error if you guard otherwise. * PrimTorch pointwise operators are updated to use empty_permuted, to avoid guarding on unbacked SymInt from empty_strided (tested in `test_dynamic_pointwise_scalar`) * Convolution is updated to skip backend selection if batch is unbacked, to avoid guarding on unbacked SymInt (tested in `test_unbacked_batch_resnet`) * I kept the helper utilities like `definitely_true` for working with possibly unbacked SymInts. They're not used right now but maybe someone will find them useful. * Added `constrain_unify` to let you specify two unbacked SymInts must have the same value Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch/pytorch#95387 Approved by: https://github.com/voznesenskym

…ytorch#95387)" This reverts commit 4833e47.

…95387) This takes the strategy described in https://docs.google.com/document/d/1lFRYAJo5nrfxRhwIzGnfi2pbLpU6T4ytSRSuLJ5qebI/edit# It is essentially pytorch#95222 but squashed and with changes that are unnecessary given that we assume nonzero returns > 1. What's in the PR: * nonzero now supports meta propagation. When `capture_dynamic_output_shape_ops`, it will return a tensor with an unbacked SymInt representing the size in question. * The unbacked SymInt is UNSOUNDLY assumed to be not equal to 0/1. We will still error if you guard otherwise. * PrimTorch pointwise operators are updated to use empty_permuted, to avoid guarding on unbacked SymInt from empty_strided (tested in `test_dynamic_pointwise_scalar`) * Convolution is updated to skip backend selection if batch is unbacked, to avoid guarding on unbacked SymInt (tested in `test_unbacked_batch_resnet`) * I kept the helper utilities like `definitely_true` for working with possibly unbacked SymInts. They're not used right now but maybe someone will find them useful. * Added `constrain_unify` to let you specify two unbacked SymInts must have the same value Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch#95387 Approved by: https://github.com/voznesenskym

Add support for nonzero, some improvements to reduce guards

91a5a77

Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]

ezyang requested review from Chillee, mruberry and ngimel as code owners February 23, 2023 18:15

github-actions bot added ciflow/inductor module: dynamo labels Feb 23, 2023

github-actions bot requested review from SherlockNoMad, albanD, antoniojkim, bdhirsh, jbschlosser, miladm, voznesenskym and wconstab February 23, 2023 18:15

pytorch-bot bot added the release notes: fx release notes category label Feb 23, 2023

ezyang added the ciflow/trunk Trigger trunk jobs on your pull request label Feb 23, 2023

voznesenskym reviewed Feb 23, 2023

View reviewed changes

Chillee reviewed Feb 23, 2023

View reviewed changes

voznesenskym approved these changes Feb 23, 2023

View reviewed changes

ezyang mentioned this pull request Feb 23, 2023

Memoize repeated nonzero calls to the same fake tensor #95399

Closed

ezyang added a commit that referenced this pull request Feb 23, 2023

Add support for nonzero, some improvements to reduce guards

71850a8

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 1fd1492 Pull Request resolved: #95387

voznesenskym reviewed Feb 23, 2023

View reviewed changes

pytorchmergebot added the Merged label Feb 24, 2023

pytorchmergebot closed this in 4833e47 Feb 24, 2023

msaroufim mentioned this pull request Mar 3, 2023

Remove mention of dynamo.optimize() in docs #96002

Closed

pruthvistony added a commit to ROCm/pytorch that referenced this pull request May 2, 2023

Revert "Add support for nonzero, some improvements to reduce guards (p…

4d111ea

…ytorch#95387)" This reverts commit 4833e47.

facebook-github-bot deleted the gh/ezyang/1844/head branch June 8, 2023 16:52

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for nonzero, some improvements to reduce guards#95387

Add support for nonzero, some improvements to reduce guards#95387
ezyang wants to merge 4 commits intogh/ezyang/1844/basefrom
gh/ezyang/1844/head

ezyang commented Feb 23, 2023 •

edited

Loading

Uh oh!

pytorch-bot bot commented Feb 23, 2023 •

edited

Loading

Uh oh!

voznesenskym Feb 23, 2023

Uh oh!

ezyang Feb 23, 2023

Uh oh!

voznesenskym Feb 23, 2023

Uh oh!

ezyang Feb 23, 2023

Uh oh!

Chillee left a comment

Uh oh!

Chillee Feb 23, 2023

Uh oh!

ezyang Feb 23, 2023

Uh oh!

voznesenskym Feb 23, 2023

Uh oh!

voznesenskym left a comment

Uh oh!

voznesenskym Feb 23, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		lower = 2 if self.specialize_zero_one else 0
		self.var_to_range[sympy_expr] = ValueRanges(lower, sympy.oo)

Conversation

ezyang commented Feb 23, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Feb 23, 2023 • 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/95387

❌ 1 Failures

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!

Chillee left a comment

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!

voznesenskym left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ezyang commented Feb 23, 2023 •

edited

Loading

pytorch-bot bot commented Feb 23, 2023 •

edited

Loading