在 C++ 中註冊已調度的運算符

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

警告

自 PyTorch 2.4 起，本教學已被棄用。請參閱 PyTorch 自訂運算符 以取得最新的 PyTorch 擴充自訂運算符指南。

調度器是 PyTorch 的內部元件，負責確定在您呼叫 torch::add 之類函式時實際應執行的程式碼。這可能並非易事，因為 PyTorch 運算需要處理許多「分層」於彼此之上的跨領域問題。以下是一些它處理的事情範例

• 在運算符的 CPU 和 CUDA 實作之間切換，取決於輸入張量的裝置。

• 在運算符的 autograd 和後端實作之間切換，取決於是否需要 autograd 處理。

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

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

• 如果您正在追蹤模型以進行匯出，則追蹤運算的執行。

如果在您的自訂運算符程式碼中，您發現自己手動編寫 if 語句來處理這些情況，則調度器 API 可以幫助組織您的程式碼。（相反地，如果您的自訂運算符非常簡單並且僅用於 CPU 推論，則您可能不需要使用調度器，只需使用基本 API 即可。）

在本教學中，我們將描述如何架構自訂運算符註冊，以使用調度器來組織各種元件。我們將假設您熟悉如何註冊運算符以及如何編寫自訂 autograd 函式。

定義模式和後端實作

調度器背後的總體原則是將運算符的實作分為多個核心，每個核心都為特定的調度鍵實作功能，例如 CPU、CUDA。調度器決定在您呼叫運算符時哪個是最高優先順序的調度鍵（這是透過查看張量引數以及一些執行緒本機狀態來完成的），並將控制權轉移到該調度鍵的核心。最終結果是，當您呼叫運算符時，我們首先執行 Autograd 核心，然後我們根據傳入張量的裝置類型重新調度到後端核心。

讓我們看看使這一切發生的各種部分。首先，我們必須定義有問題的運算符的模式。與簡單的 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 張量！ （在內部，我們將這些稱為「捕獲所有」核心，因為它們捕獲所有情況。）為了確保 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. 每個 dispatch key（例如 CPU 或 CUDA）對應一個 TORCH_LIBRARY_IMPL，用於註冊該 key 的實作。 如果您喜歡，您可以將 TORCH_LIBRARY_IMPL 程式碼區塊進一步細分為每個運算子一個區塊。 如果您每個運算子實作都有一個獨立的檔案，但不想在標頭檔中公開運算子，這會很方便；您可以直接將註冊放在定義運算子的 cpp 檔案中。

注意

您知道您也可以為 PyTorch 中現有的核心運算子編寫 TORCH_LIBRARY_IMPL 程式碼區塊嗎？ 這就是 PyTorch 的 XLA 支援的實作方式：torch_xla 程式庫包含一個 TORCH_LIBRARY_IMPL，它為 XLA dispatch key 上的所有基本運算子提供實作。

對於不需要 autograd 的運算子

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

在下一節中，我們將討論如何為運算子新增 autograd 支援。 但對於不需要 autograd 支援的運算，應註冊以下核心，以提高可用性並使您的運算子的行為與 PyTorch 的內建運算子類似。

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

上面的程式碼行註冊了一個 Autograd 核心，它在正向傳播時附加一個虛擬的 NotImplemented 節點（保留輸入的 require_grad 屬性）。 在反向傳播時，NotImplemented 節點會引發錯誤。 這對於在較大的模型中進行偵錯很有幫助，因為以前很難準確地找出在正向傳播期間 requires_grad 屬性遺失的位置。

In-place 或 view 運算

為了確保正確性和最佳效能，如果您的運算子以 in-place 方式修改輸入，或者傳回與其中一個輸入別名的張量，則應採取兩個額外的步驟

1. 除了上面的 Autograd 核心之外，還需註冊一個 ADInplaceOrView 核心。 此核心處理必要的簿記工作，以確保 in-place 或 view 運算的正確性。 重要的是要注意，此 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 boxed 核心依賴於其邏輯中的運算子 schema 資訊。 如果您的運算子以 in-place 方式修改輸入，或者傳回與其中一個輸入別名的張量，則務必確保您的 schema 正確反映了這一點。 有關如何註釋 schema 的更多資訊，請參閱 這裡。

新增 autograd 支援

此時，我們有一個具有 CPU 和 CUDA 實作的運算子。 我們如何為其新增 autograd 支援？ 您可能猜到了，我們將註冊一個 autograd 核心（類似於 自訂 autograd 函數 教學中描述的內容）！ 但是，有一個轉折：與 CPU 和 CUDA 核心不同，autograd 核心需要重新分派：它需要回調到 dispatcher 以取得推論核心，例如 CPU 或 CUDA 實作。

因此，在我們編寫 autograd 核心之前，讓我們先編寫一個分派函數，該函數會呼叫 dispatcher 以找到適合您運算子的正確核心。 此函數構成您運算子的公共 C++ API — 事實上，PyTorch 的 C++ API 中的所有張量函數實際上都是以相同的方式在底層呼叫 dispatcher。 以下是分派函數的樣子：

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

讓我們分解一下：

• 在第一行中，我們從 dispatcher 中查找與我們要分派到的運算子對應的類型化運算子句柄。findSchemaOrThrow 接受兩個參數：運算子的（命名空間限定的）名稱和運算子的重載名稱（通常只是空字串）。typed 將動態類型化的句柄轉換為靜態類型化的句柄（執行運行時測試以確保您給出了正確的 C++ 類型），以便我們可以對其執行正常的 C++ 呼叫。我們將 decltype(myadd) 傳遞給它，因為分派函數的類型與註冊到 dispatcher 的底層核心的類型相同。

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

