Python Bindings for SymInts #78135

Krovatkin · 2022-05-23T23:02:16Z

This PR adds support for SymInts in python. Namely,

THPVariable_size now returns sym_sizes()
python arg parser is modified to parse PyObjects into ints and SymbolicIntNodes
pybind11 bindings for SymbolicIntNode are added, so size expressions can be traced
a large number of tests added to demonstrate how to implement python symints.

facebook-github-bot · 2022-05-23T23:02:22Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/78135
📄 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

✅ No Failures (0 Pending)

As of commit 50810d1 (more details on the Dr. CI page):

Expand to see more

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

ezyang · 2022-06-01T20:36:46Z

c10/core/SymbolicIntNode.h

+  virtual SymbolicIntNode* wrap(int64_t num) { TORCH_CHECK(false, "NYI"); };
+  virtual bool bool_() { TORCH_CHECK(false, "NYI"); };
+  virtual int64_t int_() { TORCH_CHECK(false, "NYI"); }
+  virtual std::string str() { TORCH_CHECK(false, "NYI"); };
  virtual std::ostream& operator<<(std::ostream& os) {


does this still need to be virtual

I think it does the LTC impl won't have a python object to redispatch or fallback to

but they'll implement str instead?

OH, just realized what you meant ... operator<< does NOT to be virtual me thinks :)

ezyang · 2022-06-01T20:38:23Z

c10/core/SymbolicIntNode.h

+  virtual SymbolicIntNode* gt(SymbolicIntNode* other) { TORCH_CHECK(false, "NYI"); };
+  virtual SymbolicIntNode* lt(SymbolicIntNode* other) { TORCH_CHECK(false, "NYI"); };
+  virtual SymbolicIntNode* wrap(int64_t num) { TORCH_CHECK(false, "NYI"); };
+  virtual bool bool_() { TORCH_CHECK(false, "NYI"); };


Hmm, not sure why we need this on top of the int conversion. Is the problem that you also want a SymbolicBoolNode as well?

@horace overrode bool in his PoC, so I added it as well. I can provide a default implementation static_cast<boo>(this->int_()) or just let implementers implement via int_.

ezyang · 2022-06-01T20:39:56Z

c10/core/TensorImpl.h

@@ -1306,6 +1306,10 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
    return numel() == 0;
  }

+  // if we are going to use sym sizes, we should be setting sym strides at the same time,
+  // otherwise it's very easy to misuse this API 
+  virtual void set_sym_sizes_and_strides(c10::SymIntArrayRef sizes, c10::SymIntArrayRef strides);


it's not clear to me why this needs to be virtual

agreed, @suo mentioned he would help to formalize this API. I'll remove virtual in the meantime.

ezyang · 2022-06-01T20:41:12Z

c10/core/TensorImpl.h

@@ -2330,6 +2335,8 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
  void set_custom_device(bool custom_device) {
    custom_device_ = custom_device;
  }
+ protected:


ezyang · 2022-06-01T20:41:31Z

c10/core/TensorImpl.h

@@ -2321,6 +2325,7 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
    //
    // Can override: strides(), is_contiguous(), sizes(), dim(), numel()
    CustomSizes = 2,
+    CustomSymSizes = 3,


Skeptical about this. I'll read through the uses first

ezyang · 2022-06-01T20:45:10Z

requirements.txt

@@ -10,3 +10,5 @@ setuptools
 six
 types-dataclasses
 typing_extensions
+dataclasses; python_version<"3.7"


we're py3.7 and up only, so this really shouldn't be needed

ah this is a bad merge.. will fix..

ezyang · 2022-06-01T20:52:13Z

torch/csrc/Size.cpp

@@ -40,6 +45,29 @@ PyObject * THPSize_NewFromSizes(int dim, const int64_t *sizes)
  return self.release();
 }

