Support large negative SymInt #99157

ezyang · 2023-04-14T15:13:42Z

Stack from ghstack (oldest at bottom):

The strategy is that we will heap allocate a LargeNegativeIntSymNodeImpl whenever we have a large negative int, so that we can keep the old is_symbolic test (now called is_heap_allocated) on SymInt. Whenever we need to do something with these ints, though, we convert them back into a plain int64_t (and then, e.g., wrap it in whatever user specificed SymNodeImpl they need.) We cannot wrap directly in the user specified SymNodeImpl as we generally do not know what the "tracing context" is from C++. We expect large negative ints to be rare, so we don't apply optimizations like singleton-ifying INT_MIN. Here's the order to review:

c10/core/SymInt.h and cpp
- is_symbolic renamed to is_heap_allocated as I needed to audit all use sites: the old is_symbolic test would return true for large negative int, but it would be wrong to then try to dispatch on the LargeNegativeIntSymNodeImpl which supports very few operations. In this file, I had to update expect_int,
- If you pass in a large negative integer, we instead heap allocate it in promote_to_negative. The function is written in a funny way to keep compact constructor code for SymInt (the heap allocation happens out of line)
- clone is now moved out-of-line
- New method maybe_as_int which will give you a constant int if it is possible, either because it's stored inline or in LargeNegativeIntSymNodeImpl. This is the preferred replacement for previous use of is_symbolic() and then as_int_unchecked().
- Rename toSymNodeImpl to toSymNode, which is more correct (since it returns a SymNode)
- Complete rewrite of normalize_symints.cpp to use new maybe_as_int. Cannot easily use the old code structure, so it's now done doing a macro and typing out each case manually (it's actually not that bad.)
- Reimplementations of all the unary operators by hand to use maybe_as_int, relatively simple.
c10/core/LargeNegativeIntSymNodeImpl.h - Just stores a int64_t value, but it has to be big and negative. Most methods are not implemented, since we will rewrap the large negative int in the real SymNodeImpl subclass before doing operations with it
The rest of the files are just rewriting code to use maybe_as_int. There is a nontrivial comment in c10/core/SymIntArrayRef.h

Very minor test adjustment in c10/test/core/SymInt_test.cpp . Plan to exercise this properly in next PR.

Companion XLA PR: pytorch/xla#4882

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

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

pytorch-bot · 2023-04-14T15:13:45Z

🔗 Helpful Links

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

📄 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 75267ac:
💚 Looks good so far! There are no failures yet. 💚

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

