之江编译XPU算子

中文版本 | English Version

仓库简介

该仓库包含 之江AI编译器 针对 XPU 硬件特性进行编译优化的算子代码扩展模块 xpu_extension， 旨在最大限度地减少性能损失，实现能在国产芯片上高效运行的算子。

项目基于之江跨平台算子库，使用 triton-community 编译器生成可在 XPU 上运行的算子代码。这些代码随后被注册到 xpu_extension 模块中，用户通过 xpu_extension 模块调用算子，实现模型在国产芯片上的高效运行。

xpu_extension 模块安装、注册和使用

xpu_extension 安装

1. 确保已安装PyTorch：

   pip install torch

2. 编译安装：

   cd emit-xpu/src/xpu_extension
   python setup.py install

   安装成功log如下：

   Installed /root/miniconda/envs/python38_torch201_cuda/lib/python3.8/site-packages/xpu_ext-0.1-py3.8-linux-x86_64.egg
   Processing dependencies for xpu-ext==0.1
   Finished processing dependencies for xpu-ext==0.1
   

3. 注册到torch框架：

   cd emit-xpu
   pip install .

   注册成功log如下：

   Successfully built xpu_extension
   Installing collected packages: xpu_extension
   Attempting uninstall: xpu_extension
       Found existing installation: xpu_extension 0.1
       Uninstalling xpu_extension-0.1:
       Successfully uninstalled xpu_extension-0.1
   Successfully installed xpu_extension-0.1
   

4. 使用：

   import torch
   import xpu_extension
   
   # 创建在 'cuda' 设备上的输入
   a = torch.ones((2, 2)).cuda()
   b = torch.ones((2, 2)).cuda()
   # 开启xpu_extension注册的算子
   with xpu_extension.use_open_triton():   
       output = torch.add(a, b)

如何编译基于XPU的算子并注册到torch框架

基于triton-community和Emitc编译生成xpu算子

算子接入Python接口

1. src/xpu_extension/kernels.xpu 文件中加入xpu算子实现。

   例如

   __attribute__((global)) void xpu_ext_add_kernel(float* v1, float* v2, float v3, float* v4, int32_t v5, int32_t v6) {
       // 算子实现
       ...
   }

2. 在 src/xpu_extension/utils.cpp 文件中，定义 op kernel，并实现启动 op kernel。 例如：

   #include "xpu/kernel/xtdk.h"
   #include "xpu/kernel/xtdk_io.h"
   #include "xpu/runtime.h"
   #include <iostream>
   #include <torch/extension.h>
   
   // 声明 op kernel
   __attribute__((global)) void xpu_ext_add_kernel(float* v1, float* v2, float v3, float* v4, int32_t v5, int32_t v6);
   
   torch::Tensor xpu_ext_add(torch::Tensor a, torch::Tensor b, float alpha) {
       
       // 输入数据类型校验
       ...
       // 设置kernel 启动的卡号，保证与数据在同一个卡上
       TORCH_CHECK(!xpu_set_device(a.device().index()), "set xpu device failed");
       // 计算启动的 cluster 和 core 数量
       const int size_per_cluster = 4096;
       int len = a.numel();
       int cluster_num = (len + size_per_cluster - 1) / size_per_cluster;
       cluster_num = cluster_num > CLUSTER_NUM ? CLUSTER_NUM : cluster_num;
   
       torch::Tensor out = torch::empty_like(a);
       // 启动 kernel
       xpu_ext_add_kernel<<<cluster_num, CORE_NUM>>>(
           (float*)a.data_ptr(), (float*)b.data_ptr(), alpha, 
           (float*)out.data_ptr(), 1, len);
   
       return out;
   }

3. 在 src/xpu_extension/xpu_ext.cpp 文件中处理导出接口的计算逻辑。

   本例子中使用 utils.cpp 和 xpu_ext.cpp 两个文件实现cpp扩展，utils.cpp 负责启动算子kernel逻辑，xpu_ext.cpp 侧重输入数据的前处理。

   #include <torch/extension.h>
   #include <vector>
   
   torch::Tensor xpu_extension_add(torch::Tensor a, 
                                   torch::Tensor b, 
                                   float alpha) {
       // 输入数据预处理或其他计算逻辑
       ...
   
       // 调用 src/xpu_extension/utils.cpp 定义的接口
       return xpu_ext_add(a, tensor_b, alpha);
   }
   
   // 这里使用 pybind11 把函数指针绑定到 xpu-ext 的模块下
   PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
       m.def("add", &xpu_extension_add, "XPU_EXT add",
           py::arg("a"), py::arg("b"), py::arg("alpha")=1.0);
   }

4. 编译 xpu_ext 扩展包

   python setup.py install

   编译成功后，可以通过如下方式使用 xpu_ext 扩展包

   import torch  # 先 import torch, 加载 torch 相关的符号到内存中
   import xpu_ext
   
   xpu_ext.add
   <built-in method add of PyCapsule object at 0x7f656a65dd80>  # 表示函数注册成功

   import torch
   import xpu_ext
   
   # 创建在 'cuda' 设备上的输入
   a = torch.ones((2, 2)).cuda()
   b = torch.ones((2, 2)).cuda()
   # 调用 xpu_ext 扩展包的 算子
   output = xpu_ext.add(a, b)

注册到torch框架（xpu_extension）

