Add API for open registration between operators and subclasses (and modes) #130064

zou3519 · 2024-07-03T21:34:08Z

Stack from ghstack (oldest at bottom):

We add torch.library.Library._register_torch_dispatch_rule. Here, a user
can provide us a specific rule to run for a specific
(torch_dispatch_class, operator) pair. The motivation is that a user
might want to extend a subclass/mode but may not have access to the
source code of the subclass/mode.

I'll make this public in a follow-up PR if we think the approach and API
is good.

Keep in mind that many subclasses will likely deliver their own open
registration solution (DTensor has register_sharding_prop_rule and NJT
has register_jagged_op); _register_torch_dispatch_rule is meant as a
catch-all open registration mechanism for when the subclass hasn't
provided anything more specific.

Test Plan:

new tests

…odes) We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests [ghstack-poisoned]

pytorch-bot · 2024-07-03T21:34:11Z

🔗 Helpful Links

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

📄 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 73bc42c with merge base 9c1ba5a ():
💚 Looks good so far! There are no failures yet. 💚

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

…sses (and modes)" We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests [ghstack-poisoned]

…odes) We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests ghstack-source-id: 7675dbc Pull Request resolved: #130064

albanD

Sounds great!

torch/_library/simple_registry.py

torch/csrc/utils/python_arg_parser.cpp

…sses (and modes)" We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests [ghstack-poisoned]

zou3519 · 2024-07-08T22:06:15Z

@pytorchbot merge

pytorchmergebot · 2024-07-08T22:07:45Z

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

huydhn · 2024-07-09T01:46:25Z

@pytorchbot revert -m 'Sorry for reverting your change but test_profiler_tree is failing in trunk after this lands https://hud.pytorch.org/pytorch/pytorch/commit/922d2737d5e0ad22ee1dcf91c48ab09d641de840, maybe a landrace' -c landrace

pytorchmergebot · 2024-07-09T01:48:29Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2024-07-09T01:48:40Z

@zou3519 your PR has been successfully reverted.

…s (and modes) (#130064)" This reverts commit 922d273. Reverted #130064 on behalf of https://github.com/huydhn due to Sorry for reverting your change but test_profiler_tree is failing in trunk after this lands https://hud.pytorch.org/pytorch/pytorch/commit/922d2737d5e0ad22ee1dcf91c48ab09d641de840, maybe a landrace ([comment](#130064 (comment)))

…sses (and modes)" We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests [ghstack-poisoned]

zou3519 · 2024-07-09T14:28:29Z

@huydhn afaict that test was added in 2022, is something else wrong with CI? I did confirm that the test failed locally.

This is the API for defining the interaction between a torch_dispatch class and a custom op. Taking API bikeshedding. Test Plan: - new tests Pull Request resolved: #130261 Approved by: https://github.com/albanD ghstack dependencies: #130064

…odes) (pytorch#130064) We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests Pull Request resolved: pytorch#130064 Approved by: https://github.com/albanD

…0261) This is the API for defining the interaction between a torch_dispatch class and a custom op. Taking API bikeshedding. Test Plan: - new tests Pull Request resolved: pytorch#130261 Approved by: https://github.com/albanD ghstack dependencies: pytorch#130064

This reverts commit bef085b. Reverted #130079 on behalf of https://github.com/izaitsevfb due to depends on #130064 which needs to be reverted ([comment](#130079 (comment)))

Fixes #129617 Pull Request resolved: #130079 Approved by: https://github.com/zou3519

…30261)" This reverts commit bb9a73f. Reverted #130261 on behalf of https://github.com/izaitsevfb due to depends on #130064 which needs to be reverted ([comment](#130261 (comment)))

izaitsevfb · 2024-07-10T21:47:48Z

@pytorchbot revert -m "fails internal builds, see D59553526" -c ghfirst

