advanced/dispatcher

在 Google Colab 中執行

Colab

下載 Notebook

Notebook

在 GitHub 上檢視

GitHub

在 C++ 中註冊分派的運算子

建立日期：2020 年 7 月 22 日 | 最後更新日期：2024 年 7 月 22 日 | 最後驗證日期：2024 年 11 月 5 日

警告

此教程在 PyTorch 2.4 版本後已棄用。請參閱 PyTorch 自定義運算子 以獲取關於使用自定義運算子擴充套件 PyTorch 的最新指南。

分派器是 PyTorch 的一個內部元件，負責確定在呼叫像 torch::add 這樣的函式時，實際應該執行哪些程式碼。這可能並不簡單，因為 PyTorch 運算子需要處理許多相互疊加的橫切關注點。以下是它處理的一些示例：

• 根據輸入張量的裝置，在運算子的 CPU 和 CUDA 實現之間切換。

• 根據是否需要自動梯度處理，在運算子的自動梯度和後端實現之間切換。

• 在必要時應用自動轉換以實現自動混合精度。

• 當運算子在 vmap 呼叫下執行時，應用批處理規則。

• 如果正在為匯出追蹤模型，則追蹤操作的執行。

如果在您的 自定義運算子程式碼 中發現您手動編寫 if 語句來處理這些情況，分派器 API 可以幫助組織您的程式碼。（反之，如果您的自定義運算子非常簡單，僅用於 CPU 推理，您可能不需要使用分派器，只需使用基本 API。）

在本教程中，我們將描述如何構造自定義運算子註冊以使用分派器來組織各種元件。我們將假設您熟悉如何 註冊運算子 以及如何編寫 自定義自動梯度函式。

定義模式和後端實現

分派器的通用原理是，它將運算子的實現劃分為多個核心，每個核心為特定的分派鍵（例如，CPU、CUDA）實現功能。分派器確定在呼叫運算子時優先順序最高的_分派鍵_（透過檢視張量引數以及一些執行緒本地狀態來完成），並將控制權轉移到該_分派鍵_的核心。最終效果是，當您呼叫運算子時，我們首先執行自動梯度核心，然後根據傳入張量的裝置型別重新分派到後端核心。

讓我們看看實現這一點的各種部分。首先，我們必須定義相關運算子的模式。與簡單的 pybind11 風格的運算子註冊不同，我們實際上此時不提供運算子的實現；我們只提供一個模式字串，指定運算子的型別簽名，所有其他核心都將遵循該簽名。

TORCH_LIBRARY(myops, m) {
  m.def("myadd(Tensor self, Tensor other) -> Tensor");
}

接下來，我們需要實際提供該運算子的一些實現。為了具體起見，以下是 CPU 上加法的一個非常簡單的實現：

Tensor myadd_cpu(const Tensor& self_, const Tensor& other_) {
  TORCH_CHECK(self_.sizes() == other_.sizes());
  TORCH_INTERNAL_ASSERT(self_.device().type() == DeviceType::CPU);
  TORCH_INTERNAL_ASSERT(other_.device().type() == DeviceType::CPU);
  Tensor self = self_.contiguous();
  Tensor other = other_.contiguous();
  Tensor result = torch::empty(self.sizes(), self.options());
  const float* self_ptr = self.data_ptr<float>();
  const float* other_ptr = other.data_ptr<float>();
  float* result_ptr = result.data_ptr<float>();
  for (int64_t i = 0; i < result.numel(); i++) {
    result_ptr[i] = self_ptr[i] + other_ptr[i];
  }
  return result;
}

我們想將此函式註冊為 myops::myadd 的實現。但是，註冊它的簡單方法（def("myadd", myadd_cpu)）會將核心註冊為在所有情況下執行，即使張量不是 CPU 張量！（在內部，我們將這些稱為“catch-all”核心，因為它們會捕獲所有情況。）為了確保 myadd_cpu 僅為 CPU 張量執行，我們可以使用 TORCH_LIBRARY_IMPL 宏。

TORCH_LIBRARY_IMPL(myops, CPU, m) {
  m.impl("myadd", myadd_cpu);
}

