Completely redo how ShapeEnv guards are generated #90528

ezyang · 2022-12-09T05:17:11Z

Stack from ghstack (oldest at bottom):

Instead of inferring shape mappings from a bunch of data structures that were plumbed in InstructionTranslator, we instead work out mappings by just iterating over the GraphArgs and mapping symbols to arguments as they show up. If multiple argument sizes/strides/offset map to the same symbol, this means they are duck sized, so we also generate extra equality tests that they must be equal. Finally, we generate 0/1 specialization guards. The resulting code is much shorter, and I think also easier to understand.

TODO: Delete all the tensor ref tracking code, it's unnecessary

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

cc @gujinghui @PenghuiCheng @XiaobingSuper @jianyuh @jgong5 @mingfeima @sanchitintel @ashokei @jingxu10 @min-jean-cho @yanbing-j @Guobing-Chen @Xia-Weiwen @mlazos @soumith @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @chunyuan-w @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @desertfire

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

pytorch-bot · 2022-12-09T05:17:13Z

🔗 Helpful Links

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

📄 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 our office hours

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

❌ 2 Failures

As of commit 8fc5d5a:

The following jobs have failed:

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

Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: f7afdb2 Pull Request resolved: #90528

ezyang · 2022-12-09T05:21:28Z

torch/fx/experimental/symbolic_shapes.py

+            global COUNTER
+            sympy_expr = Symbol(f"s{COUNTER}", positive=True, integer=True)
+            COUNTER += 1
+            sympy_expr.shape_env = self


I need to fix this up: the problem is that Sympy deduplicates symbols that have the same name, and this is very confusing when there are multiple ShapeEnvs over different generations. Need to figure out proper way to impede this.

Can we just initialize a sympy with a starting int? Sympy then increments this as it works, and we read it out as we wrap up the frame.

Then, when you have "generation"-al systems like dynamo, the lifecycle is:

Enter frame

Make shape_env(counter_pos)

do stuff (shape_env increments counter_pos as it goes)

Exit frame, record counter_pos + 1 at shape_env end time

Enter frame

Make shape_env(counter_pos)
...
etc

My main problem is that I want the number to reset every fresh frame, because I generally want my compilation to be indifferent to what other computation went on. So I can't have a generation counter either.

Sorry, I think I misunderstood you, what is the desired behavior?

I get to number from 0 every iteration, but I have a distinct Symbol.

ezyang · 2022-12-09T05:21:51Z

torch/fx/experimental/symbolic_shapes.py

@@ -428,8 +428,16 @@ def wrapper(self, *args, **kwargs):

 # This stub exists so we can easily add metadata to sympy symbols
 class Symbol(sympy.Symbol):
-    __slots__: List[str] = []
+    __slots__: List[str] = ['snames', 'shape_env']
+    snames: List[str]


In the end, it was pretty useful for debugging purposes to see what snames produced a symbol!

ezyang · 2022-12-09T05:23:52Z

