关于源码分析:OneFlow源码解析OpKernel与解释器

撰文｜郑建华
更新｜赵露阳

Op与Kernel的注册

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

注册信息保留在公有的map变量中。UserOpRegistryMgr的头文件
（https://github.com/Oneflow-In...）中定义了3个宏，REGISTER_USER_OP、REGISTER_USER_OP_GRAD、REGISTER_USER_KERNEL别离用于注册op、grad_op、kernel。

1.1 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（https://github.com/Oneflow-In...）会创立一个OpRegistry（https://github.com/Oneflow-In...）对象，这个类和UserOpRegisterTrigger（https://github.com/Oneflow-In...）类一样，只是为结构OpRegistryResult（https://github.com/Oneflow-In...）用的两头类型。

OpRegistry会暂存两头后果并在Finish中设置一些默认推导逻辑。UserOpRegisterTrigger的构造函数会调用注册逻辑。动态变量就是为了触发构造函数从而调用注册逻辑，将结构好的OpRegistryResult保留到UserOpRegistryMgr（https://github.com/Oneflow-In...）（key是op_type，如relu）。

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

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

1.2 ReluKernel的注册

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

static UserOpRegisterTrigger<OpKernelRegistry> g_register_trigger0 =  UserOpRegistryMgr::Get().    CheckAndGetOpKernelRegistry("relu").    .SetCreateFn(...)    .SetIsMatchedHob(UnaryPrimitiveExists(ep::primitive::UnaryOp::kRelu, "y", "x"))    .SetInplaceProposalFn([](const user_op::InferContext&,                             const user_op::AddInplaceArgPair& AddInplaceArgPairFn) -> Maybe<void> {      OF_RETURN_IF_ERROR(AddInplaceArgPairFn("y", 0, "x", 0, true));      return Maybe<void>::Ok();    });

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

[]() {    return user_op::NewOpKernel<UnaryPrimitiveKernel>(        "y", "x", [](user_op::KernelComputeContext* ctx) {            const user_op::TensorDesc* src = ctx->TensorDesc4ArgNameAndIndex("x", 0);            const user_op::TensorDesc* dst = ctx->TensorDesc4ArgNameAndIndex("y", 0);            return ep::primitive::NewPrimitive<ep::primitive::ElementwiseUnaryFactory>(                ctx->device_type(), ep::primitive::UnaryOp::kRelu, src->data_type(),                dst->data_type());        });}

对于relu来说，NewOpKernel就是new一个UnaryPrimitiveKernel对象并返回函数指针。最终注册的后果，会把OpKernelRegistryResult保留到UserOpRegistryMgr（key是op_type_name，如"relu"）。 1.3 Op和Kernel注册相干的类关系图

UserOpExpr的结构

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

ReluFunctor
（https://github.com/Oneflow-In...）的外围代码如下：

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
（https://github.com/Oneflow-In...）的构造函数中，次要是结构UserOpExpr（https://github.com/Oneflow-In...）。

每一个user op通过OpBuilder的Build()后，都会生成相应的UserOpExpr，用于存储属性、类型/shape/设施等推导办法，用于接下来op/kernel的理论计算。UserOpExpr蕴含以下成员：

base_attrs_
tensor_desc_infer_fn_
dtype_infer_fn_
device_and_stream_infer_fn_

它们别离用于存储该user op相干attrs属性、input/output tensor shape推导办法、数据类型data type推导办法、设施及计算流推导办法等。除了罕用的UserOpExpr、还有一些用于零碎op的BuiltinOpExpr。

OpBuilder的Input/Output调用次要是操作UserOpConf的proto对象，Build函数内会批改UserOpConf对象，比方依据OpRegistryResult::op_def补充默认值到attr。

之后结构UserOpExpr对象，UserOpConf对象被保留到UserOpExpr的父类BuiltinOpExprImpl<UserOpConf>的op_proto_字段，对于relu来说，op_proto_次要保留input, output等信息。UserOpExpr初始化时会从OpRegistryResult拷贝函数变量。

Functor的执行

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

整个链路很长，本篇笔记只以Eager Local Mode下，对次要执行流程做一些阐明。

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

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

LazyInterpreter: 用于lazy mode下的分布式动态图执行模式
EagerLocalInterpreter: 用于eager local mode本地单卡执行模式（和pytorch单卡或DDP对齐）
EagerGlobalInterpreter: 用于eager global mode，的分布式动态图执行模式

各个Interpreter的关系如下：

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

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

3.2 Apply

通过下面咱们晓得，EagerLocalInterpreter、EagerGlobalnterpreter和LazyInterpreter 都将为其包裹上AutogradInterpreter的壳，通过AutogradInterpreter触发Apply的调用。顾名思义，AutogradInterpreter的作用次要是和autograd相干，其次要为eager mode下前向的op节点插入对应的，用于反向计算grad的节点。

上面以最罕用的（Eager Mode）模式，解说Apply的执行办法。在Eager Mode（无论是eager local还是eager consistent）模式下，理论都会走到EagerInterpreter的Apply（https://github.com/Oneflow-In...）办法：