TORCH_LIBRARY_IMPL 允許我們為一個特定分派鍵（在此例中為 CPU）的運算子註冊實現。每次呼叫 impl 都會將一個 CPU 核心與相應的運算子（我們之前在 TORCH_LIBRARY 塊中定義）相關聯。如果我們還有一個 CUDA 實現 myadd_cuda，我們可以在單獨的 TORCH_LIBRARY_IMPL 塊中註冊它。

TORCH_LIBRARY_IMPL(myops, CUDA, m) {
  m.impl("myadd", myadd_cuda);
}

這些註冊可以分佈在檔案甚至庫邊界之間；因此，例如，您可以將這兩個 TORCH_LIBRARY_IMPL 塊編譯到單獨的 myops_cpu 和 myops_cuda 動態庫中。總的來說，您的註冊結構將如下所示：

1. 一個單一的 TORCH_LIBRARY，在一箇中心位置列出您名稱空間中的所有自定義運算子。

2. 每個分派鍵一個 TORCH_LIBRARY_IMPL，它為該鍵（例如，CPU 或 CUDA）註冊實現。如果您願意，您可以將 TORCH_LIBRARY_IMPL 塊進一步細分為每個運算子一個塊。如果您為每個實現檔案單獨存放一個檔案，但不想在標頭檔案中公開這些運算子，這會很方便；您可以只將註冊放在定義運算子的 cpp 檔案中。

注意

您知道您也可以為 PyTorch 中現有的核心運算子編寫 TORCH_LIBRARY_IMPL 塊嗎？XLA 對 PyTorch 的支援就是這樣實現的：torch_xla 庫包含一個 TORCH_LIBRARY_IMPL，它在 XLA 分派鍵上為所有基本運算子提供實現。

對於不需要自動梯度的運算子

注意：本節僅適用於 PyTorch 版本 >= 1.10。

在下一節中，我們將討論如何為運算子新增自動梯度支援。但對於不需要自動梯度支援的運算子，應註冊以下核心以提高可用性並使您的運算子表現得像 PyTorch 的內建運算子。

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl(op, autogradNotImplementedFallback());
}

上面的行註冊了一個 Autograd 核心，它在前向傳播時附加一個虛擬的 NotImplemented 節點（保留輸入的 require_grad 屬性）。在後向傳播時，NotImplemented 節點會引發錯誤。這對於在較大的模型中除錯可能很有幫助，因為在以前很難精確地確定 requires_grad 屬性在前向傳播過程中丟失的位置。

原地或檢視運算子

為確保正確性和最佳效能，如果您的運算子原地修改輸入或返回一個與輸入之一別名的張量，應採取兩個額外步驟：

1. 除了上述 Autograd 核心外，還註冊一個 ADInplaceOrView 核心。此核心處理必要的簿記，以確保原地或檢視操作的正確性。重要的是要注意，此 ADInplaceOrView 核心只能與 autogradNotImplementedFallback 一起使用。

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl(op, autogradNotImplementedFallback());
}
TORCH_LIBRARY_IMPL(myops, ADInplaceOrView, m) {
  m.impl(op, autogradNotImplementedInplaceOrViewFallback());
}

2. 上面註冊的 Autograd 或 ADInplaceOrView 框裝核心在其邏輯中依賴於運算子模式資訊。如果您的運算子原地修改輸入或返回一個與輸入之一別名的張量，確保您的模式能夠正確反映這一點非常重要。有關如何註釋模式的更多資訊，請參閱 此處。

新增自動梯度支援

此時，我們有一個同時具有 CPU 和 CUDA 實現的運算子。如何為其新增自動梯度支援？正如您可能猜到的，我們將註冊一個自動梯度核心（類似於在 自定義自動梯度函式 教程中描述的）！但是，有一個竅門：與 CPU 和 CUDA 核心不同，自動梯度核心需要_重新分派_：它需要回調分派器以訪問推理核心，例如 CPU 或 CUDA 實現。

因此，在我們編寫自動梯度核心之前，讓我們編寫一個_分派函式_，它呼叫分派器以查詢您運算子的正確核心。此函式構成了您運算子的公共 C++ API — 事實上，PyTorch C++ API 中的所有張量函式在底層都以相同的方式呼叫分派器。以下是分派函式的外觀：

Tensor myadd(const Tensor& self, const Tensor& other) {
  static auto op = torch::Dispatcher::singleton()
    .findSchemaOrThrow("myops::myadd", "")
    .typed<decltype(myadd)>();
  return op.call(self, other);
}

