当前位置：网站首页>如何在OneFlow中新增算子
如何在OneFlow中新增算子
2022-07-23 19:37:00 【OneFlow深度学习框架】
本文将以开发一个 leaky_relu（准确说是 leaky_relu_yzh op，因为 master 分支的 leaky_relu 组合了其它知识点）为例介绍如何在 OneFlow 中新增算子（https://github.com/Oneflow-Inc/oneflow/pull/8350）。
1
背景
op 与 kernel
op 与 kernel 是两个有关联的概念。op 是逻辑上的算子，包含 OneFlow Compiler 在构建计算图时所需要的必要信息，如输入、输出形状，哪些张量需要自动求导等信息。有了 op 中的信息，OneFlow Compiler 就可以构建计算图并依据计算图做资源申请、构建等操作（如根据张量的输入输出大小申请内存）， 但是 op 中不包含具体的处理数据的逻辑。
在真正需要处理数据时，OneFlow Runtime 会启动 kernel 完成计算，所以 kernel 中包含了具体处理数据的逻辑。对于一个逻辑上的 op，OneFlow Runtime 会根据数据类型、硬件设备（比如是 CPU 还是 CUDA）的具体情况，选择启动不同的 kernel。
OneFlow 中的系统 op 与 user op
定义 op 实现 kernel 计算逻辑 导出 functional 接口 实现用于求导的反向逻辑
op class 输入 input 输出 output 属性 attrs 导出并实现推导接口
def OneFlow_LeakyReluYZHOp : OneFlow_BaseOp<"leaky_relu_yzh", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {//...}
NoSideEffect 表示该算子无副作用（即不会改变内存、网络、管道、磁盘等的系统状态），这个特性可以指导某些优化操作 NoGrad 表示该算子在数学上没有梯度（不可导） CpuOnly 表示该算子只支持在 CPU 设备上执行 SupportNonContiguous 表示该算子是否支持 NonContiguous 张量（关于 Contiguous Tensor 的概念，可以参考 PyTorch Internals 中的相关内容 ）
// 一个输入 xlet input = (insOneFlow_Tensor:$x);
OneFlow_Tensor Variadic<OneFlow_Tensor>：指可变 tensor，比如 concat op，支持 concat 可变个数的 tensor。 Optional<OneFlow_Tensor>：表示这个 tensor 是可选的，既可以有也可以没有，比如 conv op 中的 add_output。
// 两个输入：a, blet input = (insOneFlow_Tensor:$a,OneFlow_Tensor:$b);
let output = (outsOneFlow_Tensor:$out0,OneFlow_Tensor:$out1);
let attrs = (insDefaultValuedAttr<F32Attr, "0.">:$rate);
let attrs = (insF32Attr:$rate);
let has_check_fn = 1; // 生成属性检查接口let has_logical_tensor_desc_infer_fn = 1; // 生成 logical shape 推导接口let has_physical_tensor_desc_infer_fn = 1; // 生成 physical shape 推导接口let has_get_sbp_fn = 1; // 生成 get sbp 接口let has_sbp_signature_infer_fn = 1; // 生成 sbp signature 推导接口，未来会移除，推荐使用 has_nd_sbp_infer_fnlet has_data_type_infer_fn = 1; // 生成 data type 推导接口let has_device_and_stream_infer_fn = 1; // 生成 device 推导接口let has_input_arg_modify_fn = 1; // 生成输入 modify 接口，比如设置 is_mutable、requires_grad（用于Lazy）等let has_output_arg_modify_fn = 1; // 生成输出 modify 接口，比如设置 is_mutable、requires_grad（用于Lazy）等let has_output_blob_time_shape_infer_fn = 1; // 生成输出 time shape 推导接口let has_nd_sbp_infer_fn = 1; // 生成 nd sbp 推导接口
let has_logical_tensor_desc_infer_fn = 1;let has_physical_tensor_desc_infer_fn = 1;let has_data_type_infer_fn = 1;let has_get_sbp_fn = 1;
op_generated.h：由解析 .td 文件生成的 op C++ 类 op_generated.cpp：由解析 .td 文件生成的 op 注册代码（包含调用 REGISTER_USER_OP 宏的代码）
Leaky ReLU CPU Kernel （https://github.com/Oneflow-Inc/oneflow/blob/7ab4b0f08c86a6f8af08b44daa510725942288fb/oneflow/user/kernels/leaky_relu_yzh_kernel.cpp） Leaky ReLU GPU KernelCPU （https://github.com/Oneflow-Inc/oneflow/blob/7ab4b0f08c86a6f8af08b44daa510725942288fb/oneflow/user/kernels/leaky_relu_yzh_kernel.cu）
template<typename T>class CpuLeakyReluYZHKernel final : public user_op::OpKernel {public:CpuLeakyReluYZHKernel() = default;~CpuLeakyReluYZHKernel() = default;private:void Compute(user_op::KernelComputeContext* ctx) const override {const user_op::Tensor* x = ctx->Tensor4ArgNameAndIndex("x", 0);user_op::Tensor* y = ctx->Tensor4ArgNameAndIndex("y", 0);const int32_t elem_cnt = x->shape().elem_cnt();const T* x_ptr = x->dptr<T>();T* y_ptr = y->mut_dptr<T>();const auto alpha = ctx->Attr<float>("alpha");FOR_RANGE(int32_t, i, 0, elem_cnt) { y_ptr[i] = x_ptr[i] > 0 ? x_ptr[i] : alpha * x_ptr[i]; }}bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }};
Compute 必须重写，在其中实现具体的运算逻辑 AlwaysComputeWhenAllOutputsEmpty 必须重写，对于绝大多数 op 而言，直接返回 false 即可。对于极少数内部需要维护状态，即使输出为空也需要调用 kernel 进行计算的 op 而言，应该返回 true
首先取得 "x"，"y" 两个 Tensor。传入Tensor4ArgNameAndIndex的字符串要和之前在OneFlowUserOps.td设置的名称一致 获取 x 的元素个数，以便后续用于 for 循环进行计算 获取属性 alpha 进入次数为 elem_cnt 的 for 循环，将结果写入
#define REGISTER_CPU_LEAKY_RELU_YZH_KERNEL(dtype) \REGISTER_USER_KERNEL("leaky_relu_yzh") \.SetCreateFn<CpuLeakyReluYZHKernel<dtype() \.SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCPU) \&& (user_op::HobDataType("y", 0) == GetDataType<dtype>::value));
op type name：为哪个 op 注册 kernel SetCreateFn<T>()：该模板方法的模板参数 T，就是我们实现的 kernel 类，OneFlow Runtime 将使用它创建 kernel 对象。 SetIsMatchedHob：因为一个 op 可能有多个 kernel，要想根据物理设备及数据格式的不同而选择不同的 kernel 进行计算，就需要调用 SetIsMatchedHob 进行设置。该方法接受一个表达式，表达式为 true 时，OneFlow 将调用该 kernel 完成计算。以上代码的匹配逻辑是：当硬件设备为 cpu，且 y 的数据类型和 dtype 一致时，选择调用注册的 kernel 类（CpuLeakyReluYZHKernel<dtype>）。
视频：CUDA 的由来（https://www.bilibili.com/video/BV1Mb4y1p7BG） 视频：CUDA 的入门小程序（https://www.bilibili.com/video/BV1bF411s76k） 视频：线程层级（https://www.bilibili.com/video/BV1MZ4y127Sq）
template<typename T>__global__ void LeakyReluForwardGpu(const int n, const float alpha, const T* x, T* y) {CUDA_1D_KERNEL_LOOP(i, n) { y[i] = x[i] > 0 ? x[i] : x[i] * alpha; }}
#define RUN_CUDA_KERNEL(func, device_ctx_ptr, thread_num, ...) \func<<<SMBlocksNum4ThreadsNum(thread_num), kCudaThreadsNumPerBlock, 0, \(device_ctx_ptr)->cuda_stream()>>>(__VA_ARGS__)
第一个参数是核函数名字 第二个参数是 device context，后续获取对应的 cuda_stream 第三个参数是要启动的线程数量，会根据线程数量来计算所需的 Block 数目。
可以看到不同设备类的 Compute 中大部分代码是重复的。一种更优的代码组织方式是用一个 .cpp 文件完成 kernel 和注册的逻辑，.cu 文件编写 GPU Kernel 函数和 GPU 模板特化的代码，.h 文件用于定义和编写注册宏。可参考 dim_gather_kernel_* （https://github.com/Oneflow-Inc/oneflow/tree/master/oneflow/user/kernels）中的代码。
OneFlow 为了适配多种设备，还提供了 Primitive 组件，可以参考：Primitive PR（ https://github.com/Oneflow-Inc/oneflow/pull/6234）
┌─────────────┐│ Module ││ (Python) │├─────────────┤│ ││ Functional │├─────────────┤│ ││ Op/Kernels ││ (C++) │└─────────────┘
实现对应的 functor 并注册 在 oneflow/core/functional/functional_api.yaml 中添加接口描述
class LeakyReluYZHFunctor {public:LeakyReluYZHFunctor() {op_ = CHECK_JUST(one::OpBuilder("leaky_relu_yzh").Input("x").Output("y").Build());}Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const float& alpha) const {MutableAttrMap attrs;JUST(attrs.SetAttr<float>("alpha", alpha));return OpInterpUtil::Dispatch<one::Tensor>(*op_, {x}, attrs);}private:std::shared_ptr<OpExpr> op_;};
在构造函数里，构造了 leaky_relu 这个op 实现 operator() 重载运算符，通过 Dispatch 调用构造好的 op，并分别传入输入，属性
m.add_functor<impl::LeakyReluYZHFunctor>("LeakyReluYZH"); // 注意最后字符串中的名字在后续的 functional_api.yaml 中会用到
在 functional_api.yaml 中添加接口描述
- name: "leaky_relu_yzh"signature: "Tensor (Tensor x, Float alpha) => LeakyReluYZH"bind_python: True
其中 name 表示导出到 Python 接口后函数的名字，比如导出后在 Python 下使用就是
flow._C.leaky_relu_yzh(...)signature 用于描述接口原型及 C++ 代码的对应关系。=> 左边的为原型；=> 右边为对应的 Functional Library 中的名字。这里LeakyRelu 和前面导出时指定的字符串是一致的。 bind_python，表示这个接口是否需要绑定 Python 接口 。比如下面的 leaky_relu_grad，我们不会在 Python 层用到（但会在 C++ 层求导使用），所以设置为 False。
import oneflow as flowimport numpy as npx_tensor = flow.Tensor(np.random.randn(3, 3))out = flow._C.leaky_relu_yzh(x_tensor, alpha=0.2)
- name: "leaky_relu_yzh_grad"signature: "Tensor (Tensor x, Tensor dy, Float alpha) => LeakyReluYZHGrad"bind_python: False实现用于求导的反向逻辑
正向过程中的输入、输出 正向过程的 attr 反向过程中上一层（正向过程中的下一层）传递过来的正向输出的梯度
为 Eager 模式注册反向
LeakyReluYZHCaptureState ：用于存储数据的结构体
struct LeakyReluYZHCaptureState : public AutoGradCaptureState {bool requires_grad; // 输入x是否需要梯度float alpha=0.0; // 输入的参数alpha};
LeakyReluYZH 类：继承自 OpExprGradFunction 的类。需要重写三个函数：Init、Capture、Apply。
class LeakyReluYZH : public OpExprGradFunction<LeakyReluYZHCaptureState> {public:Maybe<void> Init(const OpExpr& op) override {//...}Maybe<void> Capture(LeakyReluYZHCaptureState* ctx, const TensorTuple& inputs,const TensorTuple& outputs, const AttrMap& attrs) const override {//...}Maybe<void> Apply(const LeakyReluYZHCaptureState* ctx, const TensorTuple& out_grads,TensorTuple* in_grads) const override {//...}};
Init：做的是一些初始化的工作，可以根据前向 op 的配置，来初始化属性。
Maybe<void> Init(const OpExpr& op) override {const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);CHECK_NOTNULL_OR_RETURN(fw_op_expr);base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());return Maybe<void>::Ok();}
Capture：用于捕捉前向的 Tensor，属性，用于后续的求导。
Maybe<void> Capture(LeakyReluYZHCaptureState* ctx, const TensorTuple& inputs,const TensorTuple& outputs, const AttrMap& attrs) const override {CHECK_EQ_OR_RETURN(inputs.size(), 1); // 判断输入个数是否为1ctx->requires_grad = inputs.at(0)->requires_grad(); // 判断输入是否需要梯度if (!ctx->requires_grad) { return Maybe<void>::Ok(); } // 如果不需要梯度，也就不需要求导了，直接返回 Maybe<void>::Ok()ComposedAttrMap composed_attrs(attrs, base_attrs_);ctx->alpha = JUST(composed_attrs.GetAttr<float>("alpha")); // 获取 alpha，并存入 ctx->alpha 中ctx->SaveTensorForBackward(inputs.at(0)); // 调用 SaveTensorForBackward 方法，保存输入的 Tensorreturn Maybe<void>::Ok();}
Apply：实际计算梯度的函数，我们可以拿到先前的 Tensor，并调用 functional 接口下注册的 LeakyReluGrad，求得梯度，并返回
Maybe<void> Apply(const LeakyReluYZHCaptureState* ctx, const TensorTuple& out_grads,TensorTuple* in_grads) const override {CHECK_EQ_OR_RETURN(out_grads.size(), 1); // 检查梯度 Tensor 个数是否为 1in_grads->resize(1); // 因为输入只有一个，所以我们 resize(1)if (ctx->requires_grad) {const auto& x = ctx->SavedTensors().at(0); // 调用 SavedTensors 接口，拿到先前通过 SaveTensorForBackward 接口保存的 Tensorin_grads->at(0) = JUST(functional::LeakyReluYZHGrad(x, out_grads.at(0), ctx->alpha)); // 拿到 x，dy，alpha，传入给 LeakyReluYZHGrad 计算，并将梯度返回给 in_grads->at(0)}return Maybe<void>::Ok();}
REGISTER_OP_EXPR_GRAD_FUNCTION("leaky_relu_yzh", LeakyReluYZH); // 第二个参数是用于求导的类名
为 Graph 模式注册反向
REGISTER_USER_OP_GRAD("leaky_relu_yzh").SetBackwardOpConfGenFn([](user_op::BackwardOpConfContext* ctx) -> Maybe<void> {// 根据前向的 op type name，拼凑出一个 leaky_relu_yzh_grad_op_name （leaky_relu_yzh_grad）const std::string leaky_relu_yzh_grad_op_name = ctx->FwOp().op_name() + "_grad";ctx->DefineOp(leaky_relu_yzh_grad_op_name, [&ctx](user_op::BackwardOpBuilder& builder) {// 构建一个 op（op type name 为 leaky_relu_yzh_grad）// 把前向输出 y 的梯度，作为 leaky_relu_yzh_grad 的输入 dy// 把前向的 x 作为 leaky_relu_yzh_grad 的输入 x// 输出为 dx// attr alpha 同前向一样return builder.OpTypeName("leaky_relu_yzh_grad").InputBind("dy", ctx->FwOp().output_grad("y", 0)).InputBind("x", ctx->FwOp().input("x", 0)).Attr("alpha", ctx->FwOp().attr<float>("alpha")).Output("dx").Build();});// 把 leaky_relu_yzh_grad_op_name 算子的输出 dx 的结果// 绑定到前向输入 x 的反向梯度上// 即：// leaky_relu_yzh 的输入 x 的梯度 = leaky_relu_yzh_grad 的输出 dxctx->FwOp().InputGradBind(user_op::OpArg("x", 0),[&ctx, &leaky_relu_yzh_grad_op_name]() -> const std::string& {return ctx->GetOp(leaky_relu_yzh_grad_op_name).output("dx", 0);});return Maybe<void>::Ok();});3 测试与文档
本文分享自微信公众号 - OneFlow（OneFlowTechnology）。
如有侵权，请联系 [email protected] 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。
边栏推荐
- Relevant interfaces of [asp.net core] option mode
- 千呼万唤，5G双卡双通到底有多重要？
- How to reset the computer system? The method is actually very simple
- I deliberately leave a loophole in the code. Is it illegal?
- Odrive application 6 encoder
- Leetcode 151. 颠倒字符串中的单词
- Osgearth uses sundog's Triton ocean and silverlining cloud effects
- QT with OpenGL (frame cache)
- Non local mean filtering / attention mechanism
- 2022/7/21训练总结
猜你喜欢
Complex data processing of multi subsystem and multi business modules -- project development practice based on instruction set Internet of things operating system
梅科尔工作室-小熊派开发笔记3
使用Jmeter和VisualVW进行压测准备
Energy principle and variational method note 18: virtual force principle
Failure after reinstalling the system (error: Reboot and select proper boot device or insert boot media in selected boot device)
梅科尔工作室-小熊派开发笔记2
Osgearth2.8 compiling silvering cloud effect
Top ten shrinking industries in China in the first half of 2022
安装Win11找不到固态硬盘如何解决？
深入浅出边缘云 | 1. 概述
随机推荐
Parity rearrangement of Bm14 linked list
2022上半年中国十大收缩行业
Energy principle and variational method note 14: summary + problem solving
How to reset the computer system? The method is actually very simple
shell脚本中$＃、$*、[email protected]、$?、$0等含义一文搞懂
Uncover the working principle of solid state disk
梅科爾工作室-小熊派開發筆記2
【无标题】
solidworkd学习笔记：草图几何关系及编辑
Atelier macoll - notes de développement de la secte de l'ours 2
Leetcode 238. product of arrays other than itself
Osgearth uses sundog's Triton ocean and silverlining cloud effects
Energy principle and variational method note 12: minimum potential energy principle
设置ASP.NET MVC站点默认页为指定页面
Leetcode 209. 长度最小的子数组
Dokcer image understanding
Top ten shrinking industries in China in the first half of 2022
Leetcode 151. 颠倒字符串中的单词
When using polymorphism, two ideas to judge whether it can be transformed downward
Energy principle and variational method note 15: solution of differential element method