OneFlow学习笔记：从Functor到OpExprInterpreter

撰文｜月踏

更新｜赵露阳

此前写过的《OneFlow学习笔记：python到C++调用过程分析》，从Python代码追到了Functor这一层，本文从Functor开始继续往下追，后面就是OpExprInterpreter。

1

Functor回顾

Functor层作为OneFlow的基础设施，为Python端和C++端提供了op操作的统一入口，这在《python到C++调用过程分析》中有详细分析，其中使用了Relu作为示例，这是为了尽可能的减小理解成本，本文继续以Relu作为示例来往下追代码，前文已经列过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 {
    ...
    return OpInterpUtil::Dispatch<Tensor>(*op_, {x});
  }
 private:
  std::shared_ptr<OpExpr> op_;
};

代码很简单，可以分成三部分来看：

• 定义了数据结构：也就是类成员变量op_，它是OpExpr类型，这是下面第二节主要讲的部分

• 构造函数：使用OpBuilder这个辅助类对op_进行了初始化，主要还是在最后调用Build()的时候，内部调用了第二节讲到的UserOpExpr中的静态函数New来进行创建

• 函数调用运算符重载函数：这里通过一个Dispatch函数来把具体的计算做调度，最终会在某个具体的设备上来真正进行计算，这里面的细节太多了，本文的第三节先讲一部分的内容，完整的链条后续会再继续总结出来

2

OpExpr

算子在OneFlow的框架中用OpExpr来抽象表示，除了表示算子之外，它还可以表示一些其它的操作，先看一下OpExpr的继承体系：

图1

算子所对应的OpExpr一般是上面图1中的橙色继承链条底端的UserOpExpr，代码定义位于oneflow/core/framework/op_expr.h，其它的这些OpExpr我目前也了解很少，以后有所了解之后再做总结，在橙色的继承链条中，每一个类的主要数据结构如下所述：

1.OpExpr是虚基类，无数据成员

2.BuiltinOpExpr是一个比较高层且重要的基类，主要维护了op_name、input arg、output arg信息：

class BuiltinOpExpr : public OpExpr {
  std::string op_name_;
  std::shared_ptr<const ArgTuple> input_arg_tuple_;
  std::shared_ptr<const ArgTuple> output_arg_tuple_;
};

3.BuiltinOpExprImpl主要维护了op proto和grad func的信息，子类通过前文《C/C++杂谈：CRTP》介绍过的CRTP的方式来使用这个类，主要是为了复用接口，这里的模板参数类型主要是由proto文件生成的类型，这也是这里叫做ProtoType的原因，以图1中的橙色继承链条为例，使用的UserOpConf来做的实例化，它是由oneflow/core/framework/user_op_conf.proto自动生成的一个数据结构，下面一同展示一下BuiltinOpExprImpl和user_op_conf.proto的主要内容：

template<typename ProtoType>
class BuiltinOpExprImpl : public BuiltinOpExpr {
  ProtoType op_proto_;
  mutable std::shared_ptr<OpExprGradFunctionIf> op_grad_func_;
};


// oneflow/core/framework/user_op_conf.proto
message UserOpConf {
  message ListString { repeated string s = 1; }
  required string op_type_name = 1;
  map<string, ListString> input = 2;
  map<string, ListString> output = 3;
  map<string, AttrValue> attr = 4;
  repeated string input_order = 5;
  repeated string output_order = 6;
}

4.最后是UserOpExpr，它维护了一些op的attrs、shape的infer function、dtype的infer function等信息：

class UserOpExpr final : public BuiltinOpExprImpl<UserOpConf> {
  AttrMap base_attrs_;
  user_op::TensorDescInferFn shape_infer_fn_;
  user_op::DataTypeInferFn dtype_infer_fn_;
  user_op::DeviceInferFn device_infer_fn_;
  mutable HashMap<Symbol<Device>, std::shared_ptr<StatefulLocalOpKernel>> device2kernel_;
  std::shared_ptr<ConsistentTensorInferCache> consistent_tensor_infer_cache_;


public:
  static Maybe<UserOpExpr> New(const std::string& op_name,