OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity
b5c4f46bb9
pytorch/torch
History
Han Qi b5c4f46bb9 Add functions to setup PrivateUse1 as a python backend device. (#157859) 
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.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/157859
Approved by: https://github.com/albanD
2025-10-01 21:32:59 +00:00
..
_awaits
_C Add functions to setup PrivateUse1 as a python backend device. (#157859) 2025-10-01 21:32:59 +00:00
_C_flatbuffer
_custom_op [BE]: ruff PLC0207 - use maxsplit kwarg (#160107) 2025-08-08 03:14:59 +00:00
_decomp [dynamic shapes] unbacked-safe slicing (#161414) 2025-09-30 01:15:19 +00:00
_dispatch
_dynamo [export] Preserve nn_module_stack for aliased nn modules (#164311) 2025-10-01 18:04:33 +00:00
_export [3.14] make unbacked_sym[int/float]_counter integers (#163920) 2025-09-30 17:42:55 +00:00
_functorch [1/N] Fix ruff warnings (#164333) 2025-10-01 16:48:32 +00:00
_higher_order_ops Fix invalid f-strings (#164112) 2025-09-30 04:17:13 +00:00
_inductor [inductor] add a runtime assert for triton shapes (#164242) 2025-10-01 18:55:33 +00:00
_lazy
_library [ignore][codex-test] Add typing to simple library registry (#161367) 2025-09-23 02:08:55 +00:00
_logging Add compile_id: Optional[CompileID] to torch._logging._internal.trace_structured_artifact (#160440) 2025-08-13 06:28:23 +00:00
_numpy Remove unnecessary list comprehensions (#164103) 2025-09-30 03:56:54 +00:00
_prims [2/N] Simplify "in" operation for containers of a single item (#164323) 2025-10-01 05:39:11 +00:00
_prims_common [BE][Easy]: Add prims common TypeGuard (#164263) 2025-10-01 03:13:10 +00:00
_refs [export] support unbacked stack (#163867) 2025-10-01 16:48:46 +00:00
_strobelight
_subclasses Add functions to setup PrivateUse1 as a python backend device. (#157859) 2025-10-01 21:32:59 +00:00
_vendor
accelerator Add unified memory APIs for torch.accelerator (#152932) 2025-08-08 17:41:22 +00:00
amp [Easy][AMP] Refactor the AMP logic for getting dtype (#162796) 2025-09-21 06:32:35 +00:00
ao [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
autograd [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
backends Add SVE128 ISA (#158932) 2025-09-29 14:49:19 +00:00
compiler Simplify PrecompileContext to no longer be a CacheArtifactManager (#162886) 2025-09-20 01:24:37 +00:00
contrib
cpu [AMP] Add deprecated decorator for torch.xxx.amp.autocast class (#163654) 2025-09-27 14:37:12 +00:00
csrc Add functions to setup PrivateUse1 as a python backend device. (#157859) 2025-10-01 21:32:59 +00:00
cuda [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
distributed [Replicate][Pipeline Parallelism] integration of new replicate function with pipeline parallelism (#164031) 2025-10-01 18:01:16 +00:00
distributions
export [export] Explicitly passing requires_grad to nn.Parameter() in deserialization (#164290) 2025-10-01 02:55:20 +00:00
fft
func
futures
fx [1/N] Fix ruff warnings (#164333) 2025-10-01 16:48:32 +00:00
headeronly Migrate DeviceType to torch/headeronly (#163999) 2025-09-30 23:13:27 +00:00
jit [BE] Delete all pre py-3.10 checks (#163653) 2025-09-23 23:22:53 +00:00
legacy
lib
linalg Revert "Add __init__.pyi to torch/linalg (#160750)" 2025-09-02 16:53:55 +00:00
masked [1/N] Simplify "in" operation for containers of a single item (#164224) 2025-09-30 19:59:43 +00:00
monitor
mps Add type annotations to MPS profiler utilities (#163486) 2025-09-27 23:00:53 +00:00
mtia [BE] Add Documentation for Device APIs (#162834) 2025-09-16 17:01:06 +00:00
multiprocessing Allow parallel start NUMA binding (#161576) 2025-08-28 01:15:58 +00:00
nativert Add SVE128 ISA (#158932) 2025-09-29 14:49:19 +00:00
nested Add NestedTensor dispatch for _is_any_true/_is_all_true (#162096) 2025-09-22 20:22:44 +00:00
nn [1/N] Fix ruff warnings (#164333) 2025-10-01 16:48:32 +00:00
numa Allow parallel start NUMA binding (#161576) 2025-08-28 01:15:58 +00:00
onnx [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
optim [optim] prevent unintended aliasing in lr_scheduler; update type annotations/docs (#163120) 2025-09-25 06:58:58 +00:00
package [2/N] Apply ruff UP035 check in torch files (#164054) 2025-09-29 03:35:32 +00:00
profiler removed duplicate imports (#161685) 2025-08-31 16:21:49 +00:00
quantization [BE][PYFMT] migrate PYFMT for torch/[p-z]*/ to ruff format (#144552) 2025-08-07 00:09:56 +00:00
signal [BE][PYFMT] migrate PYFMT for torch/[p-z]*/ to ruff format (#144552) 2025-08-07 00:09:56 +00:00
sparse Use computed buffer sizes of torch for cusparseLt metadata (#163125) 2025-09-19 22:12:40 +00:00
special [BE][PYFMT] migrate PYFMT for torch/[p-z]*/ to ruff format (#144552) 2025-08-07 00:09:56 +00:00
testing [1/N] Fix ruff warnings (#164333) 2025-10-01 16:48:32 +00:00
utils Add functions to setup PrivateUse1 as a python backend device. (#157859) 2025-10-01 21:32:59 +00:00
xpu [2/N] Apply ruff UP035 check in torch files (#164054) 2025-09-29 03:35:32 +00:00
__config__.py
__future__.py
__init__.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_appdirs.py
_classes.py
_compile.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_custom_ops.py
_environment.py
_guards.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_jit_internal.py [Code Clean] Remove deadcodes about Python3.9 [7/N] (#163646) 2025-09-24 07:30:50 +00:00
_linalg_utils.py
_lobpcg.py
_lowrank.py
_meta_registrations.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_namedtensor_internals.py
_ops.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_python_dispatcher.py
_size_docs.py
_sources.py
_storage_docs.py
_streambase.py
_tensor_docs.py
_tensor_str.py [1/N] Fix ruff warnings (#164333) 2025-10-01 16:48:32 +00:00
_tensor.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_thread_safe_fork.py
_torch_docs.py Update documentation for torch.index_select (#163616) 2025-09-25 18:29:17 +00:00
_utils_internal.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_utils.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_VF.py
_vmap_internals.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_weights_only_unpickler.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
CMakeLists.txt Revert "[RELAND] Always build USE_DISTRIBUTED (#160449) and Make distributed modules importable even when backend not built (#159889) (#162594)" 2025-09-25 13:47:46 +00:00
custom_class_detail.h
custom_class.h
extension.h
functional.py unify broadcast_shapes functions and avoid duplicates (#160251) 2025-08-16 00:54:32 +00:00
header_only_apis.txt Migrate DeviceType to torch/headeronly (#163999) 2025-09-30 23:13:27 +00:00
hub.py [Fix] Adding missing f prefixes to formatted strings [1/N] (#164065) 2025-09-29 04:53:00 +00:00
library.h Using std::make_unique<T>() instead of unique<T>(new T()) (#160723) 2025-08-19 10:25:47 +00:00
library.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
overrides.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
py.typed
quasirandom.py
random.py
return_types.py
script.h
serialization.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
storage.py
torch_version.py
types.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
version.py.tpl