+PyObject * THPSize_NewFromSymSizes(const at::Tensor& self_)
+{
+  HANDLE_TH_ERRORS


you shouldn't need this macro here; this function is not directly bound to python

ezyang · 2022-06-01T21:05:52Z

torch/csrc/utils/python_arg_parser.h

@@ -389,9 +390,52 @@ inline std::vector<int64_t> PythonArgs::intlist(int i) {
  return intlistWithDefault(i, signature.params[i].default_intlist);
 }

+TORCH_API bool is_symint_node(py::handle obj);


Why not just include the header?

yeah let me inline it back.

ezyang · 2022-06-01T21:15:32Z

torch/csrc/utils/python_arg_parser.h

+  const auto size1 = signature.params[i].size;
+  if (size1 > 0 && THPUtils_checkLong(args[i])) {
+    return std::vector<c10::SymInt>(size1, c10::SymInt(THPUtils_unpackIndex(args[i])));
+  }


You need to replicate this logic for a solitary symint arg as well

yes, a good catch, thank you!

ezyang · 2022-06-01T21:35:49Z

torch/csrc/jit/python/init.cpp

+        // we need to clear SymIntTable until we have python
+        // otherwise python classes are already deregistered
+
+        //c10::getSymIntTable().clear();


this is dead now

as a dinosaur

ezyang · 2022-06-01T21:35:59Z

c10/core/SymIntTable.cpp

@@ -14,6 +14,11 @@ std::shared_ptr<SymbolicIntNode> SymIntTable::getNode(size_t index) {
  return nodes_[index];
 }

+void SymIntTable::clear() {
+  std::lock_guard<std::mutex> lock(mutex_);
+  nodes_.clear();


this is dead now

ezyang · 2022-06-01T21:37:57Z

c10/core/TensorImpl.cpp

@@ -792,6 +792,13 @@ void TensorImpl::ShareExternalPointer(
  }
 }

+void TensorImpl::set_sym_sizes_and_strides(c10::SymIntArrayRef sizes, c10::SymIntArrayRef strides) {
+    has_symbolic_sizes_strides_ = true;
+    sizes_strides_policy_ = static_cast<uint8_t>(SizesStridesPolicy::CustomSizes);


This seems to me like CustomSymSizes isn't actually being used!

We are setting the CustomSizes policy for python tensors (i.e. made via make_wrapper_class) so calls to sizes() would throw for those. Unfortunately, it means that sym_sizes() also throws. We actually would like to just run the default implementation in this case hence CustomSymSizes which is indeed overridden by LTC. I'm open to how we can make this cleaner.

@suo any suggestions?

I think the easiest thing is to just call into python if python key is set and you have a custom sizes policy.

ezyang · 2022-06-01T21:42:06Z

torch/csrc/autograd/python_variable.cpp

@@ -622,9 +622,6 @@ static PyObject* THPVariable_make_wrapper_subclass(PyObject*, PyObject* args, Py
  if (r.toBool(10)) {
    data.unsafeGetTensorImpl()->set_sizes_strides_policy(c10::TensorImpl::SizesStridesPolicy::CustomStrides);
  }
-  if (r.toBool(11)) {
-    data.unsafeGetTensorImpl()->set_custom_device(true);
-  }


oh crap, a bad merge :(

ezyang · 2022-06-01T21:44:47Z

torch/csrc/autograd/python_variable.cpp

+  // NB: pin_memory doesn't actually do anything
+  // TODO: strides variant?
+  static PythonArgParser parser({
+    "_make_wrapper_subclass(PyObject* cls, SymIntArrayRef size, SymIntArrayRef strides, int64_t? storage_offset=None, *, MemoryFormat? memory_format=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)",


How about keeping only one _make_wrapper_subclass and just having a second overload for PythonArgParser? Having it as an overload should also help reduce duplication in this variant of the function.

mkkk... it's already a pretty branch and long function

It being long is a good reason not to copy paste right ;)

ezyang · 2022-06-01T21:46:31Z

torch/csrc/jit/python/init.cpp

+      pyobj_ = std::make_shared<c10::SafePyObject>(pyobj.release().ptr(), getPyInterpreter());
+    };
+
+  virtual SymbolicIntNode* wrap(int64_t num) {


remind me again why we are doing raw pointer memory management here

ezyang · 2022-06-01T21:47:06Z

torch/csrc/jit/python/init.cpp

+
+  virtual bool bool_() {
+    py::gil_scoped_acquire acquire;
+    return py::str(getPyObj().attr("__bool__")()).is(py::str(Py_True));


why are you doing a string comparison to test what the bool result is?

lemme try simplifying it a bit. This what SO recommends to do, but I do agree it's convoluted. There's no py::cast to bool but maybe we don't have to do py::str on both sides.

When taking questionable advice from stack overflow I highly recommend leaving a link to the URL of the question

cleaned up this part. Now it should make more sense.

ezyang · 2022-06-01T21:50:18Z

torch/csrc/jit/python/init.cpp

+  virtual SymbolicIntNode* dispatch_common_(const char* fname, SymbolicIntNode* other) {
+    auto pother = dynamic_cast<PythonSymbolicIntNode*>(other);
+    TORCH_CHECK(pother);
+    auto magic_fname = std::string("__") + fname + std::string("__");


I'd much rather you had taken the magic_fname as argument lol. With a macro you could paste the __ together with a string constant without having to do a string concat every function call (which is wasteful)

haha sorry, I was going to do but didn't do it before your review.

ezyang · 2022-06-01T23:03:55Z

torch/csrc/jit/python/init.cpp

+    .def_static("isinstance", [](py::object obj, bool convert) -> bool {
+      return pybind11::detail::type_caster<std::shared_ptr<c10::SymbolicIntNode>>().load(obj, convert);
+      //return false;
+    })


This method is pretty weird

dead code. sorry..

ezyang · 2022-06-01T23:07:26Z

torch/csrc/jit/python/init.cpp

+      if (torch::is_symint_node(b)) {
+        return std::shared_ptr<c10::SymbolicIntNode>(a->add(b.cast<c10::SymbolicIntNode*>()));
+      } else {
+        return std::shared_ptr<c10::SymbolicIntNode> (a->add(a->wrap(b.cast<int64_t>())));


This here feels like a helper function would help a bit here. But it's also not entirely clear you want to wrap integers into symbolic int nodes (that denote plain integers); it seems like it would be more user friendly if these showed up at dispatch site as plain integers. It might be a bit easier here to make the add method accept an IValue instead of a SymbolicIntNode, so you can pass in either an int or symbolic int without needing to unconditionally accept a SymbolicIntNode.

it looks way nicer rn.
I don't want to introduce a dependency on IValue into SymbolicIntNode :( . It seems more complex architecturally and possibly less user friendly since both LTC and AOTAutograd will need to parse IValues explicitly. Both LTC and AOTAutograd already wrap ints into sympy.Integer or prim::Constant.

Krovatkin · 2022-06-14T02:10:22Z

@pytorchbot merge this

pytorch-bot · 2022-06-14T02:10:24Z

❌ 🤖 pytorchbot command failed:

@pytorchbot: error: unrecognized arguments: this

usage: @pytorchbot [-h] {merge,revert,rebase} ...

Try @pytorchbot help for more info.

ezyang · 2022-06-14T02:16:48Z

@pytorchbot merge

pytorchmergebot · 2022-06-14T02:17:55Z

@pytorchbot successfully started a merge job. Check the current status here

github-actions · 2022-06-14T02:18:37Z

Hey @Krovatkin.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

bdhirsh · 2022-06-14T15:35:45Z

c10/core/TensorImpl.h

-            static_cast<uint8_t>(SizesStridesPolicy::CustomSizes))) {
-      return sym_sizes_custom();
-    }
+  virtual c10::SymIntArrayRef sym_sizes() const {


Why do we want this to be virtual now, instead of doing the policy thing that we do with all of our other de-virtualized methods?

(I need to fix it up for functionalization, which will be pretty easy - just curious on the reasoning)

And do we even need a sym_sizes_custom() anymore if sym_sizes() is virtual?

@bdhirsh python tensor subclasses and LTC want to do different things when policy is set to CustomSizes and there's no easy way to implement both via _custom so we had to make sym_sizes() virtual for now

datumbox · 2022-06-15T11:37:27Z

We suspect that this PR broke TorchVision's tests. Information available here: pytorch/vision#6166 (comment)

ezyang · 2022-06-15T13:50:40Z

@pytorchbot revert -m "broke torchvision tests" -c weird

pytorchmergebot · 2022-06-15T13:52:11Z

@pytorchbot successfully started a revert job. Check the current status here

This reverts commit d332724. Reverted #78135 on behalf of https://github.com/ezyang due to broke torchvision tests

This reverts commit b8db0a0. [ghstack-poisoned]

This reverts commit b8db0a0. ghstack-source-id: 602ffd6 Pull Request resolved: #79608

This reverts commit b8db0a0. ghstack-source-id: 602ffd6 Pull Request resolved: pytorch#79608

This reverts commit b8db0a0. ghstack-source-id: 602ffd6 Pull Request resolved: #79608

Krovatkin requested review from albanD and soulitzer as code owners May 23, 2022 23:02

facebook-github-bot added the cla signed label May 23, 2022

facebook-github-bot added the oncall: jit Add this issue/PR to JIT oncall triage queue label May 23, 2022

This was referenced May 31, 2022

[WIP] Parallel version of binding symint into Python #77173

Closed

[WIP] SymInt Python part #76971

Closed

sym sizes #76773

Closed

Chillee mentioned this pull request May 31, 2022

[WIP] Dynamic shape poc #78587

Closed

7 tasks

Krovatkin changed the title ~~Python SymInts collab [WIP]~~ Python Bindings for SymInts Jun 1, 2022

Krovatkin force-pushed the krovatkin/pybind_symint branch from e43d362 to 8bdcc93 Compare June 1, 2022 17:59

ezyang reviewed Jun 1, 2022

View reviewed changes

Krovatkin added 4 commits June 13, 2022 12:06

fix python docs error

8c59c65

switch to using temporarily

b6f75a7

fix lints

fac7500

missing gil in wrap

4e3fa42

Krovatkin force-pushed the krovatkin/pybind_symint branch from b539c7d to 4e3fa42 Compare June 13, 2022 19:06

make overrides test pass

50810d1

Krovatkin mentioned this pull request Jun 13, 2022

Adding python bindings for SymInt #79460

Closed

pytorchmergebot added the Merged label Jun 14, 2022

pytorchmergebot closed this in d332724 Jun 14, 2022

bdhirsh reviewed Jun 14, 2022

View reviewed changes

vfdev-5 mentioned this pull request Jun 15, 2022

CI broken due to Core issue pytorch/vision#6166

Closed

pytorchmergebot added the Reverted label Jun 15, 2022

pytorchmergebot added a commit that referenced this pull request Jun 15, 2022

Revert "Python Bindings for SymInts (#78135)"

b8db0a0

This reverts commit d332724. Reverted #78135 on behalf of https://github.com/ezyang due to broke torchvision tests

ezyang added a commit that referenced this pull request Jun 15, 2022

Revert "Revert "Python Bindings for SymInts (#78135)""

e0d33b0

This reverts commit b8db0a0. [ghstack-poisoned]

ezyang added a commit that referenced this pull request Jun 15, 2022

Update on "Revert "Revert "Python Bindings for SymInts (#78135)"""

06ebcd9

This reverts commit b8db0a0. [ghstack-poisoned]

ezyang added a commit that referenced this pull request Jun 15, 2022

Revert "Revert "Python Bindings for SymInts (#78135)""

dde2adf

This reverts commit b8db0a0. ghstack-source-id: 602ffd6 Pull Request resolved: #79608

Krovatkin pushed a commit to Krovatkin/pytorch that referenced this pull request Jun 15, 2022

Revert "Revert "Python Bindings for SymInts (pytorch#78135)""

8a0fb1b

This reverts commit b8db0a0. ghstack-source-id: 602ffd6 Pull Request resolved: pytorch#79608

wconstab pushed a commit that referenced this pull request Jun 17, 2022

Revert "Revert "Python Bindings for SymInts (#78135)""

4c6b824

This reverts commit b8db0a0. ghstack-source-id: 602ffd6 Pull Request resolved: #79608

zengk95 mentioned this pull request Jun 22, 2022

[Meta] CI Revert Tracker #66178

Closed

github-actions bot deleted the krovatkin/pybind_symint branch February 17, 2024 01:50

Python Bindings for SymInts #78135

Python Bindings for SymInts #78135

Uh oh!

Conversation

Uh oh!

Uh oh!

Uh oh!

🔗 Helpful links

✅ No Failures (0 Pending)

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!

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!

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!