評價此頁
★ ★ ★ ★ ★
advanced/extend_dispatcher

使用 C++ 為新後端擴充套件 Dispatcher#

創建於: 2021年02月01日 | 最後更新: 2024年09月23日 | 最後驗證: 2024年11月05日

在本教程中，我們將逐步介紹將 Dispatcher 擴充套件到 PyTorch 倉庫外部新增新裝置所需的所有必要步驟，並對其進行維護以與原生 PyTorch 裝置保持同步。在此，我們假設您熟悉如何在 C++ 中 註冊已分派的運算子 以及如何編寫 自定義自動微分函式。

注意

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

什麼是新後端？#

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

  • 如果您有現有 PyTorch 運算子的新演算法，請向 PyTorch 傳送 PR。

  • 如果您想提出新的運算子，請向 PyTorch 傳送功能請求/PR。

  • 如果您想新增對新裝置/硬體（如 Google TPU 和定製晶片）的支援，這通常需要使用硬體特定的 API 來編寫核心，請遵循本教程並向 PyTorch 新增一個“out-of-tree”（樹外）後端。

  • 如果您想為現有運算子新增支援，但採用不同的 Tensor 佈局/表示（如稀疏和量化），這會強制您的核心以一種更有效的方式編寫，考慮到佈局/表示的限制，請遵循本教程並向 PyTorch 新增一個“out-of-tree”（樹外）後端。

在本教程中，我們將主要關注新增一個新的“out-of-tree”（樹外）裝置。為不同的張量佈局新增“out-of-tree”（樹外）支援可能與裝置共享許多通用步驟，但我們尚未看到此類整合的示例，因此可能需要 PyTorch 額外的工作來支援它。

為您的後端獲取 Dispatch Key#

PyTorch 運算子是用 C++ 實現的，並透過 Python 繫結在 Python 前端中可用。PyTorch Dispatcher 將運算子的實現劃分為多個核心，每個核心都與特定的 Dispatch Key 相關聯。在 PyTorch 中支援新後端本質上意味著用 C++ 為每個 PyTorch 運算子編寫一個核心，然後將它們註冊到一個代表您定製後端的 Dispatch Key 中。

Dispatch Key 是您在 Dispatcher 系統中的識別符號。Dispatcher 檢視輸入 Tensor 上攜帶的 Dispatch Key，並相應地呼叫正確的核心。PyTorch 提供了三個保留的 Dispatch Key（及其對應的 Autograd Key），用於原型化“out-of-tree”（樹外）後端擴充套件：

  • PrivateUse1/AutogradPrivateUse1

  • PrivateUse2/AutogradPrivateUse2

  • PrivateUse3/AutogradPrivateUse3

您可以選擇以上任何一個 Key 來原型化您的定製後端。要建立一個 Tensor 在 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 類假定您的 Tensor 由類似 CPU/CUDA 的儲存支援。我們還為沒有儲存的後端提供了 OpaqueTensorImpl。您可能需要調整/覆蓋某些方法以適應您的定製硬體。PyTorch 倉庫中的一個示例是 Vulkan TensorImpl。

注意

一旦原型完成，並且您計劃為後端擴充套件進行定期釋出，請隨時提交 PR 到 pytorch/pytorch 以為您的後端保留專用的 Dispatch Key。

獲取 PyTorch 運算子的完整列表#

PyTorch 在生成的 C++ 檔案 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++ 簽名相比的別名和突變（mutation）註釋。這是 Dispatcher 用來查詢運算子的唯一識別符號。

  • dispatch 和 default 是布林欄位，提供了有關原生 PyTorch 核心可以執行的操作的資訊，從而暗示後端擴充套件者是否需要實現核心。更多詳細資訊請參閱 為新後端註冊核心。

為新後端註冊核心#

要將您的核心註冊到 PyTorch Dispatcher，您可以使用 在 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 中，這些運算子在其伴隨註釋的元資料中具有 dispatch 設定為 True *並且* default 設定為 False。

  • 可選註冊：後端擴充套件者可以跳過對這些運算子的註冊，而不會犧牲任何支援。但是，如果後端擴充套件者想覆蓋 PyTorch 提供的預設核心，他們仍然可以向其後端註冊定製核心，Dispatcher 將僅為您的後端使用它。例如，PyTorch 當前的 max_pool2d 實現將其 indices 作為前向輸出的一部分返回，這在 torch_xla 中會產生開銷，因此 torch_xla 會註冊自己的 max_pool2d 核心。

    • 在 RegistrationDeclarations.h 中，這些運算子在其伴隨註釋的元資料中具有 dispatch 設定為 False *或* default 設定為 True。

