You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Paddle/paddle/fluid/operators/jit/README.md

3.8 KiB

JIT Kernel

结合函数模板和JIT生成需要的kernel函数。 这里的kernel是比Operator中kernel更小级别的算子单元更侧重的是在不同硬件上的性能。可以有多重第三方库的实现每种实现有自己的UseMe函数负责什么条件下可以被调用。 这里实现的函数可以非常细粒度的函数方法比如Vector MUL 也可以是一个复杂的逻辑比如LSTM等。复杂的逻辑也可以由自己的底层函数拼接而成。 目前仅支持CPU上的高性能计算。

目录结构

PaddlePaddle/Paddle/paddle/fluid/
├── ...
└── operators/
    ├── .../
    └── jit/
        ├── ...
        ├── gen/
        │   └── ...
        |── more/
        │   ├── ...
        │   ├── mkl/
        │   │   └── ...
        │   ├── mkldnn/
        │   │   └── ...
        │   ├── mix/
        │   │   └── ...
        │   ├── intrinsic/
        │   │   └── ...
        │   └── openblas/
        │       └── ...
        └── refer/
            └── ...

基本类的定义都放在根目录下根目录下包括gen,more和refer三个目录。每个目录下都是一种或者多种实现每种kernel算子都需要有reference的实现用作单元测试的基准其他的实现都是可选的。

  • gen: 代表使用jit生成的code需要依赖xbyak库。该实现最关心的就是性能。
  • refer: 代表reference的实现每种kernel算子都需要有在CPU上的reference的实现他主要关心的算法逻辑的正确性。
  • more: 下面可以放入跟多实现可以包括mklmkldnnintrinsicopenblas等也可以是自身已有的kernel组合。

动态获取

提供一个jit::Get方法根据kernel类别获取每种实现都有自己的使用范围根据范围动态和当前条件选择需要的kernel函数。

测试

  • 逻辑测试 所有实现都要与refer的code对比需要满足精度要求 包括float和double的数据类型
  • 性能测试 所有实现的性能对比,并且与最终的jit::Get方法对比,该方法拿到的性能需要在各种条件下都是最好的。

如何添加新的算子

  • KernelType 中添加 your_key .
  • 实现Reference 的逻辑这个是必须是在CPU上的实现并且不能依赖任何第三方库。实现后在refer/CmakeLists.txt中添加USE_JITKERNEL_REFER(your_key)来使用该kernel.
  • (optional) 实现更多的算法在more目录下可以依赖mklintrinsic或者mkldnn等第三方库。
  • (optional) 实现基于Xbyak的生成codegen目下。 jitcode需要实现自己的JitCodeCreator并注册在与refer相同的KernelType上。
  • 必要时可以添加新的KernelTuples,可以参考XYZNTuples新加的Attr类型需要特例化JitCodeKey方法。
  • test.cc中添加unit test至少需要测试floatdouble两种数据类型,如有必要需要支持额外的数据类型,比如int8的相关函数。
  • benchmark.cc中添加相应的性能对比同一种kernel需要对比所有实现并且确保jit::Get得到的实现一直是速度最快的。

优点

  • 统一的Get方法接口简单。
  • 同一套逻辑可以有多套实现,可以依赖多套第三方库,互不影响。
  • 目录结构清晰,不会在某个文件中有多个宏定义,导致的可读性差问题。
  • 优化方便,可以直接针对某种属性针对性优化,并不影响其他属性下的性能。
  • 可以支持多种平台包括LinuxMac 和 Windows至少可以保证每种平台都可以正常work。后期也可以针对不同平台有针对的优化。框架层面可以使用统一接口不必关心底层实现。