熱文：OneFlow源碼解析：算子簽名的自動推斷

撰文 | 鄭建華

(資料圖片僅供參考)

OneFlow是一個原生支持分布式訓練的、高性能的深度學習框架。最近讀了一些OneFlow的源碼、架構設計和代碼實現的文章，簡單梳理一下自己的理解。主要通過圖形展示調用過程和類之間的關系，只對部分重要的代碼作一下分析。

深度學習框架是一個復雜的系統，而用戶使用最多的就是算子（op）。用戶通過op構造模型，進行訓練、預測。這個筆記就從op入手，看看從Python前端到C++底層，OneFlow如何執行算子的計算邏輯。

具體地說，以比較簡單的Relu算子為例，分析如下代碼怎么執行：

# import會觸發一系列初始化工作，暫時忽略import oneflow as flow# tensor的實現其實很復雜，因為要融合local和分布式的global tensort = flow.tensor([-1, 0, 1])r = flow.relu(t)

1

編譯環境

在開始分析之前，需要搭建環境編譯OneFlow的源碼，因為有些代碼是在編譯構建過程中自動生成的。在分析的過程中，這些自動生成的代碼也是必要的環節。

OneFlow提供了官方的編譯鏡像（https://hub.docker.com/r/oneflowinc/manylinux2014_x86_64_cuda11.2）。用這個鏡像可以非常方便地搭建編譯環境（https://github.com/Oneflow-Inc/oneflow#option-2-build-in-docker-container-recommended）。

我使用的OneFlow版本是v0.7.0。本地編譯環境目錄結構如下，build是

cmake的構建目錄，oneflow是源碼目錄。

.├── build└── oneflow

編譯比較耗時，可以把兩個目錄mount到容器，便于后續查看build目錄中生成的文件。

在cmake配置、構建過程中，會下載很多第三方源碼包，如果網絡狀況不好容易超時，直接重試cmake/make即可。

# docker run -itd -v $PWD/oneflow:/mnt/oneflow -v $PWD/build:/mnt/build \# manylinux2014_x86_64_cuda11.2 bashcd /mnt/buildcmake -S /mnt/oneflowcmake --build . # --parallel 8cd ../oneflow/pythonpython3 setup.py bdist_wheelpip install ./dist/oneflow-0.7.0+cpu-cp38-cp38-linux_x86_64.whl

用GDB追蹤OneFlow的執行過程

王益：Use GDB to Walkthrough OneFlow Source Code（https://quip.com/JuQ0AuodVJn4）

CMAKE_BUILD_TYPE=Debug cmake -S /mnt/oneflowcmake --build . --parallel 8source /mnt/build/source.shgdb python3b oneflow::one::MakeLocalTensorFromDatarunimport oneflow as flowflow.Tensor([[1,2,3],[4,5,6]])

2

Python Binding

OneFlow底層是C++實現，通過pybind11實現Python Binding。月踏在《從Python到C++調用過程分析》對相關內容做了講解。

2.1 Relu的Python包路徑

# python/oneflow/__init__.pyfrom oneflow._C import relu# python/oneflow/_C/__init__.pyfrom oneflow._oneflow_internal._C import *

2.2 module處理邏輯的注冊

Python代碼主要在python/oneflow目錄，C++實現的包主要在_oneflow_internal下，pybind11的綁定代碼位于init.cpp（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/init.cpp）：

