PR #157859: Add functions to setup PrivateUse1 as a python backend device.

• Repository: pytorch/pytorch
• Author: qihqi
• State: closed | merged: False
• Link: https://github.com/pytorch/pytorch/pull/157859

Description (problem / solution / changelog)

Fixes #156052 and #156444.

This PR setup the privateuseone key in Python to be used as a python backend for pytorch. Meaning that, after calling setup_privateuseone_for_python_backend('npy'), one can use a subclass to with that device to hold arbitrary python data as "device data" and use torch.library to register ops that takes that Tensor.

Changes done in this PR:

1. Register an vanilla Device Guard: I extended NoOpDeviceGuard to have allow device index of 0 and to not raise errors when event related functions are accessed. If I don't do those, when calling backward I would get errors. (CPU backend uses NoOpDeviceGuard just fine, although there seems to be special treatment of CPU in the autograd engine.
2. Tensor subclass allows not having __torch_dispatch__ if the device is not CUDA or CPU. The comment of the check suggests it was to avoid segfault when calling into ops that expects a storage. Here we have a different device so will not call into those ops.
3. python function that invokes the other incantations to setup the privateusekey backend.

This took inspiration of https://github.com/bdhirsh/pytorch_open_registration_example and https://github.com/tinygrad/tinygrad/blob/master/extra/torch_backend/wrapped_tensor.cpp; great thanks to @bdhirsh and @geohot.

Changed files

• build_variables.bzl (modified, +1/-0)
• c10/core/impl/DeviceGuardImplInterface.cpp (modified, +8/-2)
• c10/core/impl/DeviceGuardImplInterface.h (modified, +3/-0)
• test/dynamo_skips/PrivateUse1BackendTest.test_backend_simple (added, +0/-0)
• test/run_test.py (modified, +1/-0)
• test/test_privateuseone_python_backend.py (added, +147/-0)
• torch/_C/__init__.pyi.in (modified, +1/-0)
• torch/_C/_acc/__init__.pyi (added, +15/-0)
• torch/_subclasses/meta_utils.py (modified, +7/-5)
• torch/csrc/Module.cpp (modified, +2/-0)
• torch/csrc/acc/Module.cpp (added, +196/-0)
• torch/csrc/acc/Module.h (added, +8/-0)
• torch/utils/backend_registration.py (modified, +79/-1)