讓我們詳細分析一下：

• 在第一行，我們從分派器中查詢一個型別化的運算子控制代碼，該控制代碼對應於我們要分派的運算子。findSchemaOrThrow 接受兩個引數：運算子的（名稱空間限定）名稱，以及運算子的過載名稱（通常只是空字串）。typed 將動態型別控制代碼轉換為靜態型別控制代碼（進行執行時測試以確保您提供了正確的 C++ 型別），以便我們可以對其進行常規的 C++ 呼叫。我們傳遞 decltype(myadd)，因為分派函式的型別與註冊到分派器的底層核心的型別相同。

  為了效能，此計算在靜態變數中完成，因此我們只需要執行一次（緩慢的）查詢。如果您輸入的運算子名稱有誤，在第一次呼叫此函式時，此查詢將出錯。

• 在第二行，我們只需使用傳遞給分派函式的所有引數來call運算子控制代碼。這將實際呼叫分派器，最終控制權將轉移到適合此次呼叫的任何核心。

有了分派函式，我們現在可以編寫自動梯度核心了：

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];
}

自動梯度函式像往常一樣使用 torch::autograd::Function 編寫，除了不是直接在 forward() 中編寫實現，我們：

1. 使用 at::AutoNonVariableTypeMode RAII 保護停用自動梯度處理，然後

2. 呼叫分派函式 myadd 以回撥分派器。

沒有 (1)，您的呼叫將無限迴圈（並堆疊溢位），因為 myadd 會將您送回此函式（因為優先順序最高的_分派鍵_仍然是自動梯度）。使用 (1)，自動梯度將從考慮的_分派鍵_集合中排除，我們將轉到下一個處理器，它將是 CPU 和 CUDA。

我們現在可以以與註冊 CPU/CUDA 函式相同的方式註冊此函式：

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl("myadd", myadd_autograd);
}

注意

在此示例中，我們將核心註冊到 Autograd，它將其作為所有後端的自動梯度核心安裝。您也可以透過使用相應的後端特定_分派鍵_來註冊特定後端的最佳化核心 — 例如，AutogradCPU 或 AutogradCUDA。要更詳細地探索這些以及其他_分派鍵_選項，請檢視 torch/_python_dispatcher.py 中提供的 PythonDispatcher 工具。

超越自動梯度

在某種意義上，分派器並沒有做太多事情：它所做的只是實現一個華而不實的 if 語句，大致如下：

class MyAddFunction : ... {
public:
  static Tensor forward(
    AutogradContext *ctx, torch::Tensor self, torch::Tensor other) {

    if (self.device().type() == DeviceType::CPU) {
      return add_cpu(self, other);
    } else if (self.device().type() == DeviceType::CUDA) {
      return add_cuda(self, other);
    } else {
      TORCH_CHECK(0, "Unsupported device ", self.device().type());
    }
  }
  ...
}

那麼為什麼還要使用分派器呢？有幾個原因：

1. 它是分散式的。您無需編寫一個引用所有這些內容的集中式 if 語句，即可組合運算子的所有部分（CPU、CUDA、自動梯度）。重要的是，第三方可以註冊額外的實現以處理其他方面，而無需修補運算子的原始定義。我們將在 在 C++ 中為新後端擴充套件分派器 中更詳細地討論擴充套件分派器。

2. 它支援比 CPU、CUDA 和自動梯度更多的_分派鍵_。您可以在 c10/core/DispatchKey.h 中檢視 PyTorch 中當前實現的_分派鍵_的完整列表。這些_分派鍵_實現了各種可選的運算子功能，如果您決定希望您的自定義運算子支援此功能，只需為您相應的鍵註冊一個核心。

3. 分派器實現了對盒裝回退函式（boxed fallback functions）的支援，這些函式可以實現一次並應用於系統中的所有運算子。盒裝回退可用於為_分派鍵_提供預設行為；如果您使用分派器來實現您的運算子，您也就選擇了所有這些操作的回退。

以下是一些您可能需要為其定義運算子的特定_分派鍵_：

自動轉換

