Keep track of source name on all allocated SymInts #90295

ezyang · 2022-12-06T16:55:38Z

Stack from ghstack (oldest at bottom):

Wow, I had to sweat so much to get this PR out lol.

This PR enforces the invariant that whenever we allocate SymInts as part of fakeification, the SymInt is associated with a Source, and in fact we store the string source name on SymbolWithSourceName. We use 'sname' as the shorthand for source name, as 'name' is already used by sympy to name symbols.

In order to store source names, we have to plumb source names from Dynamo to PyTorch. This made doing this PR a bit bone crushing, because there are many points in the Dynamo codebase where we are improperly converting intermediate tensors into fake tensors, where there is no source (and there cannot be, because it's a frickin' intermediate tensor). I've fixed all of the really awful cases in earlier PRs in the stack. This PR is just plumbing in source names from places where we do have it.

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

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

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

pytorch-bot · 2022-12-06T16:55:41Z

🔗 Helpful Links

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

📄 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 89e7c71:

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> cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 desertfire [ghstack-poisoned]

Wow, I had to sweat so much to get this PR out lol. This PR enforces the invariant that whenever we allocate SymInts as part of fakeification, the SymInt is associated with a Source, and in fact we store the string source name on SymbolWithSourceName. We use 'sname' as the shorthand for source name, as 'name' is already used by sympy to name symbols. In order to store source names, we have to plumb source names from Dynamo to PyTorch. This made doing this PR a bit bone crushing, because there are many points in the Dynamo codebase where we are improperly converting intermediate tensors into fake tensors, where there is no source (and there cannot be, because it's a frickin' intermediate tensor). I've fixed all of the really awful cases in earlier PRs in the stack. This PR is just plumbing in source names from places where we do have it. 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 peterbell10 desertfire [ghstack-poisoned]

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

voznesenskym · 2022-12-07T17:20:51Z

