工程開發
at-dispatch-v2 avatar

at-dispatch-v2

將 PyTorch ATen C++ 程式碼從舊版 AT_DISPATCH 巨集轉換為現代 AT_DISPATCH_V2 API 的自動化工具。

簡介

此技能為將 PyTorch ATen C++ 運算子實作從舊版分派巨集遷移至現代化 AT_DISPATCH_V2 架構提供了系統化框架。專為 PyTorch 核心開發人員與系統工程師設計,透過 Dispatch_v2.h 標頭檔中定義的標準化轉換規則、參數重新排序邏輯與巨集展開指引,簡化移植流程。目標針對 aten/src/ATen/native 目錄下的檔案,確保舊有程式碼庫能持續與演進中的 PyTorch 分派需求保持相容。

  • 自動化處理舊版模式(如 AT_DISPATCH_ALL_TYPES_AND* 與 AT_DISPATCH_FLOATING_TYPES_AND*)至統一 AT_DISPATCH_V2 格式的遷移。

  • 提供類型群組展開的清晰對應表,確保正確使用 AT_EXPAND() 處理如 AT_ALL_TYPES、AT_INTEGRAL_TYPES 與 AT_FLOATING_TYPES 等基礎類型集。

  • 使用 AT_WRAP 封裝 Lambda 表達式,以處理包含內部逗號或多行邏輯的複雜 C++ 程式碼區塊。

  • 將 #include <ATen/Dispatch_v2.h> 與現有分派相依項目整合,標準化標頭檔管理。

  • 支援對類型列表的細粒度控制,允許開發人員在應用廣泛類型群組巨集後,附加如 kBFloat16、kHalf 或 kBool 等特定類型。

  • 在核心實作中識別舊版分派巨集,包含 ATen 原生運算子與 CUDA 核心。

  • 透過重新排序參數(將 scalar_type 與運算子名稱置於呼叫開頭)來系統化替換巨集。

  • 使用 AT_WRAP 封裝 Lambda 本體,以維持前處理期間正確的 C++ 巨集解析。

  • 針對涉及的特定類型類別應用正確的 AT_EXPAND 巨集,例如 AT_ALL_TYPES_AND_COMPLEX 或 AT_BAREBONES_UNSIGNED_TYPES。

  • 根據必要的簽章驗證最終轉換結果:AT_DISPATCH_V2(scalar_type, name, lambda, type_groups, individual_types)。

  • 確保個別類型遵循 API 規範,置於類型群組之後且不被 AT_EXPAND 包覆。

倉庫統計

Star 數
181
Fork 數
24
Open Issue 數
4
主要語言
Python
預設分支
main
同步狀態
閒置
最近同步時間
2026年4月29日 下午02:12
在 GitHub 查看