在 C++ 中為新後端擴充套件排程器¶

創建於：2021 年 2 月 1 日 | 最後更新：2024 年 9 月 23 日 | 最後驗證：2024 年 11 月 5 日

在本教程中，我們將逐步講解擴充套件排程器以新增一個存在於 pytorch/pytorch 倉庫之外的新裝置並維護它以使其與原生 PyTorch 裝置同步的所有必要步驟。我們假設你熟悉如何在 C++ 中註冊排程運算元以及如何編寫自定義 autograd 函式。

注意

本教程涉及許多 PyTorch 內部元件，這些元件正在積極改進中。如果你決定按照本教程操作，請預期 API 會有所變化。我們將及時更新本教程，使其與最新 API 保持同步。

什麼是新後端？¶

向 PyTorch 新增新後端需要後端擴充套件者投入大量開發和維護工作。在新增新後端之前，讓我們首先考慮一些常見的用例和推薦的解決方案。

如果你對現有 PyTorch 運算元有新的演算法，請向 PyTorch 提交 PR。
如果你想提出新的運算元，請向 PyTorch 提交功能請求/PR。
如果你想為 Google TPU 和定製晶片等新裝置/硬體新增支援，這通常需要使用特定硬體的 API 來編寫核函式，請按照本教程為 PyTorch 新增一個樹外 (out-of-tree) 後端。
如果你想為現有運算元新增支援，但使用不同的張量佈局/表示形式，例如稀疏和量化，這要求你的核函式在給定佈局/表示限制的情況下以更高效的方式編寫，請按照本教程為 PyTorch 新增一個樹外 (out-of-tree) 後端。

在本教程中，我們將主要關注如何新增一個新的樹外裝置。為不同的張量佈局新增樹外支援可能與裝置有很多共同步驟，但我們尚未見到此類整合的示例，因此 PyTorch 可能需要額外的工作來支援它。

為你的後端獲取排程鍵¶

PyTorch 運算元是用 C++ 實現的，並透過 Python 繫結在 Python 前端中提供。PyTorch 排程器將一個運算元的實現劃分為多個核函式，每個核函式都與一個特定的排程鍵關聯。在 PyTorch 中支援新後端本質上意味著用 C++ 為每個 PyTorch 運算元編寫一個核函式，然後將它們註冊到排程器中代表你定製後端的排程鍵。

排程鍵是你在排程器系統中的識別符號。排程器檢視輸入張量攜帶的排程鍵，並據此呼叫正確的核函式。PyTorch 提供了三個保留的排程鍵（及其對應的 Autograd 鍵）用於原型化樹外後端擴充套件：

PrivateUse1/AutogradPrivateUse1
PrivateUse2/AutogradPrivateUse2
PrivateUse3/AutogradPrivateUse3

你可以選擇上面任何一個鍵來原型化你的定製後端。要在 PrivateUse1 後端上建立張量，你需要在 TensorImpl 建構函式中設定排程鍵。

/* Example TensorImpl constructor */
TensorImpl(
    Storage&& storage,
    DispatchKeySet ks,
    const caffe2::TypeMeta data_type);

// To create a TensorImpl on PrivateUse1 backend, pass in the following ks to TensorImpl creation.
DispatchKeySet ks = c10::DispatchKeySet{c10::DispatchKey::PrivateUse1, c10::DispatchKey::AutogradPrivateUse1};

注意，上面的 TensorImpl 類假設你的張量是由 CPU/CUDA 等儲存支援的。對於沒有儲存的後端，我們也提供了 OpaqueTensorImpl。你可能需要調整/重寫某些方法以適應你的定製硬體。pytorch 倉庫中的一個例子是 Vulkan TensorImpl。

注意

原型完成後，如果你計劃為你的後端擴充套件定期釋出版本，請隨時向 pytorch/pytorch 提交 PR，為你的後端保留一個專用的排程鍵。

獲取 PyTorch 運算元完整列表¶

PyTorch 在生成檔案 build/aten/src/ATen/RegistrationDeclarations.h 中提供了可擴充套件 C++ 運算元的完整列表。此檔案僅在從原始碼構建 PyTorch 後可用。以下是該檔案的一個片段：