ERROR    ...
ERROR    reason: Action failed: fbcode//caffe2:_C_impl_cuda (fbcode//buck2/platform/execution:linux-x86_64#22f53947965bda2d) (cxx_compile torch/csrc/utils/python_arg_parser.cpp (pic))
ERROR    Remote command returned non-zero exit code 1
ERROR    Remote action, reproduce with: `frecli cas download-action 9c60783056764f8cb03d98ef26e521ad27a68f14daa5ed7ef3c9d9790a5e70cb:145`
ERROR    Stdout: <empty>
ERROR    Stderr:
ERROR    fbcode/caffe2/torch/csrc/utils/python_arg_parser.cpp:265:3: error: unknown type name 'PYBIND11_CONSTINIT'
ERROR      PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
ERROR      ^
ERROR    fbcode/caffe2/torch/csrc/utils/python_arg_parser.cpp:265:33: error: definition or redeclaration of 'gil_safe_call_once_and_store' not allowed inside a function
ERROR      PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
ERROR                                ~~~~^
ERROR    fbcode/caffe2/torch/csrc/utils/python_arg_parser.cpp:265:33: error: no member named 'gil_safe_call_once_and_store' in namespace 'pybind11'
ERROR      PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
ERROR                                ~~~~^
ERROR    fbcode/caffe2/torch/csrc/utils/python_arg_parser.cpp:265:61: error: expected ';' at end of declaration
ERROR      PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
ERROR                                                                ^
ERROR                                                                ;
ERROR    fbcode/caffe2/torch/csrc/utils/python_arg_parser.cpp:268:7: error: use of undeclared identifier 'storage'
ERROR          storage
ERROR          ^
ERROR    5 errors generated.
ERROR 
ERROR      --> fbcode/buck2/prelude/python/sourcedb/build.bxl:25:9
ERROR       |
ERROR    25 |         fail(error_message)
ERROR       |         ^^^^^^^^^^^^^^^^^^^
ERROR       |
ERROR

pytorchmergebot · 2024-07-10T21:50:21Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2024-07-10T21:50:34Z

@zou3519 your PR has been successfully reverted.

…s (and modes) (#130064)" This reverts commit c23d103. Reverted #130064 on behalf of https://github.com/izaitsevfb due to fails internal builds, see [D59553526](https://www.internalfb.com/diff/D59553526) ([comment](#130064 (comment)))

…sses (and modes)" We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests [ghstack-poisoned]

XuehaiPan · 2024-07-11T16:49:32Z

torch/csrc/utils/python_arg_parser.cpp

+  static auto find_torch_dispatch_rule =
+      py::module_::import("torch._library.simple_registry")
+          .attr("find_torch_dispatch_rule");


Suggested change

static auto find_torch_dispatch_rule =

py::module_::import("torch._library.simple_registry")

.attr("find_torch_dispatch_rule");

static const py::handle find_torch_dispatch_rule =

py::module_::import("torch._library.simple_registry")

.attr("find_torch_dispatch_rule").release();

Use py::handle instead of py::object.

What's the difference?
Also I can fix this up (go back to the original approach you suggested) when we update pybind11 internally (which will hopefully be soon); I want to merge this PR sooner than later because it may have performance implications.

py::handle is equivalent to PyObject *. It will not do Py_DECREF(obj.m_ptr) on C++ instance destruction where py::object will. handle = object.release() will leak the instance to keep it alive on Python side.

…sses (and modes)" We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests [ghstack-poisoned]

This is the API for defining the interaction between a torch_dispatch class and a custom op. Taking API bikeshedding. Test Plan: - new tests Pull Request resolved: #130261 Approved by: https://github.com/albanD ghstack dependencies: #130064

…odes) (pytorch#130064) We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests Pull Request resolved: pytorch#130064 Approved by: https://github.com/albanD

…0261) This is the API for defining the interaction between a torch_dispatch class and a custom op. Taking API bikeshedding. Test Plan: - new tests Pull Request resolved: pytorch#130261 Approved by: https://github.com/albanD ghstack dependencies: pytorch#130064

This reverts commit bef085b. Reverted pytorch#130079 on behalf of https://github.com/izaitsevfb due to depends on pytorch#130064 which needs to be reverted ([comment](pytorch#130079 (comment)))

…torch#130261)" This reverts commit bb9a73f. Reverted pytorch#130261 on behalf of https://github.com/izaitsevfb due to depends on pytorch#130064 which needs to be reverted ([comment](pytorch#130261 (comment)))

