关于机器学习:OneFlow源码阅读2OpKernel与解释器

Op与Kernel的注册

持续追踪执行流程会发现，ReluFunctor在结构UserOpExpr时会用到UserOpRegistryMgr治理的Op与Kernel。Op示意算子的形容信息，Kernel实现在不同设施上的计算。

注册信息保留在公有的map变量中。UserOpRegistryMgr的头文件中定义了3个宏，别离用于注册op、grad_op、kernel。

ReluOp的注册

REGISTER_USER_OP负责UserOp的注册。通过检索代码能够找到这个宏的应用场景。ReluOp相干的源代码在这3个文件中：

class定义: build/oneflow/core/framework/op_generated.h
注册op、op的局部实现: build/oneflow/core/framework/op_generated.cpp
次要实现: oneflow/oneflow/user/ops/relu_op.cpp

REGISTER_USER_OP宏在op_generated.cpp中开展后代码如下：

static UserOpRegisterTrigger<OpRegistry> g_register_trigger715 =  ::oneflow::user_op::UserOpRegistryMgr::Get()  .CheckAndGetOpRegistry("relu")  .Input("x")  .Output("y")  .SetGetSbpFn(&ReluOp::GetSbp)  .SetLogicalTensorDescInferFn(&ReluOp::InferLogicalTensorDesc)  .SetPhysicalTensorDescInferFn(&ReluOp::InferPhysicalTensorDesc)  .SetDataTypeInferFn(&ReluOp::InferDataType);

调用流程如下：

CheckAndGetOpRegistry会创立一个OpRegistry对象，这个类和UserOpRegisterTrigger类一样，只是为结构OpRegistryResult用的两头类型。OpRegistry会暂存两头后果并在Finish中设置一些默认推导逻辑。UserOpRegisterTrigger的构造函数会调用注册逻辑。动态变量就是为了触发构造函数从而调用注册逻辑，将结构好的OpRegistryResult保留到UserOpRegistryMgr（key是op_type，如relu）。

ReluOp示意一个具体的op_type，负责为OpRegistryResult提供Op特有的办法。

OpRegistryResult把不同的Op形象为一个通用的构造（便于对立注册治理），次要蕴含形容信息，保留了op的输入输出形容，以及数据类型、sbp等的推导逻辑函数。对于relu来说，次要是记录了几个推导函数要调用ReluOp的静态方法；op_def次要蕴含input/output的名字。

ReluKernel的注册

ReluKernel在relu_kernel.cpp中注册，过程和Op的注册相似。REGISTER_USER_KERNEL宏产开后如下所示：

static UserOpRegisterTrigger<OpKernelRegistry> g_register_trigger0 =  UserOpRegistryMgr::Get().    CheckAndGetOpKernelRegistry("relu").    // 通过模版参数指定kernel为ReluKernel    SetCreateFn<ReluKernel>().    // 参数不是bool表达式，应该是一个高阶表达式对象    SetIsMatchedHob(ReluPrimitiveExists() == true);

留神SetCreateFn只是把一个如下的lambda表达式赋值给result_.create_fn，这个字段很重要，后续执行就是通过它获取kernel。

[] () -> const OpKernel* { return NewOpKernel<T>(); }

对于relu来说，NewOpKernel就是new一个ReluKernel对象并返回指针。

最终注册的后果，会把OpKernelRegistryResult保留到UserOpRegistryMgr（key是op_type，如relu）。

Op和Kernel注册相干的类关系图

UserOpExpr的结构

上一篇提到，functional_api.yaml.cpp中的functional::Relu函数通过find("Relu")获取事后注册的PackedFunctor<impl::ReluFunctor>，调用其call办法会执行impl::ReluFunctor。

ReluFunctor的外围代码如下：

