diff --git a/docs/design/phi/kernel_migrate_cn.md b/docs/design/phi/kernel_migrate_cn.md new file mode 100644 index 00000000000..4cc3ec06005 --- /dev/null +++ b/docs/design/phi/kernel_migrate_cn.md @@ -0,0 +1,328 @@ +# OpKernel 迁移指南 + +## 一、迁移背景 + +由于 Paddle fluid 的算子范式存在诸多问题，我们设计实现了新的 PHI 算子范式，后续 PHI 算子范式将逐渐统一并替代原有 fluid 算子范式，诸多执行调度的适配将仅围绕 PHI 算子范式进行（例如，新动态图执行引擎仅在 PHI 算子范式上才能充分发挥其性能优势），fluid 算子范式将逐渐被废弃并移除。详细的问题背景及设计概要请参考 [《飞桨高可复用算子库 PHI 设计文档》](./design_cn.md)。 + +因此，我们正在逐渐将 Python 2.x API 体系中运算类 API 关联的 Op 迁移至 `paddle/phi` 目录，目前我们已经基本完成了必要算子的 InferShape 函数与 CPU&GPU OpKernel 的迁移工作，本文档旨在帮助剩余异构设备 XPU/MKLDNN/NPU/MLU 等 OpKernel 的迁移。 + +## 二、迁移内容 + +迁移工作是将原先在 `paddle/fluid/operators` 目录下实现的相应异构设备或者第三方库的 OpKernel 实现，改写为 PHI 形式的函数式 Kernel，并放置到 `paddle/phi/kernels` 目录或者外部 CustomDevice repo 的 kernels 目录。具体地，以 `log_softmax` Op 为例，对于不同设备后端来讲，迁移的对象及放置位置如下： + +- XPU：`paddle/fluid/operators/log_softmax_op_xpu.cc` 修改并迁移至 `paddle/phi/kernels/xpu/log_softmax_kernel.cc & log_softmax_grad_kernel.cc` +- MKLDNN：`paddle/fluid/operators/mkldnn/`log_softmax_mkldnn_op`.cc` 修改并迁移至 `paddle/phi/kernels/onednn/log_softmax_kernel.cc` +- NPU：`paddle/fluid/operators/log_softmax_op_npu.cc` 修改并迁移至外部 CustomDevice repo `PaddleCustomDevice/backends/npu/kernels/log_softmax_kernel.cc` +- MLU：`paddle/fluid/operators/log_softmax_op_mlu.cc` 修改并迁移至外部 MLU 对应的插件式硬件适配 repo 中 + +## 三、迁移步骤 + +迁移异构 OpKernel 可以按照以下几个步骤进行： + +1. 创建 kernel.cc 文件 +2. 迁移并改写原先的 OpKernel 实现，完成注册 +3. 编译并通过单元测试 +4. 移除原 OpKernel 文件 + +迁移参考 PR： + +- where_index xpu kernel 迁移 PR：[#40255](https://github.com/PaddlePaddle/Paddle/pull/40255) +- log_softmax mkldnn kernel 迁移 PR：[#43941](https://github.com/PaddlePaddle/Paddle/pull/43941) + +### 3.1 创建 kernel.cc 文件 + +首先，根据当前迁移的 Kernel 所属的设备，确定迁移后 Kernel 文件所应该放置的子目录，例如 xpu kernel 迁移后放置到 `paddle/phi/kernels/xpu` 子目录，mkldnn kernel 迁移后放置到 `paddle/phi/kernels/onednn` 子目录。 + +然后，在相应子目录中创建 `xxx_kernel.cc` 和 `xxx_grad_kernel.cc` 文件，创建文件后在文件头部添加 License，include 必要的头文件，一般来讲至少需要 include 以下两个头文件： + +> 注意：在 phi/kernels 子目录下，前向和反向的 kernel 是分不同文件放置的，此处和迁移前不同 + +```c++ +#include "paddle/phi/kernels/log_softmax_kernel.h" + +#include "paddle/phi/core/kernel_registry.h" +``` + +一个是相应 kernel 的声明头文件，另一个是 kernel 注册所需的头文件，其他必要的头文件，按需添加即可，原则上不添加非必要的头文件。 + +最后，添加命名空间和所迁移 kernel 的函数定义，这里函数定义需要和 xxx_kernel.h 中的函数声明一致，直接从头文件中拷贝至 .cc 中即可： + +```c++ +namespace phi { + +template +void LogSoftmaxKernel(const Context& dev_ctx, + const DenseTensor& x, + int axis, + DenseTensor* out) { +} + +} // namespace phi +``` + +> 注意： +> 1. Kernel 函数定义和 xxx_kernel.h 中的函数声明需要一致，如果不一致，则当前异构 Kernel 使用了非标准的参数输入，属于不合规的行为，需要先规范 Kernel 的实现 +> 2. 如果在 phi/kernels 目录下找不到相应 Kernel 的函数定义，则暂时不需要迁移该 Kernel + +### 3.2 迁移并改写原先 OpKernel 实现，完成注册 + +完成文件创建与函数定义创建后，即可将原先的 OpKernel 函数实现拷贝至相应的 Kernel 函数体中。然后按照如下步骤完善代码： + +1. 移除 ExecutionContext 相关逻辑 +2. 替换对象类型或函数 +3. 迁移依赖的工具函数 +4. 完成注册 + +#### 3.2.1 移除 ExecutionContext 相关逻辑 + +第一步，以 `log_softmax_mkldnn_op` 为例，先将原先 LogSoftmaxMKLDNNKernel 的实现拷贝过来，删除掉其中从 ExecutionContext 获取参数的代码。 + +```c++ +template +void LogSoftmaxKernel(const Context& dev_ctx, + const DenseTensor& x, + int axis, + DenseTensor* out) { + // auto& dev_ctx = + // ctx.template device_context(); // delete this line + const auto& mkldnn_engine = dev_ctx.GetEngine(); + + // const Tensor* x = ctx.Input("X"); // delete this line + // Tensor* out = ctx.Output("Out"); // delete this line + + // int axis = ctx.Attr("axis"); // delete this line + axis = axis >= 0 ? axis : x->dims().size() + axis; + + LogSoftmaxMKLDNNHandler handler(mkldnn_engine, ctx.GetPlace(), x, axis); + + auto src_memory_p = handler.AcquireSrcMemory(x); + auto dst_memory_p = handler.AcquireDstMemory(out); + + auto logsoftmax_p = handler.AcquireForwardPrimitive(); + + auto& astream = platform::MKLDNNDeviceContext::tls().get_stream(); + logsoftmax_p->execute( + astream, + {{DNNL_ARG_SRC, *src_memory_p}, {DNNL_ARG_DST, *dst_memory_p}}); + astream.wait(); + + out->set_mem_desc(dst_memory_p->get_desc()); +} +``` + +移除时，需要确认函数传入的参数命名，和函数内实际使用的参数命名是一致的，如果不一致，需要将函数实现中的参数名替换为函数参数列表中的命名。在当前示例中，参数名刚好是一致的。 + +#### 3.2.2 替换对象类型或函数 + +第二步，需要对其他的函数实现写法进行更新。 + +首先，需要根据参数类型的变化，更新参数成员访问的方式，最主要的差别是，原先通过 ExecutionContext 取出的输入 Tensor 一般是指针，例如：`const Tensor* x = ctx.Input("X");`，但是迁移之后，输入 Tensor 统一改为了引用的写法，例如：`const DenseTensor& x`，因此参数的访问方式会有所变化。继续以 `log_softmax_mkldnn_op` 为例，这里需要改动的地方如下方注释所示： + +```c++ +template +void LogSoftmaxKernel(const Context& dev_ctx, + const DenseTensor& x, + int axis, + DenseTensor* out) { + const auto& mkldnn_engine = dev_ctx.GetEngine(); + // axis = axis >= 0 ? axis : x->dims().size() + axis; + axis = axis >= 0 ? axis : x.dims().size() + axis; + + // LogSoftmaxMKLDNNHandler handler(mkldnn_engine, ctx.GetPlace(), x, axis); + LogSoftmaxMKLDNNHandler handler(mkldnn_engine, ctx.GetPlace(), &x, axis); + + // auto src_memory_p = handler.AcquireSrcMemory(x); + auto src_memory_p = handler.AcquireSrcMemory(&x); + auto dst_memory_p = handler.AcquireDstMemory(out); + + auto logsoftmax_p = handler.AcquireForwardPrimitive(); + + auto& astream = platform::MKLDNNDeviceContext::tls().get_stream(); + logsoftmax_p->execute( + astream, + {{DNNL_ARG_SRC, *src_memory_p}, {DNNL_ARG_DST, *dst_memory_p}}); + astream.wait(); + + out->set_mem_desc(dst_memory_p->get_desc()); +} +``` + +需要将原先函数实现中部分仅在 fluid 中使用的类型或函数，替换为 PHI 中对应的类型或函数，一些替换的映射关系如下： + +| fluid 写法 | phi 写法 | +|---|---| +| `farmework::Tensor` | `DenseTensor` | +| `farmework::LoDTensor` | `DenseTensor` | +| 模板参数 `DeviceContext` | 模板参数 `Context` | +| `platform::XXXDeviceContext` | `XXXContext` | +| `out->mutbale_data(ctx.GetPlace()/place)` | `dev_ctx.template Alloc(out)` | +| `auto* ptr = out->mutbale_data()` | `auto* ptr = out->data()` | +| `out->mutbale_data(dims, place)` | `out->Resize(dims); dev_ctx.template Alloc(out)` | +| `out->mutbale_data(place, dtype)` | `dev_ctx.Alloc(out, dtype)` | +| `platform::erros::XXX` | `phi::erros::XXX` | +| `platform::float16/bfloat16/complex64/complex128` | `dtype::float16/bfloat16/complex64/complex128` | +| `framework::Eigen***` | `Eigen***` | +| `platform::XXXPlace` | `phi::XXXPlace` | +| `framework::DefaultCPUGenerator()` | `dev_ctx.GetGenerator()->GetCPUEngine()` | +| `framework::LoD` | `phi::LoD` | +| `framework::TensorCopy/TensorCopySync` | `phi::Copy` | +| `platform::is_xxx_place` | `place.GetType() == phi::AllocationType::XXX` | + +> 注意：PHI 最终会作为独立的库进行编译，服务于 fluid、infrt、自定义算子等上层组件，因此原则上 PHI 中的文件不能 include fluid 中的头文件，迁移时注意尽可能不要 include 非必要的 fluid 头文件 + +继续以 `log_softmax_mkldnn_op` 为例，这里需要将 `platform::MKLDNNDeviceContext` 替换为 `OneDNNContext` ，改动的地方如下方注释所示： + +```c++ +template +void LogSoftmaxKernel(const Context& dev_ctx, + const DenseTensor& x, + int axis, + DenseTensor* out) { + const auto& mkldnn_engine = dev_ctx.GetEngine(); + axis = axis >= 0 ? axis : x.dims().size() + axis; + + LogSoftmaxMKLDNNHandler handler( + mkldnn_engine, dev_ctx.GetPlace(), x, axis); + + auto src_memory_p = handler.AcquireSrcMemory(&x); + auto dst_memory_p = handler.AcquireDstMemory(out); + + auto logsoftmax_p = handler.AcquireForwardPrimitive(); + + // auto& astream = platform::MKLDNNDeviceContext::tls().get_stream(); + auto& astream = OneDNNContext::tls().get_stream(); + logsoftmax_p->execute( + astream, {{DNNL_ARG_SRC, *src_memory_p}, {DNNL_ARG_DST, *dst_memory_p}}); + astream.wait(); + + out->set_mem_desc(dst_memory_p->get_desc()); +} +``` + +至此，Kernel 函数迁移与改写完成。 + +#### 3.2.3 迁移依赖的函数 + +第三步，迁移 Kernel 依赖的函数。 + +迁移 Kernel 时，Kernel 调用的相关 function 及 functor 也需要一并迁移到 phi，根据所依赖 function 或者 functor 使用场景不同，可以分为以下几种情况放置： + +1. 仅有当前所迁移 Kernel 使用的辅助函数（具体到设备，比如 log_softmax 的 mkldnn Kernel），一律和 Kernel 实现放到同一个设备文件夹中 + - 如果辅助函数相关代码较少，就直接和 Kernel 实现放到同一个 .cc 文件中 + - 如果辅助函数相关代码较多，就在 Kernel 所在的设备目录创建相应的 xxx_utils.h 管理代码 +2. 其他情况，在 phi/kernels/funcs 目录下创建 .h/cc 文件管理代码，如果当前依赖的辅助函数可以直接归类到 phi/kernels/funcs 目录下已有的文件中，则直接放过去，不必创建新的文件 + +继续以 `log_softmax_mkldnn_op` 为例，它的 Kernel 使用了 LogSoftmaxMKLDNNHandler，并且只有当前 Kernel 使用，符合上述情况 1，因此将其直接迁移到 log_softmax_kernel.cc 文件中，放到 Kernel 函数上方即可。 + +```c++ +template +class LogSoftmaxMKLDNNHandler + : public paddle::platform:: + MKLDNNHandlerNoCachingT { + public: + LogSoftmaxMKLDNNHandler(const dnnl::engine mkldnn_engine, + Place cpu_place, + const DenseTensor& x, + const int axis) + : paddle::platform::MKLDNNHandlerNoCachingT( + mkldnn_engine, cpu_place) { + this->AcquireForwardPrimitiveDescriptor( + dnnl::prop_kind::forward_inference, x.mem_desc(), axis); + } +}; + +template +void LogSoftmaxKernel(const Context& dev_ctx, + const DenseTensor& x, + int axis, + DenseTensor* out) { + const auto& mkldnn_engine = dev_ctx.GetEngine(); + axis = axis >= 0 ? axis : x.dims().size() + axis; + + LogSoftmaxMKLDNNHandler handler( + mkldnn_engine, dev_ctx.GetPlace(), x, axis); + + auto src_memory_p = handler.AcquireSrcMemory(&x); + auto dst_memory_p = handler.AcquireDstMemory(out); + + auto logsoftmax_p = handler.AcquireForwardPrimitive(); + + auto& astream = OneDNNContext::tls().get_stream(); + logsoftmax_p->execute( + astream, {{DNNL_ARG_SRC, *src_memory_p}, {DNNL_ARG_DST, *dst_memory_p}}); + astream.wait(); + + out->set_mem_desc(dst_memory_p->get_desc()); +} +``` + +> 注意： +> 1. 由于 MKLDNNHandlerNoCachingT 涉及代码较多，我们暂时未将其迁移至 phi 目录，因此这里仍然 include 了 fluid/platform 目录下的头文件，后续该部分代码也需要迁移至 phi 目录 +> 2. 根据编程风格指南，输入参数一般是 const & 的形式，因此迁移 LogSoftmaxMKLDNNHandler 时将其输入参数 x 由指针改为 const & 形式，相应 LogSoftmaxKernel 内也改为直接传递 x 值的形式 + + +#### 3.2.4 完成注册 + +第四步，添加 PHI Kernel 的注册代码。 + +注册 Kernel 的方式比较简单，在 paddle 内部直接使用 PD_REGISTER_KERNEL 宏注册即可。继续以 `log_softmax_mkldnn_op` 为例，注册写法如下： + +```c++ +PD_REGISTER_KERNEL(log_softmax, + OneDNN, + ALL_LAYOUT, + phi::LogSoftmaxKernel, + float, + phi::dtype::bfloat16) {} +``` + +字段说明： +1. log_softmax: Kernel 名称，和已有 CPU&GPU 设备注册的 Kernel 名称一致（注意不一定和原先 Op 的名称一致，需要参考 `paddle/phi/kernels/cpu/log_softmax_kernel.cc|cu` 文件末尾的注册的 Kernel 名称） +2. OneDNN: Backend 名称，目前在 paddle 内部迁移异构 Kernel 会遇到的包括 XPU，OneDNN +3. phi::LogSoftmaxKernel: Kernel 函数的名称，记得带上 namespace phi +4. 剩余的均为数据类型，注册的数据类型对齐原有 Kernel 即可。 + +> 注意 +> 1. phi Kernel 的注册宏末尾是函数体 { }，不是直接加分号，此处与旧的注册宏有区别 +> 2. 注册 Kernel 的宏声明需要在 global namespace + +对于在外部 CustomDevice 注册 Kernel，注册写法略有不同，以 log_softmax_op_npu(ascend) 为例，注册写法如下： + +```c++ +PD_REGISTER_PLUGIN_KERNEL(log_softmax, + ascend, + ALL_LAYOUT, + custom_kernel::LogSoftmaxKernel, + float, + phi::dtype::float16, + double) {} +``` + +不同之处包括： +1. 注册宏的名称不一样，此处是 PD_REGISTER_PLUGIN_KERNEL +2. Backend 的名称是用户注册的 CustomDevice 的名称，此处是 ascend。 + +### 3.3 编译并通过单元测试 + +至此，一个 Kernel 的主体代码迁移工作已经完成，接下来，需要先重新 cmake & make 确认下是否有语法错误或者其他编译错误，根据错误的提示进行问题的修复。 + +PHI Kernel 的优先级高于 fluid OpKernel ，因此编译通过后，可以先通过执行相应的单元测试，验证迁移后 Kernel 的正确性，继续以 `log_softmax_mkldnn_op` 为例，这里可以验证 test_log_softmax_mkldnn_op.py 是否可以执行通过，通过意味着 Kernel 迁移正确。 + +`ctest -R test_log_softmax_mkldnn_op -V` + + +### 3.4 移除原 OpKernel 文件 + +前序步骤完成之后，需要移除原先 OpKernel 实现以及注册代码。一般来讲，可以直接删除原先 fluid operators 目录以下相应文件，继续以 `log_softmax_mkldnn_op` 为例，需要移除 `paddle/fluid/operators/mkldnn/`log_softmax_mkldnn_op`.cc` ，不维护冗余的代码。 + +> 注意，如果出现找不到符号的报错，可能需要将部分 C++单测中的 `USE_OP(fluid_op_type)` 手动改为 `USE_OP_ITSELF(fluid_op_type)` ，以及将 `USE_OP_DEVICE_KERNEL(fluid_op_type, MKLDNN)` 改为 `PD_DECLARE_KERNEL(phi_kernel_name, OneDNN, ALL_LAYOUT)`。 + + +## 四、注意事项 + +1. 通过宏 REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE 注册的 Kernel 暂时不支持迁移至 phi，后续我们会扩展机制支持 +2. Kernel 注册时 Layout 字段目前不太准确，例如对于 OneDNN 来讲，它 Kernel 的 Layout 应该是 ONEDNN，但是目前仍然声明是 ALL_LAYOUT，后续我们也会扩展这里的机制 +3. 迁移的时候，尽可能避免对原 Kernel 实现的逻辑改动，如果觉得它原来写得不好，想要优化，可以拆分 PR 进行（担心出现性能变化，CI 又发现不了，后续导致 CE 模型性能下降） +4. 按照编程风格指南，迁移后 kernel.cc 文件中对应的 kernel.h 需要在最前面 include +5. 注意迁移 Kernel 之后，DenseTensor*的输出参数，仍然是返回值的定位，所以在 kernel 内注意只能写该参数的成员值，而不能读它的值用作逻辑判断，可能会出错 +6. 异构设备 Kernel 迁移时，一些写法不确定的地方可以参考已有的 CPU&GPU Kernel