Autocast _分派鍵_實現了對 自動混合精度 (AMP) 的支援。Autocast 包裝器核心通常會在執行 op 之前將傳入的 float16 或 float32 CUDA 張量轉換為首選精度。例如，浮點 CUDA 張量上的 matmuls 和 convolutions 通常在 float16 中執行速度更快，記憶體佔用更少，而不會損害收斂性。Autocast 包裝器僅在 啟用自動轉換的上下文中 有效。

這是一個假設的自定義 matmul 的自動轉換包裝器及其註冊：

// Autocast-specific helper functions
#include <ATen/autocast_mode.h>

Tensor mymatmul_autocast(const Tensor& self, const Tensor& other) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
  return mymatmul(at::autocast::cached_cast(at::kHalf, self),
                  at::autocast::cached_cast(at::kHalf, other));
}

TORCH_LIBRARY_IMPL(myops, Autocast, m) {
  m.impl("mymatmul", mymatmul_autocast);
}

cached_cast(kHalf, tensor) 將 tensor 轉換為 float16（如果 tensor 是 CUDA 且為 float32），否則，它將 tensor 保持不變（參見 原生支援自動轉換的 op 的合格策略）。這確保瞭如果網路在任何 float16 和 float32 CUDA 張量的混合體上呼叫 mymatmul，mymatmul 將以 float16 執行。同時，對非 CUDA、整數型別或 float64 輸入的 mymatmul 呼叫不受影響。建議使用 cached_cast 來遵循您自己的自動轉換包裝器中的原生合格策略，但這並非強制要求。例如，如果您想強制所有輸入型別的 float16 執行，您可以 return mymatmul(self.half(), other.half()); 而不是使用 cached_cast。

請注意，與我們的自動梯度核心一樣，我們在重新分派之前排除了 Autocast 鍵。

預設情況下，如果未提供自動轉換包裝器，我們將直接回退到常規運算子實現（不進行自動轉換）。（我們沒有為此示例使用 myadd，因為逐點加法不需要自動轉換，應該直接回退。）

何時應註冊自動轉換包裝器？不幸的是，關於 op 的首選精度沒有明確的規則。您可以透過檢視 轉換列表 來了解一些原生 op 的首選精度。一般指導：

• 執行歸約的 op 應該在 float32 中執行，

• 任何在內部執行卷積或 gemm 的 op 應該在 float16 中執行，並且

• 具有多個浮點張量輸入的其他 op 應將其標準化為通用精度（除非實現支援不同精度的輸入）。

如果您的自定義 op 屬於第三類，promote_type 模板有助於確定輸入張量中最寬的浮點型別，這是執行型別的最安全選擇。

#include <ATen/autocast_mode.h>

Tensor my_multiple_input_op_autocast(const Tensor& t0, const Tensor& t1) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
  // The required at::kHalf argument is an optimistic initial guess.
  auto exec_type = at::autocast::promote_type(at::kHalf, t0, t1);
  return my_multiple_input_op(at::autocast::cached_cast(exec_type, t0),
                              at::autocast::cached_cast(exec_type, t1));
}

如果您的自定義 op 已 啟用自動梯度，您只需為主自動梯度包裝器註冊的相同名稱編寫並註冊一個自動轉換包裝器。例如，如果您想要上面自動梯度部分所示的 myadd 函式的自動轉換包裝器，您只需要：

Tensor myadd_autocast(const Tensor& self, const Tensor& other) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
  return myadd(at::autocast::cached_cast(<desired dtype>, self),
               at::autocast::cached_cast(<desired dtype>, other));
}

TORCH_LIBRARY_IMPL(myops, Autocast, m) {
  m.impl("myadd", myadd_autocast);
}

沒有單獨的技巧可以使後向方法相容自動轉換。但是，您自定義自動梯度函式中定義的後向方法將在與自動轉換為前向方法設定的 dtype 相同的 dtype 中執行，因此您應該選擇適合您的前向和後向方法的 <desired dtype>。

批處理

批處理張量允許您以每個示例的方式編寫程式碼，然後在 vmap 呼叫下執行時自動將其批處理。編寫批處理規則的 API 目前正在開發中，但一旦穩定，您就可以透過在 Batched _分派鍵_處註冊一個核心來為 vmap 新增支援。

追蹤器

Tracer _分派鍵_支援在執行 torch.jit.trace 時將運算子的呼叫記錄到跟蹤中。我們打算提供一個盒裝回退，它將實現任意操作的跟蹤，請參閱 issue #41478 以跟蹤進度。