ONEFLOW source code parsing: automatic inference of operator signature

writing  |  Zheng Jianhua

OneFlow Is a native support for distributed training 、 High performance deep learning framework . Recently I read some OneFlow Source code 、 Architecture design and code implementation , Simply sort out your understanding . It mainly shows the relationship between calling process and class through graphics , Just analyze some important code .

Deep learning framework is a complex system , And the operator that users use most is （op）. User pass op Build a model , Training 、 forecast . This note starts from op Starting with , Let's see Python Front end to C++ Bottom ,OneFlow How to execute the operator's computational logic .

To be specific , In a simpler way Relu For example, operators , Analyze how the following code executes ：

# import It will trigger a series of initialization , Temporary neglect
import oneflow as flow
# tensor The implementation of is actually very complicated , Because we need to integrate local And distributed global tensor
t = flow.tensor([-1, 0, 1])
r = flow.relu(t)

1、 Compile environment

Before starting the analysis , You need to build an environment to compile OneFlow Source code , Because some code is generated automatically during the compilation and construction process . In the course of analysis , These automatically generated codes are also necessary .

OneFlow The official compilation image is provided （

https://hub.docker.com/r/oneflowinc/manylinux2014_x86_64_cuda11.2

）. Using this image, you can easily build a compilation environment （

https://github.com/Oneflow-Inc/oneflow#option-2-build-in-docker-container-recommended

）.

I use OneFlow The version is v0.7.0. The directory structure of the local compilation environment is as follows ,build yes cmake Build directory ,oneflow It's the source directory .

.
├── build
└── oneflow

Compilation is time consuming , You can put two directories mount To the container , For subsequent viewing build Files generated in the directory .

stay cmake To configure 、 During construction , Many third-party source code packages will be downloaded , If the network condition is not good, it is easy to time out , Just try again cmake/make that will do .

# docker run -itd -v $PWD/oneflow:/mnt/oneflow -v $PWD/build:/mnt/build \
# manylinux2014_x86_64_cuda11.2 bash
cd /mnt/build
cmake -S /mnt/oneflow
cmake --build . # --parallel 8
cd ../oneflow/python
python3 setup.py bdist_wheel
pip install ./dist/oneflow-0.7.0+cpu-cp38-cp38-linux_x86_64.whl

use GDB track OneFlow Implementation process of

Wang Yi ：Use GDB to Walkthrough OneFlow Source Code（

https://quip.com/JuQ0AuodVJn4

）

CMAKE_BUILD_TYPE=Debug cmake -S /mnt/oneflow
cmake --build . --parallel 8

source /mnt/build/source.sh
gdb python3

b oneflow::one::MakeLocalTensorFromData

run

import oneflow as flow

flow.Tensor([[1,2,3],[4,5,6]])

2、Python Binding

OneFlow The bottom is C++ Realization , adopt pybind11 Realization Python Binding. Moon treads on 《 from Python To C++ Call process analysis 》 Explained the relevant contents .

2.1 Relu Of Python Package path

# python/oneflow/__init__.py
from oneflow._C import relu

# python/oneflow/_C/__init__.py
from oneflow._oneflow_internal._C import *

2.2 module Handle the registration of logic

Python The code is mainly in python/oneflow Catalog ,C++ The implemented packages are mainly in _oneflow_internal Next ,pybind11 The binding code for is located at 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, &quot;CfgMessage&quot;);
 ::oneflow::cfg::Pybind11ModuleRegistry().ImportAll(m);
 ::oneflow::OneflowModuleRegistry().ImportAll(m);
}

among OneflowModuleRegistry（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/init.cpp#L106

） It is the binding of operator and other modules ;Pybind11ModuleRegistry（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/init.cpp#L105

） It should be custom 、 similar protobuf Configuration data structure binding .

from OneflowModuleRegistry The detailed calling process is as follows ：

Put the code together （

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/of_api_registry.cpp

）：

using SubModuleMap = std::map<std::string, std::vector<std::function<void(pybind11::module&)>>>;

SubModuleMap* GetSubModuleMap() {
 static SubModuleMap sub_module_map;
 return &sub_module_map;
}

