enable dispatch stub for backend PrivateUse1 #99611

caizhi-mt · 2023-04-20T07:42:38Z

When expanding the new backend of pytorch in the form of out ot tree, Privateuse1 will be reused. So we also need to support PrivateUse1 in the dispatch stub module

cc @bdhirsh

pytorch-bot · 2023-04-20T07:42:42Z

🔗 Helpful Links

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

📄 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 bf93485:
💚 Looks good so far! There are no failures yet. 💚

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

caizhi-mt · 2023-04-23T02:34:17Z

@bdhirsh Please help me review this commit, thank you

bdhirsh · 2023-04-24T17:18:46Z

aten/src/ATen/native/DispatchStub.h

@@ -170,6 +172,10 @@ struct DispatchStub<rT (*)(Args...), T> {
    impl.mps_dispatch_ptr = reinterpret_cast<void*>(fn_ptr);
  }

+  void set_private_use1_dispatch_ptr(FnPtr fn_ptr) {


nit: can you make it set_privateuse1_dispatch_ptr (for greppability)

bdhirsh · 2023-04-24T17:22:24Z

One comment: I don't think getting DispatchStub to work with externally backends, is strictly necessary, since for any ATen kernels that are implemented with it, you can always register your own dedicated kernel to the PrivateUse1 dispatch key for that operator.

Is the idea that you want to be able to re-use some kernels from ATen that are implemented in terms of dispatch stub? That sounds pretty useful, I just wanted to confirm.

If so, it would be good to add a test. Do you think you can add a DispatchStub handler in test/cpp_extensions/open_registration_extension.cpp, and exercise that logic in test/test_cpp_extensions_open_device_registration.py?

caizhi-mt · 2023-04-25T14:33:01Z

yes,that's exactly what I want to do, "re-use some kernels" through dispatch stub。
A simple unit test has been added, which re-use CPU Abs kernel through "abs_stub"。

bdhirsh · 2023-04-26T20:58:03Z

test/cpp_extensions/open_registration_extension.cpp

+  }
+}
+
+REGISTER_PRIVATEUSE1_DISPATCH(abs_stub, &abs_kernel);


it would be nice to be able to test that we actually ran this kernel (right now the "test" just calls abs() and checks that the outputs match).

One way to do it would be the way it's done in custom_add_Tensor below - add a global counter for whether our custom abs kernel was called, that we increment inside of the kernel. Then expose a custom_abs_called helper to python that you can assert in the python code.

Thank you for the suggestion. I have modified the unit test according to your suggestion, added custom_abs_called and abs_counter in test.

bdhirsh · 2023-05-02T21:46:37Z

Can you take a look at the failing CI?

It looks like the custom cpp extension is failing to be compiled:

2023-05-01T10:03:59.2852263Z  const at::Tensor& custom_resize_(const at::Tensor& self, at::IntArrayRef size,
2023-05-01T10:03:59.2852610Z                    ^~~~~~~~~~~~~~
2023-05-01T10:03:59.2853410Z /var/lib/jenkins/workspace/test/cpp_extensions/open_registration_extension.cpp:63:19: note: ‘const at::Tensor& custom_resize_(const at::Tensor&, c10::IntArrayRef, c10::optional<c10::MemoryFormat>)’ previously defined here
2023-05-01T10:03:59.2854072Z  const at::Tensor& custom_resize_(
2023-05-01T10:03:59.2854350Z                    ^~~~~~~~~~~~~~
2023-05-01T10:03:59.2854649Z ninja: build stopped: subcommand failed.

bdhirsh · 2023-05-02T21:47:19Z

test/cpp_extensions/open_registration_extension.cpp

+    c10::IntArrayRef size,
+    c10::optional<c10::MemoryFormat> optional_memory_format) {
+  // Since this custom device is just for testing, not bothering to implement kernels.
+  return self;


it looks like it's failing because custom_resize_() is already defined further down.

You can just use the existing one instead of defining this new one.

ok, I have updated the code

all checks have passed, could this PR be merged?

bdhirsh · 2023-05-05T14:10:10Z

@pytorchbot merge

pytorch-bot · 2023-05-05T14:10:14Z

This PR needs to be approved by an authorized maintainer before merge.

caizhi-mt · 2023-05-06T06:05:02Z

This PR need to be approved by an authorized maintainer before merge. Could you approve this PR? @bdhirsh

caizhi-mt · 2023-05-08T00:27:41Z

@bdhirsh Could you approve this PR? Thank you

caizhi-mt · 2023-05-09T06:13:45Z

@ezyang Could you review and approve this PR? Thank you

ezyang · 2023-05-09T15:11:15Z

Sorry, y'all have already gone a review on this, but I want to put the brakes on this. The dispatch stub API is not intended for external users, and for the most part we've expunged all the cases where we registered kernels that call into stubs so that they don't CompositeExplicitAutograd anymore. Concretely, what operators do you think you need this for?

caizhi-mt · 2023-05-10T02:46:27Z

Actually I want to extend a new hardware backend based on "PrivateUse1" key. And I want to use the dispatchstub mechanism to call my kernels. such as: https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/UnaryOps.cpp#L446
This is why I commit this PR。

ezyang · 2023-05-10T14:19:56Z

Are you really using TensorIterator to implement your kernels though? Seems... unlikely.

caizhi-mt · 2023-05-11T02:50:32Z

yes, I actually used TensorIterator.

I am doing some cuda-compatible work using Privateuse1 key base on our hardware which will reuse cuda kernels.
We will refer to cuda kernels to implement our own kernels using TensorIterator

PrivateUse1 key is just like "CPU" key, "CUDA" key. I think it is necessary to support privateuse1 in dispatchstub.

ezyang

ok fine

ezyang · 2023-05-11T21:18:25Z

@pytorchbot merge

pytorchmergebot · 2023-05-11T21:20:58Z

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 · 2023-05-11T21:20:59Z

Merge failed

Reason: This PR is too stale; the last push date was more than 3 days ago. Please rebase and try again. You can rebase and merge by leaving the following comment on this PR:
@pytorchbot merge -r
Or just rebase by leaving @pytorchbot rebase comment

Details for Dev Infra team

Raised by workflow job

ezyang · 2023-05-12T00:30:17Z

@pytorchbot merge -r

pytorchmergebot · 2023-05-12T00:32:13Z

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

pytorchmergebot · 2023-05-12T00:32:19Z

Successfully rebased enable_privateuse1_dispatch onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout enable_privateuse1_dispatch && git pull --rebase)

pytorchmergebot · 2023-05-12T00:33:27Z

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

caizhi-mt · 2023-05-12T03:59:19Z

@pytorchbot merge

pytorchmergebot · 2023-05-12T03:59:38Z

The merge job was canceled. If you believe this is a mistake,then you can re trigger it through pytorch-bot.

pytorchmergebot · 2023-05-12T04:02:05Z

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

pytorchbot added the open source label Apr 20, 2023

bdhirsh reviewed Apr 24, 2023

View reviewed changes

bdhirsh added the module: backend non-standard backend support label Apr 24, 2023

bdhirsh mentioned this pull request Apr 24, 2023

Contribute to the privateuse1 backend. #98406

Open

caizhi-mt force-pushed the enable_privateuse1_dispatch branch 2 times, most recently from d217ca1 to a287040 Compare April 25, 2023 14:26

mikaylagawarecki added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Apr 25, 2023

bdhirsh reviewed Apr 26, 2023

View reviewed changes

caizhi-mt force-pushed the enable_privateuse1_dispatch branch 2 times, most recently from bde88cf to effa7de Compare May 1, 2023 07:49

bdhirsh reviewed May 2, 2023

View reviewed changes

caizhi-mt force-pushed the enable_privateuse1_dispatch branch 3 times, most recently from 9253ca6 to ce72f84 Compare May 5, 2023 01:39

bdhirsh added the release notes: composability release notes category label May 5, 2023

caizhi-mt force-pushed the enable_privateuse1_dispatch branch 2 times, most recently from 6099fdc to aba8401 Compare May 6, 2023 02:27

ezyang approved these changes May 11, 2023

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label May 11, 2023

pytorchmergebot added the merging label May 11, 2023

pytorchmergebot removed the merging label May 11, 2023

caizhi-mt added 3 commits May 12, 2023 00:32

enable dispatch stub for backend PrivateUse1

611be82

rename "private_use1" to "privateuse1" and add unit tests

64b48f6

modify unit test for custom dispatch_stub module

bf93485

pytorchmergebot force-pushed the enable_privateuse1_dispatch branch from aba8401 to bf93485 Compare May 12, 2023 00:32

pytorchmergebot added the merging label May 12, 2023

pytorchmergebot added Merged and removed merging labels May 12, 2023

pytorchmergebot closed this in bf50180 May 12, 2023

neeldug mentioned this pull request Jun 10, 2023

torchgen/gen_backend_stubs.py compatibility with DispatchStubs #103370

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

enable dispatch stub for backend PrivateUse1 #99611

enable dispatch stub for backend PrivateUse1 #99611

caizhi-mt commented Apr 20, 2023 •

edited by pytorch-bot bot

pytorch-bot bot commented Apr 20, 2023 •

edited

caizhi-mt commented Apr 23, 2023

bdhirsh Apr 24, 2023

caizhi-mt Apr 25, 2023

bdhirsh commented Apr 24, 2023

caizhi-mt commented Apr 25, 2023

bdhirsh Apr 26, 2023

caizhi-mt May 1, 2023

bdhirsh commented May 2, 2023

bdhirsh May 2, 2023

caizhi-mt May 3, 2023

caizhi-mt May 5, 2023

bdhirsh commented May 5, 2023

pytorch-bot bot commented May 5, 2023

caizhi-mt commented May 6, 2023

caizhi-mt commented May 8, 2023

caizhi-mt commented May 9, 2023

ezyang commented May 9, 2023 •

edited

caizhi-mt commented May 10, 2023

ezyang commented May 10, 2023

caizhi-mt commented May 11, 2023

ezyang left a comment

ezyang commented May 11, 2023

pytorchmergebot commented May 11, 2023

pytorchmergebot commented May 11, 2023

ezyang commented May 12, 2023

pytorchmergebot commented May 12, 2023

pytorchmergebot commented May 12, 2023

pytorchmergebot commented May 12, 2023

caizhi-mt commented May 12, 2023

pytorchmergebot commented May 12, 2023

pytorchmergebot commented May 12, 2023

enable dispatch stub for backend PrivateUse1 #99611

enable dispatch stub for backend PrivateUse1 #99611

Conversation

caizhi-mt commented Apr 20, 2023 • edited by pytorch-bot bot

pytorch-bot bot commented Apr 20, 2023 • edited

🔗 Helpful Links

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

✅ No Failures

caizhi-mt commented Apr 23, 2023

bdhirsh Apr 24, 2023

Choose a reason for hiding this comment

caizhi-mt Apr 25, 2023

Choose a reason for hiding this comment

bdhirsh commented Apr 24, 2023

caizhi-mt commented Apr 25, 2023

bdhirsh Apr 26, 2023

Choose a reason for hiding this comment

caizhi-mt May 1, 2023

Choose a reason for hiding this comment

bdhirsh commented May 2, 2023

bdhirsh May 2, 2023

Choose a reason for hiding this comment

caizhi-mt May 3, 2023

Choose a reason for hiding this comment

caizhi-mt May 5, 2023

Choose a reason for hiding this comment

bdhirsh commented May 5, 2023

pytorch-bot bot commented May 5, 2023

caizhi-mt commented May 6, 2023

caizhi-mt commented May 8, 2023

caizhi-mt commented May 9, 2023

ezyang commented May 9, 2023 • edited

caizhi-mt commented May 10, 2023

ezyang commented May 10, 2023

caizhi-mt commented May 11, 2023

ezyang left a comment

Choose a reason for hiding this comment

ezyang commented May 11, 2023

pytorchmergebot commented May 11, 2023

Merge started

pytorchmergebot commented May 11, 2023

Merge failed

ezyang commented May 12, 2023

pytorchmergebot commented May 12, 2023

pytorchmergebot commented May 12, 2023

pytorchmergebot commented May 12, 2023

Merge started

caizhi-mt commented May 12, 2023

pytorchmergebot commented May 12, 2023

pytorchmergebot commented May 12, 2023

Merge started

caizhi-mt commented Apr 20, 2023 •

edited by pytorch-bot bot

pytorch-bot bot commented Apr 20, 2023 •

edited

ezyang commented May 9, 2023 •

edited