aten/src/ATen/native/TensorShape.cpp

 Tensor as_strided_tensorimpl_meta_symint(const Tensor& self, SymIntArrayRef sym_size, SymIntArrayRef sym_stride, optional<c10::SymInt> sym_storage_offset_) {
  auto sym_storage_offset = sym_storage_offset_.value_or(self.sym_storage_offset());
  auto result = at::detail::make_tensor<TensorImpl>(
      c10::TensorImpl::VIEW, Storage(self.storage()), self.key_set(), self.dtype());
-  setStrided(result, sym_size, sym_stride, sym_storage_offset);
+  setStridedUnchecked(result, sym_size, sym_stride, sym_storage_offset);


as_strided calls are the usual reason we generate a guard on a base symbol. Removing all the tests is probably the wrong thing to do, but we're also generating these guards for stupid reasons, e.g., check out this godforsaken guard

File "/data/users/ezyang/a/pytorch/torch/fx/experimental/proxy_tensor.py", line 185, in wrap_with_proxy set_meta(proxy, e) File "/data/users/ezyang/a/pytorch/torch/fx/experimental/proxy_tensor.py", line 131, in set_meta proxy.node.meta['val'] = snapshot_fake(val) File "/data/users/ezyang/a/pytorch/torch/fx/experimental/proxy_tensor.py", line 121, in snapshot_fake return val.detach() File "/data/users/ezyang/a/pytorch/torch/_subclasses/fake_tensor.py", line 896, in __torch_dispatch__ r = func(*args, **kwargs) File "/data/users/ezyang/a/pytorch/torch/_ops.py", line 285, in __call__ return self._op(*args, **kwargs or {}) File "/data/users/ezyang/a/pytorch/torch/_decomp/decompositions.py", line 1572, in nop_decomposition return aten.alias(x) File "/data/users/ezyang/a/pytorch/torch/_ops.py", line 500, in __call__ return self._op(*args, **kwargs or {}) File "/data/users/ezyang/a/pytorch/torch/_meta_registrations.py", line 1258, in meta_alias return self.view(self.shape) File "/data/users/ezyang/a/pytorch/torch/_refs/__init__.py", line 3935, in view return _reshape_view_helper(a, *shape, allow_copy=False) File "/data/users/ezyang/a/pytorch/torch/_refs/__init__.py", line 3140, in _reshape_view_helper return prims.view_of(a) File "/data/users/ezyang/a/pytorch/torch/_ops.py", line 285, in __call__ return self._op(*args, **kwargs or {}) File "/data/users/ezyang/a/pytorch/torch/_prims/__init__.py", line 1782, in _view_of_meta return a.as_strided(a.shape, a.stride(), a.storage_offset())

I can't even... what the fuck lol

voznesenskym · 2022-12-09T09:27:58Z

torch/_dynamo/guards.py

-        if not config.dynamic_shapes:
-            return None
-
-        expr_to_tensor_ref: Dict[sympy.Symbol, Dict[TensorReference, None]] = {}


do you even need TensorReference anymore?

no, i'm going to delete it

voznesenskym · 2022-12-09T09:30:16Z

torch/_dynamo/guards.py

+                except Exception:
+                    # TODO: this is getting suppressed smh
+                    logging.warning(f"failing guard allocated at {tb}")
+                    raise


holdup - today we got it to where we have no missing symbols, why are suppressing something?

the warning is not showing up for some reason when I trigger this locally haha

voznesenskym · 2022-12-09T09:45:34Z

torch/_dynamo/guards.py

+                expr_as_str = " and ".join(exprs)
+                code_parts.append(expr_as_str)
+                verbose_code_parts.append(expr_as_str)


nit move these down, we want these last, I think, as invoking sympy is costly.

not only is it costly, but it is wrong to access size before we know we have a tensor

oh yes, that too, I was trying to remember why I initially had it here.

Instead of inferring shape mappings from a bunch of data structures that were plumbed in InstructionTranslator, we instead work out mappings by just iterating over the GraphArgs and mapping symbols to arguments as they show up. If multiple argument sizes/strides/offset map to the same symbol, this means they are duck sized, so we also generate extra equality tests that they must be equal. Finally, we generate 0/1 specialization guards. The resulting code is much shorter, and I think also easier to understand. TODO: Delete all the tensor ref tracking code, it's unnecessary Signed-off-by: Edward Z. Yang <ezyangfb.com> cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: 5e07c30 Pull Request resolved: #90528

Instead of inferring shape mappings from a bunch of data structures that were plumbed in InstructionTranslator, we instead work out mappings by just iterating over the GraphArgs and mapping symbols to arguments as they show up. If multiple argument sizes/strides/offset map to the same symbol, this means they are duck sized, so we also generate extra equality tests that they must be equal. Finally, we generate 0/1 specialization guards. The resulting code is much shorter, and I think also easier to understand. TODO: Delete all the tensor ref tracking code, it's unnecessary Signed-off-by: Edward Z. Yang <ezyangfb.com> cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

Instead of inferring shape mappings from a bunch of data structures that were plumbed in InstructionTranslator, we instead work out mappings by just iterating over the GraphArgs and mapping symbols to arguments as they show up. If multiple argument sizes/strides/offset map to the same symbol, this means they are duck sized, so we also generate extra equality tests that they must be equal. Finally, we generate 0/1 specialization guards. The resulting code is much shorter, and I think also easier to understand. TODO: Delete all the tensor ref tracking code, it's unnecessary Signed-off-by: Edward Z. Yang <ezyangfb.com> cc gujinghui PenghuiCheng XiaobingSuper jianyuh jgong5 mingfeima sanchitintel ashokei jingxu10 min-jean-cho yanbing-j Guobing-Chen Xia-Weiwen mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang chunyuan-w zhuhaozhe blzheng wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: fc7f6c9 Pull Request resolved: #90528

Instead of inferring shape mappings from a bunch of data structures that were plumbed in InstructionTranslator, we instead work out mappings by just iterating over the GraphArgs and mapping symbols to arguments as they show up. If multiple argument sizes/strides/offset map to the same symbol, this means they are duck sized, so we also generate extra equality tests that they must be equal. Finally, we generate 0/1 specialization guards. The resulting code is much shorter, and I think also easier to understand. TODO: Delete all the tensor ref tracking code, it's unnecessary Signed-off-by: Edward Z. Yang <ezyangfb.com> cc gujinghui PenghuiCheng XiaobingSuper jianyuh jgong5 mingfeima sanchitintel ashokei jingxu10 min-jean-cho yanbing-j Guobing-Chen Xia-Weiwen mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang chunyuan-w zhuhaozhe blzheng wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: f838a96 Pull Request resolved: #90528

ezyang · 2022-12-10T13:31:09Z

@pytorchbot merge -f "previous ci was good, lint fix only"

pytorchmergebot · 2022-12-10T13:34:59Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

@ezyang

Fixes a minor I noticed in #90528 also a follow up to #89000. @ezyang Pull Request resolved: #90630 Approved by: https://github.com/ezyang

Signed-off-by: Eli Uriegas <eliuriegas@meta.com> Follow up to #90528 Fixes #90696 Pull Request resolved: #90704 Approved by: https://github.com/weiwangmeta, https://github.com/atalman, https://github.com/malfet

Completely redo how ShapeEnv guards are generated

a74d339

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

ezyang mentioned this pull request Dec 9, 2022

Keep track of source name on all allocated SymInts #90295

Closed

ezyang mentioned this pull request Dec 9, 2022

Revert guaranteed symint allocation #90381

Closed

pytorch-bot bot added the release notes: fx release notes category label Dec 9, 2022

github-actions bot requested review from albanD, anjali411, antoniojkim, bdhirsh, Chillee, miladm, SherlockNoMad, voznesenskym and wconstab December 9, 2022 05:17

github-actions bot added ciflow/inductor module: dynamo labels Dec 9, 2022

ezyang added a commit that referenced this pull request Dec 9, 2022

Completely redo how ShapeEnv guards are generated

97a72d5

Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: f7afdb2 Pull Request resolved: #90528

ezyang commented Dec 9, 2022

View reviewed changes

voznesenskym reviewed Dec 9, 2022

View reviewed changes

ezyang added the ciflow/trunk Trigger trunk jobs on your pull request label Dec 10, 2022

ezyang added a commit that referenced this pull request Dec 10, 2022

Completely redo how ShapeEnv guards are generated

5b679d5

Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: 5e07c30 Pull Request resolved: #90528

voznesenskym approved these changes Dec 10, 2022

View reviewed changes

github-actions bot added the module: mkldnn Related to Intel IDEEP or oneDNN (a.k.a. mkldnn) integration label Dec 10, 2022

ezyang added a commit that referenced this pull request Dec 10, 2022

Completely redo how ShapeEnv guards are generated

453408c

Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: fc7f6c9 Pull Request resolved: #90528

ezyang added a commit that referenced this pull request Dec 10, 2022

Completely redo how ShapeEnv guards are generated

c668b42

Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: f838a96 Pull Request resolved: #90528

pytorchmergebot added the Merged label Dec 10, 2022

pytorchmergebot closed this in 45109ec Dec 10, 2022

Skylion007 added a commit to Skylion007/pytorch that referenced this pull request Dec 10, 2022

Fix perf bug in pytorch#90528

9eca984

Skylion007 mentioned this pull request Dec 10, 2022

Fix perf bug in #90528 #90630

Closed

pytorchmergebot pushed a commit that referenced this pull request Dec 11, 2022

Fix perf bug in #90528 (#90630)

184f6b5

Fixes a minor I noticed in #90528 also a follow up to #89000. @ezyang Pull Request resolved: #90630 Approved by: https://github.com/ezyang

seemethere mentioned this pull request Dec 12, 2022

Guard Symbol and ShapeGuardPrinter behind HAS_SYMPY #90704

Closed

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

Completely redo how ShapeEnv guards are generated #90528

Completely redo how ShapeEnv guards are generated #90528

Uh oh!

Conversation

ezyang commented Dec 9, 2022 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Dec 9, 2022 • 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/90528

❌ 2 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!

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!

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!

ezyang commented Dec 10, 2022

Uh oh!

pytorchmergebot commented Dec 10, 2022

Merge started

Uh oh!

Uh oh!

ezyang commented Dec 9, 2022 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Dec 9, 2022 •

edited

Loading