Any simple example about the new way to register custom device through PrivateUse1

I know in pytorch 2.1, it is more convient for user to register their own device in pytorch,
Facilitating New Backend Integration by PrivateUse1 — PyTorch Tutorials 2.1.0+cu121 documentation,
it includes ddp amp sparce or anything else, but I have not seen any example code for how to utilize this new feature,
GitHub - Ascend/pytorch: Ascend PyTorch adapter (torch_npu). Mirror of https://gitee.com/ascend/pytorch is a good example but it is too complicated, so does the official Pytorch plan to provide a complete routine to explain how to use this new feature?

Hey!

I’m not sure what you mean here? The tutorial covers how to register your implementation onto the PrivateUse1 device and rename it.
Then you can use it as you use any other device. Nothing special is needed.

Hello, I mean any complete example instead of scattered code

Hey!

In that case no, I don’t think there is any plan beyond the tutorial, doc, tests and end to end example in a separate repo I’m afraid.

find, thanks :grinning:

For completeness I would add to the discussion the particular link to the relevant file with tests, which was very useful for me: pytorch/test/cpp_extensions/open_registration_extension.cpp at 7fd3b6988681557a4dfc35dc99807cdd78e805c9 · pytorch/pytorch · GitHub
It gives minimal complete standalone example of new device secretly storing data on CPU RAM, and showcasing most of the things developer of new backend might need.

It is compiled dynamically in this file pytorch/test/test_cpp_extensions_open_device_registration.py at 7fd3b6988681557a4dfc35dc99807cdd78e805c9 · pytorch/pytorch · GitHub
however can be also compiled as a static python library in the way explained in this tutorial Extending dispatcher for a new backend in C++ — PyTorch Tutorials 2.4.0+cu121 documentation

1 Like