在 C++ 中為新的後端擴展 dispatcher¶

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

在本教學中，我們將逐步介紹所有必要的步驟，以擴展 dispatcher，以新增一個位於 pytorch/pytorch 儲存庫之外的裝置，並維護它以與原生 PyTorch 裝置保持同步。在這裡，我們假設您熟悉如何在 C++ 中註冊已分派的運算子以及如何編寫自訂 autograd 函式。

注意

本教學涉及 PyTorch 內部的許多內部元件，這些元件正在積極改進中，如果您決定遵循本教學，請預期 API 會發生變更。我們將使本教學與最新的 API 保持同步。

什麼是新的後端？¶

向 PyTorch 新增新的後端需要後端擴充程式進行大量的開發和維護。在新增新的後端之前，讓我們先考慮一些常見的用例和推薦的解決方案

如果您有現有 PyTorch 運算子的新演算法，請向 PyTorch 發送 PR。
如果您想提出一個新的運算子，請向 PyTorch 發送功能請求/PR。
如果您想新增對新裝置/硬體（例如 Google TPU 和自訂晶片）的支援，這通常需要使用特定於硬體的 API 來編寫核心，請遵循本教學並向 PyTorch 新增一個樹外後端。
如果您想新增對現有運算子的支援，但使用不同的張量佈局/表示形式（例如稀疏和量化），這會強制您以更有效的方式編寫核心，因為佈局/表示形式的限制，請遵循本教學並向 PyTorch 新增一個樹外後端。

在本教學中，我們將主要關注在下方新增一個新的樹外裝置。為不同的張量佈局新增樹外支援可能與裝置共享許多常見步驟，但我們尚未看到此類整合的範例，因此可能需要 PyTorch 進行額外的工作才能支援它。

取得後端的 dispatch key¶

PyTorch 運算子是在 C++ 中實現的，並透過 Python 綁定在 Python 前端中提供。 PyTorch dispatcher 將運算子的實現劃分為多個核心，每個核心都與特定的 dispatch key 關聯。在 PyTorch 中支援新的後端本質上意味著用 C++ 為每個 PyTorch 運算子編寫一個核心，然後將它們註冊到 dispatcher 中代表您自訂後端的 dispatch key。

Dispatch key 是您在 dispatcher 系統中的識別碼。 dispatcher 會查看輸入張量上攜帶的 dispatch key，並相應地呼叫正確的核心。 PyTorch 提供了三個保留的 dispatch key（及其對應的 Autograd key），用於原型化樹外後端擴展

PrivateUse1/AutogradPrivateUse1
PrivateUse2/AutogradPrivateUse2
PrivateUse3/AutogradPrivateUse3

您可以選擇上面的任何一個 key 來原型化您的自訂後端。若要在 PrivateUse1 後端上建立張量，您需要在 TensorImpl 建構函式中設定 dispatch key。

