写点什么

一个算子在深度学习框架中的旅程

作者:OneFlow
  • 2022 年 6 月 14 日
  • 本文字数:11496 字

    阅读完需:约 38 分钟

一个算子在深度学习框架中的旅程

撰文|赵露阳


算子即 Operator,这里简称 op。op 是深度学习的基础操作,任意深度学习框架中都包含了数百个 op,这些 op 用于各种类型的数值、tensor 运算。


在深度学习中,通过 nn.Module 这样搭积木的方式搭建网络,而 op 就是更基础的,用于制作积木的配方和原材料。


譬如如下的一个 demo 网络:


import oneflow as torch                  class TinyModel(torch.nn.Module):def __init__(self):    super(TinyModel, self).__init__()
self.linear1 = torch.nn.Linear(100, 200) self.activation = torch.nn.ReLU() self.linear2 = torch.nn.Linear(200, 10) self.softmax = torch.nn.Softmax()
def forward(self, x): x = self.linear1(x) x = self.activation(x) x = self.linear2(x) x = self.softmax(x) return xtinymodel = TinyModel() print('The model:') print(tinymodel)
复制代码


从结构来看,这个网络是由各种 nn.Module 如 Linear、ReLU、Softmax 搭建而成,但从本质上,这些 nn.Module 则是由一个个基础 op 拼接,从而完成功能的。这其中就包含了 Matmul、Relu、Softmax 等 op。 在 OneFlow 中,对于一个已有 op,是如何完成从 Python 层->C++层的调用、流转和执行过程?本文将以


output = flow.relu(input)


为例,梳理一个 op 从 Python -> C++执行的完整过程。


首先,这里给出一个流程示意图:



下面,将分别详细从源码角度跟踪其各个环节。


1、Binding


这里,binding 是指 Python 和 C++代码的绑定。通常,我们用 Python 搭建网络,训练模型,调用函数完成各种操作。实际上,这些函数通常在 Python 层只是一层 wrapper,底层实现还是通过 C++代码完成的,那么 Python -> C++是如何调用的?这就需要用到 Python 和 C++的绑定。


在深度学习框架的实现中,即可以用 Python 原生的 C API,也可以通过 pybind11 来完成函数绑定,在 OneFlow 中,二者均有使用,譬如:


  • oneflow/api/python/framework/tensor.cpp

  • oneflow/api/python/framework/tensor_functions.cpp


中涉及到的 tensor.xxx 方法都是通过 Python C API 完成了函数绑定;


  • oneflow/core/functional/functional_api.yaml


中定义的诸多 flow.xxx 方法则是通过 pybind 实现的绑定。这里关于 Python C API 和 pybind 不做过多介绍,具体用法可以参考相应文档:


  • https://docs.python.org/zh-cn/3.8/c-api/index.html

  • https://pybind11.readthedocs.io/en/stable/index.html


下面我们回到 flow.relu 方法,我们在 Python 层调用的 flow.relu 实际是调用了在


python/oneflow/init.py


中定义的 oneflow._C.relu。 _C 表示其实现位于底层 C++。和 PyTorch 类似,我们也基于.yaml 定义了一套接口导出及 code gen 的规则,譬如在 functional_api.yaml 中,我们可以看到 Relu 的导出接口的函数签名:


- name: "relu"  signature: "Tensor (Tensor x, Bool inplace=False) => Relu"  bind_python: True
复制代码


从 yaml 定义可以看出,flow._C.relu 接收两个参数,tensor 和一个 bool 值,其绑定了 C++的 Relu 方法,函数返回值也是 tensor。实际上,在 OneFlow 编译时,会通过执行


tools/functional/generate_functional_api.py


这个文件,对 functional_api.yaml 进行解析和代码生成,动态生成 C++的.h 和.cpp 文件。


  • build/oneflow/core/functional/functional_api.yaml.h

  • build/oneflow/core/functional/functional_api.yaml.cpp


