Add utility to get computed kernel in torch.library #158393

mikaylagawarecki · 2025-07-15T22:35:13Z

Adds OperatorEntry::getComputedKernelForDispatchKey which returns the KernelFunction corresponding to OperatorEntry.dispatchTable_[dispatch_ix] for a given dispatch key

Specifically it returns a SafeKernelFunction that holds a KernelToken. This KernelToken is registered to the KernelFunction in OperatorEntry.kernels_ and will be invalidated when the KernelFunction is destructed (i.e. when the AnnotatedKernel that holds this KernelFunction is removed from kernels_, which happens when the corresponding impl is deregistered).
SafeKernelFunction can be called via callBoxed, the validity of the token will be checked before this happens
SafeKernelFunction is pybinded and getComputedKernelForDispatchKey is exposed to the frontend ia torch.library.get_kernel

Related to #155330

Stack from ghstack (oldest at bottom):

-> Add utility to get computed kernel in torch.library #158393

cc @albanD

[ghstack-poisoned]

pytorch-bot · 2025-07-15T22:35:17Z

🔗 Helpful Links

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

📄 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 172857e with merge base 34ec5ed ():
💚 Looks good so far! There are no failures yet. 💚

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

ghstack-source-id: c7793b0 Pull Request resolved: #158393

[ghstack-poisoned]

ghstack-source-id: 439450f Pull Request resolved: #158393

Related to #155330 [ghstack-poisoned]

ghstack-source-id: e876743 Pull Request resolved: #158393

test/test_custom_ops.py

Related to #155330 [ghstack-poisoned]

ghstack-source-id: 762f9b6 Pull Request resolved: #158393

aten/src/ATen/core/boxing/KernelFunction.h

Related to #155330 cc albanD [ghstack-poisoned]

ghstack-source-id: 5925282 Pull Request resolved: #158393

albanD

Sounds pretty good!
Only small questions!

albanD · 2025-08-01T14:28:42Z

torch/library.py

+        op = op._name
+
+    if isinstance(dispatch_key, str):
+        dispatch_key = torch._C.DispatchKey.__members__[dispatch_key]


What is the error you get when passing a wrong dispatch key here?

oops, added proper error handling here

torch/library.py

test/test_custom_ops.py

aten/src/ATen/core/dispatch/OperatorEntry.cpp

albanD · 2025-08-01T15:52:37Z

aten/src/ATen/core/dispatch/OperatorEntry.cpp

+  auto [annotatedKernel, _] =
+      computeDispatchTableEntryWithDebug(c10::Dispatcher::singleton(), k);
+
+  return SafeKernelFunction(&annotatedKernel.kernel);


I think it would be nice to grab the debug string here and add that to the __repr__ we get from python?

This gives something like

SafeKernelFunction(debug='registered at /data/users/mg1998/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutograd_0.cpp:2309')

Do you think that is meaningful enough or should we add more info

Yes, it is 100% super useful. These error messages saved me a few times for multiple-registration errors!
And for python users, it should point to their code directly. Which is even better so they know which function this is!

aten/src/ATen/core/boxing/KernelFunction_impl.h

albanD · 2025-08-01T15:57:47Z

aten/src/ATen/core/boxing/KernelFunction.h

+
+  // List of tokens that need to be invalidated when this KernelFunction is
+  // destroyed
+  mutable std::vector<std::weak_ptr<KernelToken>> tokens_;


Why mutable?

Also why weak_ptr and not shared_ptr?

why mutable

I think this is necessary in order to make registerToken const, which was in turn needed to allow SafeKernelFunction to take in const KernelFunction*, removing this would necessitate const_cast-ing the annotatedKernel.kernel in getComputedKernelForDispatchKey, wdyt

why weak_ptr and not shared_ptr

What's the benefit of shared_ptr over weak_ptr here? if we use weak_ptr, the KernelToken dies with the SafeKernelFunction, which I think is what we want to achieve here (?)

aten/src/ATen/core/boxing/KernelFunction.h

Adds `OperatorEntry::getComputedKernelForDispatchKey` which returns the KernelFunction corresponding to `OperatorEntry.dispatchTable_[dispatch_ix]` for a given dispatch key - Specifically it returns a `SafeKernelFunction` that holds a `KernelToken`. This `KernelToken` is registered to the `KernelFunction` in `OperatorEntry.kernels_` and will be invalidated when the `KernelFunction` is destructed (i.e. when the `AnnotatedKernel` that holds this `KernelFunction` is removed from `kernels_`, which happens when the corresponding impl is deregistered). - `SafeKernelFunction` can be called via `callBoxed`, the validity of the token will be checked before this happens - `SafeKernelFunction` is pybinded and `getComputedKernelForDispatchKey` is exposed to the frontend ia `torch.library.get_kernel` Related to #155330 cc albanD [ghstack-poisoned]