In a very hacky prototype PR, (#90370) I store the name associations on a guard env structure instead of the fake tensor itself. An alternative to your PR here would be to have a structure that is just name<>symint at a higher level so as not to have str/name notions on symints.

ezyang · 2022-12-07T19:10:20Z

I'm not really sure I want to plumb the snames all the way anymore. But half of this PR is useful, which is making sure we don't call fakeify when we don't have a source in hand.

Wow, I had to sweat so much to get this PR out lol. This PR enforces the invariant that whenever we allocate SymInts as part of fakeification, the SymInt is associated with a Source, and in fact we store the string source name on SymbolWithSourceName. We use 'sname' as the shorthand for source name, as 'name' is already used by sympy to name symbols. In order to store source names, we have to plumb source names from Dynamo to PyTorch. This made doing this PR a bit bone crushing, because there are many points in the Dynamo codebase where we are improperly converting intermediate tensors into fake tensors, where there is no source (and there cannot be, because it's a frickin' intermediate tensor). I've fixed all of the really awful cases in earlier PRs in the stack. This PR is just plumbing in source names from places where we do have it. 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 peterbell10 desertfire [ghstack-poisoned]

voznesenskym · 2022-12-07T23:58:35Z

torch/_dynamo/utils.py

-def wrap_to_fake_tensor_and_record(e, tx, ignore_subclass=False):
-    # The not fake tensor check here is annoying - ideally, fake tensors never call this during wrapping.
-    # However, get_fake_value takes args and passes them through this, which may include fake tensors.
-    # see tree_map(fake_wrapper, args) in get_fake_value.


voznesenskym · 2022-12-07T23:59:21Z

torch/_dynamo/utils.py

-                e, tx.fake_mode, static_shapes, tx, ignore_subclass=ignore_subclass
+                e,
+                tx.fake_mode,
+                static_shapes,
+                tx,
+                ignore_subclass=ignore_subclass,
+                sname=sname,


Nit: this is getting large enough to warrant its own struct

voznesenskym · 2022-12-08T00:00:31Z

torch/_dynamo/utils.py

@@ -1053,7 +1049,11 @@ def get_fake_value(node, tx):
    from .exc import TorchRuntimeError, unimplemented, Unsupported

    op = node.op
-    fake_wrapper = functools.partial(wrap_to_fake_tensor_and_record, tx=tx)
+
+    def fake_wrapper(e):


this isn't a fake_wrapper anymore, imo worth rename

voznesenskym · 2022-12-08T00:01:04Z

torch/_dynamo/variables/builder.py

+            assert "source" in options
+            if options["source"] is None:
+                kwargs["static_shapes"] = True
+                kwargs["sname"] = "__constant_illegal_sname"


is this ever checked anywhere? Why do we need a __constant_illegal_sname sname?

voznesenskym · 2022-12-08T00:02:01Z

torch/_subclasses/meta_utils.py

+        if sname is None:
+            sname = f"__unknown_tensor{len(self.tensor_memo)}"


should we have a prefix pattern so we know where the unknowns are coming from?

voznesenskym · 2022-12-08T00:03:38Z

torch/fx/experimental/symbolic_shapes.py

-            self.replacements[s] = stride_expr
+            s: sympy.Symbol
+            if isinstance(stride_expr, SymbolWithSourceName):
+                s = stride_expr


is this used anywhere? did you mean for self.replacements[s] = stride_expr to be tabbed over?

oh for sym_stride.append(self.create_symintnode(s))

Tad hard to follow, maybe give the name a little length?

voznesenskym · 2022-12-08T00:04:19Z

torch/fx/experimental/symbolic_shapes.py

@@ -530,15 +545,20 @@ def create_symintnode(self, sym: "sympy.Expr"):
    # This is guaranteed to return a symbol or its negation is a sympy.Symbol,
    # but there may be a replacement that allows it to be immediately
    # simplified
-    def create_symbol(self, val: int, *, simplify: bool = True) -> "sympy.Expr":
+    def create_symbol(self, val: int, *, simplify: bool = True, sname: str) -> "sympy.Expr":
+        assert isinstance(sname, str), f"{type(sname)} {sname}"


nit: type is enough probably, if sname is a tensor this print is gonna be ugly

voznesenskym · 2022-12-08T00:05:09Z

torch/fx/experimental/symbolic_shapes.py

+    # with it; e.g., suppose we already have a tensor of size 3 in scope,
+    # which was assigned s3, then shape_env.duck_int(3) we will get back s3.


hyper-ultra-nit: change one of the digits in the ex so folks dont think s3 is 3, s2 is 2, etc.

Wow, I had to sweat so much to get this PR out lol. This PR enforces the invariant that whenever we allocate SymInts as part of fakeification, the SymInt is associated with a Source, and in fact we store the string source name on SymbolWithSourceName. We use 'sname' as the shorthand for source name, as 'name' is already used by sympy to name symbols. In order to store source names, we have to plumb source names from Dynamo to PyTorch. This made doing this PR a bit bone crushing, because there are many points in the Dynamo codebase where we are improperly converting intermediate tensors into fake tensors, where there is no source (and there cannot be, because it's a frickin' intermediate tensor). I've fixed all of the really awful cases in earlier PRs in the stack. This PR is just plumbing in source names from places where we do have it. 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 peterbell10 desertfire [ghstack-poisoned]

facebook-github-bot · 2022-12-10T13:17:46Z

This pull request has been merged in b68dead.

Keep track of source name on all allocated SymInts

4f222be

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

This was referenced Dec 6, 2022

ShapeEnv.create_symbolic_sizes_strides_storage_offset #89962

Closed

Ensure that we fakeify tensor subclasses when they are initially tracked #90009

Closed

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

github-actions bot requested review from albanD, anjali411, antoniojkim, bdhirsh, Chillee, miladm, SherlockNoMad, voznesenskym and wconstab December 6, 2022 16:55

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

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

Keep track of source name on all allocated SymInts

320f1ff

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

ezyang mentioned this pull request Dec 7, 2022

Revert guaranteed symint allocation #90381

Closed

ezyang added 3 commits December 7, 2022 11:10

voznesenskym reviewed Dec 7, 2022

View reviewed changes

voznesenskym reviewed Dec 8, 2022

View reviewed changes

voznesenskym approved these changes Dec 8, 2022

View reviewed changes

ezyang mentioned this pull request Dec 9, 2022

Completely redo how ShapeEnv guards are generated #90528

Closed

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

pytorchmergebot closed this in b68dead Dec 10, 2022

facebook-github-bot added the Merged label Dec 10, 2022

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Keep track of source name on all allocated SymInts #90295

Keep track of source name on all allocated SymInts #90295

Uh oh!

ezyang commented Dec 6, 2022 •

edited

Loading

Uh oh!

pytorch-bot bot commented Dec 6, 2022 •

edited

Loading

Uh oh!

voznesenskym commented Dec 7, 2022

Uh oh!

ezyang commented Dec 7, 2022

Uh oh!

voznesenskym Dec 7, 2022

Uh oh!

voznesenskym Dec 7, 2022

Uh oh!

voznesenskym Dec 8, 2022

Uh oh!

voznesenskym Dec 8, 2022

Uh oh!

voznesenskym Dec 8, 2022

Uh oh!

voznesenskym Dec 8, 2022

Uh oh!

voznesenskym Dec 8, 2022

Uh oh!

voznesenskym Dec 8, 2022

Uh oh!

voznesenskym Dec 8, 2022

Uh oh!

facebook-github-bot commented Dec 10, 2022

Uh oh!

Uh oh!

		if sname is None:
		sname = f"__unknown_tensor{len(self.tensor_memo)}"

		# with it; e.g., suppose we already have a tensor of size 3 in scope,
		# which was assigned s3, then shape_env.duck_int(3) we will get back s3.

Keep track of source name on all allocated SymInts #90295

Keep track of source name on all allocated SymInts #90295

Uh oh!

Conversation

ezyang commented Dec 6, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Dec 6, 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/90295

❌ 2 Failures

Uh oh!

voznesenskym commented Dec 7, 2022

Uh oh!

ezyang commented Dec 7, 2022

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!

facebook-github-bot commented Dec 10, 2022

Uh oh!

Uh oh!

ezyang commented Dec 6, 2022 •

edited

Loading

pytorch-bot bot commented Dec 6, 2022 •

edited

Loading