并在.cpp 文件中调用相应的 functor 完成 C++层面的函数调用。这里,还是以 flow._C.relu 为例,其对应的 functor 定义位于 oneflow/core/functional/impl/activation_functor.cpp:


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 {    ...  }
private: std::shared_ptr<OpExpr> op_;};
复制代码


ReluFunctor 通过


ONEFLOW_FUNCTION_LIBRARY(m) {  m.add_functor<impl::ReluFunctor>("Relu");  ...}
复制代码


完成 functor 的注册,注册成 functional 接口后,在 Python 层 flow._C.relu 就完成了和“Relu”的绑定。同时,这个函数在 C++中也可以通过 functional::Relu 直接调用。


2、Functor


Functor 不仅是 Python -> C++交互的核心,也是 op 调用、输入参数推导和检查的第一站。通常,各种 op 在 functor 层需要完成对输入 tensor 的 shape、dtype、维度、元素个数等各种 check,以及对 op 特有的逻辑进行解析和处理。Relu Functor 代码如下:


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 {    if (inplace) {      JUST(CheckInplaceValid(x));      std::shared_ptr<TensorTuple> outputs = std::make_shared<TensorTuple>(1);      outputs->at(0) = x;      JUST(OpInterpUtil::Dispatch(*op_, {x}, outputs.get(), AttrMap{}));      return outputs->at(0);    } else {      return OpInterpUtil::Dispatch<Tensor>(*op_, {x});    }  }
private: std::shared_ptr<OpExpr> op_;};
复制代码


可以看见,ReluFunctor 是比较简单的,其定义了一个私有变量


std::shared_ptr<OpExpr> op_;


这个 op_即需要执行的 Relu op,通过 OpBuilder 进行构建;functor 的 operator()内部,根据是否 inplace 走到 2 个不同分支,并最终通过 OpInterpUtil::Dispatch()将 op、输入 tensor 和参数派发至 Interpreter 处理。


3、Dispatch


各种 op 在 functor 中完成 check 和逻辑处理后,大多需要通过 OpInterpUtil::Dispatch() 进行派发,其目的地是 Interpreter。在 Interpreter 中,将会对 op 进行更进一步的处理。在 oneflow/core/framework/op_interpreter/op_interpreter_util.h 中,我们可以看见多种重载的 Dispatch 模板代码:


class OpInterpUtil { public:  template<typename T>  static Maybe<T> Dispatch(const OpExpr& op_expr, const TensorTuple& inputs, const AttrMap& attrs) {    return Dispatch<T>(op_expr, inputs, OpExprInterpContext(attrs));  }
template<typename T> static Maybe<T> Dispatch(const OpExpr& op_expr, const TensorTuple& inputs) { return Dispatch<T>(op_expr, inputs, OpExprInterpContext(AttrMap{})); }
template<typename T> static Maybe<T> Dispatch(const OpExpr& op_expr, const TensorTuple& inputs, const OpExprInterpContext& ctx);
static Maybe<void> Dispatch(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs, const AttrMap& attrs) { return Dispatch(op_expr, inputs, outputs, OpExprInterpContext(attrs)); }
static Maybe<void> Dispatch(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs) { return Dispatch(op_expr, inputs, outputs, OpExprInterpContext(AttrMap{})); }
static Maybe<void> Dispatch(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs, const OpExprInterpContext& ctx);
复制代码


这些重载,是为了应对不同的输入、输出以及 OpExprInterpContext 的情况。譬如这个 OpExprInterpContext 是 op 在 Interpreter 中所需的上下文,可能携带 op 计算所需要的属性(如 conv2d op 所需要的 kernel_size、padding 等)、device、sbp、parallel 等描述信息。这些重载的 Dispatch 最终都会走到:


/* static */ 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);}
复制代码


Dispatch 至此,剩下的就要交给 Interpreter 了。


4、Interpreter


Get Interpreter


这里先看看 GetInterpreter,这里其实就是获取所需的 Interpreter,来负责 op 接下来的执行。省略 check 相关的逻辑,主要代码如下:oneflow/core/framework/op_interpreter/op_interpreter_util.cpp


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;      }    } else {      if (inputs.at(0)->is_consistent()) {        ...        return g_eager_consistent_interpreter;      } else {        ...        return g_eager_mirrored_interpreter;      }    }    UNIMPLEMENTED_THEN_RETURN();  }  return g_lazy_interpreter;}
复制代码


通过上面的逻辑可以看出,Interpreter 大体上分为 Eager Interpteter 和 Lazy Interpreter;其中 Eager Interpteter 又根据 Eager Mirrored 和 Eager Consistent 有所区别。具体就是以下 3 种子类实现:


  • EagerMirroredInterpreter

  • EagerConsistentInterpreter

  • LazyInterpreter


普通的 Eager mode 下(无论是单卡还是 DDP 的情况)都会走到 EagerMirroredInterpreter 的逻辑;在普通 Eager Mode 之外,为输入 tensor 设置了 sbp、placement 则会进入到 EagerConsistentInterpreter 的逻辑;在 Lazy Mode 时(使用 nn.Graph),则会进入到 LazyInterpreter


下面,我们看下这 3 种 Interpreter 的构建:


std::shared_ptr<AutogradInterpreter> BuildEagerInterpreter(const bool& is_mirrored) {  std::shared_ptr<OpExprInterpreter> internal;  if (is_mirrored) {    internal = std::make_shared<EagerMirroredInterpreter>();  } else {    internal = std::make_shared<EagerConsistentInterpreter>();  }  return std::make_shared<AutogradInterpreter>(internal);}
std::shared_ptr<AutogradInterpreter> BuildLazyInterpreter() { auto internal = std::make_shared<LazyInterpreter>(); return std::make_shared<AutogradInterpreter>(internal);}
复制代码


可见,这 3 种 Interpreter 构建完成后,都会以私有变量 internal 的形式,参与 AutogradInterpreter 的构建,并最终返回 AutogradInterpreter


class AutogradInterpreter { public:  AutogradInterpreter() = delete;  AutogradInterpreter(const std::shared_ptr<OpExprInterpreter>& internal) : internal_(internal) {}
virtual ~AutogradInterpreter() = default;
Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs, const AttrMap& attrs) const { return Apply(op_expr, inputs, outputs, OpExprInterpContext(attrs)); }
Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs) const { return Apply(op_expr, inputs, outputs, OpExprInterpContext(AttrMap{})); }
Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs, const OpExprInterpContext& ctx) const;
private: std::shared_ptr<OpExprInterpreter> internal_;};
复制代码


Apply()


通过上面我们知道,EagerMirroredInterpreterEagerConsistentInterpreter LazyInterpreter 都将为其包裹上 AutogradInterpreter 的壳,通过 AutogradInterpreter 触发 Apply 的调用。顾名思义,AutogradInterpreter 的作用主要是和 autograd 相关,其主要为 eager mode 下前向的 op 节点插入对应的用于反向计算 grad 的节点。


我们看看这部分代码,关键部分的作用在注释里给出:


