Port OpSchema.__post_init__ and OpSchema._recompute_comparison_key to C++ by swolchok · Pull Request #161695 · pytorch/pytorch

swolchok · 2025-08-28T06:59:34Z

Stack from ghstack (oldest at bottom):

I initially didn't see good results porting this, but it was apparently because of pybind11 function calling overhead. (pybind11's object-handling primitives seem fine enough.) I'm interested in setting up nanobind, but this demonstrates it's not blocking.

Differential Revision: D81530102

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @ezyang @msaroufim @dcci

… does make it faster I initially didn't see good results porting this, but it was apparently because of pybind11 function calling overhead. (pybind11's object-handling primitives seem fine enough.) I'm interested in setting up nanobind, but this demonstrates it's not blocking. [ghstack-poisoned]

pytorch-bot · 2025-08-28T06:59:38Z

🔗 Helpful Links

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

📄 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

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 0f67393 with merge base a63221a ():
💚 Looks good so far! There are no failures yet. 💚

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

… does make it faster I initially didn't see good results porting this, but it was apparently because of pybind11 function calling overhead. (pybind11's object-handling primitives seem fine enough.) I'm interested in setting up nanobind, but this demonstrates it's not blocking. ghstack-source-id: cfd236c Pull Request resolved: #161695

ezyang · 2025-09-18T03:51:08Z

torch/csrc/autograd/python_variable.cpp