Tensor abs(const Tensor & self); // {"schema": "aten::abs(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}
Tensor & abs_(Tensor & self); // {"schema": "aten::abs_(Tensor(a!) self) -> Tensor(a!)", "dispatch": "True", "default": "True"}
Tensor & abs_out(Tensor & out, const Tensor & self); // {"schema": "aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "True", "default": "False"}
Tensor absolute(const Tensor & self); // {"schema": "aten::absolute(Tensor self) -> Tensor", "dispatch": "False", "default": "False"}
Tensor & absolute_(Tensor & self); // {"schema": "aten::absolute_(Tensor(a!) self) -> Tensor(a!)", "dispatch": "False", "default": "False"}
Tensor & absolute_out(Tensor & out, const Tensor & self); // {"schema": "aten::absolute.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "False", "default": "False"}
Tensor angle(const Tensor & self); // {"schema": "aten::angle(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}
Tensor & angle_out(Tensor & out, const Tensor & self); // {"schema": "aten::angle.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "True", "default": "False"}
Tensor sgn(const Tensor & self); // {"schema": "aten::sgn(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}

一個運算元關聯著多個欄位。讓我們以 abs_out 為例進行分解：

Tensor & abs_out(Tensor & out, const Tensor & self); 是該運算元的 C++ 簽名，你的 C++ 核函式應該與此簽名完全匹配。
aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) 是代表該運算元的唯一 schema，與 C++ 簽名相比，它還包含別名和修改註釋。這是排程器用來查詢運算元的唯一識別符號。
dispatch 和 default 是布林欄位，提供了關於原生 PyTorch 核函式功能的資訊，從而暗示了後端擴充套件者是否需要實現該核函式。更多詳情請參見為新後端註冊核函式。

為新後端註冊核函式¶

要將你的核函式註冊到 PyTorch 排程器，可以使用在 C++ 中註冊排程運算元中描述的 TORCH_LIBRARY_IMPL API。

TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl(<schema_my_op1>, &my_op1);
  m.impl(<schema_my_op2>, &my_op2);
  m.impl(<schema_my_op2_backward>, &my_op2_backward);
}

現在讓我們深入瞭解，哪些運算元需要定製後端的核函式，以及這些核函式內部究竟是什麼。

PyTorch 目前有超過 1600 個運算元，並且仍在增長。後端擴充套件要跟上這個速度是不現實的。即使對於 CPU 或 CUDA 等原生後端，為每個新運算元編寫專用核函式通常也需要大量工作。

幸運的是，一些原生 PyTorch 核函式編寫方式使得它們可以分解為幾個已知運算元的組合。換句話說，你只需實現一組已知運算元（下文需要註冊的運算元），而不是所有 PyTorch 運算元。

PyTorch 運算元可以分為兩類：

需要註冊的運算元：這些運算元的 PyTorch 原生實現是後端特定的，因此需要為定製後端提供核函式。否則，在定製後端上呼叫此類運算元將會出錯。
- 在 RegistrationDeclarations.h 中，這些運算元在其 accompanying comments 中的元資料裡，dispatch 設定為 True 且 default 設定為 False。
註冊是可選的：後端擴充套件者可以跳過註冊這些運算元而不犧牲任何支援。然而，如果後端擴充套件者想要覆蓋 PyTorch 提供的預設核函式，他們仍然可以將自己定製的核函式註冊到其後端，排程器將僅為你的後端使用它。例如，PyTorch 當前的 max_pool2d 實現將 indices 作為前向輸出的一部分返回，這在 torch_xla 中會產生開銷，因此 torch_xla 轉而為 max_pool2d 註冊了自己的核函式。
- 在 RegistrationDeclarations.h 中，這些運算元在其 accompanying comments 中的元資料裡，dispatch 設定為 False 或 default 設定為 True。

新後端的 Autograd 支援¶

梯度公式大多是純數學的，因此對所有後端都通用。PyTorch 通常會將核函式註冊到別名排程鍵 Autograd，這意味著它可以被所有後端使用。

對於這些運算元，你無需擔心它們的導數公式，只需編寫 RegistrationDeclarations.h 中運算元的前向定義，PyTorch 會自動為你處理反向傳播。

Tensor my_op1(const Tensor& self, const Tensor& other) {
  // call your backend-specific APIs to implement my_op so that
  // it matches PyTorch's native behavior
}
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl(<schema_my_op1>, &my_op);
}