ghstack-source-id: 1c15180 Pull Request resolved: #158393

Adds `OperatorEntry::getComputedKernelForDispatchKey` which returns the KernelFunction corresponding to `OperatorEntry.dispatchTable_[dispatch_ix]` for a given dispatch key - Specifically it returns a `SafeKernelFunction` that holds a `KernelToken`. This `KernelToken` is registered to the `KernelFunction` in `OperatorEntry.kernels_` and will be invalidated when the `KernelFunction` is destructed (i.e. when the `AnnotatedKernel` that holds this `KernelFunction` is removed from `kernels_`, which happens when the corresponding impl is deregistered). - `SafeKernelFunction` can be called via `callBoxed`, the validity of the token will be checked before this happens - `SafeKernelFunction` is pybinded and `getComputedKernelForDispatchKey` is exposed to the frontend ia `torch.library.get_kernel` Related to #155330 cc albanD [ghstack-poisoned]

ghstack-source-id: ec2351a Pull Request resolved: #158393

Adds `OperatorEntry::getComputedKernelForDispatchKey` which returns the KernelFunction corresponding to `OperatorEntry.dispatchTable_[dispatch_ix]` for a given dispatch key - Specifically it returns a `SafeKernelFunction` that holds a `KernelToken`. This `KernelToken` is registered to the `KernelFunction` in `OperatorEntry.kernels_` and will be invalidated when the `KernelFunction` is destructed (i.e. when the `AnnotatedKernel` that holds this `KernelFunction` is removed from `kernels_`, which happens when the corresponding impl is deregistered). - `SafeKernelFunction` can be called via `callBoxed`, the validity of the token will be checked before this happens - `SafeKernelFunction` is pybinded and `getComputedKernelForDispatchKey` is exposed to the frontend ia `torch.library.get_kernel` Related to #155330 cc albanD [ghstack-poisoned]

ghstack-source-id: 72c37fb Pull Request resolved: #158393

Adds `OperatorEntry::getComputedKernelForDispatchKey` which returns the KernelFunction corresponding to `OperatorEntry.dispatchTable_[dispatch_ix]` for a given dispatch key - Specifically it returns a `SafeKernelFunction` that holds a `KernelToken`. This `KernelToken` is registered to the `KernelFunction` in `OperatorEntry.kernels_` and will be invalidated when the `KernelFunction` is destructed (i.e. when the `AnnotatedKernel` that holds this `KernelFunction` is removed from `kernels_`, which happens when the corresponding impl is deregistered). - `SafeKernelFunction` can be called via `callBoxed`, the validity of the token will be checked before this happens - `SafeKernelFunction` is pybinded and `getComputedKernelForDispatchKey` is exposed to the frontend ia `torch.library.get_kernel` Related to #155330 cc albanD [ghstack-poisoned]

ghstack-source-id: a4e17da Pull Request resolved: #158393

albanD

Awesome!
Thanks!

mikaylagawarecki · 2025-08-13T16:01:10Z

@pytorchbot merge -r

pytorchmergebot · 2025-08-13T16:02:43Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[ghstack-poisoned]

pytorchmergebot · 2025-08-13T16:02:56Z

Successfully rebased gh/mikaylagawarecki/320/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/158393)

ghstack-source-id: 8f833e2 Pull Request resolved: #158393

pytorchmergebot · 2025-08-13T16:04:18Z

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

pytorchmergebot · 2025-08-13T16:35:54Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-jammy-py3.10-clang18-asan / build

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

mikaylagawarecki · 2025-08-13T18:57:31Z

@pytorchbot merge

pytorchmergebot · 2025-08-13T18:59:17Z

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

Adds `OperatorEntry::getComputedKernelForDispatchKey` which returns the KernelFunction corresponding to `OperatorEntry.dispatchTable_[dispatch_ix]` for a given dispatch key - Specifically it returns a `SafeKernelFunction` that holds a `KernelToken`. This `KernelToken` is registered to the `KernelFunction` in `OperatorEntry.kernels_` and will be invalidated when the `KernelFunction` is destructed (i.e. when the `AnnotatedKernel` that holds this `KernelFunction` is removed from `kernels_`, which happens when the corresponding impl is deregistered). - `SafeKernelFunction` can be called via `callBoxed`, the validity of the token will be checked before this happens - `SafeKernelFunction` is pybinded and `getComputedKernelForDispatchKey` is exposed to the frontend ia `torch.library.get_kernel` Related to #155330 Pull Request resolved: #158393 Approved by: https://github.com/albanD