The strategy is that we will heap allocate a LargeNegativeIntSymNodeImpl whenever we have a large negative int, so that we can keep the old `is_symbolic` test (now called `is_heap_allocated`) on SymInt. Whenever we need to do something with these ints, though, we convert them back into a plain `int64_t` (and then, e.g., wrap it in whatever user specificed SymNodeImpl they need.) We cannot wrap directly in the user specified SymNodeImpl as we generally do not know what the "tracing context" is from C++. We expect large negative ints to be rare, so we don't apply optimizations like singleton-ifying INT_MIN. Here's the order to review: * c10/core/SymInt.h and cpp * `is_symbolic` renamed to `is_heap_allocated` as I needed to audit all use sites: the old `is_symbolic` test would return true for large negative int, but it would be wrong to then try to dispatch on the LargeNegativeIntSymNodeImpl which supports very few operations. In this file, I had to update expect_int, * If you pass in a large negative integer, we instead heap allocate it in `promote_to_negative`. The function is written in a funny way to keep compact constructor code for SymInt (the heap allocation happens out of line) * clone is now moved out-of-line * New method maybe_as_int which will give you a constant int if it is possible, either because it's stored inline or in LargeNegativeIntSymNodeImpl. This is the preferred replacement for previous use of is_symbolic() and then as_int_unchecked(). * Rename toSymNodeImpl to toSymNode, which is more correct (since it returns a SymNode) * Complete rewrite of `normalize_symints.cpp` to use new `maybe_as_int`. Cannot easily use the old code structure, so it's now done doing a macro and typing out each case manually (it's actually not that bad.) * Reimplementations of all the unary operators by hand to use `maybe_as_int`, relatively simple. * c10/core/LargeNegativeIntSymNodeImpl.h - Just stores a int64_t value, but it has to be big and negative. Most methods are not implemented, since we will rewrap the large negative int in the real SymNodeImpl subclass before doing operations with it * The rest of the files are just rewriting code to use `maybe_as_int`. There is a nontrivial comment in c10/core/SymIntArrayRef.h Very minor test adjustment in c10/test/core/SymInt_test.cpp . Plan to exercise this properly in next PR. Signed-off-by: Edward Z. Yang <ezyangmeta.com> [ghstack-poisoned]

kit1980 · 2023-04-14T18:40:13Z

.github/ci_commit_pins/xla.txt

@@ -1 +1 @@
-27fc8b28ee0a6dc5a223555b980be9fad4c697da
+6577a152c9e102d2046a73b46079d588a059118d


Did you mean to update xla hash?

albanD

Very cool! The new maybe_as_int makes a lot of these use case easier to read I think

(accept conditional on the revert of the xla pin change!)

albanD · 2023-04-14T20:41:12Z

c10/core/SymInt.h

@@ -207,22 +205,42 @@ class C10_API SymInt {

  operator SymFloat() const;

+  // Don't use this.  Prefer maybe_as_int instead


Do we still have "valid" use cases for this?

albanD · 2023-04-14T20:51:24Z

c10/core/SymInt.cpp

+      if (auto mb = sci.maybe_as_int()) {                            \
+        return RET(OP(*ma, *mb));                                    \
+      } else {                                                       \
+        auto b = sci.toSymNode();                                    \


Why not toSymNodeImplUnowned when dealing with sci?

Because the argument takes a const SymNode& so I have to manufacture a SymNode anyway. Could optimize if we change the signature here

albanD · 2023-04-14T20:58:53Z

c10/core/SymNodeImpl.h

@@ -163,6 +164,9 @@ class C10_API SymNodeImpl : public c10::intrusive_ptr_target {
  virtual std::string str() {
    TORCH_CHECK(false, "NYI");
  };
+  virtual int64_t large_negative_int() {
+    return 0; // not a large negative int!


nit: make it explicit that this function will return a large negative int if one is stored in this SymNodeImpl and 0 in all other cases (meaning a large negative int is not stored in here).
I guess we don't go with optional as we can't inline it?

albanD · 2023-04-14T21:01:59Z

torch/csrc/jit/serialization/unpickler.cpp

@@ -88,7 +88,8 @@ void restoreAccurateTypeTags(const IValue& root, const TypePtr& type_tag) {
        // no op, there is nothing to tag
        break;
      case c10::SymIntType::Kind:
-        TORCH_CHECK(!w.value.toSymInt().is_symbolic());
+        // TODO: Can this really show up though? :think:
+        TORCH_CHECK(!w.value.toSymInt().is_heap_allocated());


Is TS serializing default values of functions?

I'm not sure haha

ezyang · 2023-04-15T01:49:28Z

Blah, I fucked something up and it only shows up as a timeout on distributed lol

The strategy is that we will heap allocate a LargeNegativeIntSymNodeImpl whenever we have a large negative int, so that we can keep the old `is_symbolic` test (now called `is_heap_allocated`) on SymInt. Whenever we need to do something with these ints, though, we convert them back into a plain `int64_t` (and then, e.g., wrap it in whatever user specificed SymNodeImpl they need.) We cannot wrap directly in the user specified SymNodeImpl as we generally do not know what the "tracing context" is from C++. We expect large negative ints to be rare, so we don't apply optimizations like singleton-ifying INT_MIN. Here's the order to review: * c10/core/SymInt.h and cpp * `is_symbolic` renamed to `is_heap_allocated` as I needed to audit all use sites: the old `is_symbolic` test would return true for large negative int, but it would be wrong to then try to dispatch on the LargeNegativeIntSymNodeImpl which supports very few operations. In this file, I had to update expect_int, * If you pass in a large negative integer, we instead heap allocate it in `promote_to_negative`. The function is written in a funny way to keep compact constructor code for SymInt (the heap allocation happens out of line) * clone is now moved out-of-line * New method maybe_as_int which will give you a constant int if it is possible, either because it's stored inline or in LargeNegativeIntSymNodeImpl. This is the preferred replacement for previous use of is_symbolic() and then as_int_unchecked(). * Rename toSymNodeImpl to toSymNode, which is more correct (since it returns a SymNode) * Complete rewrite of `normalize_symints.cpp` to use new `maybe_as_int`. Cannot easily use the old code structure, so it's now done doing a macro and typing out each case manually (it's actually not that bad.) * Reimplementations of all the unary operators by hand to use `maybe_as_int`, relatively simple. * c10/core/LargeNegativeIntSymNodeImpl.h - Just stores a int64_t value, but it has to be big and negative. Most methods are not implemented, since we will rewrap the large negative int in the real SymNodeImpl subclass before doing operations with it * The rest of the files are just rewriting code to use `maybe_as_int`. There is a nontrivial comment in c10/core/SymIntArrayRef.h Very minor test adjustment in c10/test/core/SymInt_test.cpp . Plan to exercise this properly in next PR. Signed-off-by: Edward Z. Yang <ezyangmeta.com> [ghstack-poisoned]

ezyang · 2023-04-15T12:09:35Z

No, it's not me!! It fails on base commit. Rebasing to master.

The strategy is that we will heap allocate a LargeNegativeIntSymNodeImpl whenever we have a large negative int, so that we can keep the old `is_symbolic` test (now called `is_heap_allocated`) on SymInt. Whenever we need to do something with these ints, though, we convert them back into a plain `int64_t` (and then, e.g., wrap it in whatever user specificed SymNodeImpl they need.) We cannot wrap directly in the user specified SymNodeImpl as we generally do not know what the "tracing context" is from C++. We expect large negative ints to be rare, so we don't apply optimizations like singleton-ifying INT_MIN. Here's the order to review: * c10/core/SymInt.h and cpp * `is_symbolic` renamed to `is_heap_allocated` as I needed to audit all use sites: the old `is_symbolic` test would return true for large negative int, but it would be wrong to then try to dispatch on the LargeNegativeIntSymNodeImpl which supports very few operations. In this file, I had to update expect_int, * If you pass in a large negative integer, we instead heap allocate it in `promote_to_negative`. The function is written in a funny way to keep compact constructor code for SymInt (the heap allocation happens out of line) * clone is now moved out-of-line * New method maybe_as_int which will give you a constant int if it is possible, either because it's stored inline or in LargeNegativeIntSymNodeImpl. This is the preferred replacement for previous use of is_symbolic() and then as_int_unchecked(). * Rename toSymNodeImpl to toSymNode, which is more correct (since it returns a SymNode) * Complete rewrite of `normalize_symints.cpp` to use new `maybe_as_int`. Cannot easily use the old code structure, so it's now done doing a macro and typing out each case manually (it's actually not that bad.) * Reimplementations of all the unary operators by hand to use `maybe_as_int`, relatively simple. * c10/core/LargeNegativeIntSymNodeImpl.h - Just stores a int64_t value, but it has to be big and negative. Most methods are not implemented, since we will rewrap the large negative int in the real SymNodeImpl subclass before doing operations with it * The rest of the files are just rewriting code to use `maybe_as_int`. There is a nontrivial comment in c10/core/SymIntArrayRef.h Very minor test adjustment in c10/test/core/SymInt_test.cpp . Plan to exercise this properly in next PR. Companion XLA PR: pytorch/xla#4882 Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #99157 Approved by: https://github.com/albanD

Support large negative SymInt

c581c21

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

ezyang requested review from albanD and soulitzer as code owners April 14, 2023 15:13

pytorch-bot bot added the release notes: jit release notes category label Apr 14, 2023

ezyang mentioned this pull request Apr 14, 2023

SymIntify randint and randperm #98968

Closed

github-actions bot requested review from antoniojkim, bdhirsh, jbschlosser, miladm, SherlockNoMad, voznesenskym and wconstab April 14, 2023 15:13

ezyang requested a review from a team as a code owner April 14, 2023 17:54

kit1980 reviewed Apr 14, 2023

View reviewed changes

albanD approved these changes Apr 14, 2023

View reviewed changes

ezyang mentioned this pull request Apr 15, 2023

Update for SymInt API changes pytorch/xla#4882

Merged

ezyang added the suppress-api-compatibility-check Suppresses the failures of API backward-compatibility linter (Lint/bc_linter) label Apr 15, 2023

pytorchmergebot added the Merged label Apr 15, 2023

pytorchmergebot closed this in 756a86d Apr 15, 2023

facebook-github-bot deleted the gh/ezyang/1996/head branch June 8, 2023 17:01

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support large negative SymInt #99157

Support large negative SymInt #99157

ezyang commented Apr 14, 2023 •

edited

pytorch-bot bot commented Apr 14, 2023 •

edited

kit1980 Apr 14, 2023

ezyang Apr 15, 2023

albanD left a comment •

edited

albanD Apr 14, 2023

albanD Apr 14, 2023

ezyang Apr 15, 2023

albanD Apr 14, 2023

albanD Apr 14, 2023

ezyang Apr 15, 2023

ezyang commented Apr 15, 2023

ezyang commented Apr 15, 2023

		@@ -1 +1 @@
		27fc8b28ee0a6dc5a223555b980be9fad4c697da
		6577a152c9e102d2046a73b46079d588a059118d

		@@ -207,22 +205,42 @@ class C10_API SymInt {

		operator SymFloat() const;

		// Don't use this. Prefer maybe_as_int instead

Support large negative SymInt #99157

Support large negative SymInt #99157

Conversation

ezyang commented Apr 14, 2023 • edited

pytorch-bot bot commented Apr 14, 2023 • edited

🔗 Helpful Links

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

✅ No Failures

Choose a reason for hiding this comment

Choose a reason for hiding this comment

albanD left a comment • edited

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ezyang commented Apr 15, 2023

ezyang commented Apr 15, 2023

ezyang commented Apr 14, 2023 •

edited

pytorch-bot bot commented Apr 14, 2023 •

edited

albanD left a comment •

edited