PYBIND11_MODULE(_oneflow_internal, m) { // ... py::class_<::oneflow::cfg::Message, std::shared_ptr<::oneflow::cfg::Message>>(m, "CfgMessage"); ::oneflow::cfg::Pybind11ModuleRegistry().ImportAll(m); ::oneflow::OneflowModuleRegistry().ImportAll(m);}

其中OneflowModuleRegistry（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/init.cpp#L106）是算子等模塊的綁定；Pybind11ModuleRegistry（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/init.cpp#L105）應該是自定義的、類似protobuf的配置數據結構的綁定。

從OneflowModuleRegistry開始的詳細調用流程如下：

把代碼放到一起看看（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/of_api_registry.cpp）：

using SubModuleMap = std::map>>;SubModuleMap* GetSubModuleMap() { static SubModuleMap sub_module_map; return &sub_module_map;}// 修改map，執行注冊void OneflowModuleRegistry::Register(std::string module_path, std::function BuildModule) { (*GetSubModuleMap())[module_path].emplace_back(BuildModule);}void OneflowModuleRegistry::ImportAll(pybind11::module& m) { for (const auto& pair : (*GetSubModuleMap())) { for (const auto& BuildModule : pair.second) { BuildSubModule(pair.first, m, BuildModule); } }}void OneflowModuleRegistry::BuildSubModule( const std::string& module_path, pybind11::module& m, const std::function& BuildModule) { // ... BuildModule(m); // ...}

從這段代碼可以看出，python module的注冊邏輯都保存在SubModuleMap中。它的key是module name；value是一組函數，BuildSubModule中調用這些函數、執行module注冊邏輯。

GetSubModuleMap中保存map單例，Register函數設置map的值，of_api_registry.h（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/of_api_registry.h）中的宏ONEFLOW_API_PYBIND11_MODULE調用Register函數處理module注冊邏輯。搜索一下可以知道Relu的注冊邏輯在build/oneflow/api/python/functional/functional_api.yaml.pybind.cpp中，這個文件中注冊了很多算子（user_op）。以Relu和pow為例，這個宏展開后的核心代碼如下：

static void OneflowApiPythonModule9623(pybind11::module&);namespace { struct OfApiRegistryInit { OfApiRegistryInit() { ::oneflow::OneflowModuleRegistry().Register("_C", &OneflowApiPythonModule9623); } }; OfApiRegistryInit of_api_registry_init;}static void OneflowApiPythonModule9623(pybind11::module & m) { m.def("relu", &functional::PyFunction); m.def("pow", &functional::PyFunction< functional::PowSchema_TTT, functional::ScalarPowSchema_TTScB, functional::ScalarPowSchema_TTSc, functional::ScalarReversePowSchema_TScT >);}

這段代碼中的類似注冊技巧，在OneFlow中的很多地方都被用到。

module注冊邏輯在函數OneflowApiPythonModule9623中（9623來自宏定義中的LINE以避免名字沖突），OfApiRegistryInit在構造對象時將這個函數注冊到SubModuleMap，匿名空間中的變量of_api_registry_init就是為了通過構造對象、在構造函數中調用注冊邏輯（而這個對象不占用任何空間）。這樣在系統加載時就通過靜態對象的初始化實現了module處理邏輯的注冊，再通過pybind11的調用完成對Python Binding的定義。

3

多個接口簽名的自動推斷

從以上代碼可以看到，Relu算子被綁定到PyFunction（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#L120）這個函數執行計算邏輯，每次調用算子都會執行PyFunction這個函數。

從簽名看，PyFunction是一個模版函數，給Python前端返回py::object作為算子執行結果。

Relu只有一個模版參數，pow有4個模版參數。每個模版參數表示算子支持的一種調用接口簽名。OneFlow可以根據Python傳過來的arguments類型，自動推斷合適的簽名，調用相關函數。

例如下面的代碼，算子pow的指數參數既支持標量，也支持tensor：

import oneflow as flowr = flow.randn(1, 10)flow.pow(r, 2)flow.pow(r, flow.ones(1, 10))

下面就來看看OneFlow是怎么實現這個功能的。

Relu算子的簽名Schema如下所示：

struct ReluSchema_TTB { using FType = Maybe (const std::shared_ptr& x, bool inplace); using R = Maybe; static constexpr FType* func = &functional::Relu; static constexpr size_t max_args = 2; static constexpr size_t max_pos_args = 2; static constexpr char const* signature = "Tensor (Tensor x, Bool inplace=False)"; static FunctionDef function_def;};

先看一下從PyFunction開始的的調用順序：

PyFunction相關的代碼如下（刪掉了一些與核心邏輯無關的內容）。

// SchemaT如 ReluSchema_TTBtemplateclass PyFunctionDispatcher { public: // schema_t是第I個簽名 template using schema_t = typename std::tuple_element>::type; // schema_size_是簽名個數，比如relu是1，pow是4 PyFunctionDispatcher() : schema_size_(sizeof...(SchemaT)) { signatures_.resize(schema_size_); InitSignatures(std::make_index_sequence{}); } template py::object call(const py::args& args, const py::kwargs& kwargs, std::index_sequence) const { // T是當前檢查的簽名，比如 ReluSchema_TTB using T = schema_t; std::vector parsed_args(T::max_args); if (ParseArgs(args, kwargs, &parsed_args, T::function_def, T::max_pos_args, /*raise_exception*/ schema_size_ == 1)) { return detail::unpack_call(*T::func, parsed_args); } return call(args, kwargs, std::index_sequence{}); } py::object call(const py::args& args, const py::kwargs& kwargs, std::index_sequence<>) const { // throw error ... return py::none(); } private: template void InitSignatures(std::index_sequence) { __attribute__((__unused__)) int dummy[] = { ((void)(signatures_[I] = schema_t::signature), 0)...}; } private: size_t schema_size_; std::vector signatures_;};// SchemaT如 ReluSchema_TTBtemplateinline py::object PyFunction(const py::args& args, const py::kwargs& kwargs) { static PyFunctionDispatcher dispatcher; return dispatcher.call(args, kwargs, std::make_index_sequence{});}// py module注冊static void OneflowApiPythonModule9623(pybind11::module & m) { m.def("relu", &functional::PyFunction); m.def("pow", &functional::PyFunction< functional::PowSchema_TTT, functional::ScalarPowSchema_TTScB, functional::ScalarPowSchema_TTSc, functional::ScalarReversePowSchema_TScT >);}

3.1 dispatcher: 算子接口簽名的自動推斷

PyFunction是一個模版函數，每個模版參數表示算子的一個接口簽名。

PyFunction及其后續執行鏈路的最重要的功能，就是實現這些簽名的自動篩選。自動篩選的實質，就是通過index_sequence逐個檢查簽名與PyFunction的參數args/kwargs是否匹配。函數內的靜態變量dispatcher實現了這個自動篩選功能。

每個算子都會特化一個PyFunction和PyFunctionDispatcher實例，也有一個算子自己的dispatcher變量。PyFunction直接將請求轉發給dispatcher.call，順帶加上一個index_sequence模版參數，正是依靠這個模版參數實現了簽名的自動篩選。

在call函數中，先確定當前檢查的簽名類型T（例如ReluSchema_TTB），然后通過ParseArgs檢查Python傳過來的參數args/kwargs與簽名T是否匹配。如果不匹配，就去掉當前簽名T，將剩余的簽名類型作為模版參數、繼續遞歸調用call函數。

如果算子只有一個簽名，就通過schema_size_ == 1通知ParseArgs（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.cpp#L48），校驗失敗時直接拋出錯誤信息。

3.2 ParseArgs: 簽名與參數的匹配

Python的keyword arguments是類似map的結構，在C++中不方便直接用，需要轉為positional arguments，同時按順序保存到parsed_args中供后續執行使用。而這個順序只能是簽名指定的順序，所以ParseArgs中只能按function_def的順序循環校驗。

函數的參數可能是各種類型，ParseArgs統一轉為PythonArg類型，并通過PyObject*類型的成員讀取Python的變量值。

參數校驗不一致的情況主要包括：

positional與keyword參數類型沖突

簽名中的keyword參數名在kwargs中不存在且不接受默認值

參數類型不符合PythonArgCheck規定的內部類型檢查要求

kwargs包含function_def中未定義的參數

3.3 unpack_call: 展開算子函數的參數

在call函數中確定算子簽名的Schema之后，直接調用unpack_call（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/unpack_call.h#L69）函數。這時已經可以確定具體的算子執行函數了，對于Relu來說就是functional::Relu，同時將Python傳過來的參數都整理到args中。

unpack_call的模版參數是函數類型，例如functional::Relu，在函數體內利用function_traits推導出函數的參數個數和返回值類型。

unpack_call_dispatcher內主要是調用f，也就是functional::Relu。但還不能直接調用這個函數。因為每個算子對應函數的簽名都不一樣，又不能把vector args直接傳給這些函數。

OneFlow通過如下步驟完成模版的特化適配：

將args展開為各個PythonArg元素，通過index_sequence和變長模版參數包的展開實現；

利用function_traits推導得到函數參數類型列表ArgsType；

As函數調用可簡化為As>()...核心是拿到各個參數的實際類型并交給As處理，最終調用ObjectAs實現各種內部數據類型的轉換。

unpack_call_dispatcher返回的是C++內部數據類型，最后要通過CastToPyObject轉為pybind11::object，主要是調用pybind11::cast函數。

class PythonArg { template T As() const { return ObjectAsHelper>()(this).GetOrThrow(); }};templatestruct unpack_call_dispatcher { template static R apply(const F& f, const std::vector& args, std::index_sequence) { // 這里適當改寫了一下，把ArgsType抽出來 using ArgsType = function_traits::args_type; return f(args[I] .As::type>>()...); }};templatepy::object unpack_call(const F& f, const std::vector& args) { constexpr size_t nargs = function_traits::nargs; using R = typename function_traits::return_type; return CastToPyObject( unpack_call_dispatcher::apply(f, args, std::make_index_sequence{}));}

3.4 簽名都無效時的錯誤處理

以上只是討論了Python參數合法，可以找到匹配的函數簽名的情況。如果傳過來的參數是非法的，根據args/kwargs找不到匹配的簽名怎么辦？

如之前的討論，PyFunctionDispatcher::call（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#L58c）是遞歸模版參數，如果當前簽名不匹配，就嘗試下一個簽名。如果所有簽名都不匹配，就會進入call的模版參數列表為空的特化版本（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#L69）。這個函數會記錄詳細的錯誤信息。

例如，flow.pow("abc", 123)會輸出如下錯誤信息：

File ".../oneflow/api/python/functional/py_function.h", line 76, in call TypeError: pow(): received an invalid combination of arguments. The valid signatures are: *0: Tensor (Tensor input, Tensor exponent) *1: Tensor (Tensor input, Scalar exponent, *, Bool inplace=False) *2: Tensor (Tensor input, Scalar exponent) *3: Tensor (Scalar exponent, Tensor input)

而Relu這種只支持一個簽名的算子，如下面看到的，參數類型錯誤時的提示信息體現了單個簽名的特點。如上所述，這是由schema_size_ == 1提示給ParseArgs的。

flow.relu(1)TypeException: File ".../oneflow/api/python/functional/py_function.cpp", line 98, in ParseArgs TypeError: relu(): argument "x" must be tensor, not int

3.5 yaml cpp的生成

functional_api.yaml的相關代碼是在cmake構建過程中生成的，對應的cmake腳本是cmake/functional.cmake。

3.6 小結

總結一下上述幾個主要組件的作用：

PyFunction是pybind11的def定義的入口函數，并為算子保存一個dispatcher對象用于推斷合適的簽名；

PyFunctionDispatcher通過模版函數的遞歸調用實現了簽名的自動篩選，通過成員變量為參數校驗和異常提示保存必要的信息；

unpack_call在編譯期就確定了具體執行的算子函數類型，這一點在PyFunctionDispatcher中是無法做到的；

unpack_call_dispatcher的作用是將vector展開為多個元素、作為調用算子函數的參數，這在unpack_call中也是無法做到的；

PythonArg是Python與C++類型轉換的橋梁，同時承擔類型檢查的職能；

基于yaml生成的2組文件，yaml.pybind.cpp中調用pybind11的m.def指定模塊調用的函數，并定義了函數簽名的Schema結構作為PyFunction的模版參數。yaml.cpp中則定義了具體的執行函數，如Relu。將二者銜接起來的就是Schema的字段func，對于Relu算子來說，簽名Schema的func字段就是函數functional:Relu。

核心是實現簽名的自動校驗推斷，參數的統一處理以及參數的合并、展開。整個過程環環相扣、自然流暢。

4

算子Functor的注冊與執行

4.1 算子Functor的注冊

追蹤一下functional::Relu（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#L40）的調用鏈路，容易發現最終會用到FunctionLibrary的靜態map變量。先看看這個map是怎么初始化的。它在add_functor_creator（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#L93）中被添加元素，后者被add_functor（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#L63）間接調用。

搜索一下add_functor和Relu，發現在activation_functor.cpp中調用宏ONEFLOW_FUNCTION_LIBRARY（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/impl/activation_functor.cpp#L444）。宏展開后代碼如下，通過定義一個靜態變量來實現調用注冊函數的目的。

static void _oneflow_function_library_0(FunctionLibrary & m);// 以定義一個靜態變量的方式調用注冊函數static int _oneflow_function_library_dummy_0 = []() { FunctionLibrary* library = FunctionLibrary::Global(); _oneflow_function_library_0(*library); return 0; }();void _oneflow_function_library_0(FunctionLibrary & m) { m.add_functor("Relu");};

稍微梳理一下就可以發現，FunctionLibrary的map中的value是類似下面這樣的lambda：

[=]() { // Func如 impl::ReluFunctor Func func; // func_name來自lambda綁定，如Relu return PackedFunctorMaker::make(func_name, func);}

注冊的調用順序如下：

那么，add_functor的模版參數為何是變長的，內部又要展開呢？是因為ScalarAdd等名字對應多個Functor。

4.2 算子Functor的執行

接下來看看functional_api.yaml.cpp中的functional::Relu函數。代碼經過整理后如下所示。

Maybe Relu(const std::shared_ptr& x, bool inplace) { static thread_local const auto& __op = CHECK_JUST( FunctionLibrary::Global()->find < Maybe, const std::shared_ptr&, bool > ("Relu")); return __op->call(x, inplace);}

核心邏輯就是func_lib.find("Relu").call(x, inplace)。

獲取__op并執行的調用順序如下（忽略op的靜態屬性）：

根據上面的討論以及調用鏈路容易發現，PackedFuncCreatorMap::Get內的靜態map變量（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#L40），其value實際是一個類似如下的lambda表達式：

[=]() { // Func如 impl::ReluFunctor Func func; // func_name來自lambda綁定，如Relu return PackedFunctorMaker::make(func_name, func);}

find返回的是it->second()，也就是調用這個lambda表達式的返回值，即PackedFunctorMaker::make的返回值，類型是PackedFunctor，這就是op__的類型。其中模版參數F的類型如decltype(ReluFunctor::operator())。

PackedFunctor構造時接受如下的lambda表達式，并保存到變量impl_中：

// func是一個函數變量，類型如 impl::ReluFunctor[func](const remove_cvref_t&... args) -> R { return func(std::forward&>(args)...);}

所以__op->call(...)就是PackedFunctor::call(...)，最終相當于調用impl::ReluFunctor::operator()(args)。

也就是說，Relu的操作就由impl::ReluFunctor執行。

需要注意的是，這里整個鏈路的分析，最關鍵的是模版參數的梳理和推導。模版參數確定后，整個邏輯還是比較清楚的。

4.3 小結

同一個名字可能對應多個Functor。所以不能只用名字作為Functor的key，需要結合簽名。

FunctionLibrary負責管理所有的Functor。但是單例不適合作為模版類，所以通過內嵌的PackedFuncCreatorMap保存簽名各異的Functor。

每種簽名都會特化一個PackedFuncCreatorMap模版類，再通過名字區分不同的Functor。

那么，PackedFunctor類的作用是什么？或者換個角度，如果沒有這個類，能否實現需求？答案是不能。

首先，yaml生成的2個cpp文件，都沒有Functor信息，只有Relu這個名字、以及Functor的簽名信息。Functor是在各個模塊根據名字注冊的。yaml與FunctionLibrary通過名字和簽名進行交互。

其次，FunctionLibrary::find返回的PackedFunctor是帶模版參數的（參數就是Functor簽名）。find能否直接返回Functor對象呢？主要是map不便存儲不同類型的Functor。即使Functor都有共同的虛基類、map的value存儲指針，但不能要求所有Functor的執行接口是一致的，虛函數不滿足這個場景的需求。所以find不能直接返回Functor對象。

PackedFunctor的作用就在于，它把真正的Functor包在自己的結構里面；它的模版參數與Functor的調用接口一致；它的call方法將Op的所有入參通過lambda轉發給Functor。

Functor能直接作為PackedFunctor的成員變量嗎？應該是可以的。PackedFunctorMaker::make的模版參數也包含Functor。但是這樣每個Functor都要特化一個PackedFunctor，編譯后的可執行程序容易膨脹。而現在的實現，PackedFunctor只根據Functor執行函數簽名特化，代價是要做一次調用轉發（編譯器有優化空間？）。

參考資料

從Python到C++調用過程分析

https://github.com/Oneflow-Inc/oneflow/tree/release/0.7.0

（本文經授權后發布，原文：https://segmentfault.com/a/1190000041843994）

其他人都在看

深度學習概述

一個算子在深度學習框架中的旅程

手把手推導分布式矩陣乘的最優并行策略

訓練千億參數大模型，離不開四種并行策略

解讀Pathways（二）：向前一步是OneFlow

關于并發和并行，Go和Erlang之父都弄錯了？

OneFlow v0.7.0發布：全新分布式接口，LiBai、Serving等一應俱全

歡迎體驗OneFlow v0.7.0：GitHub - Oneflow-Inc/oneflow: OneFlow is a performance-centered and open-source deep learning framework.OneFlow is a performance-centered and open-source deep learning framework. - GitHub - Oneflow-Inc/oneflow: OneFlow is a performance-centered and open-source deep learning framework.https://github.com/Oneflow-Inc/oneflow/