//  modify map, Enforcement registration
void OneflowModuleRegistry::Register(std::string module_path,
 std::function<void(pybind11::module&)> 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<void(pybind11::module&)>& BuildModule) {
 // ...
 BuildModule(m);
 // ...
}

As you can see from this code ,python module The registration logic is stored in SubModuleMap in . its key yes module name;value Is a set of functions ,BuildSubModule These functions are called in 、 perform module Registration logic .

GetSubModuleMap Kept in map Single case ,Register Function settings map Value ,of_api_registry.h（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/of_api_registry.h

） Macro in ONEFLOW_API_PYBIND11_MODULE call Register Function processing module Registration logic . A search will reveal Relu The registration logic for is in build/oneflow/api/python/functional/functional_api.yaml.pybind.cpp in , Many operators are registered in this file （user_op）. With Relu and pow For example , The core code after the grand exhibition is as follows ：

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(&quot;relu&quot;, &functional::PyFunction<functional::ReluSchema_TTB>);
 m.def(&quot;pow&quot;, &functional::PyFunction<
 functional::PowSchema_TTT, functional::ScalarPowSchema_TTScB,
 functional::ScalarPowSchema_TTSc, functional::ScalarReversePowSchema_TScT
 >);
}

Similar registration techniques in this code , stay OneFlow A lot of places are used in .

module Register logic in function OneflowApiPythonModule9623 in （9623 From the macro definition LINE To avoid name conflicts ）,OfApiRegistryInit Register this function to when constructing the object SubModuleMap, Variables in anonymous space of_api_registry_init To construct objects 、 Call the registration logic in the constructor （ And this object doesn't take up any space ）. In this way, the initialization of static objects is implemented when the system is loaded module Handle the registration of logic , Re pass pybind11 Complete the call to Python Binding The definition of .

3、 Automatic inference of multiple interface signatures

As you can see from the above code ,Relu The operator is bound to PyFunction（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#L120

） This function performs computational logic , Every time the operator is called PyFunction This function .

Judging from the signature ,PyFunction Is a template function , to Python Front end return py::object As a result of operator execution .

Relu There is only one template parameter ,pow Yes 4 Template parameters . Each template parameter represents a calling interface signature supported by the operator .OneFlow According to Python From here arguments type , Automatically infer appropriate signatures , Call related functions .

For example, the following code , operator pow The exponential parameter of supports both scalars , Also support tensor：

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

So let's see OneFlow How to realize this function .

Relu Operator signature Schema As shown below ：

struct ReluSchema_TTB {
 using FType = Maybe<one::Tensor> (const std::shared_ptr<one::Tensor>& x, bool inplace);
 using R = Maybe<one::Tensor>;
 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 = &quot;Tensor (Tensor x, Bool inplace=False)&quot;;
 static FunctionDef function_def;
};

First take a look from PyFunction The starting call sequence ：

PyFunction The relevant code is as follows （ Some content irrelevant to the core logic has been deleted ）.

// SchemaT Such as  ReluSchema_TTB
template<typename... SchemaT>
class PyFunctionDispatcher {
 public:
 // schema_t It's No I A signature
 template<size_t I>
 using schema_t = typename std::tuple_element<I, std::tuple<SchemaT...>>::type;

 // schema_size_ Is the number of signatures , such as relu yes 1,pow yes 4
 PyFunctionDispatcher() : schema_size_(sizeof...(SchemaT)) {
 signatures_.resize(schema_size_);
 InitSignatures(std::make_index_sequence<sizeof...(SchemaT)>{});
 }

 template<size_t I0, size_t... I>
 py::object call(const py::args& args, const py::kwargs& kwargs,
 std::index_sequence<I0, I...>) const {
 // T Is the signature of the current check , such as  ReluSchema_TTB
 using T = schema_t<I0>;
 std::vector<PythonArg> 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<I...>{});
 }

 py::object call(const py::args& args, const py::kwargs& kwargs, std::index_sequence<>) const {
 // throw error ...
 return py::none();
 }

 private:
 template<size_t... I>
 void InitSignatures(std::index_sequence<I...>) {
 __attribute__((__unused__)) int dummy[] = {
 ((void)(signatures_[I] = schema_t<I>::signature), 0)...};
 }

 private:
 size_t schema_size_;
 std::vector<const char*> signatures_;
};

