Change SymIntNode into an intrusive pointer (#82548)

This will make the pointer type a single word, which is important
for packing it into an int64_t

This time, this diff doesn't segfault when you build with DEBUG mode; more details at pybind/pybind11#4099

Signed-off-by: Edward Z. Yang <ezyangfb.com>
Pull Request resolved: #82548
Approved by: https://github.com/albanD

Author: ezyang

.github/ci_commit_pins/xla.txt

@@ -1 +1 @@
-b3342319e96a0becd139019620d8665605b78475
+1693c1c728c05a3afda946415acc9dfd6adf421d

c10/core/SymInt.cpp

@@ -6,8 +6,11 @@
 namespace c10 {
 
 std::array<SymIntNode, 2> normalize_symints(SymInt a_, SymInt b_) {
-  SymIntNode a = a_.is_symbolic() ? a_.toSymIntNodeImpl() : nullptr;
-  SymIntNode b = b_.is_symbolic() ? b_.toSymIntNodeImpl() : nullptr;
+  SymIntNode a, b;
+  if (a_.is_symbolic())
+    a = a_.toSymIntNodeImpl();
+  if (b_.is_symbolic())
+    b = b_.toSymIntNodeImpl();
 
   SymIntNodeImpl* common = a ? a.get() : b.get();
   // TODO: technically we need to check that the classes match

c10/core/SymInt.h

@@ -2,13 +2,14 @@
 
 #include <c10/macros/Macros.h>
 #include <c10/util/Exception.h>
+#include <c10/util/intrusive_ptr.h>
 
 #include <memory>
 
 namespace c10 {
 
 class SymIntNodeImpl;
-using SymIntNode = std::shared_ptr<SymIntNodeImpl>;
+using SymIntNode = c10::intrusive_ptr<SymIntNodeImpl>;
 
 // `SymInt` is a C++ wrapper class around int64_t data_ which  and is used to
 // represent concrete dimension values.

c10/core/SymIntNodeImpl.h

@@ -3,20 +3,26 @@
 #include <c10/core/SymInt.h>
 #include <c10/macros/Macros.h>
 #include <c10/util/Exception.h>
+#include <c10/util/intrusive_ptr.h>
 #include <memory>
 #include <mutex>
 #include <vector>
 
 namespace c10 {
 
 class SymIntNodeImpl;
-using SymIntNode = std::shared_ptr<SymIntNodeImpl>;
+using SymIntNode = c10::intrusive_ptr<SymIntNodeImpl>;
 
-class C10_API SymIntNodeImpl
-    : public std::enable_shared_from_this<SymIntNodeImpl> {
+class C10_API SymIntNodeImpl : public c10::intrusive_ptr_target {
  public:
   c10::SymInt toSymInt();
   virtual ~SymIntNodeImpl(){};
+
+  template <typename T>
+  c10::intrusive_ptr<T> dyn_cast() const {
+    return c10::intrusive_ptr<T>::reclaim_copy(dynamic_cast<T*>(this));
+  }
+
   // these could be pure virtual when we implement LTC versions
   virtual SymIntNode add(const SymIntNode& other) {
     TORCH_CHECK(false, "NYI");

c10/core/SymIntTable.cpp

@@ -5,7 +5,7 @@ namespace c10 {
 uint64_t SymIntTable::addNode(SymIntNode sin) {
   std::lock_guard<std::mutex> lock(mutex_);
   auto index = nodes_.size();
-  nodes_.push_back(sin);
+  nodes_.push_back(std::move(sin));
   return index;
 }
 SymIntNode SymIntTable::getNode(size_t index) {
@@ -17,7 +17,7 @@ SymIntNode SymIntTable::getNode(size_t index) {
 c10::SymInt SymIntNodeImpl::toSymInt() {
   // We will need to figure out a way
   // to dedup nodes
-  auto sit_sp = this->shared_from_this();
+  auto sit_sp = SymIntNode::reclaim_copy(this);
   return SymInt::toSymInt(sit_sp);
 }
 

c10/test/core/SymInt_test.cpp

@@ -21,7 +21,7 @@ TEST(SymIntTest, ConcreteInts) {
 }
 
 TEST(SymIntTest, AddNode) {
-  auto n = std::make_shared<SymIntNodeImpl>();
+  auto n = c10::make_intrusive<SymIntNodeImpl>();
   auto i = n->toSymInt();
   EXPECT_TRUE(i.is_symbolic());
 }

test/cpp/lazy/test_lazy_ops.cpp

@@ -91,7 +91,7 @@ TEST(LazyDynamicOpsTest, NarrowCopy) {
   auto y = torch::rand({Y_DIM}).to(kLazy);
   auto ly = torch::lazy::TryGetLtcTensor(y);
   auto dim_node = MakeNode<SizeNode>(ly->GetIrValue(), 0);
-  auto lmn = std::make_shared<torch::lazy::SymIntNodeImpl>(dim_node);
+  auto lmn = c10::make_intrusive<torch::lazy::SymIntNodeImpl>(dim_node);
   auto z = x.narrow_copy_symint(X_DIM_INDEX, 0, lmn->toSymInt());
   AllClose(z.cpu(), x.cpu().narrow_copy(X_DIM_INDEX, 0, Y_DIM));
 }

tools/autograd/templates/python_functions.cpp

@@ -11,6 +11,7 @@
 #include <torch/csrc/autograd/python_variable.h>
 #include <torch/csrc/autograd/saved_variable.h>
 #include <pybind11/pybind11.h>
+#include <torch/csrc/utils/pybind.h>
 
 // NOTE: See [Sharded File] comment in VariableType
 

torch/csrc/Size.cpp

@@ -11,6 +11,7 @@
 
 #include <torch/csrc/autograd/python_variable.h>
 #include <torch/csrc/jit/frontend/tracer.h>
+#include <torch/csrc/utils/pybind.h>
 
 struct THPSize {
   PyTupleObject tuple;
@@ -58,6 +59,8 @@ PyObject* THPSize_NewFromSymSizes(const at::Tensor& self_) {
           !torch::jit::tracer::isTracing(),
           "JIT Tracing of SymInts isn't supported");
       auto py_symint = py::cast(si.toSymIntNodeImpl()).release().ptr();
+      if (!py_symint)
+        throw python_error();
       PyTuple_SET_ITEM(ret.get(), i, py_symint);
     } else {
       if (torch::jit::tracer::isTracing()) {

torch/csrc/jit/python/init.cpp

@@ -145,7 +145,7 @@ class PythonSymIntNodeImpl : public c10::SymIntNodeImpl {
   virtual SymIntNode wrap(int64_t num) override {
     py::gil_scoped_acquire acquire;
     auto r = getPyObj().attr("wrap")(num);
-    return std::make_shared<PythonSymIntNodeImpl>(r);
+    return c10::make_intrusive<PythonSymIntNodeImpl>(r);
   }
 
   virtual bool bool_() override {
@@ -166,11 +166,11 @@ class PythonSymIntNodeImpl : public c10::SymIntNodeImpl {
   virtual SymIntNode dispatch_common_(
       const char* fname,
       const SymIntNode& other) {
-    auto pother = std::dynamic_pointer_cast<PythonSymIntNodeImpl>(other);
+    auto pother = dynamic_cast<PythonSymIntNodeImpl*>(other.get());
     TORCH_CHECK(pother);
     py::gil_scoped_acquire acquire;
     auto r = getPyObj().attr(fname)(pother->getPyObj());
-    return std::make_shared<PythonSymIntNodeImpl>(r);
+    return c10::make_intrusive<PythonSymIntNodeImpl>(r);
   }
 
   virtual SymIntNode add(const SymIntNode& other) override {
@@ -1182,12 +1182,12 @@ void initJITBindings(PyObject* module) {
       .def_static(
           "new_symint",
           [](py::object obj) -> c10::SymIntNode {
-            return std::make_shared<PythonSymIntNodeImpl>(obj);
+            return c10::make_intrusive<PythonSymIntNodeImpl>(obj);
           })
       .def(
           "get_pyobj",
           [](c10::SymIntNode a) -> py::object {
-            if (auto psn = std::dynamic_pointer_cast<PythonSymIntNodeImpl>(a)) {
+            if (auto* psn = dynamic_cast<PythonSymIntNodeImpl*>(a.get())) {
               return py::reinterpret_borrow<py::object>(psn->getPyObj());
             }
             return py::none();

torch/csrc/lazy/core/shape_inference.cpp

@@ -452,9 +452,9 @@ std::vector<Shape> compute_shape_expand(
   for (const auto idx : c10::irange(_sizes.size())) {
     if (_sizes[idx].is_symbolic()) {
       c10::SymIntNode symbolicIntNode = _sizes[idx].toSymIntNodeImpl();
-      auto lazySymIntNode =
-          std::dynamic_pointer_cast<torch::lazy::SymIntNodeImpl>(
-              symbolicIntNode);
+      auto* lazySymIntNode =
+          dynamic_cast<torch::lazy::SymIntNodeImpl*>(symbolicIntNode.get());
+      TORCH_INTERNAL_ASSERT(lazySymIntNode);
       auto size_node = lazySymIntNode->node_;
       auto static_value =
           std::dynamic_pointer_cast<torch::lazy::DimensionNode>(size_node)

torch/csrc/lazy/core/tensor_impl.cpp

@@ -95,7 +95,7 @@ LTCTensorImpl::LTCTensorImpl(LazyTensor&& tensor)
   for (auto i : c10::irange(rank)) {
     auto dim_node = getBackend()->GetIrBuilder()->MakeSizeNode(
         this->tensor_->GetIrValue(), i);
-    auto sn = std::make_shared<torch::lazy::SymIntNodeImpl>(dim_node);
+    auto sn = c10::make_intrusive<torch::lazy::SymIntNodeImpl>(dim_node);
     sym_sizes_.push_back(sn->toSymInt());
   }
 }

torch/csrc/utils/python_arg_parser.cpp

@@ -335,6 +335,9 @@ auto handle_torch_function_no_python_arg_parser(
       // NOLINTNEXTLINE(clang-diagnostic-writable-strings)
       py::object torch_function =
           PyObject_FastGetAttrString(arg.ptr(), torch_function_name_str);
+      if (!torch_function) {
+        TORCH_INTERNAL_ASSERT(0);
+      }
 
       // See https://github.com/pytorch/pytorch/issues/63767
       if (PyObject_FastGetAttrString(torch_function.ptr(), "__self__")

torch/csrc/utils/python_arg_parser.h

@@ -487,7 +487,9 @@ inline bool is_symint_node(py::handle obj) {
 
 inline PyObject* toPyObject(c10::SymInt symint) {
   if (symint.is_symbolic()) {
-    return py::cast(symint.toSymIntNodeImpl()).release().ptr();
+    auto r = py::cast(symint.toSymIntNodeImpl()).release().ptr();
+    TORCH_INTERNAL_ASSERT(r);
+    return r;
   } else {
     return THPUtils_packInt64(symint.data());
   }

torchgen/dest/lazy_ir.py

@@ -48,10 +48,7 @@ def node_ctor_arg_rvalue_string(arg: LazyArgument) -> str:
                 return f"lazy_{arg.name}_tensorlist"
             elif arg.is_symint_or_list:
                 cpp_type = arg.lazy_type.cpp_type()
-                return (
-                    f"{cpp_type}(std::dynamic_pointer_cast<torch::lazy::SymIntNodeImpl>"
-                    f"({arg.name}.toSymIntNodeImpl())->node_, 0)"
-                )
+                return f"{cpp_type}(dynamic_cast<torch::lazy::SymIntNodeImpl*>({arg.name}.toSymIntNodeImpl().get())->node_, 0)"
             return f"lazy_{arg.name}->GetIrValue()"
         elif isinstance(arg.lazy_type, OptionalCType):
             if arg.is_wrapped_scalar: