投稿
  1. 首页
  2. C

OneFlow学习笔记：从Functor到OpExprInterpreter

code • 2022-4-25 • C • 阅读 109

OneFlow学习笔记：从Functor到OpExprInterpreter,第1张

撰文｜月踏

更新｜赵露阳

此前写过的《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 operator()(const std::shared_ptr& x, bool inplace) const {
    ...
    return OpInterpUtil::Dispatch(*op_, {x});
  }
 private:
  std::shared_ptr 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 input_arg_tuple_;
  std::shared_ptr 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
class BuiltinOpExprImpl : public BuiltinOpExpr {
  ProtoType op_proto_;
  mutable std::shared_ptr 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 input = 2;
  map output = 3;
  map 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 {
  AttrMap base_attrs_;
  user_op::TensorDescInferFn shape_infer_fn_;
  user_op::DataTypeInferFn dtype_infer_fn_;
  user_op::DeviceInferFn device_infer_fn_;
  mutable HashMap, std::shared_ptr> device2kernel_;
  std::shared_ptr consistent_tensor_infer_cache_;


public:
  static Maybe New(const std::string& op_name, ...);
};
 
 
这些类的接口部分基本和数据结构对应，大家可以自行脑补，上面仅列出了一个UserOpExpr的静态New接口，它用来创建一个UserOpExpr对象，前面的one::OpBuilder("relu")最终就会调到这个函数来创建OpExpr对象。
 
 
3
 
 
OpExprInterpreter
 
 
简单来讲，OpExprInterpreter用来根据OpExpr的不同类型来做分发，也就是后面接不通的处理流程，这在OneFlow中被称为不同的执行模式，目前OneFlow支持的执行模式有eager和lazy，其中eager又可以被继续细分为mirror和consistent（注：OneFlow v0.7.0版本之后统称“global”），如下图所示：
 
 
 
 
图2
 
 
显而易见，上面的OpExprInterpreter总共派生出前面所说的mirror、consistent、lazy三种interpreter，除此之外，图2中还有一个标为橙色的AutogradInterpreter类，它和OpExprInterpreter之间是has-a的关系，并提供一个Appy接口来做三种执行模式的选择，下面是简化后的代码：
 
 
 
 
class AutogradInterpreter {
  std::shared_ptr internal_;
public:
  Maybe Apply(const OpExpr& op_expr, ...) const { ... }
};
 
 
我们先从文章开头贴的ReluFunctor代码中调用的OpInterpUtil::Dispatch来开始追，这里调用的Dispatch的定义在oneflow/core/framework/op_interpreter/op_interpreter_util.h，这是一系列的重载函数，可以简单把它们看作一堆helper function，不管调用的是哪个重载版本的dispatch，最终都会汇入下面这个重载版本的Dispatch中，位于oneflow/core/framework/op_interpreter/op_interpreter_util.cpp+142：
 
 
 
 
Maybe OpInterpUtil::Dispatch(
      const OpExpr& op_expr, 
      const TensorTuple& inputs,
      TensorTuple* outputs,
      const OpExprInterpContext& ctx) {
  return JUST(GetInterpreter(inputs, ctx, op_expr))->Apply(op_expr, inputs, outputs, ctx);
}
 
 
先看这里的几个参数，op_expr是前面创建的UserOpExpr类型的对象，TensorTuple可以简单认为是vector，inputs/outputs也就是相应的输入输出Tensor，OneFlow中的Tensor细节可以参考前文《Global View的相关概念和实现》中的第三节，最后一个参数是OpExprInterpContext类型，主要用于保存op的attributes信息，定义于oneflow/core/framework/op_interpreter.h+36，下面是主要的数据结构：
 
 
 
 
struct OpExprInterpContext {
  ...
  AttrMap attrs;
  Optional> device;
  Optional> parallel_desc;
  Optional> nd_sbp;
  std::shared_ptr state;
};
 
 
再继续看OpInterpUtil::Dispatch中的GetInterpreter()调用，它会根据提供的上下文信息来创建前面图2所示的AutogradInterpreter对象：
 
 
 
 
Maybe GetInterpreter(const TensorTuple& inputs, const OpExprInterpContext& ctx,
                                          const OpExpr& op_expr) {
  static const auto& g_lazy_interpreter = BuildLazyInterpreter();
  static const auto& g_eager_consistent_interpreter = BuildEagerInterpreter(/*is_mirrored=*/false);
  static const auto& g_eager_mirrored_interpreter = BuildEagerInterpreter(/*is_mirrored=*/true);
  if (!LazyMode::is_enabled()) {
    if (inputs.empty()) {
      if (ctx.parallel_desc.has_value()) {
        JUST(ctx.nd_sbp);
        CHECK_OR_RETURN(!ctx.device.has_value());
        return g_eager_consistent_interpreter;
      } else {
        CHECK_OR_RETURN(!ctx.nd_sbp.has_value());
        return g_eager_mirrored_interpreter;
      }
    }
...
 
 
再然后用创建的AutogradInterpreter对象调用了AutogradInterpreter的Apply接口来做三种执行模式的选择，它的实现位于oneflow/core/framework/op_interpreter/op_interpreter.cpp+86：
 
 
 
 
Maybe AutogradInterpreter::Apply(const OpExpr& op_expr, const TensorTuple& inputs,
                                       TensorTuple* outputs, const OpExprInterpContext& ctx) const {
  bool requires_grad = false;
  if (autograd::GradMode::is_enabled() && !JUST(op_expr.IsGradDisabled())) {
    requires_grad =
        std::any_of(inputs.begin(), inputs.end(),
                    [](const std::shared_ptr& tensor) { return tensor->requires_grad(); });
  }
  {
    autograd::AutoGradMode mode(false);
    JUST(internal_->Apply(op_expr, inputs, outputs, ctx));
  }
  // Lazy mode will construct backward compute graph in passes, so disable autograd if lazy mode.
  std::shared_ptr grad_closure(nullptr);
  if (requires_grad && !LazyMode::is_enabled()) {
    grad_closure = JUST(op_expr.GetOrCreateOpGradClosure());
    auto backward_fn =
        std::make_shared(const TensorTuple&, TensorTuple*, bool)>>(
            [=](const TensorTuple& out_grads, TensorTuple* in_grads,
                bool create_graph) -> Maybe {
              autograd::AutoGradMode mode(create_graph);
              JUST(grad_closure->Apply(out_grads, in_grads));
              return Maybe::Ok();
            });
    JUST(GetThreadLocalAutogradEngine()->AddBackwardFuncPtr(op_expr.op_type_name() + "_backward",
                                                            backward_fn, inputs, outputs));
  }
  // Update outputs autograd meta
  // Note: if requires_grad is True, we will create a new autograd meta for each output
  // in `AddBackwardFuncPtr` to support inplace operation, so the update should after
  // `AddBackwardFuncPtr`
  for (auto& output : *outputs) {
    output->set_is_leaf(inputs.size() == 0 || !requires_grad);
    if (!output->requires_grad()) {
      JUST(output->set_requires_grad(
          requires_grad && IsSupportRequireGradDataType(output->dtype()->data_type())));
    }
  }
  if (requires_grad && !LazyMode::is_enabled()) {
    // Capture inputs and outputs after `AddBackwardFuncPtr` because of that grad function
    // node has been attached to them.
    JUST(grad_closure->Capture(inputs, *outputs, ctx));
  }
  return Maybe::Ok();
}
 
 
这里主要看JUST(internal_->Apply(op_expr, inputs, outputs, ctx));（后面列的代码都和backward相关，本文只关注forward这条主线），它其实调用了所持有的OpExprInterpreter的Apply函数，下面根据前面图2中的三种Interpreter来看下后面的流程。
 
 
3.1 Mirror mode
 
 
如果我们选择的是mirror的执行模式，internal_->Apply实际会调用到EagerMirroredInterpreter的基类EagerInterpreter中的Apply，位于oneflow/core/framework/op_interpreter/op_interpreter.cpp+51：
 
 
 
 
Maybe EagerInterpreter::Apply(const OpExpr& op_expr, ...) const {
#define APPLY_IF(op_type)                                              \
  if (const auto* op = dynamic_cast(&op_expr)) { \
    return ApplyImpl(*op, inputs, outputs, ctx);                       \
  }


  APPLY_IF(UserOp);
  APPLY_IF(VariableOp);
  APPLY_IF(CastToMirroredOp);
  ...
}
 
 
这里其实又使用dynamic_cast来根据OpExpr的实际类型做了一次动态分发，也可以结合下面这张图来辅助理解：
 
 
 
 
图3
 
 
我们是从ReluFunctor中过来的，创建的是UserOpExpr，所以这里会调用EagerMirroredInterpreter中的下面这个ApplyImpl函数，位于oneflow/core/framework/op_interpreter/eager_mirrored_op_interpreter.cpp+191：
 
 
 
 
Maybe EagerMirroredInterpreter::ApplyImpl(const UserOpExpr& op_expr,
                                                const TensorTuple& inputs, TensorTuple* outputs,
                                                const OpExprInterpContext& ctx) const {
  return NaiveInterpret(op_expr, inputs, outputs, ctx);
}
 
 
这里又继续调用同一个文件中的NaiveInterpret函数，这个函数很长，主要在做进入OneFlow虚拟机之前的准备工作，其中最重要的准备工作是根据输入输出的Tensor对象来创建虚拟机需要的vm::EagerBlobObject对象，它的定义位于oneflow/core/eager/eager_blob_object.h+83，主要数据成员如下：
 
 
 
 
class EagerBlobObject final : public BlobObject {
  std::unique_ptr blob_;
  std::unique_ptr header_buffer_;
  std::shared_ptr tensor_storage_;
  std::atomic is_shape_synced_;
  int64_t storage_offset_;
  intrusive::shared_ptr compute_local_dep_object_;
};
 
 
在EagerBlobObject的数据成员中，Blob和TensorStorage维护了真正的数据存储空间，另外，从上面代码可见，EagerBlobObject也有继承关系，总结如下图：
 
 
 
 
图4
 
 
关于NaiveInterpret，内容比较多，主要是在为进入虚拟机做准备，下面展示最后一段代码，它是进入OneFlow虚拟机的入口：
 
 
 
 
Maybe NaiveInterpret(const UserOpExpr& user_op_expr, ...) {
  ...
  JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe {
    return builder->LocalCallOpKernel(
        kernel, 
        input_eager_blob_objects, 
        output_eager_blob_objects,
        ctx, 
        op_device);
  }));
  return Maybe::Ok();
}
 
 
虚拟机的内容不在本文范畴，以后抽时间继续学习。
 
 
3.2 Global mode
 
 
关于Global的概念，前文《Global View的相关概念和实现》中有详细的分析，这里就直接使用其中的概念了。如果我们选择的是Global的执行模式，internal_->Apply实际和mirror模式一样会调用到EagerInterpreter中的Apply，位于oneflow/core/framework/op_interpreter/op_interpreter.cpp+51：
 
 
 
 
Maybe EagerInterpreter::Apply(const OpExpr& op_expr, ...) const {
#define APPLY_IF(op_type)                                              \
  if (const auto* op = dynamic_cast(&op_expr)) { \
    return ApplyImpl(*op, inputs, outputs, ctx);                       \
  }


  APPLY_IF(UserOp);
  APPLY_IF(VariableOp);
  APPLY_IF(CastToMirroredOp);
  ...
}
 
 
这里使用dynamic_cast来根据OpExpr的实际类型动态（本文示例是UserOpExpr这个类型）分发到了EagerConsistentInterpreter中的ApplyImpl函数，定义位于oneflow/core/framework/op_interpreter/eager_consistent_op_interpreter.cpp+194：
 
 
 
 
Maybe EagerConsistentInterpreter::ApplyImpl(const UserOpExpr& op_expr,
                                                  const TensorTuple& inputs, TensorTuple* outputs,
                                                  const OpExprInterpContext& ctx) const {
  return InterpretThenInitConsistentId(op_expr, inputs, outputs, ctx);
}
 
 
这里InterpretThenInitConsistentId是一个函数指针，指向了用NonRecursiveInitConsistentId作为装饰器来包装Interpret这个函数的函数，简单来看下装饰器这部分代码，先看DECORATE宏，位于oneflow/core/common/decorator.h+39：
 
 
 
 
template class Decorator>
struct WithDecorator final {
  template
  struct Decorate;
  template
  struct Decorate final {
    template
    static T Call(Args... args) {
      return Decorator::template Call(args...);
    }
  };
};


#define DECORATE(fn_ptr, decorator) \
  (&WithDecorator::Decorate::Call)
 
 
其中WithDecorator算是一个装饰器包装器，Decorator是它的模板的模板类型参数，表示实际的装饰器，然后调用实际的装饰器中的Call函数，在本例中WithDecorator使用NonRecursiveInitConsistentId作为Decorator来实例化，NonRecursiveInitConsistentId定义位于oneflow/core/framework/tensor_consistent_id.h+35：
 
 
 
 
template
struct NonRecursiveInitConsistentId, Arg0, Arg1, TensorTuple*, Args...> {
  template (*func)(Arg0, Arg1, TensorTuple*, Args...)>
  static Maybe Call(Arg0 arg0, Arg1 arg1, TensorTuple* outputs, Args... args) {
    auto* recursive_depth = MutThreadLocalConsistentIdDepth();
    ++*recursive_depth;
    Maybe ret = func(arg0, arg1, outputs, args...);
    --*recursive_depth;
    if (*recursive_depth == 0 && ret.IsOk()) { JUST(InitConsistentId(outputs)); }
    return ret;
  }
};
 
 
从上面可以看出NonRecursiveInitConsistentId这个Decorator的作用是用来保证InitConsistentId只被执行一次。继续看eager模式的主线，也就是被这个装饰器所装饰的Interpret这个函数，位于oneflow/core/framework/op_interpreter/eager_consistent_op_interpreter.cpp+112，这个函数内容也稍多，总结一下主要做了下面几件事：
 
 
  •  
    创建前文《Global View的相关概念和实现》第三节中讲到的ConsistentTensorMeta信息，存于ConsistentTensorInferResult这个数据结构中
     