Maybe<void> AutogradInterpreter::Apply(const OpExpr& op_expr, const TensorTuple& inputs,                                       TensorTuple* outputs, const OpExprInterpContext& ctx) const {  // 判断是否需要计算梯度,如果处于GradMode的作用域切改op注册时没有禁用梯度  // 则requires_grad的值根据输入tensor的requires_grad属性判断  // any of input tensors requires_grad==True,则表示需要计算梯度  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(); });  }// 这一坨逻辑比较丑陋,是因为近期支持了oneflow系统中支持了stride&&view机制// 而大部分op尚未注册stride推导、尚未支持non-contiguous的输入tensor// 所以需要在这对这部分op的输入进行强制转换,将其变为contiguous的// NOTE: if this op not support stride, then need to tensor->contiguous()#define HANDLE_NON_CONTIGUOUS_INPUT(tensor_tuple_ptr)                                       \  TensorTuple tmp_inputs;                                                                   \  if (!LazyMode::is_enabled() && !JUST(op_expr.SupportNonContiguous())) {                   \    tmp_inputs.resize(inputs.size());                                                       \    for (size_t i = 0; i < inputs.size(); i++) { tmp_inputs[i] = inputs[i]->contiguous(); } \    tensor_tuple_ptr = &tmp_inputs;                                                         \  }
const TensorTuple* inputs_ptr = &inputs; HANDLE_NON_CONTIGUOUS_INPUT(inputs_ptr);
// 这里是进行实际Interpreter执行的主要过程 { autograd::AutoGradMode mode(false); JUST(internal_->Apply(op_expr, *inputs_ptr, outputs, ctx)); }
// 这里主要是为了eager mode下,且requires_grad==True的op, // 插入反向节点(AddNode)用于autograd,该节点包含反向梯度计算的方法(backward_fn) // 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<BackwardFunction>(); backward_fn->body = [=](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(); }; backward_fn->status = [=]() { return grad_closure->state()->SavedTensors().size() > 0; }; JUST(GetThreadLocalAutogradEngine()->AddNode(op_expr.op_type_name() + "_backward", backward_fn, *inputs_ptr, 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_ptr->size() == 0 || !requires_grad); ... if (!output->requires_grad()) { JUST(output->set_requires_grad( requires_grad && IsSupportRequireGradDataType(output->dtype()->data_type()))); } } // 捕获前向的inputs outputs,反向计算时可能用到 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_ptr, *outputs, ctx)); } return Maybe<void>::Ok();}
复制代码


上面一坨逻辑有点多,让我们看一下重点,对于简单的 Relu op,我们只需关注这部分代码:


// 这里是进行实际Interpreter执行的主要过程  {    autograd::AutoGradMode mode(false);    JUST(internal_->Apply(op_expr, *inputs_ptr, outputs, ctx));  }
复制代码


这里,还是以上面的 flow.relu 为例,由于是简单的 Eager Mode,所以实际会走到 EagerInterpreter 的 Apply 方法:


Maybe<void> EagerInterpreter::Apply(const OpExpr& op_expr, const TensorTuple& inputs,                                    TensorTuple* outputs, const OpExprInterpContext& ctx) 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); APPLY_IF(CastFromMirroredOp); APPLY_IF(ConsistentToConsistentOp); APPLY_IF(CastToConsistentOp); APPLY_IF(CastFromConsistentOp); APPLY_IF(DistributeSplitOp); APPLY_IF(DistributeCloneOp); APPLY_IF(DistributeConcatOp); APPLY_IF(DistributeAddOp); APPLY_IF(FunctionOp); APPLY_IF(SelectTopNOp)#undef APPLY_IF
OF_UNIMPLEMENTED() << "The type " << op_expr.op_type_name() << " has not been supported in EagerInterpreter::Apply.";}
复制代码


这里,通过宏定义 APPLY_IF,增加了对不同类型 op 的分支处理。对于大多数用户来说,用到的 op 都是 UserOp 类型,所以这里实际上会走到这个分支中:


if (const auto* op = dynamic_cast<const UserOpExpr*>(&op_expr)) {    return ApplyImpl(*op, inputs, outputs, ctx);  }
复制代码


再看看 EagerMirroredInterpreter::ApplyImpl,位于 oneflow/core/framework/op_interpreter/eager_mirrored_op_interpreter.cpp:


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。


NaiveInterpret


NaiveInterpret 简单来说,主要用于做以下几件事:


  • check input tensor 的 device 是否一致

  • 生成 output tensor

  • 为 output tensor 推导和检查 shape/stride/dtype

  • 构建 op 执行指令,并派发至 vm


简化版的代码如下:


