pytorch/test/cpp_extensions/open_registration_extension

PyTorch OpenReg

This folder contains a self-contained example of a PyTorch out-of-tree backend leveraging the "PrivateUse1" backend from core.

How to use

Install as standalone with python setup.py develop (or install) from this folder. You can run test via python {PYTORCH_ROOT_PATH}/test/test_openreg.py.

Design principles

For simplicity anything that can be implemented from python is done so. A real implementation will most likely want to call these different APIs from c++ directly.

The current version sends everything back to python and contains enough implementation to run basic model, transfer host/device and printing.

The codebase is split as follows:

• pytorch_openreg/__init__.py
  • imports torch to get core state initialized.
  • imports ._aten_impl to register our aten op implementations to torch.
  • imports .C to load our c++ extension that registers more ops, allocator and hooks.
  • renames the PrivateUse1 backend and register our python-side module.
• pytorch_openreg/_aten_impl.py
  • Define a new torch.Library that registers a fallback that will be called whenever a backend kernel for PrivateUse1 is called. It contains the logic to handle all kind of native functions, computing the output metadata, allocating it and only calling into the device daemon to perform computation.
• pytorch_openreg/_device_daemon.py
  • contains the Allocator (responsible for allocating memory on the device side and host side, as int8 buffers).
  • contains Driver, which as user-process driver to deal with some information needed to be done in driver.
  • contains Executor, which as device-process exector to do something related device logic.
• pytorch_openreg/_meta_parser.py mainly contain utilities to send objects over the wire from the user process to the device process.
  • The main class there is OpenRegTensorMeta that contains all the metadata sent to the device which should be enough for it to populate the output Tensor.

Next steps

The main next step would be to:

• Replace the current open_registration_extension.cpp test in PyTorch CI with this.