其他
如何在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 = (ins OneFlow_Tensor:$x);OneFlow_Tensor Variadic<OneFlow_Tensor>：指可变 tensor，比如 concat op，支持 concat 可变个数的 tensor。 Optional<OneFlow_Tensor>：表示这个 tensor 是可选的，既可以有也可以没有，比如 conv op 中的 add_output。
// 两个输入：a, b let input = (ins OneFlow_Tensor:$a, OneFlow_Tensor:$b );let output = (outs OneFlow_Tensor:$out0, OneFlow_Tensor:$out1); let attrs = (ins DefaultValuedAttr<F32Attr, "0.">:$rate ); let attrs = (ins F32Attr:$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_fn let 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 flow import numpy as np
x_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); // 判断输入个数是否为1 ctx->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 方法，保存输入的 Tensor return 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 个数是否为 1 in_grads->resize(1); // 因为输入只有一个，所以我们 resize(1) if (ctx->requires_grad) { const auto& x = ctx->SavedTensors().at(0); // 调用 SavedTensors 接口，拿到先前通过 SaveTensorForBackward 接口保存的 Tensor in_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); // 第二个参数是用于求导的类名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 的输出 dx ctx->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(); });