  •  
    为output创建相应的EagerConsistentTensorImpl和ConsistentTensor
     
  •  
    根据输入输出Tensor，创建前面图3展示的vm::EagerBlobObject对象，这些对象会在OneFlow的虚拟机中被用到，这中间可能会做boxing的 *** 作，这部分目前不太熟悉，以后熟悉了再单独总结
     
  •  
    进入虚拟机，调度并执行当前的这个op
     

 
简化过的代码如下所示：
 
 
 
 
Maybe Interpret(const UserOpExpr& user_op_expr, const TensorTuple& inputs,
                      TensorTuple* outputs, const OpExprInterpContext& ctx) {
  // step 1
  const auto& infer_args = JUST(ConsistentTensorMetaInferArgs::New(ctx.attrs, inputs));
  std::shared_ptr result =
      JUST(user_op_expr.mut_consistent_tensor_infer_cache()->GetOrInfer(*infer_args));
  const auto& output_tensor_metas = result->output_tensor_metas();
  // step 2
  for (int i = 0; i < outputs->size(); ++i) {
    if (!outputs->at(i)) {
      const auto& tensor_impl = JUST(EagerConsistentTensorImpl::New(
          output_tensor_metas.at(i), tensor_device, parallel_id, false, false));
      outputs->at(i).reset(new ConsistentTensor(tensor_impl));
    }
  }
  // step 3
  for (int i = 0; i < inputs.size(); ++i) {
    const auto& local_tensor = JUST(input->cur_rank_phy_tensor());
    input_eager_blob_objects->at(i) = JUST(local_tensor->eager_blob_object());
  }
  for (int i = 0; i < outputs->size(); ++i) {
    const auto& local_tensor = JUST(outputs->at(i)->cur_rank_phy_tensor());
    output_eager_blob_objects->at(i) = JUST(local_tensor->eager_blob_object());
  }
  // step 4
  JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe {
    return builder->LocalCallOpKernel(kernel, input_eager_blob_objects, output_eager_blob_objects,
                                      result, ctx, result->op_device());
  }));
  return Maybe::Ok();
}
 
 
这就是在进入虚拟机之前EagerMirroredInterpreter的大概工作主线。
 
 
3.3 Lazy mode
 
 
这部分目前还不熟悉，熟悉了之后再单独总结。
 
 
本文主要梳理了OpExprInterpreter的主要职责和相关实现，主要参考的是OneFlow的官方代码和之前的一些相关文章，下面是相关链接：
 
 
  •  
    https://github.com/Oneflow-Inc/oneflow
     
  •  
    C/C++杂谈：CRTP
     
  •  
    Python到C++调用过程分析
     
  •  
    Global view的相关概念和实现
     

 
