Press "Enter" to skip to content

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, ...);
};

 

这些类的接口部分基本和数据结构对应,大家可以自行脑补,上面仅列出了一个UserOpExpr的静态New接口,它用来创建一个UserOpExpr对象,前面的 one:: OpBuilder( “relu” ) 最终就会调到这个函数来创建OpExpr对象。

 

3

 

OpExprInterpreter

 

简单来讲,OpExprInterpreter用来根据OpExpr的不同类型来做分发,也就是后面接不通的处理流程,这在OneFlow中被称为不同的执行模式,目前OneFlow支持的执行模式有eager和lazy,其中eager又可以被继续细分为mirror和consistent( 注:OneFlow v 0 .7. 0版本之后统称“global” ) ,如下图所示:

 

 

图2

 

显而易见,上面的OpExprInterpreter总共派生出前面所说的mirror、consistent、lazy三种interpreter,除此之外,图2中还有一个标为橙色的AutogradInterpreter类,它和OpExprInterpreter之间是has-a的关系,并提供一个Appy接口来做三种执行模式的选择,下面是简化后的代码:

 

class AutogradInterpreter {
std::shared_ptr<OpExprInterpreter> internal_;
public:
Maybe<void> 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<void> 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<Tensor>,inputs/outputs也就是相应的输入输出Tensor,OneFlow中的Tensor细节可以参考前文《 Global View的相关概念和实现 》中的第三节,最后一个参数是OpExprInterpContext类型,主要用于保存op的attributes信息,定义于oneflow/core/framework/op_interpreter.h+36,下面是主要的数据结构:

 

struct OpExprInterpContext {
...
AttrMap attrs;
Optional<Symbol<Device>> device;
Optional<Symbol<ParallelDesc>> parallel_desc;
Optional<Symbol<cfg::NdSbp>> nd_sbp;
std::shared_ptr<user_op::OpKernelState> state;
};

 

再继续看OpInterpUtil::Dispatch中的GetInterpreter()调用,它会根据提供的上下文信息来创建前面图2所示的AutogradInterpreter对象:

 

Maybe<AutogradInterpreter> 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<void> 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>& 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<OpExprGradClosure> grad_closure(nullptr);
if (requires_grad && !LazyMode::is_enabled()) {
grad_closure = JUST(op_expr.GetOrCreateOpGradClosure());
auto backward_fn =
std::make_shared<std::function<Maybe<void>(const TensorTuple&, TensorTuple*, bool)>>(
[=](const TensorTuple& out_grads, TensorTuple* in_grads,
bool create_graph) -> Maybe<void> {
autograd::AutoGradMode mode(create_graph);
JUST(grad_closure->Apply(out_grads, in_grads));
return Maybe<void>::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<void>::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<void> EagerInterpreter::Apply(const OpExpr& op_expr, ...) 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(CastToMirroredOp);
...
}

 

这里其实又使用dynamic_cast来根据OpExpr的实际类型做了一次动态分发,也可以结合下面这张图来辅助理解:

 

 

图3

 

我们是从ReluFunctor中过来的,创建的是UserOpExpr,所以这里会调用EagerMirroredInterpreter中的下面这个ApplyImpl函数,位于oneflow/core/framework/op_interpreter/eager_mirrored_op_interpreter.cpp+191:

 

Maybe<void> 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> blob_;
std::unique_ptr<char[]> header_buffer_;
std::shared_ptr<TensorStorage> tensor_storage_;
std::atomic<bool> is_shape_synced_;
int64_t storage_offset_;
intrusive::shared_ptr<LocalDepObject> compute_local_dep_object_;
};

 

在EagerBlobObject的数据成员中,Blob和TensorStorage维护了真正的数据存储空间,另外,从上面代码可见,EagerBlobObject也有继承关系,总结如下图:

 

 

图4

 

关于NaiveInterpret,内容比较多,主要是在为进入虚拟机做准备,下面展示最后一段代码,它是进入OneFlow虚拟机的入口:

 

Maybe<void> NaiveInterpret(const UserOpExpr& user_op_expr, ...) {
...
JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
return builder->LocalCallOpKernel(
kernel,
input_eager_blob_objects,
output_eager_blob_objects,
ctx,
op_device);
}));
return Maybe<void>::Ok();
}

 

虚拟机的内容不在本文范畴,以后抽时间继续学习。

 

3.2 Global mode

 

关于Global的概念,前文《 Global View的相关概念和实现 》中有详细的分析,这里就直接使用其中的概念了。如果我们选择的是Global的执行模式, internal_->Apply 实际和mirror模式一样会调用到EagerInterpreter中的Apply,位于oneflow/core/framework/op_interpreter/op_interpreter.cpp+51:

 

Maybe<void> EagerInterpreter::Apply(const OpExpr& op_expr, ...) 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(CastToMirroredOp);
...
}

 

这里使用dynamic_cast来根据OpExpr的实际类型动态(本文示例是UserOpExpr这个类型)分发到了EagerConsistentInterpreter中的ApplyImpl函数,定义位于oneflow/core/framework/op_interpreter/eager_consistent_op_interpreter.cpp+194:

 

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是一个函数指针,指向了用NonRecursiveInitConsistentId作为装饰器来包装Interpret这个函数的函数,简单来看下装饰器这部分代码,先看DECORATE宏,位于oneflow/core/common/decorator.h+39:

 

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...);
}
};
};
#define DECORATE(fn_ptr, decorator) \
(&WithDecorator<decorator>::Decorate<decltype(fn_ptr)>::Call<fn_ptr>)

 

其中WithDecorator算是一个装饰器包装器,Decorator是它的模板的模板类型参数,表示实际的装饰器,然后调用实际的装饰器中的Call函数,在本例中WithDecorator使用NonRecursiveInitConsistentId作为Decorator来实例化,NonRecursiveInitConsistentId定义位于oneflow/core/framework/tensor_consistent_id.h+35:

 

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) {
auto* recursive_depth = MutThreadLocalConsistentIdDepth();
++*recursive_depth;
Maybe<void> 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<void> 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<const ConsistentTensorInferResult> 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<void> {
return builder->LocalCallOpKernel(kernel, input_eager_blob_objects, output_eager_blob_objects,
result, ctx, result->op_device());
}));
return Maybe<void>::Ok();
}

 

这就是在进入虚拟机之前EagerMirroredInterpreter的大概工作主线。

 

3.3 Lazy mode

 

这部分目前还不熟悉,熟悉了之后再单独总结。

 

本文主要梳理了OpExprInterpreter的主要职责和相关实现,主要参考的是OneFlow的官方代码和之前的一些相关文章,下面是相关链接:

 

https://github.com/Oneflow-Inc/oneflow

 

C/C++杂谈:CRTP

 

Python到C++调用过程分析

 

Global view的相关概念和实现

 

其他人都在看

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。