ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Paddle-Lite x86 JIT Kernel 架构解析:从函数模板到运行时生成的端侧高性能算子体系

Paddle-Lite x86 JIT Kernel 架构解析:从函数模板到运行时生成的端侧高性能算子体系 Paddle-Lite x86 JIT Kernel 架构解析从函数模板到运行时生成的端侧高性能算子体系【免费下载链接】Paddle-LitePaddlePaddle High Performance Deep Learning Inference Engine for Mobile and Edge (飞桨高性能深度学习端侧推理引擎项目地址: https://gitcode.com/GitHub_Trending/pa/Paddle-LitePaddle-Lite 的lite/backends/x86/jit/目录实现了一套基于函数模板 JITJust In Time代码生成的轻量级 Kernel 体系用于在 x86 CPU 上以极致性能执行向量级运算与复杂逻辑如 LSTM、GRU、SeqPool 等组合算子。本文将以仓库文档 lite/backends/x86/jit/README.md 为主体结合该目录下的 kernel_base.h、registry.h、helper.h 等源码系统讲解其设计思想、动态获取与缓存机制、测试与基准要求并给出新增算子的完整落地步骤。读完本文你将能够理解并复用这套多实现 自动选择 缓存的算子框架为 Paddle-Lite 扩展高性能算子提供可直接落地的路径。一、什么是 JIT Kernel比 Operator 更小的性能单元在 Paddle-Lite 的语境下一个 Operator算子是框架层的计算单元而这里的JIT Kernel 是比 Operator 中 Kernel 更小级别的算子单元更侧重在不同硬件上的性能表现。它的核心特征如下多种实现并存同一个 Kernel 可以有多种第三方库实现如 MKL、MKL-DNN、OpenBLAS、Intrinsic 代码等每种实现通过自身的CanBeUsed函数声明在什么条件下可以被调用见 kernel_base.h 中的KernelMore::CanBeUsed。粒度可粗可细既可以是非常细粒度的函数例如VMul向量乘法、VAdd向量加法也可以是复杂逻辑例如 LSTM、GRU 等复杂逻辑可以完全由底层函数拼接而成。面向 CPU 高性能目前该体系仅支持 CPU 上的高性能计算文档明确说明Currently its only supported on CPU yet。从源码结构看KernelType枚举集中定义了当前支持的全部 Kernel 类别见 kernel_base.h包括kCRFDecoding、kEmbSeqPool、kGRUH1、kGRUHtPart1、kGRUHtPart2、kHSum、kHMax、kLSTMCtHt、kLSTMC1H1、kLayerNorm、kMatMul、kNCHW16CMulNC、kSeqPool、kSoftmax、kStrideASum、kStrideScal、kVAdd、kVAddBias、kVAddRelu、kVBroadcast、kVCopy、kVExp、kVIdentity、kVMul、kVRelu、kVScal、kSgd、kVSigmoid、kVSquare、kVSub、kVTanh等 30 余种基本覆盖了向量运算、规约、激活、RNN 单元、矩阵乘、SGD 更新等常见计算模式。二、目录结构与三类实现refer / gen / more基本类的定义都放在lite/backends/x86/jit/根目录下根目录下包含gen、more、refer三个子目录。每种 Kernel 算子都必须有 reference 实现作为单元测试的基准其他实现都是可选的。Paddle-Lite/lite/backends/x86/jit/ ├── ... ├── gen/ # 基于 Xbyak 的 JIT 生成代码 │ ├── act.cc/h # 激活类VRelu/VExp/VSigmoid/VTanh 等 │ ├── blas.cc/h # BLAS 相关 │ ├── embseqpool.cc/h │ ├── gru.cc/h # GRU H1 / HtPart1 / HtPart2 │ ├── hopv.cc/h # 水平规约 HSum / HMax / StrideASum │ ├── jitcode.h # JitCode 基类 │ ├── lstm.cc/h # LSTM CtHt / C1H1 │ ├── matmul.cc/h │ ├── seqpool.cc/h │ ├── sgd.cc/h │ └── vbroadcast.cc/h ├── more/ # 其他第三方库或混合实现 │ ├── intrinsic/ # SSE/AVX Intrinsic 实现crf_decoding、layer_norm │ ├── mix/ # 混合组合实现mix.cc/h │ └── mkl/ # 依赖 Intel MKL 的实现 └── refer/ # reference 参考实现refer.cc/h三类实现各自的定位文档原文要点gen代表使用 JIT 生成的 code需要依赖 Xbyak 库。该实现最关心的就是性能。从 CMakeLists.txt 可以看到gen子目录仅在WITH_XBYAK AND NOT APPLE AND NOT WIN32条件下被编译即目前主要在 Linux 平台启用 Xbyak 动态生成机器码。refer代表 reference 实现。每种 Kernel 算子都必须有在 CPU 上的 reference 实现它只关心算法逻辑的正确性不依赖任何第三方库所有其他实现的逻辑都要与它保持一致。more可以放入更多实现包括mkl、mkldnn、intrinsic、openblas等也可以是自身已有的 Kernel 组合mix。在more目录中当前仓库实际落地了三种实现来源见 more/CMakeLists.txtmklMKL 库实现、mix已有 Kernel 的组合实现、intrinsic如 crf_decoding.cc、layer_norm.cc 等手写 Intrinsic 优化。原文档目录树中的mkldnn与openblas属于可扩展位置的示意当前仓库未放置对应实现目录。目录结构的源码印证注册与选择顺序在 helper.h 的GetAllCandidateKernels实现中候选 Kernel 的收集顺序为jitcode more refer先尝试获取 JIT 生成代码仅float数据类型支持见GetJitCode的std::enable_if特化再遍历more实现池中CanBeUsed(attr)返回 true 的实现最后兜底加入 reference 实现CHECK(ref ! nullptr) Refer Kernel can not be empty.强制要求 refer 必须存在。这正是每种 kernel 都需要有 reference 实现这一设计在代码层面的强制约束。三、动态获取机制四个核心接口框架为上层调用者提供了四种获取函数指针的途径覆盖全量枚举默认最优缓存复用逻辑基准四种场景接口返回内容适用场景GetAllCandidateFuncs根据输入的 Kernel 类别获取满足要求的所有函数实现需要遍历全部实现按具体输入属性动态测试挑选最优GetDefaultBestFunc返回一个默认最优的函数实现该函数基于一些通用配置离线 tuning 后得到能覆盖大多数情况下的最优结果KernelFuncs::Cache()返回默认最优函数同时缓存函数指针属性一致时直接返回上次指针否则按属性新建并缓存GetReferFunc返回该 Kernel 最原始的逻辑函数与输入大小和属性无关每个 Kernel 有且只有一个 CPU 上的 refer 实现关键设计说明所有实现保证结果一致但速度不一致GetAllCandidateFuncs返回的每个实现逻辑等价因此可以根据具体输入属性大小做运行时 benchmark手动选择最优函数。默认最优的离线 tuning在 helper.h 的GetDefaultBestFunc中当前实现是直接返回候选列表的第一个return funcs[0]注释明确说明此处后续可做运行时 benchmark 返回最优当前第一个即最优的顺序是通过离线 tuning 保证的——即注册顺序本身就是按性能优先级排列的。缓存命中KernelFuncs::Cache()使用LITE_THREAD_LOCAL的线程局部单例以JitCodeKey(attr)作为 key用std::mapint64_t, func_type缓存函数指针operator[]与At(attr)等价可直接用下标访问。refer 的唯一性GetReferFunc不关心 attr直接查ReferKernelPool得到唯一的 CPU reference 函数它表征了该 Kernel 的原始逻辑。调用示例一从缓存中直接获取默认最优函数所有 Kernel 的调用只需在头文件中包含lite/backends/x86/jit/kernels.h该文件是编译时自动生成的由 CMakeLists.txt 通过file(WRITE/APPEND)在${PADDLE_BINARY_DIR}/lite/backends/x86/jit/kernels.h生成并注入了helper.h与registry.h的 include。using T float; jit::seq_pool_attr_t attr(width, jit::SeqPoolType::kSum); auto seqpool_func jit::KernelFuncsjit::SeqPoolTupleT, platform::CPUPlace::Cache().At(attr); seqpool_func(src_data, dst_data, attr);这里seq_pool_attr_t的定义在 kernel_base.h包含hheight构造时默认 1、wwidth、typeSeqPoolType取值kSum/kAvg/kSqrt构造参数attr(width, jit::SeqPoolType::kSum)即按宽度与池化类型构建属性。调用示例二跑一遍所有实现并输出实现类别using T float; jit::seq_pool_attr_t attr(width, jit::SeqPoolType::kSum); auto funcs jit::GetAllCandidateFuncsWithTypesjit::SeqPoolTupleT, platform::CPUPlace(attr); for (auto f : funcs) { LOG(INFO) Kernel implementation type: f.first; f.second(src_data, dst_data, attr); }GetAllCandidateFuncsWithTypes返回std::vectorstd::pairstd::string, func_type其中f.first是实现类别名称——JIT 生成的实现其ImplType()返回JitCode见 gen_base.hreference 实现返回Refer见 kernel_base.h 中ReferKernel::ImplType其余 more 实现返回各自注册的名称。属性类型与 KernelTuple 的对应关系每个 Kernel 都由一个KernelTuple描述其数据类型 属性类型 函数签名三要素并通过DECLARE_KERNELTUPLE宏与KernelType一一对应见 kernel_base.h。以文档中反复出现的SeqPoolTuple为例template typename T struct SeqPoolTuple { static constexpr KernelType kernel_type kSeqPool; typedef T data_type; typedef seq_pool_attr_t attr_type; typedef void (*func_type)(const T*, T*, const seq_pool_attr_t*); };仓库中还预置了多种通用元组模板以简化定义XYZNTuplex,y,z,n 四参、AXYNTuplea,x,y,n、AXYNSTuplea,x,y,n,stride、XYNTuplex,y,n、XRNTuplex,返回值,n、XRNSTuplex,返回值,n,stride以及自定义结构的LSTMTuple、GRUTuple、MatMulTuple、SgdTuple、EmbSeqPoolTuple、LayerNormTuple、SoftmaxTuple等。其中 LSTM 与 GRU 的属性lstm_attr_t/gru_attr_t支持配置维度d、门控激活类型act_gate、候选激活act_candLSTM 还额外支持 cell 激活act_cell与 peepholeuse_peephole开关。四、底层三大池JitCode / Kernel / Refer 的分池管理从 kernel_pool.h 的源码可以清晰看到JIT 体系使用三个独立的单例池来管理不同性质的实现JitCodePoolKT按KernelType模板参数实例化以int64_t由JitCodeKey(attr)计算得到为 key存储std::unique_ptrGenBase即已生成的机器码对象提供Has/Insert/AllKernels接口且为LITE_THREAD_LOCAL线程局部实例保证多线程下互不干扰。KernelPool以KernelKey(KernelType, Place)为 key存储std::vectorKernelPtr管理所有more类实现。ReferKernelPool结构与KernelPool相同但独立成池专门存放 reference 实现——注释明确说明Every kernel should have refer code and it should be used in unit tests, so refer kernels should have its own independent kernel pool即 refer 独立成池是为了保证单测基准不被其他实现污染。此外还有一个JitCodeCreatorPool存储GenCreatorJitCodeCreatorAttr负责按 attr 条件创建 JIT 代码。KernelKey的哈希由type_ 8 | place组合计算见 kernel_key.h。JIT 代码的生成流程见 helper.h 的GetJitCode可概括为计算JitCodeKey(attr)→ 查JitCodePool是否已有缓存 → 若无则到JitCodeCreatorPool中按KernelKey找到 creator 列表 → 遍历调用CanBeUsed(attr)与CreateJitCode(attr)生成代码并回填缓存。同一属性只生成一次机器码后续直接复用。五、测试与基准正确性与性能的双重保障文档明确要求所有 JIT Kernel 必须通过两类测试逻辑测试Unit Test所有实现都要与 refer 代码对比满足精度要求且必须覆盖float和double两种数据类型如有必要需支持额外数据类型例如int8相关函数。性能测试Benchmark对比同一 Kernel 所有实现的性能并且与最终的jit::GetDefaultBestFunc对比保证该方法拿到的实现在各种条件下都是最好的对应 benchmark 中的回归约束。这套要求对应 README 中的 Solid Test 原则refer 是唯一且正确的逻辑基准其余任何实现JIT 生成、MKL、Intrinsic 等都必须先通过正确性对齐再进行性能择优。六、如何添加新的算子7 步完整指南结合文档与源码新增一个 JIT Kernel 的完整步骤如下在KernelType中添加your_key即在 kernel_base.h 的枚举中追加一个类别如kYourKey。实现 Reference 逻辑必需必须在 CPU 上实现且不能依赖任何第三方库。实现后在 refer/CMakeLists.txt 中添加USE_JITKERNEL_REFER_LITE(your_key)来启用该 kernelCMake 会将其追加写入自动生成的kernels.h。参照 refer.cc 中REGISTER_REFER_KERNEL(func)的写法一个 refer kernel 同时注册float与double两个数据类型的实现。可选在more目录实现更多算法可以依赖 MKL、Intrinsic 或 MKL-DNN 等第三方库每个实现需要定义自己的CanBeUsed条件继承KernelMoreKernelTuple并重写CanBeUsed。可选在gen目录实现基于 Xbyak 的生成代码JIT code 需要实现自己的JitCodeCreator派生自JitCodeCreatorAttr实现CanBeUsed、CodeSize、CreateJitCode并注册在与 refer 相同的KernelType上REGISTER_JITKERNEL_GEN_LITE宏。添加新的KernelTuple需要与KernelType一一对应是数据类型 属性类型 返回函数类型的打包参考SeqPoolTuple新加的 Attr 类型需要特例化JitCodeKey方法template typename Attr int64_t JitCodeKey(const Attr attr)见 kernel_key.h它是 JIT 代码缓存的 key。在test.cc中添加单元测试至少测试float和double两种数据类型必要时支持额外数据类型如int8。在benchmark.cc中添加性能对比同一种 kernel 需要对比所有实现并确保GetDefaultBestFunc得到的实现一直是速度最快的。文档对新增算子的注册宏也有明确要求在refer/CmakeLists.txt中使用USE_JITKERNEL_REFER_LITE(your_key)而 registry.h 中提供了完整的宏体系REGISTER_JITKERNEL_REFER_LITE注册 refer、REGISTER_KERNEL_MORE_LITE/REGISTER_JITKERNEL_MORE注册 more 实现、REGISTER_JITKERNEL_GEN_LITE注册 JIT 生成器、以及对应的USE_JITKERNEL_*_LITE系列。值得注意的是REGISTER_KERNEL_MORE_LITE和REGISTER_JITKERNEL_GEN_LITE宏内部都通过extern int LiteTouchJitKernelReg_##kernel_type##_refer_CPUPlace_()静态断言该 Kernel 已存在 refer 实现从编译期强制保证refer 必须先行。七、设计优点总结文档最后总结了这套 JIT Kernel 体系的设计优势结合源码可以归纳为接口方便灵活调用所有 Kernel 统一通过自动生成的 kernels.h 暴露调用方只需包含一个头文件即可使用全部已注册 Kernel。一套逻辑多套实现互不影响同一逻辑可有 JIT 生成、MKL、Intrinsic、mix 组合等多套实现各自独立注册、独立CanBeUsed判断依赖不同的第三方库也互不干扰。目录结构清晰gen / refer / more分层明确避免在单个文件中堆叠大量宏定义导致可读性差。优化方便可针对性优化可以直接针对某种属性如特定宽度、特定维度做针对性优化不影响其他属性下的性能表现。多平台支持至少可以保证 Linux、Mac、Windows 上每种平台都能正常 workrefer 与 more 实现平台无关gen目录在WITH_XBYAK AND NOT APPLE AND NOT WIN32下编译见 CMakeLists.txt框架层面使用统一接口不必关心底层实现后期也可针对不同平台做定向优化。这套reference 保正确、多实现竞争择优、JIT 生成保极致性能、线程局部缓存保复用的完整体系是 Paddle-Lite x86 后端性能优化的重要基础设施也为在移动端与边缘设备上实现高性能推理提供了可扩展的算子加速框架。如果你需要在 Paddle-Lite 中新增一个高频调用的 CPU 算子按照第六节的 7 步流程接入该体系即可在保证正确性的同时获得接近手写汇编的性能收益。【免费下载链接】Paddle-LitePaddlePaddle High Performance Deep Learning Inference Engine for Mobile and Edge (飞桨高性能深度学习端侧推理引擎项目地址: https://gitcode.com/GitHub_Trending/pa/Paddle-Lite创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表