其他人都在看
 
 
  •  
    资源依赖的“诅咒” 
     
  •  
    “远见者”特斯拉AI主管Karpathy
     
  •  
    我，机器学习工程师，决定跑路了
     
  •  
    对抗软件系统复杂性：恰当分层，不多不少
     
  •  
    解读Pathways（二）：向前一步是 OneFlow
     
  •  
    OneFlow v0.7.0发布：全新分布式接口，LiBai、Serving等一应俱全
     

 
欢迎下载体验OneFlow v0.7.0最新版本：
 
 
GitHub - Oneflow-Inc/oneflow: OneFlow is a performance-centered and open-source deep learning framework.https://github.com/Oneflow-Inc/oneflow/
 
 ✖					
										


					

						
欢迎分享，转载请注明来源：内存溢出
原文地址: https://outofmemory.cn/langs/717022.html

					

				

								

										
										
										
										
										
									

								

					
						
							
						 赞
						(0)
					
					

						
							 打赏
						
							
								微信扫一扫 微信扫一扫
							
							
								支付宝扫一扫 支付宝扫一扫
							
						
					

				

				

					

						

							
								code
								code
								管理员组
							
						

						

							

								
									
										
									
									0
								
								
									
										
									
									0
								
							

							
							

								
									
										
									
								
							

						

					

				

			

			
            

            	            

            
			

								

					
					

         				
							
								
							 上一篇
						
						2022-04-25
					

				

												

					
					

         				
							下一篇 
           						
							
						
						2022-04-25
					

				

							


			
			

				

					

						发表评论
					

										

						
						

							

								登录后才能评论
							

							
						

					

									

								
 评论列表（0条）

				
    
									

				
											

		


	


	







	


























    
    

    

        
            保存