Unify SymIntNode and SymFloatNode into SymNode (pytorch#87817)

This refactor was prompted by challenges handling mixed int/float
operations in C++.  A previous version of this patch
added overloads for each permutation of int/float and was unwieldy
pytorch#87722  This PR takes a different
approach.

The general outline of the patch is to combine the C++ types SymIntNode
and SymFloatNode into a single type, SymNode.  This is type erased; we
no longer know statically at C++ if we have an int/float and have to test
it with the is_int()/is_float() virtual methods.  This has a number of
knock on effects.

- We no longer have C++ classes to bind to Python.  Instead, we take an
  entirely new approach to our Python API, where we have a SymInt/SymFloat
  class defined entirely in Python, which hold a SymNode (which corresponds
  to the C++ SymNode).  However, SymNode is not pybind11-bound; instead,
  it lives as-is in Python, and is wrapped into C++ SymNode using PythonSymNode
  when it goes into C++.  This implies a userland rename.

  In principle, it is also possible for the canonical implementation of SymNode
  to be written in C++, and then bound to Python with pybind11 (we have
  this code, although it is commented out.)  However, I did not implement
  this as we currently have no C++ implementations of SymNode.

  Because we do return SymInt/SymFloat from C++ bindings, the C++ binding
  code needs to know how to find these classes.  Currently, this is done
  just by manually importing torch and getting the attributes.

- Because SymInt/SymFloat are easy Python wrappers, __sym_dispatch__ now
  takes SymInt/SymFloat, rather than SymNode, bringing it in line with how
  __torch_dispatch__ works.

Some miscellaneous improvements:

- SymInt now has a constructor that takes SymNode.  Note that this
  constructor is ambiguous if you pass in a subclass of SymNode,
  so an explicit downcast is necessary.  This means toSymFloat/toSymInt
  are no more.  This is a mild optimization as it means rvalue reference
  works automatically.

- We uniformly use the caster for c10::SymInt/SymFloat, rather than
  going the long way via the SymIntNode/SymFloatNode.

- Removed some unnecessary toSymInt/toSymFloat calls in normalize_*
  functions, pretty sure this doesn't do anything.

- guard_int is now a free function, since to guard on an int you cannot
  assume the method exists.  A function can handle both int and SymInt
  inputs.

- We clean up the magic method definition code for SymInt/SymFloat/SymNode.
  ONLY the user classes (SymInt/SymFloat) get magic methods; SymNode gets
  plain methods; this is to help avoid confusion between the two types.

Signed-off-by: Edward Z. Yang <[email protected]>

cc @jansel @mlazos @soumith @voznesenskym @yanboliang @penguinwu @anijain2305
Pull Request resolved: pytorch#87817
Approved by: https://github.com/albanD, https://github.com/anjali411

Author: ezyang

.github/ci_commit_pins/xla.txt

@@ -1 +1 @@
-095ee628212f0235ad0d6908bdd514123639fc86
+1e9b8bdc75114ac6c16305c970be37a1cd2fdb1c

.lintrunner.toml

@@ -439,7 +439,7 @@ command = [
     """--error-description=\
         This line has an isinstance call that directly refers to \
         int or float.  This is error-prone because you may also \
-        have wanted to allow SymIntNode or SymFloatNode in your test.  \
+        have wanted to allow SymInt or SymFloat in your test.  \
         To suppress this lint, use an appropriate type alias defined \
         in torch._prims_common; use IntLike/FloatLike when you would accept \
         both regular and symbolic numbers, Dim for ints representing \

aten/src/ATen/FunctionalStorageImpl.cpp

@@ -95,7 +95,7 @@ c10::SymInt get_nbytes(const Tensor& value) {
   if (value.unsafeGetTensorImpl()->has_symbolic_sizes_strides()) {
     // Today, the two implementations of SymInt are in Python (proxy tensor),
     // and lazy tensor (LTC/XLA).
-    // LTC hasn't implemented SymInt support yet though (torch::lazy::SymIntNodeImpl).
+    // LTC hasn't implemented SymInt support yet though
     // Once it does, we should remove this check.
     if (value.key_set().has(c10::DispatchKey::Python)) {
       return value.storage().sym_nbytes();

aten/src/ATen/core/ivalue.h

@@ -562,7 +562,7 @@ struct TORCH_API IValue final {
   IValue(c10::SymInt i) {
     if (i.is_symbolic()) {
       tag = Tag::SymInt;
-      payload.u.as_intrusive_ptr = i.toSymIntNodeImpl().release();
+      payload.u.as_intrusive_ptr = i.toSymNodeImpl().release();
     } else {
       tag = Tag::Int;
       payload.u.as_int = i.as_int_unchecked();
@@ -578,7 +578,7 @@ struct TORCH_API IValue final {
   IValue(c10::SymFloat i) {
     if (i.is_symbolic()) {
       tag = Tag::SymFloat;
-      payload.u.as_intrusive_ptr = i.toSymFloatNodeImpl().release();
+      payload.u.as_intrusive_ptr = i.toSymNodeImpl().release();
     } else {
       tag = Tag::Double;
       payload.u.as_double = i.as_float_unchecked();
@@ -812,10 +812,10 @@ struct TORCH_API IValue final {
     // for both SymFloat and double
     if (s.isSymInt()) {
       tag = Tag::SymInt;
-      payload.u.as_intrusive_ptr = s.toSymInt().toSymIntNodeImpl().release();
+      payload.u.as_intrusive_ptr = s.toSymInt().toSymNodeImpl().release();
     } else if (s.isSymFloat()) {
       tag = Tag::SymFloat;
-      payload.u.as_intrusive_ptr = s.toSymFloat().toSymFloatNodeImpl().release();
+      payload.u.as_intrusive_ptr = s.toSymFloat().toSymNodeImpl().release();
     } else if (s.isFloatingPoint()) {
       tag = Tag::Double;
       payload.u.as_double = s.toDouble();

aten/src/ATen/core/ivalue_inl.h

@@ -219,7 +219,7 @@ inline at::Generator IValue::toGenerator() const& {
 inline c10::SymInt IValue::toSymInt() const {
   AT_ASSERT(isSymInt() || isInt(), "Expected SymInt or int but got ", tagKind());
   if (isSymInt()) {
-    return c10::SymInt::toSymInt(toIntrusivePtr<c10::SymIntNodeImpl>());
+    return c10::SymInt(toIntrusivePtr<c10::SymNodeImpl>());
   } else {
     return c10::SymInt(payload.u.as_int);
   }
@@ -228,7 +228,7 @@ inline c10::SymInt IValue::toSymInt() const {
 inline c10::SymFloat IValue::toSymFloat() const {
   AT_ASSERT(isSymFloat() || isDouble(), "Expected SymFloat or double but got ", tagKind());
   if (isSymFloat()) {
-    return c10::SymFloat::toSymFloat(toIntrusivePtr<c10::SymFloatNodeImpl>());
+    return c10::SymFloat(toIntrusivePtr<c10::SymNodeImpl>());
   } else {
     return c10::SymFloat(payload.u.as_double);
   }

aten/src/ATen/core/jit_type.h

@@ -1310,7 +1310,6 @@ struct TORCH_API SymIntType : public Type {
     return "SymInt";
   }
   std::string annotation_str_impl(TypePrinter printer = nullptr) const override {
-    // TODO: will become a Union[SymIntNodeImpl|int] in the near future
     return "int";
   }
   static const TypeKind Kind = TypeKind::SymIntType;

aten/src/ATen/test/scalar_test.cpp

@@ -194,34 +194,3 @@ TEST(TestScalar, TestFormatting) {
   ASSERT_EQ("(2,3.1)", format(Scalar(c10::complex<float>(2.0, 3.1))));
   ASSERT_EQ("4", format(Scalar(Scalar(4).toSymInt())));
 }
-
-TEST(TestSymInt, Basic) {
-  Scalar foo;
-  auto a_impl = c10::make_intrusive<c10::SymIntNodeImpl>();
-  foo = Scalar(a_impl->toSymInt());
-  ASSERT_EQ(a_impl.use_count(), 2);
-  Scalar bar{foo};
-  ASSERT_EQ(a_impl.use_count(), 3);
-  auto baz = bar;
-  ASSERT_EQ(a_impl.use_count(), 4);
-  auto foo2 = std::move(bar);
-  ASSERT_EQ(a_impl.use_count(), 4);
-  ASSERT_TRUE(foo2.isSymInt());
-  // NOLINTNEXTLINE(bugprone-use-after-move,clang-analyzer-cplusplus.Move)
-  ASSERT_TRUE(bar.isIntegral(false));
-  foo2 = SymInt(4);
-  ASSERT_FALSE(foo2.isSymInt());
-  ASSERT_EQ(foo2.toSymInt().expect_int(), 4);
-  // NOLINTNEXTLINE(clang-diagnostic-self-assign-overloaded)
-  foo2 = foo2;
-  ASSERT_FALSE(foo2.isSymInt());
-  ASSERT_EQ(foo2.toSymInt().expect_int(), 4);
-
-  ASSERT_EQ(a_impl.use_count(), 3);
-
-  ASSERT_THROW(foo.to<double>(), c10::Error);
-
-  Scalar int_s = 3;
-  TORCH_CHECK(int_s.toSymInt().expect_int(), 3);
-
-}

build_variables.bzl

@@ -958,6 +958,7 @@ libtorch_python_core_sources = [
     "torch/csrc/utils/object_ptr.cpp",
     "torch/csrc/utils/python_arg_parser.cpp",
     "torch/csrc/utils/python_dispatch.cpp",
+    "torch/csrc/utils/python_symnode.cpp",
     "torch/csrc/utils/structseq.cpp",
     "torch/csrc/utils/tensor_apply.cpp",
     "torch/csrc/utils/tensor_dtypes.cpp",

c10/core/Scalar.h

@@ -92,18 +92,17 @@ class C10_API Scalar {
 
   SymInt toSymInt() const {
     if (Tag::HAS_si == tag) {
-      return c10::SymInt::toSymInt(intrusive_ptr<SymIntNodeImpl>::reclaim_copy(
-          static_cast<SymIntNodeImpl*>(v.p)));
+      return c10::SymInt(intrusive_ptr<SymNodeImpl>::reclaim_copy(
+          static_cast<SymNodeImpl*>(v.p)));
     } else {
       return toLong();
     }
   }
 
   SymFloat toSymFloat() const {
     if (Tag::HAS_sd == tag) {
-      return c10::SymFloat::toSymFloat(
-          intrusive_ptr<SymFloatNodeImpl>::reclaim_copy(
-              static_cast<SymFloatNodeImpl*>(v.p)));
+      return c10::SymFloat(intrusive_ptr<SymNodeImpl>::reclaim_copy(
+          static_cast<SymNodeImpl*>(v.p)));
     } else {
       return toDouble();
     }

c10/core/SymFloat.cpp

@@ -1,32 +1,27 @@
 #include <c10/core/SymFloat.h>
-#include <c10/core/SymFloatNodeImpl.h>
+#include <c10/core/SymNodeImpl.h>
 #include <array>
 
 namespace c10 {
 
-SymFloatNode SymFloat::toSymFloatNodeImpl() const {
+SymNode SymFloat::toSymNodeImpl() const {
   TORCH_CHECK(is_symbolic());
-  return SymFloatNode::reclaim_copy(toSymFloatNodeImplUnowned());
+  return SymNode::reclaim_copy(toSymNodeImplUnowned());
 }
 
-static std::array<SymFloatNode, 2> normalize_symfloats(
-    SymFloat a_,
-    SymFloat b_) {
-  SymFloatNode a, b;
+static std::array<SymNode, 2> normalize_symfloats(SymFloat a_, SymFloat b_) {
+  SymNode a, b;
   if (a_.is_symbolic())
-    a = a_.toSymFloatNodeImpl();
+    a = a_.toSymNodeImpl();
   if (b_.is_symbolic())
-    b = b_.toSymFloatNodeImpl();
+    b = b_.toSymNodeImpl();
 
-  SymFloatNodeImpl* common = a ? a.get() : b.get();
-  // TODO: technically we need to check that the classes match
+  SymNodeImpl* common = a ? a.get() : b.get();
   if (!a) {
-    a = common->wrap(a_.as_float_unchecked());
-    a_.toSymFloat(a); //
+    a = common->wrap_float(a_.as_float_unchecked());
   }
   if (!b) {
-    b = common->wrap(b_.as_float_unchecked());
-    b_.toSymFloat(b);
+    b = common->wrap_float(b_.as_float_unchecked());
   }
   return {a, b};
 }
@@ -36,40 +31,36 @@ SymFloat SymFloat::operator+(SymFloat sci) const {
     return SymFloat(data_ + sci.data_);
   }
   auto res = normalize_symfloats(*this, sci);
-  return SymFloat::toSymFloat(res[0]->add(res[1]));
+  return SymFloat(res[0]->add(res[1]));
 }
 
 SymFloat SymFloat::operator-(SymFloat sci) const {
   if (!is_symbolic() && !sci.is_symbolic()) {
     return SymFloat(data_ - sci.data_);
   }
   auto res = normalize_symfloats(*this, sci);
-  return SymFloat::toSymFloat(res[0]->sub(res[1]));
+  return SymFloat(res[0]->sub(res[1]));
 }
 
 SymFloat SymFloat::operator*(SymFloat sci) const {
   if (!is_symbolic() && !sci.is_symbolic()) {
     return SymFloat(data_ * sci.data_);
   }
   auto res = normalize_symfloats(*this, sci);
-  return SymFloat::toSymFloat(res[0]->mul(res[1]));
+  return SymFloat(res[0]->mul(res[1]));
 }
 
 SymFloat SymFloat::operator/(SymFloat sci) const {
   if (!is_symbolic() && !sci.is_symbolic()) {
     return SymFloat(data_ / sci.data_);
   }
   auto res = normalize_symfloats(*this, sci);
-  return SymFloat::toSymFloat(res[0]->truediv(res[1]));
-}
-
-c10::SymFloat SymFloat::toSymFloat(SymFloatNode sin_sp) {
-  return c10::SymFloat(std::move(sin_sp));
+  return SymFloat(res[0]->truediv(res[1]));
 }
 
 std::ostream& operator<<(std::ostream& os, SymFloat s) {
   if (s.is_symbolic()) {
-    os << s.toSymFloatNodeImpl()->str();
+    os << s.toSymNodeImpl()->str();
   } else {
     os << s.as_float_unchecked();
   }

c10/core/SymFloat.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <c10/core/SymFloatNodeImpl.h>
+#include <c10/core/SymNodeImpl.h>
 #include <c10/macros/Macros.h>
 #include <c10/util/Exception.h>
 #include <c10/util/intrusive_ptr.h>
@@ -14,20 +14,21 @@ namespace c10 {
 class C10_API SymFloat {
  public:
   /*implicit*/ SymFloat(double d) : data_(d){};
-  SymFloat(SymFloatNode ptr)
-      : data_(std::numeric_limits<double>::quiet_NaN()), ptr_(std::move(ptr)){};
+  SymFloat(SymNode ptr)
+      : data_(std::numeric_limits<double>::quiet_NaN()), ptr_(std::move(ptr)) {
+    TORCH_CHECK(ptr_->is_float());
+  };
   SymFloat() : data_(0.0) {}
 
-  SymFloatNodeImpl* toSymFloatNodeImplUnowned() const {
+  SymNodeImpl* toSymNodeImplUnowned() const {
     return ptr_.get();
   }
 
-  SymFloatNodeImpl* release() && {
+  SymNodeImpl* release() && {
     return std::move(ptr_).release();
   }
 
-  SymFloatNode toSymFloatNodeImpl() const;
-  static c10::SymFloat toSymFloat(SymFloatNode sin);
+  SymNode toSymNodeImpl() const;
 
   double expect_float() const {
     TORCH_CHECK(!is_symbolic());
@@ -53,7 +54,7 @@ class C10_API SymFloat {
  private:
   // TODO: optimize to union
   double data_;
-  SymFloatNode ptr_;
+  SymNode ptr_;
 };
 
 C10_API std::ostream& operator<<(std::ostream& os, SymFloat s);