<div align="center">

Torch-MLU-Ops
===========================

![system](https://img.shields.io/badge/system-ubuntu22.04_debian10.11-blue.svg)![python](https://img.shields.io/badge/python-3.10-green)![pytorch](https://img.shields.io/badge/pytorch-2.1_2.4_2.5-green)![release](https://img.shields.io/badge/release-1.3.2-green.svg)![license](https://img.shields.io/badge/license-BSD_3-blue.svg)

<div align="center">
  <b>
    <a href="https://www.cambricon.com/docs/sdk_1.15.0/bangtransformer_0.4.0/index.html">
      <font size="4"> 📖 Torch_MLU_Ops用户手册</font>
    </a>
  </b>
  &nbsp;&nbsp;&nbsp;&nbsp;
  <b>
    <a href="https://developer.cambricon.com/">
      <font size="4"> 🌏 寒武纪开发者社区</font>
    </a>
  </b>
  &nbsp;&nbsp;&nbsp;&nbsp;
  <b>
    <a href="https://sdk.cambricon.com/download?sdk_version=V1.13.0&component_name=Basis">
      <font size="4"> 🛠️ 依赖组件获取</font>
    </a>
  </b>
</div>


---
<div align="left">

## 目录

- [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/)。
