Torch-MLU-Ops
===========================
---
## 目录
- [Torch-MLU-Ops](#torch-mlu-ops)
- [目录](#目录)
- [简介](#简介)
- [安装](#安装)
- [环境要求](#环境要求)
- [通过docker镜像启动](#通过docker镜像启动)
- [编译安装](#编译安装)
- [测试](#测试)
- [目录结构](#目录结构)
- [关键特性](#关键特性)
## 简介
Torch-MLU-Ops是寒武纪设计和开发的PyTorch第三方算子库。对于使用PyTorch框架的开发者，通过Torch-MLU-Ops，能够便捷地使用这些自定义算子，进行算子的集成、评测和业务部署。
Torch-MLU-Ops已全量覆盖LLM（Large Language Model）推理场景下的常见算子。作为后端，已支持Cambricon vLLM、Cambricon TGI、Cambricon Stable Diffusion web UI、Cambricon ComfyUI以及Cambricon Diffusers。
## 安装
### 环境要求
Torch-MLU-Ops支持的操作系统以及软件依赖如下。
不同平台支持的软件版本如下：
| 操作系统 | 编译器版本 | CMake版本 |
| ------------------------| ----------------------| -------------------------- |
| Ubuntu22.04 x86-64 | gcc 9.4.0或者更高版本 | 3.10或者更高版本 |
| Debian10.11 x86-64 | gcc 9.4.0或者更高版本 | 3.10或者更高版本 |
编译、运行Torch-MLU-Ops同时还需要配置以下依赖。
| Torch-MLU-Ops | Torch-MLU | CNNL | CNNL_Extra | CNCL | CNToolkit |
| -----------------| ------------------| ---------| --------------- | --------| -----------|
| v1.3.2 | v1.24.1 | v1.28.3 | v1.12.3 | v1.24.1 | v3.15.6 |
| v1.3.1 | v1.24.1 | v1.28.3 | v1.12.2 | v1.24.1 | v3.15.5 |
| v1.3.0 | v1.24.0 | v1.28.2 | v1.12.1 | v1.24.0 | v3.15.4 |
| v1.2.3 | v1.24.0 | v1.28.1 | v1.12.1 | v1.24.0 | v3.15.3 |
| v1.2.2 | v1.23.1 | v1.27.4 | v1.11.3 | v1.23.0 | v3.14.3 |
| v1.2.1 | v1.23.1 | v1.27.3 | v1.11.3 | v1.22.3 | v3.14.2 |
| v1.2.0 | v1.23.0 | v1.27.1 | v1.11.1 | v1.22.1 | v3.14.1 |
| v1.1.4 | v1.22.2 | v1.26.7 | v1.10.4 | v1.21.1 | v3.13.7 |
| v1.1.3 | v1.22.2 | v1.26.6 | v1.10.3 | v1.21.1 | v3.13.5 |
此外运行Torch-MLU-Ops还依赖PyTorch环境。Torch-MLU-Ops已支持的PyTorch版本请参考。
| Torch-MLU-Ops | PyTorch |
| -----------------| ------------------|
| v1.3.2 | v2.1, v2.4, v2.5 |
| v1.3.1 | v2.1, v2.4, v2.5 |
| v1.3.0 | v2.1, v2.4, v2.5 |
| v1.2.3 | v2.1, v2.4, v2.5 |
| v1.2.2 | v2.1, v2.4, v2.5 |
| v1.2.1 | v2.1, v2.4, v2.5 |
| v1.2.0 | v2.1, v2.4, v2.5 |
| v1.1.4 | v2.1, v2.3, v2.4 |
| v1.1.3 | v2.1, v2.3 |
### 通过docker镜像启动
寒武纪提供的解决方案镜像已提供了Torch-MLU-Ops所需要的依赖，请参考Cambricon PyTorch Container Image使用。您可以从[寒武纪开发者社区](https://account.cambricon.com/interaction/SMmmcGFmw837gghtZ7VR4)获取该镜像。
### 编译安装
Torch-MLU-Ops编译脚本依赖于`NEUWARE_HOME`环境变量，该环境变量指向寒武纪CNToolkit的安装路径，通常为`/usr/local/neuware`。
在运行示例程序前，您需要将`NEUWARE_HOME`中的`lib64`目录添加到`LD_LIBRARY_PATH`环境变量中，例如：
```shell
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$NEUWARE_HOME/lib64
```
如果您使用寒武纪提供的docker镜像，以上环境变量均已经被设置好了。
设置好`NEUWARE_HOME`环境变量后，您可以使用以下命令编译Torch-MLU-Ops。
从远端clone仓库，进入项目的主目录：
```shell
cd torch_mlu_ops
```
源码安装：
```shell
pip install -e .
```
wheel包安装：
```shell
python setup.py bdist_wheel
pip install dist/torch_mlu_ops*.whl
```
> _注意，如果安装过程中遇到权限问题，请获取相关权限或切换到拥有权限的user执行pip install 这步操作。_
移除Torch-MLU-Ops（Optional）：
```shell
pip uninstall torch_mlu_ops
```
### 测试
在测试前，请确保Torch-MLU-Ops已安装，您可以通过 `pip list` 命令查询Torch-MLU-Ops的安装情况。若有如下类似打印，则表明安装成功。
```shell
torch-mlu-ops 1.1.0+pt21
```
Torch-MLU-Ops模块的导入方法请参考。
```python
import torch_mlu_ops as tmo
```
在`./tests/ops_pytest`目录下，提供了`ops`级别的测试示例程序，您可以参考如下命令进行测试。
```shell
cd tests/ops_pytest/
./run_test.sh
# or 单独测试某个测例，例如
python test_flash_attention.py
```
在`./tests/kernels_pytest`目录下，提供了`kernels`级别的测试示例程序，您可以参考如下命令进行测试。
```shell
cd tests/kernels_pytest
./build.sh # 需要先编译
./run_test.sh
```
## 目录结构
```
torch_mlu_ops
|-- benchmarks ## 性能benchmarks测试代码目录
|-- csrc ## C以及BangC源码目录
| |-- common
| |-- kernels
| |-- ops
| `-- torch_api
|-- docs ## 文档目录
| |-- release_notes
| `-- user_guide
|-- tests ## 测试代码目录，包含了kernel和ops级别的测试
| |-- kernels_pytest
| `-- ops_pytest
|-- tools ## 工具目录，主要包含CI相关代码
| |-- ci
`-- torch_mlu_ops ## PyTorch第三方算子接口实现相关代码
```
## 关键特性
Torch-MLU-Ops提供的自定义算子列表以及接口说明请参考[Torch-MLU-Ops用户手册](./docs/user_guide/)。