Maybe<void> EagerInterpreter::Apply(const OpExpr& op_expr, const TensorTuple& inputs,                                    TensorTuple* outputs, const OpExprInterpContext& ctx) const {#define APPLY_IF(op_type)                                              \  if (const auto* op = dynamic_cast<const op_type##Expr*>(&op_expr)) { \    return ApplyImpl(*op, inputs, outputs, ctx);                       \  }  APPLY_IF(UserOp);  APPLY_IF(VariableOp);  APPLY_IF(CastToLocalOp);  APPLY_IF(CastFromLocalOp);  APPLY_IF(GlobalToGlobalOp);  APPLY_IF(CastToGlobalOp);  APPLY_IF(CastFromGlobalOp);  APPLY_IF(DistributeSplitOp);  APPLY_IF(DistributeCloneOp);  APPLY_IF(DistributeConcatOp);  APPLY_IF(DistributeAddOp);  APPLY_IF(FunctionOp);  APPLY_IF(SelectTopNOp)#undef APPLY_IF  OF_UNIMPLEMENTED() << "The type " << op_expr.op_type_name()                     << " has not been supported in EagerInterpreter::Apply.";}

这里通过宏定义APPLY_IF，减少了对不同类型op的分支解决，将op_expr dynamic_cast成相应子类op实现的Expr，如对于大多数用户来说，用到的op都是UserOp类型，所以这里实际上会走到这个分支中：

if (const auto* op = dynamic_cast<const UserOpExpr*>(&op_expr)) {    return ApplyImpl(*op, inputs, outputs, ctx);}

再看看EagerLocalInterpreter::ApplyImpl（https://github.com/Oneflow-In...）：

Maybe<void> EagerLocalInterpreter::ApplyImpl(const UserOpExpr& op_expr, const TensorTuple& inputs,                                             TensorTuple* outputs,                                             const OpExprInterpContext& ctx) const {  return NaiveInterpret(op_expr, inputs, outputs, ctx);}

其最终实现是NaiveInterpret（https://github.com/Oneflow-In...）。

3.3 NaiveInterpret

NaiveInterpret简略来说，次要用于做以下四件事：

check input tensor的device是否统一
生成output tensor
为output tensor推导和查看shape/stride/dtype
构建op执行指令，并派发至vm

简化版的代码如下：

Maybe<void> NaiveInterpret(const UserOpExpr& user_op_expr, const TensorTuple& inputs,                           const Symbol<Device>& default_device, TensorTuple* outputs,                           const OpExprInterpContext& ctx) {  const auto& attrs = ctx.attrs;  // 查看input tensor是否位于雷同device上  ...        // 推导outout tensor的设施类型  // Infer devices  if (!user_op_expr.has_device_and_stream_infer_fn()) {    stream = JUST(GetDefaultStreamByDevice(default_device));    for (int i = 0; i < outputs->size(); i++) {      auto* tensor_impl = JUST(TensorImpl4Tensor(outputs->at(i)));      *JUST(tensor_impl->mut_device()) = default_device;    }  } else {    need_check_mem_case = false;    stream = JUST(user_op_expr.InferDeviceAndStream(attrs, inputs, outputs));  }  // 推导outout tensor的形态、数据类型  // Infer shapes and dtypes  const auto& device_tag = stream->device()->type();  JUST(user_op_expr.InferPhysicalTensorDesc(      attrs, device_tag,      [&](int32_t i) -> const TensorMeta* {        return CHECK_JUST(TensorImpl4Tensor(inputs[i]))->mut_tensor_meta();      },      [&](int32_t i) -> TensorMeta* {        // using thread_local TensorMeta pointer if inplace.        // using tensor_impl TensorMeta pointer if not inplace.        return output_tensor_metas->at(i);      }));  // 为output tensor初始化eager_blob_object  for (int i = 0; i < output_eager_blob_objects->size(); i++) {    auto* tensor_impl = JUST(TensorImpl4Tensor(outputs->at(i)));    if (!output_eager_blob_objects->at(i)) {      if (!JUST(user_op_expr.SupportNonContiguous())) {        std::shared_ptr<Stride> stride(new Stride(*tensor_impl->shape()));        tensor_impl->mut_tensor_meta()->set_stride(stride);      }      const auto& dep_object = NewLocalDepObject();      JUST(tensor_impl->InitEagerBlobObject(dep_object));      output_eager_blob_objects->at(i) = JUST(tensor_impl->eager_blob_object());    } else {      // output i is inplaced.      // check thread_local TensorMeta and tensor_impl TensorMeta.      CHECK_OR_RETURN(tensor_impl->tensor_meta()->shape() == output_tensor_metas->at(i)->shape());      CHECK_OR_RETURN(tensor_impl->tensor_meta()->dtype() == output_tensor_metas->at(i)->dtype());    }  }  // 从user_op_expr中取出kernel  const auto& kernel = JUST(user_op_expr.MutKernel4Stream(stream));  kernel->set_need_check_mem_case(need_check_mem_case);  for (int64_t index : kernel->output_tuple_indexes4mut2_obns()) {    output_eager_blob_objects->at(index)->set_is_shape_synced(false);  }  // kernel dispatch至VM，期待后续理论的调度执行  JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {    return builder->Call(kernel, input_eager_blob_objects, output_eager_blob_objects, ctx, stream);  }));  return Maybe<void>::Ok();}

PhysicalRun承受一个lambda functor作为参数，这里即InstructionsBuilder->Call办法，该办法承受kernel、input/output的eager blob object、kernel执行的上下文作为参数。Call办法理论会实现OpCall指令的构建，并最终将其派发至vm指令列表中，期待VM理论调度执行。

参考资料

OneFlow学习笔记：Op注册
（https://mp.weixin.qq.com/s/eF...）
从Functor到OpExprInterpreter
https://github.com/Oneflow-In...
https://zhuanlan.zhihu.com/p/...

（本文经受权后公布，原文https://segmentfault.com/a/11...）

欢送下载体验 OneFlow v0.8.0 最新版本：
https://github.com/Oneflow-In...