Adds `OperatorEntry::getComputedKernelForDispatchKey` which returns the KernelFunction corresponding to `OperatorEntry.dispatchTable_[dispatch_ix]` for a given dispatch key - Specifically it returns a `SafeKernelFunction` that holds a `KernelToken`. This `KernelToken` is registered to the `KernelFunction` in `OperatorEntry.kernels_` and will be invalidated when the `KernelFunction` is destructed (i.e. when the `AnnotatedKernel` that holds this `KernelFunction` is removed from `kernels_`, which happens when the corresponding impl is deregistered). - `SafeKernelFunction` can be called via `callBoxed`, the validity of the token will be checked before this happens - `SafeKernelFunction` is pybinded and `getComputedKernelForDispatchKey` is exposed to the frontend ia `torch.library.get_kernel` Related to pytorch#155330 Pull Request resolved: pytorch#158393 Approved by: https://github.com/albanD

…158393)" This reverts commit 1196bb1.

Adds `OperatorEntry::getComputedKernelForDispatchKey` which returns the KernelFunction corresponding to `OperatorEntry.dispatchTable_[dispatch_ix]` for a given dispatch key - Specifically it returns a `SafeKernelFunction` that holds a `KernelToken`. This `KernelToken` is registered to the `KernelFunction` in `OperatorEntry.kernels_` and will be invalidated when the `KernelFunction` is destructed (i.e. when the `AnnotatedKernel` that holds this `KernelFunction` is removed from `kernels_`, which happens when the corresponding impl is deregistered). - `SafeKernelFunction` can be called via `callBoxed`, the validity of the token will be checked before this happens - `SafeKernelFunction` is pybinded and `getComputedKernelForDispatchKey` is exposed to the frontend ia `torch.library.get_kernel` Related to pytorch#155330 Pull Request resolved: pytorch#158393 Approved by: https://github.com/albanD

…158393)" This reverts commit 1196bb1.

…158393)" This reverts commit 1196bb1. (cherry picked from commit a975dfe)

…158393)" This reverts commit 1196bb1.

Add utility to get computed kernel in torch.library

fcdcdaf

[ghstack-poisoned]

mikaylagawarecki added a commit that referenced this pull request Jul 15, 2025

Add utility to get computed kernel in torch.library

3df8757

ghstack-source-id: c7793b0 Pull Request resolved: #158393

mikaylagawarecki marked this pull request as draft July 16, 2025 14:44

Update on "Add utility to get computed kernel in torch.library"

0407e30

[ghstack-poisoned]

mikaylagawarecki added a commit that referenced this pull request Jul 16, 2025

Add utility to get computed kernel in torch.library

ef89e2b

ghstack-source-id: 439450f Pull Request resolved: #158393

Update on "Add utility to get computed kernel in torch.library"

848ef21

Related to #155330 [ghstack-poisoned]

mikaylagawarecki added a commit that referenced this pull request Jul 24, 2025

Add utility to get computed kernel in torch.library

67a5dc2

ghstack-source-id: e876743 Pull Request resolved: #158393

mikaylagawarecki commented Jul 24, 2025

View reviewed changes

test/test_custom_ops.py Show resolved Hide resolved

Update on "Add utility to get computed kernel in torch.library"

51a262d

Related to #155330 [ghstack-poisoned]

mikaylagawarecki added a commit that referenced this pull request Jul 24, 2025

Add utility to get computed kernel in torch.library

7d80793

ghstack-source-id: 762f9b6 Pull Request resolved: #158393

mikaylagawarecki added module: python frontend For issues relating to PyTorch's Python frontend topic: new features topic category labels Jul 24, 2025

mikaylagawarecki requested a review from albanD July 24, 2025 21:35

mikaylagawarecki added release notes: python_frontend python frontend release notes category and removed module: python frontend For issues relating to PyTorch's Python frontend labels Jul 24, 2025

mikaylagawarecki commented Jul 24, 2025

View reviewed changes

aten/src/ATen/core/boxing/KernelFunction.h Show resolved Hide resolved

mikaylagawarecki marked this pull request as ready for review July 24, 2025 21:37

Update on "Add utility to get computed kernel in torch.library"

4df1e1d

Related to #155330 cc albanD [ghstack-poisoned]

mikaylagawarecki added a commit that referenced this pull request Jul 24, 2025

Add utility to get computed kernel in torch.library

27feb54