// SchemaT Such as  ReluSchema_TTB
template<typename... SchemaT>
inline py::object PyFunction(const py::args& args, const py::kwargs& kwargs) {
 static PyFunctionDispatcher<SchemaT...> dispatcher;
 return dispatcher.call(args, kwargs, std::make_index_sequence<sizeof...(SchemaT)>{});
}

// py module register
static void OneflowApiPythonModule9623(pybind11::module & m) {
 m.def(&quot;relu&quot;, &functional::PyFunction<functional::ReluSchema_TTB>);
 m.def(&quot;pow&quot;, &functional::PyFunction<
 functional::PowSchema_TTT, functional::ScalarPowSchema_TTScB,
 functional::ScalarPowSchema_TTSc, functional::ScalarReversePowSchema_TScT
 >);
}

3.1 dispatcher:  Automatic inference of operator interface signature

PyFunction Is a template function , Each template parameter represents an interface signature of the operator .

PyFunction And the most important functions of the subsequent execution link , It is to realize the automatic filtering of these signatures . The essence of automatic filtering , It is through index_sequence Check the signatures and one by one PyFunction Parameters of args/kwargs match . Static variables in functions dispatcher This automatic filtering function is realized .

Each operator will specialize in one PyFunction and PyFunctionDispatcher example , There is also an operator of its own dispatcher Variable .PyFunction Forward the request directly to dispatcher.call, Add one by the way index_sequence Template parameters , It is this template parameter that enables the automatic filtering of signatures .

stay call Function , First determine the signature type of the current check T（ for example ReluSchema_TTB）, And then through ParseArgs Check Python Parameters transmitted args/kwargs And signature T match . If it doesn't match , Just remove the current signature T, Use the remaining signature types as template parameters 、 Continue to call recursively call function .

If the operator has only one signature , Just through schema_size_ == 1 notice ParseArgs（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.cpp#L48

）, When the verification fails, an error message is directly thrown .

3.2 ParseArgs:  Match between signature and parameter

Python Of keyword arguments Is similar to map Structure , stay C++ It is not convenient to use directly , Need to change to positional arguments, At the same time, save in order to parsed_args For subsequent execution . And this order can only be the order specified by the signature , therefore ParseArgs You can only press function_def Sequential cyclic verification of .

The arguments to a function can be of various types ,ParseArgs Unified into PythonArg type , And pass PyObject* Members of type read Python Variable value .

Inconsistent parameter verification mainly includes ：

• positional And keyword Argument type conflict

• In signature keyword The parameter name is in kwargs Does not exist in and does not accept the default value

• Parameter type does not conform to PythonArgCheck Specified internal type inspection requirements

• kwargs contain function_def Parameters not defined in

3.3 unpack_call:  Expand the parameters of the operator function

stay call Function to determine the signature of the operator Schema after , Call directly unpack_call（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/unpack_call.h#L69

） function . At this time, the specific operator execution function can be determined , about Relu That is to say functional::Relu, At the same time Python The passed parameters are sorted to args in .

unpack_call The template parameter of is the function type , for example functional::Relu, Use... In the function body function_traits The number of parameters and return value type of the function are derived .

unpack_call_dispatcher The inner part is mainly to call f, That is to say functional::Relu. But you can't call this function directly . Because the signature of the corresponding function of each operator is different , You can't vector args Directly to these functions .

OneFlow Complete the special adaptation of the template through the following steps ：

• take args Expand to each PythonArg Elements , adopt index_sequence And variable length template parameter package ;

• utilize function_traits Derive the function parameter type list ArgsType;

• As Function calls can be simplified to As<typename tuple_element<I, typename ArgsType>>()... The core is to get the actual type of each parameter and give it to As Handle , The final call ObjectAs Realize the conversion of various internal data types .

unpack_call_dispatcher The return is C++ Internal data types , Finally, we have to pass CastToPyObject To pybind11::object, Mainly called pybind11::cast function .

class PythonArg {
 template<typename T>
 T As() const {
 return ObjectAsHelper<oneflow::detail::remove_cvref_t<T>>()(this).GetOrThrow();
 }
};

