# PyTorch 自定义算子工程

本工程展示了两种不同的 PyTorch 自定义算子实现方法，分别使用 PyBind11 和 TORCH_LIBRARY 机制。

## 工程结构

```
torch_custom_op/
├── pybind11_custom_op/     # 使用 PyBind11 实现的自定义算子
└── torch_library_custom_op/ # 使用 TORCH_LIBRARY 实现的自定义算子
```

## 子项目说明

### 1. pybind11_custom_op

使用 PyBind11 将 C++ 实现的算子暴露给 Python，并在 C++ 层面实现 CPU 与 CUDA 的手动分发逻辑。

**特点**：
- 直接使用 PyBind11 绑定 C++ 函数
- 通过检查输入张量的设备类型手动调度执行后端
- 适合需要更灵活控制的场景

**关键文件**：
- `test_pybind_expand.cpp`：C++ 源码，包含 CPU/CUDA 实现及 PyBind11 绑定
- `setup.py`：构建脚本，基于 setuptools 和 torch.utils.cpp_extension
- `test_pybind_expand.py`：测试脚本，验证算子正确性及 CUDA 支持
- `build.sh`：自动化脚本，清理、编译、安装并运行测试

### 2. torch_library_custom_op

使用 PyTorch 官方的 TORCH_LIBRARY 机制实现自定义算子，支持为不同设备提供专门的实现。

**特点**：
- 使用 PyTorch 官方推荐的 TORCH_LIBRARY 注册机制
- 支持为不同设备（CPU/CUDA）提供专门的实现
- 自动处理设备分发，无需手动检查设备类型

**关键文件**：
- `custom_op/test_torch_library_expand.cpp`：定义算子接口和初始化 Python 模块
- `custom_op/test_ops_impl.cpp`：实现算子逻辑并注册到不同设备
- `custom_op/CMakeLists.txt`：使用 CMake 配置编译
- `custom_op/setup.py`：使用 setup.py 配置编译和安装
- `custom_op/build.sh`：编译脚本

## 快速开始

### 环境准备

使用以下镜像创建容器：

```bash
docker pull harbor.sourcefind.cn:5443/dcu/admin/base/vllm:0.11.0-ubuntu22.04-dtk25.04.2-1226-das1.7-py3.10-20251226
```

### 使用方法

#### 对于 pybind11_custom_op：

```bash
cd pybind11_custom_op
bash build.sh
```

#### 对于 torch_library_custom_op：

支持两种编译方法：

1. 使用 build.sh 编译：
   ```bash
   cd torch_library_custom_op/custom_op
   bash build.sh
   ```

2. 使用 pip 安装：
   ```bash
   cd torch_library_custom_op/custom_op
   pip install --no-build-isolation .
   cd test
   python test_torch_library_expand_pip_install.py
   ```

## 两种实现方法对比

| 特性 | PyBind11 实现 | TORCH_LIBRARY 实现 |
|------|--------------|-------------------|
| 设备分发 | 手动检查设备类型 | 自动处理设备分发 |
| 注册机制 | 直接绑定 C++ 函数 | 使用 TORCH_LIBRARY 宏注册 |
| 灵活性 | 更高，可自定义绑定逻辑 | 更规范，符合 PyTorch 设计 |
| 集成度 | 与 PyTorch 生态集成度较低 | 与 PyTorch 生态完全集成 |
| 适用场景 | 需要更灵活控制的场景 | 标准 PyTorch 自定义算子开发 |

## 测试

每个子项目都包含测试脚本，用于验证算子的正确性和性能：

- `pybind11_custom_op/test_pybind_expand.py`
- `torch_library_custom_op/custom_op/test_torch_library_expand_build_sh.py`
- `torch_library_custom_op/test/test_torch_library_expand_pip_install.py`

## 总结

本工程展示了两种不同的 PyTorch 自定义算子实现方法，各有优缺点：

- **PyBind11 实现**：适合需要更灵活控制的场景，可自定义绑定逻辑，但需要手动处理设备分发。
- **TORCH_LIBRARY 实现**：符合 PyTorch 官方推荐的做法，自动处理设备分发，与 PyTorch 生态完全集成。

开发者可以根据具体需求选择适合的实现方法。