add detail

detail.md

+# PyTorch 自定义算子开发与接入指南
+
+本文档基于 `torch_custom_op` 工程，详细介绍如何使用 `TORCH_LIBRARY`、`TORCH_LIBRARY_IMPL` 和 `PyInit_【NAME】` 等技术开发和接入 PyTorch 自定义算子。
+
+## 1. 概述
+
+PyTorch 提供了强大的自定义算子开发能力，通过 `TORCH_LIBRARY` 和 `TORCH_LIBRARY_IMPL` 等宏，可以：
+- 定义新的算子接口
+- 为不同设备（CPU/CUDA）提供专门的实现
+- 无缝集成到 PyTorch 生态系统中
+
+## 2. 环境准备
+
+### 2.1 依赖项
+- PyTorch 1.7+
+- C++ 编译器（支持 C++14）
+- CMake (3.18+)
+- Python 3.6+
+
+### 2.2 项目结构
+推荐的项目结构如下：
+
+```
+my_custom_ops/
+├── setup.py            # Python 包配置
+├── build.sh            # 编译脚本
+├── my_ops.h            # 算子声明
+├── my_ops_impl.cpp     # 算子实现
+├── my_ops.cpp          # 模块初始化和算子定义
+└── test_ops.py         # 测试脚本
+```
+
+## 3. 开发步骤
+
+### 3.1 定义算子接口
+
+使用 `TORCH_LIBRARY` 宏定义算子接口，在 `.cpp` 文件中：
+
+```cpp
+#include <torch/library.h>
+#include <ATen/ATen.h>
+#include <Python.h>
+#include "my_ops.h"
+
+// 定义算子库
+TORCH_LIBRARY(my_ops, m) {
+    m.def("add_one(Tensor input) -> Tensor");
+    m.def("multiply_by_two(Tensor input) -> Tensor");
+}
+
+// Python 模块初始化函数
+extern "C" {
+PyObject *PyInit_my_ops(void) {
+    static struct PyModuleDef module_def = {
+        PyModuleDef_HEAD_INIT,
+        "my_ops",  // 模块名
+        "My custom operations module",  // 文档
+        -1,
+        NULL  // 方法定义
+    };
+    return PyModule_Create(&module_def);
+}
+}
+```
+
+### 3.2 实现算子逻辑
+
+在 `.h` 文件中声明算子函数：
+
+```cpp
+#pragma once
+#include <torch/library.h>
+#include <ATen/ATen.h>
+
+namespace my_ops_impl {
+    at::Tensor add_one(at::Tensor input);
+    at::Tensor multiply_by_two(at::Tensor input);
+}
+```
+
+在 `.cpp` 文件中实现算子逻辑并使用 `TORCH_LIBRARY_IMPL` 注册：
+
+```cpp
+#include <torch/library.h>
+#include <ATen/ATen.h>
+#include "my_ops.h"
+
+namespace my_ops_impl {
+    // 操作符的具体实现
+    at::Tensor add_one(at::Tensor input) {
+        return input + 1;
+    }
+    
+    at::Tensor multiply_by_two(at::Tensor input) {
+        return input * 2;
+    }
+}
+
+// 注册 CPU 实现
+TORCH_LIBRARY_IMPL(my_ops, CPU, m) {
+    m.impl("add_one", &my_ops_impl::add_one);
+    m.impl("multiply_by_two", &my_ops_impl::multiply_by_two);
+}
+
+// 注册 CUDA 实现（如果有 CUDA）
+#ifdef TORCH_HAS_CUDA
+TORCH_LIBRARY_IMPL(my_ops, CUDA, m) {
+    m.impl("add_one", &my_ops_impl::add_one);
+    m.impl("multiply_by_two", &my_ops_impl::multiply_by_two);
+}
+#endif
+```
+
+### 3.3 编译配置
+
+#### 3.3.1 使用 CMake 配置
+
+创建 `CMakeLists.txt` 文件：
+
+```cmake
+cmake_minimum_required(VERSION 3.18)
+project(my_ops)
+
+find_package(Torch REQUIRED)
+set(CMAKE_CXX_STANDARD 14)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+add_library(my_ops SHARED my_ops.cpp my_ops_impl.cpp)
+target_link_libraries(my_ops "${TORCH_LIBRARIES}")
+set_target_properties(my_ops PROPERTIES PREFIX "")
+```
+
+#### 3.3.2 使用 setup.py 配置
+
+创建 `setup.py` 文件，支持自动检测 CUDA 并使用相应的扩展：
+
+```python
+import os
+import torch
+from setuptools import setup, find_packages
+from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
+
+library_name = "my_ops"
+current_dir = os.path.dirname(os.path.abspath(__file__))
+
+# 源文件列表
+sources = [
+    os.path.join(current_dir, "my_ops.cpp"),
+    os.path.join(current_dir, "my_ops_impl.cpp"),
+]
+
+# 检查CUDA是否可用
+use_cuda = torch.cuda.is_available()
+extension = CUDAExtension if use_cuda else CppExtension
+
+if use_cuda:
+    import glob
+    cuda_files = glob.glob(os.path.join(current_dir, "*.cu"))
+    sources.extend(cuda_files)
+
+extra_compile_args = {
+    'cxx': ['-O2', '-std=c++17'],
+}
+
+if use_cuda:
+    extra_compile_args['nvcc'] = ['-O2']
+
+# 创建包目录和 __init__.py
+package_dir = os.path.join(current_dir, library_name)
+os.makedirs(package_dir, exist_ok=True)
+
+init_py_path = os.path.join(package_dir, "__init__.py")
+if not os.path.exists(init_py_path):
+    with open(init_py_path, "w") as f:
+        f.write("""
+from ._C import *
+
+__all__ = ['add_one', 'multiply_by_two']
+""")
+
+setup(
+    name=library_name,
+    version='0.1.0',
+    description='Custom operations for PyTorch',
+    author='Your Name',
+    
+    # 关键：指定包
+    packages=[library_name],
+    package_dir={library_name: library_name},
+    
+    # 扩展模块 - 注意命名格式
+    ext_modules=[
+        extension(
+            name=f"{library_name}._C",
+            sources=sources,
+            extra_compile_args=extra_compile_args,
+            include_dirs=[current_dir],
+        )
+    ],
+    
+    # 命令类
+    cmdclass={
+        'build_ext': BuildExtension
+    },
+    
+    # 依赖
+    install_requires=['torch>=1.10.0'],
+    
+    # 确保生成正确的 .dist-info
+    zip_safe=False,
+)
+```
+
+#### 3.3.3 编译脚本
+
+创建 `build.sh` 脚本：
+
+```bash
+#!/bin/bash
+rm -rf build
+mkdir -p build
+cd build
+cmake ..
+make -j
+cp my_ops.so ..
+cd ..
+```
+
+## 4. 编译和安装
+
+### 4.1 使用 build.sh 编译
+
+```bash
+bash build.sh
+```
+
+### 4.2 使用 pip 安装
+
+```bash
+pip install --no-build-isolation .
+```
+
+## 5. 在 Python 中使用自定义算子
+
+创建测试脚本 `test_ops.py`：
+
+```python
+import torch
+import os
+import sys
+
+# 添加当前目录到 Python 路径
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+# 加载自定义操作符
+try:
+    import my_ops
+    print("成功加载 my_ops 扩展")
+except ImportError as e:
+    print(f"加载扩展失败: {e}")
+    print("请先编译 C++ 扩展")
+    sys.exit(1)
+
+def test_add_one():
+    """测试 add_one 操作符"""
+    print("\n测试 add_one 操作符:")
+    
+    # 创建测试张量
+    x = torch.tensor([1.0, 2.0, 3.0])
+    print(f"输入: {x}")
+    
+    # 调用自定义操作符
+    y = torch.ops.my_ops.add_one(x)
+    print(f"输出: {y}")
+    
+    # 验证结果
+    expected = x + 1
+    assert torch.allclose(y, expected), f"结果不匹配: {y} vs {expected}"
+    print("✓ 测试通过")
+
+def test_multiply_by_two():
+    """测试 multiply_by_two 操作符"""
+    print("\n测试 multiply_by_two 操作符:")
+    
+    # 创建测试张量
+    x = torch.tensor([1.0, 2.0, 3.0])
+    print(f"输入: {x}")
+    
+    # 调用自定义操作符
+    y = torch.ops.my_ops.multiply_by_two(x)
+    print(f"输出: {y}")
+    
+    # 验证结果
+    expected = x * 2
+    assert torch.allclose(y, expected), f"结果不匹配: {y} vs {expected}"
+    print("✓ 测试通过")
+
+def test_cuda_support():
+    """测试 CUDA 支持（如果可用）"""
+    if torch.cuda.is_available():
+        print("\n测试 CUDA 支持:")
+        
+        # 创建 CUDA 张量
+        x = torch.tensor([1.0, 2.0, 3.0]).cuda()
+        print(f"CUDA 输入: {x}")
+        
+        # 测试 add_one
+        y = torch.ops.my_ops.add_one(x)
+        print(f"add_one 输出: {y}")
+        assert torch.allclose(y, x + 1), f"CUDA add_one 结果不匹配"
+        
+        # 测试 multiply_by_two
+        z = torch.ops.my_ops.multiply_by_two(x)
+        print(f"multiply_by_two 输出: {z}")
+        assert torch.allclose(z, x * 2), f"CUDA multiply_by_two 结果不匹配"
+        
+        print("✓ CUDA 测试通过")
+    else:
+        print("\nCUDA 不可用，跳过 CUDA 测试")
+
+if __name__ == "__main__":
+    print("测试自定义算子")
+    print("=" * 50)
+    
+    test_add_one()
+    test_multiply_by_two()
+    test_cuda_support()
+    
+    print("\n" + "=" * 50)
+    print("所有测试通过！")
+```
+
+运行测试：
+
+```bash
+python test_ops.py
+```
+
+## 6. 常见问题和解决方案
+
+### 6.1 模块初始化函数错误
+
+**错误信息**：
+```
+加载扩展失败: dynamic module does not define module export function (PyInit_my_ops)
+```
+
+**解决方案**：
+确保在 C++ 文件中正确定义了 `PyInit_my_ops` 函数，并且函数名与模块名一致。
+
+### 6.2 CUDA 支持问题
+
+**错误信息**：
+```
+undefined reference to `torch::library::Library::impl(...)`
+```
+
+**解决方案**：
+确保正确定义了 `TORCH_HAS_CUDA` 宏，并且在 CUDA 环境下编译。
+
+### 6.3 算子注册问题
+
+**错误信息**：
+```
+undefined symbol: _ZN5torch8library7Library4implERKNS_6SymbolEPFSt10unique_ptrINS_5autograd15AutogradContextESt14default_deleteIS3_EERKSt13unordered_mapISt14basic_stringIcSt11char_traitsIcESaIcEESt6vectorINS_5IValueESaIS7_EESt4hashIS9_ESt8equal_toIS9_ESaISB_EE
+```
+
+**解决方案**：
+确保使用了正确的 PyTorch 版本，并且编译时链接了正确的 PyTorch 库。
+
+## 7. 高级功能
+
+### 7.1 多设备支持
+
+可以为不同设备（CPU、CUDA）提供不同的实现：
+
+```cpp
+// CPU 实现
+TORCH_LIBRARY_IMPL(my_ops, CPU, m) {
+    m.impl("add_one", &my_ops_impl::add_one_cpu);
+}
+
+// CUDA 实现
+#ifdef TORCH_HAS_CUDA
+TORCH_LIBRARY_IMPL(my_ops, CUDA, m) {
+    m.impl("add_one", &my_ops_impl::add_one_cuda);
+}
+#endif
+```
+
+### 7.2 复杂算子定义
+
+可以定义更复杂的算子，支持多个输入和输出：
+
+```cpp
+TORCH_LIBRARY(my_ops, m) {
+    m.def("add_scalar(Tensor input, Scalar scalar) -> Tensor");
+    m.def("addmm(Tensor self, Tensor mat1, Tensor mat2, Scalar beta=1, Scalar alpha=1) -> Tensor");
+}
+```
+
+### 7.3 算子重载
+
+可以为不同类型的输入提供重载：
+
+```cpp
+TORCH_LIBRARY(my_ops, m) {
+    m.def("add(Tensor a, Tensor b) -> Tensor");
+    m.def("add(Tensor a, Scalar b) -> Tensor");
+}
+```
+
+## 8. 总结
+
+使用 `TORCH_LIBRARY` 和 `TORCH_LIBRARY_IMPL` 等技术，我们可以：
+
+1. **定义清晰的算子接口**：使用 `TORCH_LIBRARY` 宏定义算子的签名
+2. **实现设备特定的逻辑**：使用 `TORCH_LIBRARY_IMPL` 为不同设备提供专门的实现
+3. **无缝集成到 PyTorch**：通过 `PyInit_【NAME】` 函数初始化 Python 模块
+4. **灵活编译和安装**：支持使用 CMake 或 setup.py 编译和安装
+
+这种方法为 PyTorch 自定义算子开发提供了一种简洁、灵活的方式，使得我们可以轻松扩展 PyTorch 的功能，满足特定的业务需求。
+
+## 9. 代码示例
+
+完整的代码示例可以参考 `torch_custom_op` 工程中的各个部分：
+
+- **1part**：展示了在单个文件中同时定义和实现算子
+- **2part**：展示了将定义和实现分离到不同文件
+- **3part**：展示了如何打包为可安装的 Python 包
+- **4part**：展示了不同的模块初始化方式
+
+通过这些示例，您可以快速上手 PyTorch 自定义算子的开发和接入。
\ No newline at end of file