為新後端提供 Autograd 支援#

梯度公式在數學上大多是純粹的，因此對所有後端都是通用的。PyTorch 通常會向別名 Dispatch Key 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 某些運算子的梯度公式可能存在不適用於所有後端的假設。在這種情況下，後端擴充套件者可以選擇透過向相應的 Dispatch Key（例如，如果您正在為後端使用 PrivateUse1，則為 AutogradPrivateUse1）註冊一個來自 `torch::autograd::Function` 的核心來覆蓋 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` 儲存庫中的一個 示例。

構建擴充套件#

“out-of-tree”（樹外）後端透過向 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++ 核心，並將其正確地註冊到 Dispatcher 中相應的名稱空間，這與支援 PyTorch 原生運算子類似。或者，您也可以在擴充套件中新增定製 API，例如 `torch_xla.core.functions.nms`，以應對這些臨時請求。

JIT 支援#

正如我們在 在 C++ 中註冊已分派的運算子 中提到的，透過 `m.impl()` API 註冊的核心支援以未裝箱（unboxed）和裝箱（boxed）兩種方式呼叫。換句話說，您的定製後端也可以像 CPU 或 CUDA 等內部後端一樣與我們的 JIT 跟蹤/指令碼前端配合使用。您甚至可以為 JIT 圖上的後端編寫專門的最佳化通道。但由於我們尚未確定 JIT 中的整合點，因此我們在此不討論，因此當前後端支援將暫時專注於 Eager 前端。

針對原生 PyTorch 後端測試您的後端#

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

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

XLA 是一個不隨 PyTorch 分發的、使用通用裝置型別測試框架的示例裝置型別。請參閱其對通用裝置型別測試框架的 擴充套件，其中包含有關阻止列出測試、阻止列出資料型別以及覆蓋測試精度的示例。

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

向後相容性#

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

因此，我們*強烈建議*“out-of-tree”（樹外）後端擴充套件者僅同步主要的 PyTorch 版本釋出，以最大程度地減少開發中斷。PyTorch 按季度釋出。後端擴充套件者應加入 pytorch.slack.com 上的 `#announcement` 頻道，以獲取最新發布更新。

已知問題及補充說明#

  • 並非所有測試套件都已經是裝置通用的。可透過在 PyTorch 程式碼庫中搜索 `instantiate_device_type_tests` 來找到可擴充套件的測試類，例如 `TestTorchDeviceType, TestViewOps, TestTensorDeviceOps, TestTypePromotion` 等。

  • 在 C++ 中沒有為在自定義後端上序列化 Python Tensor 物件提供擴充套件點。目前，您只能透過修改 PyTorch Tensor `__reduce_ex__` 方法 或在“out-of-tree”（樹外）儲存庫中進行 monkey patching 來擴充套件它。

  • 如果您的後端不允許直接記憶體訪問，則在支援 view 運算子時應格外注意，因為它們應該共享儲存。對 view tensor 所做的更改需要傳播到其基本 tensor，反之亦然。

  • 如果您的後端無法與原生 PyTorch 最佳化器一起工作（例如，需要在後向傳播中攜帶狀態，如 torch-xla），則在 C++ 中沒有最佳化器的擴充套件點。此類用例目前只能透過在“out-of-tree”（樹外）儲存庫中新增自定義 API 或 monkey patching 來實現。

未來工作#

使 PyTorch 中的每個元件都能無縫地擴充套件到“out-of-tree”（樹外）後端需要對 PyTorch 內部進行大量更改。以下是我們正在積極進行的一些專案，它們可能會在未來改善使用者體驗：

  • 改進通用測試框架的測試覆蓋率。

  • 改進 `Math` 核心覆蓋率和更全面的測試，以確保 `Math` 核心的行為與其他後端（如 `CPU/CUDA`）匹配。

  • 重構 `RegistrationDeclarations.h` 以攜帶最少的資訊，並儘可能重用 PyTorch 的程式碼生成。

  • 支援一個後端回退核心，以自動將輸入轉換為 CPU，並將結果轉換回自定義後端。這將允許“完全”的運算子覆蓋，即使您沒有為每個運算子編寫核心。

保持聯絡#

請使用 PyTorch 開發討論 進行提問和討論。如果您有任何功能請求或錯誤報告，請在 Github 上 提交問題。

如果您有興趣參與上述任何未來工作（例如，在 C++ 中為 PyTorch 運算子新增更多 `Math` 核心），請透過 Github 或 Slack 與我們聯絡！

在此頁面上
PyTorch 庫