+#endif
+}
+
+static bool arg_type_tensor_or_tensor_list_like(py::handle arg) {


Do you have objections to keeping the old Python implementation around in comments as a reference, or do you think it will bitrot immediately?

commented-out code generally bitrots. if we want it around we would need to make it modally enabled and add tests

ezyang · 2025-09-18T03:52:48Z

How good do you think the preexisting testing here is? We are looking at some (orthogonal) parts of DTensor where the testing is definitely NOT enough and I'm wondering if this applies here or not.

ezyang · 2025-09-18T03:54:46Z

torch/csrc/autograd/python_variable.cpp

+  c10::SmallVector<py::object, 8> args_to_hash;
+  size_t idx = 0;
+  for (const auto& e : args_schema) {
+    if (idx >= static_argnum || arg_type_tensor_or_tensor_list_like(e)) {


reordered cond, nice

ezyang · 2025-09-18T03:57:33Z

torch/csrc/autograd/python_variable.cpp

+    int idx = 0;
+    auto kwargs_schema = py::reinterpret_borrow<py::dict>(raw_kwargs_schema);
+    for (const auto& k : static_kwargkey_list) {
+      PyObject* item = PyDict_GetItem(kwargs_schema.ptr(), k.ptr());


safer to call PyDict_GetItemWithError

oh wow I didn't fully realize what was going on with PyDict_GetItem

ezyang · 2025-09-18T03:59:51Z

torch/csrc/autograd/python_variable.cpp

+        kwargs_to_hash[idx++] = py::none();
+      }
+    }
+    PyObject* comparison_key = PyTuple_Pack(


Should test that PyTuple_Pack succeeded, and ditto below

ezyang · 2025-09-18T04:00:40Z

torch/csrc/autograd/python_variable.cpp

+    if (!PyList_Check(static_kwargkey.ptr())) {
+      PyErr_SetString(
+          PyExc_TypeError, "self.schema_info.static_kwargkey must be a list!");
+      return false;


It might be more uniform with the rest of the codebase if we just raised exceptions when this happens, but I guess this is not too difficult to get right.

ezyang · 2025-09-18T04:02:22Z

torch/csrc/autograd/python_variable.cpp

+  const auto dtensor_spec_class = get_dtensor_spec_class();
+  bool has_symints = false;
+  for (const auto& a : args_schema) {
+    if (Py_TYPE(a.ptr()) != (PyTypeObject*)(dtensor_spec_class.ptr()) &&


direct type test first, interesting that this helps!

yeah I briefly considered upstreaming extensions to PyFloat_Check etc. to first do the exact checks because they're way cheaper. the CPython API functions can't be inlined so there's function call overhead and they also don't themselves start with the direct test.

ezyang

I reviewed logic equivalence and CPython API usage, I did NOT review memory management

swolchok · 2025-09-18T17:24:03Z

How good do you think the preexisting testing here is? We are looking at some (orthogonal) parts of DTensor where the testing is definitely NOT enough and I'm wondering if this applies here or not.

I was feeling mildly guilty for not adding focused tests here yesterday. Can just go ahead and do it if I can't find OpSchema-specific tests. I have a historical tendency of incorrectly assuming that PyTorch code is automatically well-tested for some reason.

…mWithError, address other review feedback on "Port OpSchema.__post_init__ and OpSchema._recompute_comparison_key to C++" I initially didn't see good results porting this, but it was apparently because of pybind11 function calling overhead. (pybind11's object-handling primitives seem fine enough.) I'm interested in setting up nanobind, but this demonstrates it's not blocking. Differential Revision: [D81530102](https://our.internmc.facebook.com/intern/diff/D81530102) cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

… does make it faster I initially didn't see good results porting this, but it was apparently because of pybind11 function calling overhead. (pybind11's object-handling primitives seem fine enough.) I'm interested in setting up nanobind, but this demonstrates it's not blocking. ghstack-source-id: 829bad9 Pull Request resolved: #161695

swolchok · 2025-09-19T03:59:28Z

@pytorchbot merge

pytorchmergebot · 2025-09-19T04:01:51Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. Pull Request resolved: #162508 Approved by: https://github.com/ezyang ghstack dependencies: #161695

… C++ (pytorch#161695) I initially didn't see good results porting this, but it was apparently because of pybind11 function calling overhead. (pybind11's object-handling primitives seem fine enough.) I'm interested in setting up nanobind, but this demonstrates it's not blocking. Differential Revision: [D81530102](https://our.internmc.facebook.com/intern/diff/D81530102) Pull Request resolved: pytorch#161695 Approved by: https://github.com/ezyang

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. Pull Request resolved: pytorch#162508 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#161695

… C++ (pytorch#161695) I initially didn't see good results porting this, but it was apparently because of pybind11 function calling overhead. (pybind11's object-handling primitives seem fine enough.) I'm interested in setting up nanobind, but this demonstrates it's not blocking. Differential Revision: [D81530102](https://our.internmc.facebook.com/intern/diff/D81530102) Pull Request resolved: pytorch#161695 Approved by: https://github.com/ezyang

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. Pull Request resolved: pytorch#162508 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#161695

…nsor_info (#162968) Next PR writes a C++ implementation. Seems good to have tests first. Pull Request resolved: #162968 Approved by: https://github.com/ezyang ghstack dependencies: #161695, #162508

… C++ (pytorch#161695) I initially didn't see good results porting this, but it was apparently because of pybind11 function calling overhead. (pybind11's object-handling primitives seem fine enough.) I'm interested in setting up nanobind, but this demonstrates it's not blocking. Differential Revision: [D81530102](https://our.internmc.facebook.com/intern/diff/D81530102) Pull Request resolved: pytorch#161695 Approved by: https://github.com/ezyang

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. Pull Request resolved: pytorch#162508 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#161695

…nsor_info (pytorch#162968) Next PR writes a C++ implementation. Seems good to have tests first. Pull Request resolved: pytorch#162968 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#161695, pytorch#162508

swolchok requested review from albanD and soulitzer as code owners August 28, 2025 06:59

This was referenced Aug 28, 2025

Fix non-const reference arguments in torch/csrc/jit/python/init.cpp #161300

Closed

Fix forced copying def_property_readonly for FunctionSchema & friends #161301

Closed

Stop accessing func._schema in _python_dispatch.correct_storage_aliasing #161292

Closed

swolchok mentioned this pull request Aug 28, 2025

Use is, not ==, to check exact type matches in _python_dispatch #161304

Closed

pytorch-bot bot added ciflow/inductor oncall: distributed Add this issue/PR to distributed oncall triage queue labels Aug 28, 2025

ezyang reviewed Sep 18, 2025

View reviewed changes

ezyang approved these changes Sep 18, 2025

View reviewed changes

pytorchmergebot added the merging label Sep 19, 2025

pytorchmergebot closed this in 76a841f Sep 19, 2025

pytorchmergebot added Merged and removed merging labels Sep 19, 2025

swolchok mentioned this pull request Sep 23, 2025

Remove torch.distributed.tensor.OpSchema.has_symints #163667

Closed

github-actions bot deleted the gh/swolchok/823/head branch October 20, 2025 02:17

Conversation

swolchok commented Aug 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 28, 2025 • 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/161695

✅ No Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ezyang commented Sep 18, 2025

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!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

swolchok commented Sep 18, 2025

Uh oh!

swolchok commented Sep 19, 2025

Uh oh!

pytorchmergebot commented Sep 19, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

swolchok commented Aug 28, 2025 •

edited

Loading

pytorch-bot bot commented Aug 28, 2025 •

edited

Loading