OpenReg Background

torch_openreg is part of this RFC, and its goal is not to implement a fully functional or high-performance PyTorch backend. Instead, it aims to:

• Track the evolution of community features and provide up-to-date standardized integration implementations, serving as the official reference and code example for integration documentation.
• The goal is to cover all functional points of new device integration into PyTorch, ensuring that the integration mechanisms themselves are robust and complete.

Notes:

• Minimalism Principle: The core goal is to enable/validate all integration paths/mechanisms for new backends with PyTorch. All features follow a “just enough” strategy to ensure correctness of the relevant integration mechanisms.
• Authenticity Principle: Implemented using integration patterns of a real CUDA-like device.

Usage

import torch

if not torch.openreg.is_available():
    print("OpenReg backend is not available in this build.")
    exit()

print("OpenReg backend is available!")

device = torch.device("openreg")

x = torch.tensor([[1., 2.], [3., 4.]], device=device)
y = x + 2
print("Result y:\n", y)
print(f"Device of y: {y.device}")

z = y.cpu()
print("Result z:\n", z)
print(f"Device of z: {z.device}")

Work Items

The current preliminary scope of work is as follows. Special thanks to @albanD for the support (subject to further updates). The plan includes both functionality development and comprehensive test coverage:

• Low Level

  • Device
  • Stream: #160099 #160100
  • Event: #160099 #160100
  • Allocator
  • Memory Management
  • RNG
  • Guard: #160101
  • DeviceInterface
  • Tensor/Storage Customization
  • Documents: #158644
  • ...

• Operator

  • C++/Python

    • Operator Registration: #158644 #159640
    • Meta Registration: #158644 #159640
    • Fallback Registration: #158644 #159640
    • Fallthrough
    • torch.autograd.Function
    • Documents
    • ...

• Python binding/frontend

  • Python C++ Binding
  • Device Agnostic API
  • Documents

• Domain

  • AMP: #162050
  • Distributed
  • Compiler
    • Add testcases related to this issue later.
  • ONNX
  • Quantization
  • Autload: #158555
  • Documents
  • ...

• Others:

  • Windows/OSX Support: #159441
  • Support source distribution for OpenReg

• Bugs:

  • The whl package contains files (stub.c) that are not expected to be in it. #159845
  • Fix the display problem of sphinx tab in pytorch_sphinx_theme2(using tab-set tab-item instead)
  • All testcases in torch_openreg should support Linux, Windows and OSX platform and should not have something like skipIfWindows, skipifOSX and so on
  • Fix the issue about TestAutograd.test_autograd_init under torch.compiler
  • Fix the issue about TestTensorType.test_backend_type_methods under torch.compiler
  • Privide API of getting capalibities for Accelerator, refer to this issue for more infos and add doc as well when it is fixed.
  • Add poison fork to accelerator integration serial document

CC

@albanD and others who are interested in.
@wschin @trajepl who are using PrivateUse1.
@zeshengzong @can-gaa-hou @hipudding who will work on this.

cc @NmomoN @mengpenghui @fwenguang @cdzhan @1274085042 @PHLens @albanD