在某些情況下，PyTorch 的反向傳播核函式實現也是裝置特定的，以便最大程度地壓榨每個後端的效能。對於這些運算元，你也會在 RegistrationDeclarations.h 中看到 op_backward 顯示為 required registration。

Tensor my_op2_backward(const Tensor& self, const Tensor& other) {
  // call your backend-specific APIs to implement my_op2_backward so that
  // it matches PyTorch's native behavior
}

// Note backward kernel is still registered to PrivateUse1 instead of AutogradPrivateUse1.
// PyTorch will wrap your backward kernel with proper autograd setup and then link to it in
// my_op2's AutogradPrivateUse1 kernel.
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl(<schema_my_op2>, &my_op2);
  m.impl(<schema_my_op2_backward>, &my_op2_backward);
}

在少數罕見情況下，PyTorch 中某些運算元的梯度公式可能包含不適用於所有後端的假設。在這些情況下，後端擴充套件者可以選擇透過將 torch::autograd::Function 中的核函式註冊到相應的排程鍵（例如，如果你的後端使用 PrivateUse1，則為 AutogradPrivateUse1）來覆蓋 PyTorch Autograd 層。

class MyAddFunction : public torch::autograd::Function<MyAddFunction> {
  public:
  static Tensor forward(AutogradContext *ctx, torch::Tensor self, torch::Tensor other) {
    at::AutoNonVariableTypeMode g;
    return myadd(self, other);
  }

  static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs) {
    auto grad_output = grad_outputs[0];
    return {grad_output, grad_output};
  }
};

Tensor myadd_autograd(const Tensor& self, const Tensor& other) {
  return MyAddFunction::apply(self, other)[0];
}

// Register the autograd kernel to AutogradPrivateUse1
TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) {
  m.impl(<myadd_schema>, &myadd_autograd);
}

// Register the inference kernel to PrivateUse1
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl(<myadd_schema>, &myadd);
}

利用這個技巧，你就可以完全控制你的後端中 my_add 運算元的訓練和推理行為。以下是 pytorch/xla 倉庫中的一個示例。

構建擴充套件¶

透過向 PyTorch 新增 C++ 擴充套件來支援樹外後端。一旦你準備好核函式和註冊，你可以透過編寫一個使用 setuptools 編譯 C++ 程式碼的 setup.py 指令碼來構建 C++ 擴充套件。以下是來自 pytorch/xla 倉庫的一個簡化示例：

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(
    name='torch_xla',
    ext_modules=[
        CppExtension(
            '_XLAC',
            torch_xla_sources,
            include_dirs=include_dirs,
            extra_compile_args=extra_compile_args,
            library_dirs=library_dirs,
            extra_link_args=extra_link_args + \
                [make_relative_rpath('torch_xla/lib')],
        ),
    ],
    cmdclass={
        'build_ext': Build,  # Build is a derived class of BuildExtension
    }
    # more configs...
)

更多詳情請參見我們的 C++ 擴充套件教程。

自定義運算元支援¶

只要自定義運算元是由現有 PyTorch 運算元（你的後端已支援的運算元）組成的，你的新後端就應該能與在 python 中擴充套件的自定義運算元無縫協作，無需編寫任何新的核函式。

對於在 C++ 中擴充套件的自定義運算元，它們通常附帶一個後端特定的 C++ 核函式實現（例如 torchvision 中的 nms 核函式）以及一個定製的 Python API（例如 torch.ops.torchvision.nms）。要支援這些運算元，後端擴充套件者需要為你的後端編寫一個 C++ 核函式，並將其正確註冊到排程器中相應的名稱空間，類似於支援 PyTorch 原生運算元。另外，你也可以在你的擴充套件中新增定製的 API（例如 torch_xla.core.functions.nms）來滿足這些臨時請求。

JIT 支援¶

正如我們在在 C++ 中註冊排程運算元中提到的，透過 m.impl() API 註冊的核函式支援以 unboxed 和 boxed 兩種方式呼叫。換句話說，你的定製後端也可以像 CPU 或 CUDA 等內建 (in-tree) 後端一樣與我們的 JIT 追蹤/指令碼化前端一起工作。你甚至可以在 JIT 圖上為你的後端編寫專門的最佳化 pass。但我們在此不討論這一點，因為我們尚未最終確定 JIT 中的整合點，因此當前的後端支援將暫時側重於 eager 前端。