Maybe<void> NaiveInterpret(const UserOpExpr& user_op_expr, const TensorTuple& inputs,                           const Symbol<Device>& default_device, TensorTuple* outputs,                           const OpExprInterpContext& ctx) {  const auto& attrs = ctx.attrs;  std::shared_ptr<EagerBlobObjectList> input_eager_blob_objects =      std::make_shared<EagerBlobObjectList>(inputs.size());  // check devices  for (int i = 0; i < inputs.size(); i++) {    const auto& input_device = JUST(inputs.at(i)->device());    if (i > 0) {      CHECK_OR_RETURN(*default_device == *input_device)          << Error::RuntimeError()          << "Expected all tensors to be on the same device, but found at least two devices, "          << default_device->ToString() << " (positional 0) and " << input_device->ToString()          << " (positional " << i << ")!";    }    input_eager_blob_objects->at(i) = JUST(inputs.at(i)->eager_blob_object());  }
// make output tensors std::shared_ptr<EagerBlobObjectList> output_eager_blob_objects = std::make_shared<EagerBlobObjectList>(outputs->size()); auto* output_tensor_metas = ThreadLocalDefaultOutputMutTensorMetas(outputs->size()); for (int i = 0; i < outputs->size(); i++) { if (!outputs->at(i)) { const auto& tensor_impl = std::make_shared<EagerMirroredTensorImpl>(); outputs->at(i) = std::make_shared<MirroredTensor>(tensor_impl); output_tensor_metas->at(i) = tensor_impl->mut_tensor_meta(); } else { bool has_eager_blob_object = JUST(outputs->at(i)->has_eager_blob_object()); CHECK_OR_RETURN(has_eager_blob_object); output_eager_blob_objects->at(i) = JUST(outputs->at(i)->eager_blob_object()); } } Symbol<Stream> stream; bool need_check_mem_case = true;
// Infer devices ...
// Infer shapes strides dtype ...
// 构建op执行指令,并派发至vm JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> { return builder->LocalCallOpKernel(kernel, input_eager_blob_objects, output_eager_blob_objects, ctx, stream); })); return Maybe<void>::Ok();}
复制代码


Interpreter 的终点是虚拟机(vm)。vm 部分,是 OneFlow 比较独特的设计,内容很多,这里暂不展开了:) 可以简单理解,派发至 vm 后,此 op 将进入一个任务执行的队列,将会等待其 vm 的调度、执行。


5、Compute


在 Interpreter 将 op 执行指令派发至 vm 后,经过调度逻辑处理后,将会在


oneflow/core/eager/opkernel_instruction_type.cpp


被触发执行,核心代码如下:


static inline void OpKernelCompute(    LocalCallOpKernelPhyInstrOperand* operand,    DeviceCtx* device_ctx, user_op::OpKernelState* state,    const user_op::OpKernelCache* cache) {
auto* opkernel = operand->mut_opkernel(); auto* compute_ctx = opkernel->UpdateComputeContext(operand->inputs().get(), operand->outputs().get(), operand->consistent_tensor_infer_result().get(), device_ctx); ... operand->user_opkernel()->Compute(compute_ctx, state, cache); opkernel->UpdateComputeContext(nullptr, nullptr, nullptr, nullptr);}
复制代码


其中,


operand->user_opkernel()->Compute(compute_ctx, state, cache);


将触发 op kernel 的实际执行。通常来说,op 的 kernel 实现根据 device 的不同,会派发到不同的实现,其一般都位于:


oneflow/user/kernels/xxx_kernel.cpp



oneflow/user/kernels/xxx_kernel.cu


这里的 Relu op 相对比较特殊,是用 primitive 实现的(primitive 也是 oneflow 中一种独特的设计,有着良好的抽象和可组合性),具体这个 UnaryPrimitive 就是 elementwise unary 的模板+UnaryFunctor 的组合。其调用链如下:



UnaryPrimitiveKernel