ghstack-source-id: 5925282 Pull Request resolved: #158393

mikaylagawarecki requested a review from zou3519 July 25, 2025 14:48

albanD reviewed Aug 1, 2025

View reviewed changes

mikaylagawarecki added a commit that referenced this pull request Aug 5, 2025

Add utility to get computed kernel in torch.library

7299f1d

ghstack-source-id: 1c15180 Pull Request resolved: #158393

mikaylagawarecki requested a review from albanD August 5, 2025 19:32

mikaylagawarecki added a commit that referenced this pull request Aug 5, 2025

Add utility to get computed kernel in torch.library

d4994c6

ghstack-source-id: ec2351a Pull Request resolved: #158393

mikaylagawarecki added a commit that referenced this pull request Aug 5, 2025

Add utility to get computed kernel in torch.library

73c3d69

ghstack-source-id: 72c37fb Pull Request resolved: #158393

mikaylagawarecki added a commit that referenced this pull request Aug 8, 2025

Add utility to get computed kernel in torch.library

53d142f

ghstack-source-id: a4e17da Pull Request resolved: #158393

mikaylagawarecki added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 12, 2025

albanD approved these changes Aug 13, 2025

View reviewed changes

Update

172857e

[ghstack-poisoned]

pytorchmergebot pushed a commit that referenced this pull request Aug 13, 2025

Add utility to get computed kernel in torch.library

8c044e8

ghstack-source-id: 8f833e2 Pull Request resolved: #158393

pytorchmergebot added the merging label Aug 13, 2025

pytorchmergebot removed the merging label Aug 13, 2025

pytorchmergebot added the merging label Aug 13, 2025

pytorchmergebot added the Merged label Aug 13, 2025

pytorchmergebot closed this in 1196bb1 Aug 13, 2025

pytorchmergebot removed the merging label Aug 13, 2025

This was referenced Aug 14, 2025

Intra Kernel Dispatcher meta-pytorch/BackendBench#12

Open

IntraKernel Dispatcher #155330

Closed

github-actions bot deleted the gh/mikaylagawarecki/320/head branch September 13, 2025 02:06

pragupta added a commit to ROCm/pytorch that referenced this pull request Sep 16, 2025

Revert "Add utility to get computed kernel in torch.library (pytorch#…

7e4a926

…158393)" This reverts commit 1196bb1.

pragupta added a commit to ROCm/pytorch that referenced this pull request Sep 17, 2025

Revert "Add utility to get computed kernel in torch.library (pytorch#…

a975dfe

…158393)" This reverts commit 1196bb1.

pragupta added a commit to ROCm/pytorch that referenced this pull request Sep 24, 2025

Revert "Add utility to get computed kernel in torch.library (pytorch#…

d1c2fe4

…158393)" This reverts commit 1196bb1. (cherry picked from commit a975dfe)

jataylo added a commit to jataylo/pytorch that referenced this pull request Oct 17, 2025

Revert "Add utility to get computed kernel in torch.library (pytorch#…

be54763

…158393)" This reverts commit 1196bb1.

Add utility to get computed kernel in torch.library #158393

Add utility to get computed kernel in torch.library #158393

Uh oh!

Conversation

mikaylagawarecki commented Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 15, 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/158393

✅ No Failures

Uh oh!

Uh oh!

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

albanD Aug 1, 2025

Choose a reason for hiding this comment

Uh oh!

mikaylagawarecki Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

albanD Aug 1, 2025

Choose a reason for hiding this comment

Uh oh!

mikaylagawarecki Aug 5, 2025 • 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!

albanD Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

albanD Aug 1, 2025

Choose a reason for hiding this comment

Uh oh!

albanD Aug 1, 2025

Choose a reason for hiding this comment

Uh oh!

mikaylagawarecki Aug 5, 2025 • 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!

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

mikaylagawarecki commented Aug 13, 2025

Uh oh!

pytorchmergebot commented Aug 13, 2025

Uh oh!

pytorchmergebot commented Aug 13, 2025

Uh oh!

pytorchmergebot commented Aug 13, 2025

Merge started

Uh oh!

pytorchmergebot commented Aug 13, 2025

Merge failed

Uh oh!

mikaylagawarecki commented Aug 13, 2025

Uh oh!

pytorchmergebot commented Aug 13, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mikaylagawarecki commented Jul 15, 2025 •

edited

Loading

pytorch-bot bot commented Jul 15, 2025 •

edited

Loading

mikaylagawarecki Aug 5, 2025 •

edited

Loading

mikaylagawarecki Aug 5, 2025 •

edited

Loading