1. 将算子注册到 aten lib中

   • 在 emit-xpu/src/xpu_extension/__init__.py 文件中，将扩展算子封装注册到 enable中。

   例如，将xpu_ext.add注册到torch的add.Tensor算子

   import xpu_ext
   
   
   def enable(lib=aten_lib):
       print("\033[31m Emit-XPU enabled\033[0m")
       lib.impl("add.Tensor", xpu_ext.add, dispatch_key)

   如果需要在Python中封装一层业务逻辑，可以将 xpu_ext 模块的算子封装添加到 ./src/xpu_extension/ops/base_op 路径下，同时在 emit-xpu/src/xpu_extension/ops/base_op/__init__.py 文件的 __all__ 包含算子封装的引用。

2. 安装：

   cd emit-xpu
   pip install .

   注册成功log如下：

   Successfully built xpu_extension
   Installing collected packages: xpu_extension
   Attempting uninstall: xpu_extension
       Found existing installation: xpu_extension 0.1
       Uninstalling xpu_extension-0.1:
       Successfully uninstalled xpu_extension-0.1
   Successfully installed xpu_extension-0.1
   

3. 使用：

   import torch
   import xpu_extension
   
   # 创建在 'cuda' 设备上的输入
   a = torch.ones((2, 2)).cuda()
   b = torch.ones((2, 2)).cuda()
   # 开启xpu_extension注册的算子
   with xpu_extension.use_open_triton():   
       output = torch.add(a, b)

测试

功能性测试

1. 添加算子功能测试用例

   在 ./tests/verify/base_op 或 ./tests/verify/base_op 目录下添加算子功能测试接口。测试框架使用pytest。

2. 执行测试用例

   运行测试用例，命令如下：

   pytest -s test_add.py

   如果需要输出 kernel 执行时间，首先设置环境变量：

   export XPURT_DISPATCH_MODE=PROFILING

性能测试

1. 添加算子性能测试用例

   在 ./tests/benchmark/base_op 目录下添加算子性能测试接口。
   性能测试流程包括正确性验证、WARMUP、REPETITION，其中性能统计只记录REPETITION阶段的执行。
   请参考 test_add.py 文件，测试框架使用pytest。

   例如：test_add.py 文件对比测试 torch.add 和 xpu_ext.add 算子。

   • torch.add 对应的 kernel 名字是 xpukernel_xpu310calc_basic
   • xpu_ext.add 对应的 kernel 名字是 xpu_ext_add_kernel

   备注：torch.add 对应的 kernel 名字可以自己提前测试时从log中查出来； xpu_ext.add 对应的 kernel 名字是在 kernels.xpu 里面定义的名字。

2. 测试端到端性能

   运行命令如下：

   pytest -s test_add.py

   执行结束后，输出log如下：

   latency_xpu_torch(ms):0.098710 latency_xpu_ext(ms):0.181973 torch/ext=54.24%
   

3. 测试kernel性能

   ./tests/benchmark/base_op/test_benchmark.py 是用于捕获 kernel 运行时间的脚本模板。 该脚本会启动指定的测试文件，并解析指定算子名字的 kernel 运行时间。

   运行命令：

   python test_benchmark.py <file_path> [<op_name1>] [<op_name2>]

   参数说明

   • file_path：要执行的 pytest 测试文件的绝对路径。
   • op_name1：第一个算子的 kernel 名字，可选
   • op_name2：第二个算子的 kernel 名字，可选

   示例

   test_add.py 文件对比测试 torch.add 和 xpu_ext.add 算子。

   • torch.add 对应的 kernel 名字是 xpukernel_xpu310calc_basic ，可以自己提前测试从log中查出来，但往往torch原生的kernel名字比较难找。
   • xpu_ext.add 对应的 kernel 名字是 xpu_ext_add_kernel ，即在 kernels.xpu 文件中定义的名字。

   捕获 torch.add 和 xpu_ext.add 对应kernel的性能数据，命令如下：

   python test_benchmark.py test_add.py xpukernel_xpu310calc_basic xpu_ext_add_kernel

   执行结束后，输出log如下：

   xpukernel_xpu310calc_basic average time 3760.0ns, xpu_ext_add_kernel average time 25796.0ns, xpu_ext_add_kernel as percentage of xpukernel_xpu310calc_basic: 14.58%
   

目录结构

.
├── README.en.md
├── README.md
├── docs（设计文档和图片）
│   ├── design_docs
│   │   └── assets
│   └── test_report
│       └── assets
├── src
│   ├── dev_ops (编译算子中间结果(.ttir, .mlir, .cpp, .xpu等)和测试)
│   │   ├── base_op
│   │   │   ├── add_func
│   │   │   ├── div.Tensor
│   │   │   ├── sub_
│   │   │   ├── vec_add
│   │   │   ├── zero_
│   │   │   ├── zeros
│   │   │   └── zeros_like
│   │   └── fused_op
│   └── xpu_extension (xpu_ext模块，release给用户使用的模块)
│       ├── kernels.xpu
│       ├── setup.py
│       ├── utils.cpp
│       └── xpu_ext.cpp
├── tests (针对xpu_extension模块的测试，采用pytest测试)
│   ├── __init__.py
│   ├── benchmark（性能测试）
│   │   └── __init__.py
│   └── verify（正确性验证）
│       ├── __init__.py
│       ├── base_op
│       │   ├── __init__.py
│       │   └── test_add.py
│       └── fused_op
│           └── __init__.py
└── utils

更新日志

v0.1

• 增加10个基础算子及其测试。

兼容性

支持的平台

• Linux

支持的硬件

• ZP61

参考

有关更多信息和相关的开源项目，请参阅以下链接：