…s (and modes) (pytorch#130064)" This reverts commit c23d103. Reverted pytorch#130064 on behalf of https://github.com/izaitsevfb due to fails internal builds, see [D59553526](https://www.internalfb.com/diff/D59553526) ([comment](pytorch#130064 (comment)))

…odes) (pytorch#130064) We add torch.library.Library._register_torch_dispatch_rule. Here, a user can provide us a specific rule to run for a specific (torch_dispatch_class, operator) pair. The motivation is that a user might want to extend a subclass/mode but may not have access to the source code of the subclass/mode. I'll make this public in a follow-up PR if we think the approach and API is good. Keep in mind that many subclasses will likely deliver their own open registration solution (DTensor has register_sharding_prop_rule and NJT has register_jagged_op); _register_torch_dispatch_rule is meant as a catch-all open registration mechanism for when the subclass hasn't provided anything more specific. Test Plan: - new tests Pull Request resolved: pytorch#130064 Approved by: https://github.com/albanD

…0261) This is the API for defining the interaction between a torch_dispatch class and a custom op. Taking API bikeshedding. Test Plan: - new tests Pull Request resolved: pytorch#130261 Approved by: https://github.com/albanD ghstack dependencies: pytorch#130064

zou3519 requested review from a team, albanD and bdhirsh July 5, 2024 13:36

albanD approved these changes Jul 8, 2024

View reviewed changes

torch/_library/simple_registry.py Outdated Show resolved Hide resolved

torch/csrc/utils/python_arg_parser.cpp Outdated Show resolved Hide resolved

zou3519 mentioned this pull request Jul 8, 2024

[custom_ops] expose torch.library.register_torch_dispatch #130261

Closed

zou3519 added ciflow/trunk Trigger trunk jobs on your pull request release notes: composability release notes category labels Jul 8, 2024

albanD approved these changes Jul 8, 2024

View reviewed changes

pytorchmergebot added the merging label Jul 8, 2024

pytorchmergebot closed this in 922d273 Jul 8, 2024

pytorchmergebot added Merged and removed merging labels Jul 8, 2024

pytorchmergebot added the Reverted label Jul 9, 2024

pytorchmergebot reopened this Jul 9, 2024

clee2000 added the ci-no-td Do not run TD on this PR label Jul 9, 2024

izaitsevfb mentioned this pull request Jul 10, 2024

[reland][custom ops] infer schema #130079

Closed

pytorchmergebot referenced this pull request Jul 10, 2024

[reland][custom ops] infer schema (#130079)

bef085b

Fixes #129617 Pull Request resolved: #130079 Approved by: https://github.com/zou3519

pytorchmergebot reopened this Jul 10, 2024

zou3519 mentioned this pull request Jul 11, 2024

Fix static py::object dangling pointer with py::gil_safe_call_once_and_store #130341

Closed

XuehaiPan reviewed Jul 11, 2024

View reviewed changes

pytorchmergebot closed this in ba94176 Jul 12, 2024

zou3519 mentioned this pull request Jul 12, 2024

[custom_op] triton_op API V0 #130637

Closed

henrylhtsang mentioned this pull request Jul 17, 2024

[aoti] Unskip some aot inductor tests #130973

Closed

Add API for open registration between operators and subclasses (and modes) #130064

Add API for open registration between operators and subclasses (and modes) #130064

Uh oh!

Conversation

zou3519 commented Jul 3, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 3, 2024 • 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/130064

✅ No Failures

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

zou3519 commented Jul 8, 2024

Uh oh!

pytorchmergebot commented Jul 8, 2024

Merge started

Uh oh!

huydhn commented Jul 9, 2024

Uh oh!

pytorchmergebot commented Jul 9, 2024

Uh oh!

pytorchmergebot commented Jul 9, 2024

Uh oh!

zou3519 commented Jul 9, 2024

Uh oh!

izaitsevfb commented Jul 10, 2024

Uh oh!

pytorchmergebot commented Jul 10, 2024

Uh oh!

pytorchmergebot commented Jul 10, 2024

Uh oh!

XuehaiPan Jul 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zou3519 Jul 11, 2024

Choose a reason for hiding this comment

Uh oh!

XuehaiPan Jul 12, 2024

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

zou3519 commented Jul 3, 2024 •

edited

Loading

pytorch-bot bot commented Jul 3, 2024 •

edited

Loading

XuehaiPan Jul 11, 2024 •

edited

Loading