class UnaryPrimitiveKernel final : public user_op::OpKernel, public user_op::CudaGraphSupport { public:  OF_DISALLOW_COPY_AND_MOVE(UnaryPrimitiveKernel);  UnaryPrimitiveKernel() = default;  ~UnaryPrimitiveKernel() = default;
using PrimitiveFactoryFuncType = std::function<std::unique_ptr<ep::primitive::ElementwiseUnary>( user_op::KernelComputeContext*)>;
UnaryPrimitiveKernel(const std::string& output_name, const std::string& input_name, PrimitiveFactoryFuncType fn) : output_name_(output_name), input_name_(input_name), primitive_factory_func_(std::move(fn)) {}
private: using user_op::OpKernel::Compute; void Compute(user_op::KernelComputeContext* ctx) const override { auto primitive = primitive_factory_func_(ctx); CHECK(primitive);
const user_op::Tensor* input_tensor = ctx->Tensor4ArgNameAndIndex(input_name_, 0); ... const int64_t elem_cnt = input_shape.elem_cnt();
if (elem_cnt != 0) { primitive->Launch(ctx->stream(), input_tensor->dptr(), output_tensor->mut_dptr(), elem_cnt); } } bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
std::string output_name_; std::string input_name_; PrimitiveFactoryFuncType primitive_factory_func_;};
复制代码


ep::primitive::ElementwiseUnary


template<UnaryOp unary_op, typename Src, typename Dst>class ElementwiseUnaryImpl : public ElementwiseUnary { public:  OF_DISALLOW_COPY_AND_MOVE(ElementwiseUnaryImpl);  ElementwiseUnaryImpl(Scalar attr0, Scalar attr1) : attr0(attr0), attr1(attr1) {}  ~ElementwiseUnaryImpl() override = default;
void Launch(Stream* stream, const void* src_ptr, void* dst_ptr, size_t count) override { CpuStream* cpu_stream = stream->As<CpuStream>();
Dst* dst = reinterpret_cast<Dst*>(dst_ptr); const Src* src = reinterpret_cast<const Src*>(src_ptr); auto functor = UnaryFunctor<DeviceType::kCPU, unary_op, Dst, Src>(attr0, attr1); cpu_stream->ParallelFor(0, count, [functor, src, dst](int64_t begin, int64_t end) { for (int64_t i = begin; i < end; i++) { dst[i] = functor(src[i]); } }); }
protected: Scalar attr0, attr1;};
复制代码


UnaryFunctor


这个 UnaryFuntor 根据不同的 Unaray op 类型,特化出不同的具体 functor 实现,具体到 Relu op,其实现位于


oneflow/core/ep/common/primitive/unary_functor.h:


template<DeviceType device, typename Dst, typename Src>struct UnaryFunctor<device, UnaryOp::kRelu, Dst, Src> {  UnaryFunctor(Scalar attr0, Scalar attr1) {}
OF_DEVICE_FUNC Dst operator()(Src src) const { const Src zero_val = static_cast<Src>(0.0); if (src <= zero_val) { return static_cast<Dst>(zero_val); } else { return static_cast<Dst>(src); } }};
复制代码


至此,我们已经完成了一个 op 的 Python -> C++ 之旅。从细节上看,是相对复杂的,但从整体流程上看,其实是比较简单的,排除了 binding,vm 调度机制等细节,其主要过程其实就 4 个环节: Functor -> Dispatch -> Interpreter -> Kernel Compute。


实现/新增一个 op,通常也不需要管中间的 Dispatch 以及 Interpreter,我们只需重点关注和该 op 强相关的部分——Functor 层面的参数、op 逻辑检查,以及 Kernel Compute 部分的实际 op 运算。


(参考代码:https://github.com/Oneflow-Inc/oneflow/commit/1dbdf8faed988fa7fd1a9034a4d79d5caf18512d


其他人都在看


欢迎下载体验 OneFlow v0.7.0 最新版本:https://github.com/Oneflow-Inc/oneflow/

发布于: 刚刚阅读数: 3
用户头像

OneFlow

关注

不至于成为世界上最快的深度学习框架。 2022.03.23 加入

★ OneFlow深度学习框架:github.com/Oneflow-Inc/oneflow ★ OF云平台:oneflow.cloud

评论

发布
暂无评论
一个算子在深度学习框架中的旅程_深度学习_OneFlow_InfoQ写作社区