pytorch/torch/csrc
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
..
acc Add functions to setup PrivateUse1 as a python backend device. (#157859) 2025-10-01 21:32:59 +00:00
api [doc]: Small typos (#162982) 2025-09-16 17:42:19 +00:00
autograd Revert "[BE] Remove HermeticPyObjectTLS and Simplify PythonOpRegistrationTrampoline (#163464)" 2025-09-30 18:20:20 +00:00
cpu
cuda [PyTorch CCA] Add an API to get expandable segment sizes (#163771) 2025-10-01 02:16:58 +00:00
distributed Add option to FakeProcessGroup to raise error if comms are invoked. (#162841) 2025-10-01 17:48:19 +00:00
dynamo [dynamo] do not pop from framelocals dict in Python 3.10 (#164316) 2025-10-01 10:20:46 +00:00
export [serialization] Add pte file to archive (#162520) 2025-09-11 07:59:11 +00:00
functionalization Reapply "Make functionalization ViewMeta serializable with pickle. (#143712)" (#163769) 2025-09-25 10:27:37 +00:00
functorch Add magic TORCH_MAKE_PYBIND_ENUM_FASTER macro (#163527) 2025-09-26 17:59:22 +00:00
fx
inductor [aoti] AOTI mingw cross compilation (#163188) 2025-10-01 02:22:06 +00:00
instruction_counter [Code Clean] Replace std::runtime_error with TORCH_CHECK (#163610) 2025-09-26 04:52:48 +00:00
jit Better error handling in torch/csrc/jit/ir/* (#163757) 2025-09-28 01:18:24 +00:00
lazy [Code Clean] Replace std::runtime_error with TORCH_CHECK (#163610) 2025-09-26 04:52:48 +00:00
monitor Add magic TORCH_MAKE_PYBIND_ENUM_FASTER macro (#163527) 2025-09-26 17:59:22 +00:00
mps [MPS] Add API to query GPU core count (#160414) 2025-08-14 00:05:17 +00:00
mtia
multiprocessing
onnx Add magic TORCH_MAKE_PYBIND_ENUM_FASTER macro (#163527) 2025-09-26 17:59:22 +00:00
profiler Add magic TORCH_MAKE_PYBIND_ENUM_FASTER macro (#163527) 2025-09-26 17:59:22 +00:00
stable [STABLE ABI] Add clone method to torch::stable::Tensor (#161896) 2025-09-22 20:39:24 +00:00
tensor
utils Revert "[BE] Remove HermeticPyObjectTLS and Simplify PythonOpRegistrationTrampoline (#163464)" 2025-09-30 18:20:20 +00:00
xpu Add a new API torch.xpu.can_device_access_peer for Intel GPU (#162705) 2025-09-16 18:00:22 +00:00
copy_utils.h
CudaIPCTypes.cpp
CudaIPCTypes.h
DataLoader.cpp Replace std::runtime_error with TORCH_CHECK (#159344) 2025-09-16 09:00:06 +00:00
DataLoader.h
Device.cpp Replace std::runtime_error with TORCH_CHECK (#159344) 2025-09-16 09:00:06 +00:00
Device.h
DeviceAccelerator.cpp Revert "Generalize torch._C._set_allocator_settings to be generic (#156175)" (#161626) 2025-08-27 21:37:14 +00:00
DeviceAccelerator.h
Dtype.cpp
Dtype.h
DynamicTypes.cpp
DynamicTypes.h
empty.c
Event.cpp
Event.h
Exceptions.cpp Remove unused PyIntXXX, THPUtils_newReal_BOOL, THPQXXX macros (#164056) 2025-09-30 13:48:25 +00:00
Exceptions.h 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
Export.h
Generator.cpp
Generator.h
itt_wrapper.cpp
itt_wrapper.h
itt.cpp
itt.h
Layout.cpp
Layout.h
MemoryFormat.cpp
MemoryFormat.h
Module.cpp Add functions to setup PrivateUse1 as a python backend device. (#157859) 2025-10-01 21:32:59 +00:00
Module.h
PyInterpreter.cpp Revert "[BE] Remove HermeticPyObjectTLS and Simplify PythonOpRegistrationTrampoline (#163464)" 2025-09-30 18:20:20 +00:00
PyInterpreter.h
PyInterpreterHooks.cpp
PyInterpreterHooks.h
python_dimname.cpp
python_dimname.h
python_headers.h
QScheme.cpp
QScheme.h
README.md
serialization.cpp Using std::vector or c10::SmallVector instead of CArray (#160959) 2025-08-20 05:32:29 +00:00
serialization.h
Size.cpp
Size.h
Storage.cpp Revert "[BE] Remove HermeticPyObjectTLS and Simplify PythonOpRegistrationTrampoline (#163464)" 2025-09-30 18:20:20 +00:00
Storage.h
StorageMethods.cpp
StorageMethods.h
StorageSharing.cpp
StorageSharing.h
Stream.cpp
Stream.h
stub.c
THConcat.h
THP.h Remove unused PyIntXXX, THPUtils_newReal_BOOL, THPQXXX macros (#164056) 2025-09-30 13:48:25 +00:00
TypeInfo.cpp
TypeInfo.h
Types.h
utils.cpp Add uuid to XPU device properties (#161392) 2025-09-02 06:41:32 +00:00
utils.h Remove unused PyIntXXX, THPUtils_newReal_BOOL, THPQXXX macros (#164056) 2025-09-30 13:48:25 +00:00

csrc

The csrc directory contains all of the code concerned with integration with Python. This is in contrast to lib, which contains the Torch libraries that are Python agnostic. csrc depends on lib, but not vice versa.

There are a number of utilities for easing integration with Python which are worth knowing about, which we briefly describe here. But the most important gotchas:

  • DO NOT forget to take out the GIL with pybind11::gil_scoped_acquire before calling Python API or bringing a THPObjectPtr into scope.

  • Make sure you include Python.h first in your header files, before any system headers; otherwise, you will get error: "_XOPEN_SOURCE" redefined error. If you pay attention to warnings, you will see where you need to do this.

Notes

Note [Storage is not nullptr]

Historically, Torch supported nullptr storage, as a minor optimization to avoid having to allocate a storage object when it would be empty. However, this is actually a confusing special case to deal with, so by-in-large, PyTorch assumes that, in fact, storage is never nullptr.

One important case where this assumption is important is when tracking the CUDA device a tensor is stored in: this information is stored solely in the storage, so if a storage is nullptr, we lose this information.

Although storage is never nullptr, the data field of c10::StorageImpl may be nullptr. This mostly occurs when we want to pre-allocate an output tensor struct, but then have it be resized and filled with data by some operator: there's no point in allocating data for it in this case!

Files

Exceptions.h

Frequently when working with the Python API, you may call a function which returns an error. In this case, we want to return directly to the Python interpreter, so that this exception can be propagated accordingly; however, because the Python API is C-based, what actually will happen is it will return control to whatever C++ code called it. Similarly, if we raise a C++ exception, prior to returning to the Python interpreter, we must set the Python error flags, so it turns into a C++ exception.

Moreover, when using the following macros, the generated warnings will be converted into python warnings that can be caught by the user.

Exceptions define helpers for two main cases:

  • For code where you write the python binding by hand, HANDLE_TH_ERRORS, END_HANDLE_TH_ERRORS and an exception class python_error. You call them like this:
// Entry point from Python interpreter
PyObject* run(PyObject* arg) {
  HANDLE_TH_ERRORS
  ...
  if (!x) throw python_error();
  // From c10/Exception.h
  TORCH_CHECK(cond, "cond was false here");
  TORCH_WARN("Warning message");
  ...
  END_HANDLE_TH_ERRORS
}

The HANDLE_TH_ERRORS macro will catch all exceptions and convert them into an appropriate Python signal. python_error is a special exception which doesn't contain any info, instead it says, "An error occurred in the Python API; if you return to the interpreter, Python will raise that exception, nothing else needs to be done."

  • For code that you bind using pybind, HANDLE_TH_ERRORS and END_HANDLE_TH_ERRORS_PYBIND can be used. They will work jointly with pybind error handling to raise pytorch errors and warnings natively and let pybind handle other errors. It can be used as:
// Function given to the pybind binding
at::Tensor foo(at::Tensor x) {
  HANDLE_TH_ERRORS
  ...
  if (!x) throw python_error();
  // pybind native error
  if (!x) throw py::value_error();
  // From c10/Exception.h
  TORCH_CHECK(cond, "cond was false here");
  TORCH_WARN("Warning message");
  ...
  END_HANDLE_TH_ERRORS_PYBIND
}

GIL

Whenever you make any calls to the Python API, you must have taken out the Python GIL, as none of these calls are thread safe. pybind11::gil_scoped_acquire is a RAII struct which handles taking and releasing the GIL. Use it like this:

void iWantToUsePython() {
  pybind11::gil_scoped_acquire gil;
  ...
}

In general, the compiler will NOT warn you if you use Python functionality without taking out the GIL, so DO NOT FORGET this call.

utils/object_ptr.h

THPPointer is a smart pointer class analogous to std::shared_ptr, but which is overloaded to handle reference counting scheme of various objects which are not based on shared_ptr. The most important overloads are:

  • PyObject (so important we've aliased it as THPObjectPtr), which hooks into Python reference counting. (By the way, that means you MUST take out the GIL before bringing one of these into scope!)

  • The various TH tensor and storage types (e.g., THTensor), which hook into TH's reference counting. (TH's reference counting IS thread safe, no locks necessary.)