template<typename F, typename R>
struct unpack_call_dispatcher {
 template<size_t... I>
 static R apply(const F& f, const std::vector<PythonArg>& args, std::index_sequence<I...>) {
 //  Here is an appropriate rewriting , hold ArgsType Pull it out
 using ArgsType = function_traits<F>::args_type;
 return f(args[I]
 .As<oneflow::detail::remove_cvref_t<typename std::tuple_element<
 I, typename ArgsType>::type>>()...);
 }
};

template<typename F>
py::object unpack_call(const F& f, const std::vector<PythonArg>& args) {
 constexpr size_t nargs = function_traits<F>::nargs;
 using R = typename function_traits<F>::return_type;
 return CastToPyObject(
 unpack_call_dispatcher<F, R>::apply(f, args, std::make_index_sequence<nargs>{}));
}

3.4  Error handling when signatures are invalid

The above is just a discussion Python Parameters are legal , You can find matching function signatures . If the parameter passed is illegal , according to args/kwargs What if no matching signature is found ？

As discussed earlier ,PyFunctionDispatcher::call（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#L58c

） Is a recursive template parameter , If the current signature does not match , Just try the next signature . If all signatures do not match , Will enter call The template parameter list of is empty （

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#L69

）. This function will record detailed error information .

for example ,flow.pow(&quot;abc&quot;, 123) The following error message will be output ：

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)

and Relu This kind of operator only supports one signature , As you can see below , The prompt message when the parameter type is wrong reflects the characteristics of a single signature . As mentioned above , This is from schema_size_ == 1 Prompt to ParseArgs Of .

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 Generation

functional_api.yaml The relevant code of is in cmake Generated during the build process , Corresponding cmake The script is cmake/functional.cmake.

3.6  Summary

Summarize the functions of the above main components ：

• PyFunction yes pybind11 Of def Defined entry function , And save a for the operator dispatcher Object is used to infer the appropriate signature ;

• PyFunctionDispatcher Through the recursive call of template function, the automatic filtering of signature is realized , Save necessary information for parameter validation and exception prompt through member variables ;

• unpack_call The type of operator function to be executed is determined at compile time , This is in PyFunctionDispatcher You can't do it in ;

• unpack_call_dispatcher The role of the vector Expand to multiple elements 、 As an argument to the calling operator function , This is in unpack_call It can't be done in ;

• PythonArg yes Python And C++ Bridge of type conversion , At the same time, it undertakes the function of type inspection ;

• be based on yaml Generated 2 Group files ,yaml.pybind.cpp Call in pybind11 Of m.def Specify the function called by the module , And define the function signature Schema Structure as PyFunction Template parameters .yaml.cpp The specific execution function is defined in , Such as Relu. What connects the two is Schema Field of func, about Relu Operators , Signature Schema Of func Fields are functions functional:Relu.

The core is to realize the automatic verification and inference of signature , Unified processing and merging of parameters 、 an . The whole process is connected 、 Natural flow .

4、 operator Functor Registration and execution of

4.1  operator Functor Registration of

Track down functional::Relu（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#L40

） Call link , It is easy to find that you will eventually use FunctionLibrary Static map Variable . First look at this map How to initialize . It's in add_functor_creator（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#L93

） Added element in , The latter is add_functor（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#L63

） Indirect invocation .

Search for add_functor and Relu, Found in activation_functor.cpp Calling macros in ONEFLOW_FUNCTION_LIBRARY（

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/impl/activation_functor.cpp#L444

）. The code after Hongzhan is as follows , The purpose of calling the registration function is achieved by defining a static variable .

static void _oneflow_function_library_0(FunctionLibrary & m);

//  Call the registration function by defining a static variable
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<impl::ReluFunctor>(&quot;Relu&quot;);
};

A little comb will reveal ,FunctionLibrary Of map Medium value It is similar to the following lambda：

[=]() {
 // Func Such as  impl::ReluFunctor
 Func func;
 // func_name come from lambda binding , Such as Relu
 return PackedFunctorMaker<func_type>::make(func_name, func);
}

The registration call sequence is as follows ：