• 在第二行中，我們只需使用傳遞到分派函數的所有參數 call 運算子句柄。 這實際上會調用 dispatcher，最後控制權將轉移到適合此呼叫的任何核心。

有了分派函數，我們現在可以編寫 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];
}

autograd 函數使用 torch::autograd::Function 正常編寫，除了我們不直接在 forward() 中編寫實作之外，我們：

1. 使用 at::AutoNonVariableTypeMode RAII guard 關閉 autograd 處理，然後

2. 呼叫分派函數 myadd 以回呼到 dispatcher。

如果沒有 (1)，您的呼叫將無限迴圈（並堆疊溢位），因為 myadd 會將您送回此函數（因為最高優先順序的 dispatch key 仍然是 autograd）。 有了 (1)，autograd 會從正在考慮的 dispatch key 集中排除，我們將轉到下一個處理程序，它將是 CPU 和 CUDA。

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

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

注意

在這個範例中，我們將核心註冊到 Autograd，這會將其安裝為所有後端的 autograd 核心。您也可以使用對應的後端特定調度鍵來註冊針對特定後端優化的核心，例如 AutogradCPU 或 AutogradCUDA。若要更詳細地了解這些及其他調度鍵選項，請查看 torch/_python_dispatcher.py 中提供的 PythonDispatcher 工具。

超越 Autograd

在某種程度上，調度器並沒有做太多事情：它所做的只是實作一個經過美化的 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. 它是分散式的。您可以組裝一個運算子 (CPU、CUDA、Autograd) 的所有部分，而無需編寫一個引用所有這些部分的集中式 if 語句。重要的是，第三方可以為其他方面註冊額外的實作，而無需修補運算子的原始定義。我們將在 為新的後端擴展調度器 中詳細討論擴展調度器。

2. 它支援比 CPU、CUDA 和 Autograd 更多的調度鍵。您可以在 c10/core/DispatchKey.h 中看到 PyTorch 目前實作的所有調度鍵的完整清單。這些調度鍵為運算子實作了各種可選功能，如果您決定讓您的自訂運算子支援此功能，您只需為適當的鍵註冊一個核心。

3. 調度器實作了對 boxed fallback 函數的支援，這些函數可以實作一次並應用於系統中的所有運算子。Boxed fallback 可用於為調度鍵提供預設行為；如果您使用調度器來實作您的運算子，您也可以選擇加入所有這些運算的 fallback。

以下是一些您可能需要為其定義運算子的特定調度鍵。

Autocast

Autocast 調度鍵實作了對 自動混合精度 (AMP) 的支援。一個 autocast 包裝核心通常會在執行操作之前，將傳入的 float16 或 float32 CUDA 張量轉換為某些首選精度。例如，浮點 CUDA 張量上的 matmul 和卷積通常在 float16 中執行得更快，並且使用的記憶體更少，而不會損害收斂性。Autocast 包裝器僅在 啟用 autocast 的上下文中 才會生效。

這是一個假設的自訂 matmul 的 autocast 包裝器，以及它的註冊：

// 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 不變（參見 資格策略 以了解原生 autocast 操作）。這確保了如果網路在任何 float16 和 float32 CUDA 張量的混合上調用 mymatmul，mymatmul 將在 float16 中執行。同時，使用非 CUDA、整數類型或 float64 輸入調用 mymatmul 不會受到影響。建議在您自己的 autocast 包裝器中使用 cached_cast 來遵循原生資格策略，但不是必需的。例如，如果您想強制所有輸入類型都執行 float16，您可以 return mymatmul(self.half(), other.half()); 而不是使用 cached_cast。

請注意，與我們的 autograd 核心一樣，我們在重新調度之前從調度中排除 Autocast 鍵。

預設情況下，如果未提供 autocast 包裝器，我們會直接 fallthrough 到常規運算子實作（不發生自動轉換）。（我們沒有在此範例中使用 myadd，因為逐點加法不需要自動轉換，而應該直接 fall through。）

何時應該註冊 autocast 包裝器？不幸的是，對於運算子的首選精度沒有明確的規則。您可以查看 轉換清單，以了解一些原生運算子的首選精度。一般指導：

• 進行簡化的運算子可能應該在 float32 中執行，

• 在底層進行卷積或 gemm 的任何運算子可能應該在 float16 中執行，並且

• 具有多個浮點張量輸入的其他運算子應該將它們標準化為通用精度（除非實作支援具有不同精度的輸入）。

如果您的自訂運算子屬於第三類，則 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));
}

如果您的自訂運算子是 啟用 autograd 的，您只需要為註冊 autograd 包裝器的相同名稱編寫和註冊一個 autocast 包裝器。例如，如果您想要 autograd 部分中顯示的 myadd 函數的 autocast 包裝器，您只需要：

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

沒有單獨的技巧來使向後方法與 autocast 相容。但是，在您的自訂 autograd 函數中定義的向後方法將以與 autocast 為向前方法設定的相同 dtype 執行，因此您應該選擇適合您的向前和向後方法的 <所需的 dtype>。

Batched

Batched 張量允許您以每個範例的方式編寫程式碼，然後在 vmap 調用下運行時自動對它們進行批次處理。用於編寫批次處理規則的 API 目前正在開發中，但一旦它穩定下來，您可以透過在 Batched 調度鍵處註冊一個核心，為您的運算子新增對 vmap 的支援。

Tracer

當您執行 torch.jit.trace 時，Tracer 調度鍵支援將運算子的調用記錄到追蹤中。我們計劃提供一個封裝的回退機制，以實現對任意操作的追蹤，請參閱 issue #41478 以追蹤進度。