class ReluFunctor { public:  ReluFunctor() { op_ = CHECK_JUST(one::OpBuilder("relu").Input("x", 1).Output("y", 1).Build()); }  Maybe<Tensor> operator()(const std::shared_ptr<Tensor>& x, bool inplace) const {    // 疏忽inplace相干逻辑    return OpInterpUtil::Dispatch<Tensor>(*op_, {x});  } private:  std::shared_ptr<OpExpr> op_;};

ReluFunctor的构造函数中，次要是结构UserOpExpr。UserOpExpr能够看作user op type的子概念，relu只有一个，scalar_add等有多个UserOpExpr。

构造函数内的调用程序如下：

OpBuilder的Input/Output调用次要是结构UserOpConf对象，Build函数内会批改UserOpConf对象，比方依据OpRegistryResult::op_def补充默认值到attr。之后结构UserOpExpr对象，UserOpConf对象被保留到UserOpExpr的父类BuiltinOpExprImpl<UserOpConf>的op_proto_字段，对于relu来说，op_proto_次要保留input, output等信息。UserOpExpr初始化时会从OpRegistryResult拷贝函数变量。

Functor的执行

ReluFunctor执行的外围逻辑是调用OpInterpUtil::Dispatch。调运程序如下：

整个链路很长，本篇笔记只对前半部分的重点内容做一些阐明。

依据环境和输出抉择解释器

Dispatch调用的GetInterpreter返回的是一个AutogradInterpreter对象，这个类是在其内含的OpExprInterpreter成员变量根底之上减少了autograd的性能。GetInterpreter内理论结构的是以下3种Interpreter，在Build函数返回时转为AutogradInterpreter。

LazyInterpreter: 应该用于lazy执行模式
EagerMirroredInterpreter: 貌似是单机（或数据并行）的动态图执行模式
EagerConsistentInterpreter: 分布式的动态图执行模式

各个Interpreter的关系如下：

GetInterpreter的作用是依据输出和环境等信息，抉择一个适合的解释器。

接着在Dispatch中调用解释器的AutogradInterpreter::Apply办法，在这个办法内调用internal_->Apply(...)，也就是上述3个解释器的Apply办法。

以EagerConsistentInterpreter为例。这个类并没有定义Apply办法，理论调用的是父类办法EagerInterpreter::Apply。这个办法中调用一系列的APPLY_IF宏，就是用dynamic_cast判断op_expr的类型，类型适合才执行，对于relu会调用UserOpExpr版的ApplyImpl办法。

装璜器

EagerConsistentInterpreter::ApplyImpl的相干代码如下所示：

// Decorator的模版参数能够通过func的类型推断// oneflow/core/common/decorator.htemplate<template<typename...> class Decorator>struct WithDecorator final {  template<typename T, typename = void> struct Decorate;  template<typename T, typename... Args>  struct Decorate<T (*)(Args...)> final {    template<T (*func)(Args...)>    static T Call(Args... args) {      return Decorator<T, Args...>::template Call<func>(args...);    }  };};// oneflow/core/framework/tensor_consistent_id.htemplate<typename Arg0, typename Arg1, typename... Args>struct NonRecursiveInitConsistentId<Maybe<void>, Arg0, Arg1, TensorTuple*, Args...> {  template<Maybe<void> (*func)(Arg0, Arg1, TensorTuple*, Args...)>  static Maybe<void> Call(Arg0 arg0, Arg1 arg1, TensorTuple* outputs, Args... args) {    // ...    Maybe<void> ret = func(arg0, arg1, outputs, args...);    // ...    return ret;  }};// 宏开展// 去掉模版参数后就是 &WithDecorator::Decorate::Callauto* InterpretThenInitConsistentId =(&WithDecorator<NonRecursiveInitConsistentId>::Decorate<__decltype(&Interpret)>::Call<&Interpret>);Maybe<void> EagerConsistentInterpreter::ApplyImpl(const UserOpExpr& op_expr,                                                  const TensorTuple& inputs, TensorTuple* outputs,                                                  const OpExprInterpContext& ctx) const {  return InterpretThenInitConsistentId(op_expr, inputs, outputs, ctx);}

InterpretThenInitConsistentId是匿名命名空间中通过宏定义的一个函数指针，如果将其中的模版参数都去掉，简化后就是函数指针&WithDecorator::Decorate::Call。也就是说，ApplyImpl函数间接把工作转发给WithDecorator::Decorate::Call，再转发给NonRecursiveInitConsistentId::Call。函数Interpret的类型决定了其余模版参数的推断，它就是模版定义中的func，在NonRecursiveInitConsistentId::Call中理论调用的就是Interpret。NonRecursiveInitConsistentId在Interprete里面套了一层，次要做传输token等解决。

这是典型的Decorator模式，奇妙地通过精心设计的模版解决泛滥场景的逻辑解决。

WithDecorator的作用次要是将具体的Decorator与调用环境解耦，能够反对多个Decorator的组合。例如GetBoxingOutput就组合了2个装璜器。

Interpret的执行

Interpret外围代码放如下（略微调整以便于演示）：

// oneflow/core/framework/op_interpreter/eager_consistent_op_interpreter.cppMaybe<void> Interpret(const UserOpExpr& user_op_expr, const TensorTuple& inputs,                           const Symbol<Device>& default_device, TensorTuple* outputs,                           const OpExprInterpContext& ctx) {  // ...  std::shared_ptr<const ConsistentTensorInferResult> result =    JUST(user_op_expr.mut_consistent_tensor_infer_cache()->GetOrInfer(*infer_args));  // ...  const auto& kernel = JUST(user_op_expr.MutKernel4Stream(result->stream()));  // ...  JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {    return builder->LocalCallOpKernel(kernel, input_eager_blob_objects, output_eager_blob_objects,                                      ctx, stream);  }));}

MutKernel4Stream获取kernel的调用程序如下：

Interpret中的result对象是TensorInferResult，而不是outputs。它的获取过程比较复杂，咱们只关注它的stream字段。relu的inputs不是空的，result.stream会被设置为inputs[0]的device，默认就是CPU。这个stream变量的类型是oneflow::Stream。（和前面的指令Stream是不同类型）

Interpret中呈现的user_op_expr.MutKernel4Stream函数调用有几点须要阐明：

这个函数是UserOpExpr的成员办法，负责批改stream2kernel_成员变量。这个map为op保护从设施到StatefulLocalOpKernel的映射。
MutKernel4Stream函数中对UserOpExpr的map成员的批改不存在数据竞争问题。user_op_expr是保留在ReluFunctor中，而ReluFunctor保留在functional_api.yaml.cpp中的functional::Relu函数的动态op__变量中（作为lambda捕捉的一部分），op__是thread_local的，所以不存在数据竞争问题。
StatefulLocalOpKernel自身并没有实现kernel计算逻辑，它只是保留kernel的一些信息。前面会看到，在生成虚拟机指令时会调用它的ChooseOpKernel办法设置理论的kernel。
BuildOpConf基于UserOpExpr的proto为kernel生成配置，减少的字段次要是device_tag，这也是来自result.stream，对relu来说也就是inputs[0]的stream。

Interpret最终会结构一个lambda表达式并传给PhysicalRun，结构指令并提交虚拟机调度执行。

参考资料

OneFlow学习笔记：Op注册
从Functor到OpExprInterpreter
OneFlow