当前位置：网站首页>OneFlow學習筆記：從Functor到OpExprInterpreter

OneFlow學習筆記：從Functor到OpExprInterpreter

2022-04-23 08:40:00 【OneFlow深度學習框架】

撰文｜月踏

更新｜趙露陽

此前寫過的《OneFlow學習筆記：python到C++調用過程分析》，從Python代碼追到了Functor這一層，本文從Functor開始繼續往下追，後面就是OpExprInterpreter。

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函數來把具體的計算做調度，最終會在某個具體的設備上來真正進行計算，這裏面的細節太多了，本文的第三節先講一部分的內容，完整的鏈條後續會再繼續總結出來

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對象。

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<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的官方代碼和之前的一些相關文章，下面是相關鏈接：

其他人都在看