針對原生 PyTorch 後端測試你的後端¶

PyTorch 允許使用其通用裝置型別測試框架在多種裝置型別上執行測試。你可以找到有關測試如何使用它的詳細資訊以及有關如何新增新的裝置型別的資訊。新增後，使用通用裝置型別測試框架的 PyTorch 測試也將使用你的裝置型別執行。請參閱此 Wiki 頁面，瞭解測試如何例項化的示例。

使用你的裝置型別執行 PyTorch 現有的測試套件對於確保正確性很重要，但並非所有 PyTorch 功能都受到每種裝置型別的支援。通用裝置型別測試框架允許進行大量自定義，以便裝置型別可以選擇要執行哪些測試、支援哪些 dtypes，甚至在比較張量是否相等時使用哪些精度。

使用通用裝置型別測試框架但不隨 PyTorch 一起釋出的裝置型別示例是 XLA。請參閱其對通用裝置型別測試框架的擴充套件，其中包含阻止列表測試、阻止列表 dtypes 和覆蓋測試精度的示例。

通用裝置型別測試框架正在積極開發中。如需請求新功能，請在 PyTorch 的 Github 上提交 issue。

向後相容性¶

目前 PyTorch 無法保證註冊運算元的向後相容性。運算元及其 schema 可能會根據需要新增/修改/刪除。註冊的核函式必須與 PyTorch 版本完全相同。如果 PyTorch 為某個運算元添加了更多引數（即使帶有預設值），你舊的註冊將無法工作，直到更新以匹配 PyTorch 的新簽名。

因此，我們強烈建議樹外後端擴充套件者僅與 PyTorch 的主要版本同步，以最大程度地減少開發中斷。PyTorch 遵循季度釋出節奏。後端擴充套件者應加入 pytorch.slack.com 上的 #announcement 頻道，以獲取有關釋出的最新更新。

已知問題和附加說明¶

並非所有測試套件都已是裝置通用型。可以透過在 PyTorch 程式碼庫中搜索 instantiate_device_type_tests 來找到可擴充套件的測試類，例如 TestTorchDeviceType, TestViewOps, TestTensorDeviceOps, TestTypePromotion 等。
目前在 C++ 中沒有用於在定製後端上序列化 Python 張量物件的擴充套件點。當前，你只能透過修改 PyTorch 張量的 __reduce_ex__ 方法或在樹外倉庫中進行 monkey patching 來擴充套件它。
如果你的後端不允許直接記憶體訪問，你應該格外注意支援檢視運算元，因為它們應該共享儲存。對檢視張量的更改需要傳播到其基礎張量，反之亦然。
如果你的後端無法與原生 PyTorch 最佳化器一起工作（例如需要在反向傳播中攜帶要更新的狀態，如 torch-xla），則在 C++ 中沒有最佳化器的擴充套件點。此類用例目前只能透過在樹外倉庫中新增定製 API 或進行 monkey patching 來實現。

未來工作¶

要使 PyTorch 的每個元件都能無縫地為樹外後端提供擴充套件點，需要對 PyTorch 內部進行大量更改。以下是我們正在積極開展的一些可能在未來改善體驗的工作項：

提高通用測試框架的測試覆蓋率。
改進 Math 核函式的覆蓋率和更全面的測試，以確保 Math 核函式的行為與其他後端（如 CPU/CUDA）匹配。
重構 RegistrationDeclarations.h，使其攜帶最少資訊，並儘可能重用 PyTorch 的 codegen。
支援後端回退 (fallback) 核函式，自動將輸入轉換為 CPU，並將結果轉換回定製後端。即使你沒有為每個運算元編寫核函式，這也將實現“完全”的運算元覆蓋。

保持聯絡¶

請使用PyTorch 開發者討論進行提問和討論。如果你有任何功能請求或 bug 報告，請在 github 上提交 issue。

如果你有興趣幫助完成上述任何未來工作項（例如為 C++ 中的 PyTorch 運算元新增更多 Math 核函式），請透過 Github 或 Slack 與我們聯絡！