/* 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 以為您的後端保留專用的 dispatch key。

取得 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!) 是表示運算子的唯一模式，與 C++ 簽名相比，它還包含別名和變異註釋。這是 dispatcher 用於尋找運算子的唯一識別碼。
dispatch 和 default 是布林欄位，提供有關原生 PyTorch 核心可以做什麼的資訊，因此暗示了後端擴充程式是否需要實現核心。更多詳細資訊可以在為新的後端註冊核心中找到。

為新的後端註冊核心¶

要將你的核心註冊到 PyTorch 調度器，你可以使用 TORCH_LIBRARY_IMPL API，詳情請參閱在 C++ 中註冊調度的運算子

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 中，這些運算子的 dispatch 在其隨附的註解中的元數據中設定為 True *且* default 設定為 False。
註冊是可選的：後端擴充程式可以跳過註冊到這些運算子，而不會犧牲任何支持。但是，如果後端擴充程式想要覆蓋 PyTorch 提供的預設核心，它們仍然可以將其客製化的核心註冊到其後端，並且調度器將僅將其用於你的後端。例如，目前 PyTorch 的 max_pool2d 的實現返回 indices 作為前向輸出的部分，這會在 torch_xla 中產生額外的開銷，因此 torch_xla 會為 max_pool2d 註冊自己的核心。
- 在 RegistrationDeclarations.h 中，這些運算子的 dispatch 在其隨附的註解中的元數據中設定為 False *或* default 設定為 True。

新後端的自動微分支援¶

梯度公式大多是純數學的，因此對於所有後端都是通用的。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 作為*需要註冊*。

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 儲存庫中。

建置擴充程式¶

通過將 C++ 擴充程式添加到 PyTorch 來支持樹外後端。一旦你準備好核心和註冊，你可以通過編寫一個使用 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++ 核心實現，例如 torchvsion 中的 nms 核心以及一個自定義的 Python API，例如 torch.ops.torchvision.nms。為了支持這些運算子，後端擴充程式需要為你的後端編寫一個 C++ 核心，並將其正確註冊到調度器中相應的命名空間，這與支持 PyTorch 原生運算子類似。或者，你也可以在你的擴充程式中添加一個自定義的 API，例如 torch_xla.core.functions.nms 來滿足這些臨時請求。

JIT 支援¶

正如我們在在 C++ 中註冊調度的運算子中提到的那樣，通過 m.impl() API 註冊的核心支持以非裝箱和裝箱的方式調用。換句話說，你的客製化後端也可以像 CPU 或 CUDA 等樹內後端一樣，與我們的 JIT 追蹤/腳本編寫前端一起工作。你也可以在 JIT 圖上為你的後端編寫專門的優化傳遞。但我們不會在這裡討論它，因為我們還沒有最終確定 JIT 中的集成點，因此目前的後端支持將側重於 eager 前端。

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

PyTorch 允許使用其通用設備類型測試框架在多個設備類型上運行測試。你可以找到關於測試如何使用它的詳細信息，以及關於如何添加新設備類型的信息。添加後，使用通用設備類型測試框架的 PyTorch 測試也將使用你的設備類型運行。請參閱這個 Wiki 頁面，以獲取如何實例化測試的範例。

使用您的裝置類型執行 PyTorch 現有的測試套件，對於確保正確性非常重要，但並非所有裝置類型都支援所有 PyTorch 功能。通用裝置類型測試框架允許進行大量的自訂，以便裝置類型可以選擇要執行的測試、它們支援的資料類型，甚至是在比較張量是否相等時要使用的精度。

一個使用通用裝置類型測試框架且未隨 PyTorch 一起提供的裝置類型範例是 XLA。請參閱其對通用裝置類型測試框架的擴展，其中包含封鎖測試、封鎖資料類型和覆蓋測試精度的範例。

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

向後相容性¶

目前 PyTorch 無法保證已註冊運算子的向後相容性。運算子及其 schema 可能會根據需要新增/修改/刪除。已註冊的 kernel 必須與 PyTorch 版本完全相同。如果 PyTorch 為運算子新增更多參數（即使帶有預設值），您的舊註冊將無法運作，直到更新為符合 PyTorch 的新簽名。

因此，我們強烈建議樹外後端擴充器僅與主要的 PyTorch 版本同步，以最大程度地減少開發中的中斷。 PyTorch 採用季度發布週期。後端擴充器應加入 pytorch.slack.com 上的#announcement 頻道，以取得最新的發布更新。

已知問題 & 附加說明¶

並非所有測試套件都已通用於裝置。可擴展的測試類別可以透過在 PyTorch 程式碼庫中搜尋 instantiate_device_type_tests 來找到，例如 TestTorchDeviceType, TestViewOps, TestTensorDeviceOps, TestTypePromotion 等。
在 C++ 中沒有擴展點可用於序列化自訂後端的 Python Tensor 物件。目前，您只能透過修改 PyTorch Tensor __reduce_ex__ 方法或在樹外儲存庫中進行 monkey patching 來擴展它。
如果您的後端不允許直接記憶體存取，您應該額外注意支援 view ops，因為它們應該共享儲存空間。對 view 張量的變更需要傳播到其基本張量，反之亦然。
如果您的後端不適用於原生 PyTorch Optimizer，則在 C++ 中沒有 Optimizer 的擴展點可用，例如需要將狀態攜帶到 backward 中進行更新，如 torch-xla。目前只能透過新增自訂 API 或在樹外儲存庫中進行 monkey patching 來完成此類用例。

未來工作¶

要使 PyTorch 中的每個元件對於樹外後端無縫可擴展，需要對 PyTorch 內部進行大量變更。以下是我們正在積極進行的一些項目，可能會改善未來的體驗

改善通用測試框架的測試覆蓋率。
改善 Math kernel 的覆蓋率和更全面的測試，以確保 Math kernel 行為與其他後端（如 CPU/CUDA）相符。
重構 RegistrationDeclarations.h 以攜帶最少的資訊，並盡可能地重複使用 PyTorch 的 codegen。
支援後端回退 kernel，以自動將輸入轉換為 CPU，並將結果轉換回自訂後端。即使您沒有為每個運算子編寫 kernel，這也將允許“完整”的運算子覆蓋。

保持聯繫¶

請使用 PyTorch 開發討論區提出問題和討論。如果您有任何功能要求或錯誤報告，請在 github 上提交 issue。

如果您有興趣協助以上任何未來工作項目（例如，在 C++ 中為 PyTorch 運算子新增更多 Math kernel），請透過 Github 或 Slack 與我們聯繫！