# pybind_custom_op
PyTorch Custom Operator with PyBind11 (Manual Dispatch)
这是一个 PyTorch 自定义算子（Custom Operator）的最小化示例项目。

本项目展示了如何使用 PyBind11 将 C++ 实现的算子暴露给 Python，并在 C++ 层面实现 CPU 与 CUDA 的手动分发（Manual Dispatch） 逻辑。

#### 项目简介
与标准的 TORCH_LIBRARY 注册机制不同，本项目使用 pybind11 直接绑定 C++ 函数。由于 pybind11 本身不具备 PyTorch 的自动设备分发功能，我们在 C++ 代码中通过检查输入张量的设备类型 (input.device().is_cuda()) 来手动调度执行后端。

#### 文件结构
```text
.
├── test_pybind_expand.cpp    # C++ 源码：包含 CPU/CUDA 实现及 PyBind11 绑定
├── setup.py                  # 构建脚本：基于 setuptools 和 torch.utils.cpp_extension
├── test_pybind_expand.py     # 测试脚本：验证算子正确性及 CUDA 支持
├── build.sh                  # 自动化脚本：清理、编译、安装并运行测试
└── README.md                 # 项目说明文档
```

#### 快速开始

使用以下镜像创建容器:

```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
```

1.  编译并运行
在项目根目录下执行：
bash build.sh

2.  脚本执行流程：
清理旧的 build/ 目录和 .so 文件。
运行 python setup.py install 编译 C++ 扩展。
将编译生成的 .so 动态库复制到当前目录。
运行 python test_pybind_expand.py 进行验证。