关于机器学习: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.h
template<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.h
template<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::Call
auto* 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.cpp
Maybe<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