that ,add_functor Why is the template parameter of variable length , The interior has to be expanded ？ Because ScalarAdd Such names correspond to multiple Functor.

4.2  operator Functor Implementation

So let's see functional_api.yaml.cpp Medium functional::Relu function . The code is sorted out as follows .

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

The core logic is func_lib.find(&quot;Relu&quot;).call(x, inplace).

obtain __op And execute the call sequence as follows （ Ignore op Static properties of ）：

According to the above discussion and the call link, it is easy to find ,PackedFuncCreatorMap::Get Static in map Variable （

https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#L40

）, Its value In fact, it is similar to the following lambda expression ：

[=]() {
 // Func Such as  impl::ReluFunctor
 Func func;
 // func_name come from lambda binding , Such as Relu
 return PackedFunctorMaker<func_type>::make(func_name, func);
}

find The return is it->second(), That is to call this lambda Return value of expression , namely PackedFunctorMaker::make The return value of , The type is PackedFunctor<F>, This is it. op__ The type of . Template parameters F The type of decltype(ReluFunctor::operator()).

PackedFunctor The following are accepted during construction lambda expression , And save it to the variable impl_ in ：

// func It's a function variable , Type as  impl::ReluFunctor
[func](const remove_cvref_t<Args>&... args) -> R {
 return func(std::forward<const remove_cvref_t<Args>&>(args)...);
}

therefore __op->call(...) Namely PackedFunctor<Func>::call(...), In the end, it is equivalent to calling impl::ReluFunctor::operator()(args).

in other words ,Relu The operation of impl::ReluFunctor perform .

It should be noted that , Here is the analysis of the whole link , The key is to sort out and derive template parameters . After the template parameters are determined , The whole logic is quite clear .

4.3  Summary

• The same name may correspond to more than one Functor. So you can't just use your name as Functor Of key, Need to combine signatures .

• FunctionLibrary Manage all Functor. But singletons are not suitable for template classes , So through the embedded PackedFuncCreatorMap Save different signatures Functor.

• Each signature will specialize in one PackedFuncCreatorMap Template class , Then distinguish different by name Functor.

that ,PackedFunctor What is the function of class ？ Or change the angle , Without this class , Whether the requirements can be realized ？ The answer is no .

• First ,yaml Generated 2 individual cpp file , None Functor Information , Only Relu The name 、 as well as Functor Signature information for .Functor It is registered by name in each module .yaml And FunctionLibrary Interact by name and signature .

• secondly ,FunctionLibrary::find Back to PackedFunctor With template parameters （ The parameter is Functor Signature ）.find Can I return directly to Functor What about objects? ？ Mainly map It is inconvenient to store different types of Functor. Even if Functor Have a common virtual base class 、map Of value Store pointer , But not all Functor The implementation interface of is consistent , Virtual functions do not meet the needs of this scenario . therefore find Cannot return directly Functor object .

• PackedFunctor The role of , It puts the real Functor Wrapped in its own structure ; Its template parameters are the same as Functor The calling interface of is consistent ; its call Methods will Op All the input parameters of the pass lambda Forward to Functor.

• Functor Can act directly as PackedFunctor Member variables of ？ It should be possible .PackedFunctorMaker::make The template parameters of also contain Functor. But this way everyone Functor We should specialize in one PackedFunctor, Compiled executable programs tend to swell . And now the implementation ,PackedFunctor Based on Functor Perform function signature specialization , The cost is to do a call forwarding （ Compiler has optimization space ？）.

Reference material ：

• from Python To C++ Call process analysis

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

（ This document is issued after authorization , original text ：https://segmentfault.com/a/1190000041843994）

Everyone else is watching

• Overview of deep learning

• The journey of an operator in the framework of deep learning

• The optimal parallel strategy of distributed matrix multiplication is derived by hand

• Train a large model with hundreds of billions of parameters , Four parallel strategies are indispensable

• Reading Pathways（ Two ）： The next step forward is OneFlow

• About concurrency and parallelism ,Go and Erlang My father is mistaken ？

• OneFlow v0.7.0 Release ：
  New distributed interface ,LiBai、Serving Everything

Welcome to download experience  OneFlow v0.7.0  The latest version ：

https://github.com/Oneflow-Inc/oneflow/