1  # ONNX Runtime Custom CUDA Operator
  2  
  3  本模块演示如何为 ONNX Runtime 开发自定义 CUDA 算子，实现完整的 Python → C++ → CUDA 集成流程。
  4  
  5  ## 学习目标
  6  
  7  1. 理解 ONNX Runtime 自定义算子架构
  8  2. 掌握 CUDA Kernel 开发与执行流语义
  9  3. 学习 C++ 胶水层实现 (`CustomOpBase`)
 10  4. 实现 Python 绑定和动态库注册
 11  
 12  ## 文件结构
 13  
 14  ```
 15  02_ORT_Custom_CUDA_Op/
 16  ├── README.md                    # 模块说明
 17  ├── CMakeLists.txt               # 构建配置
 18  ├── requirements.txt             # Python 依赖
 19  ├── src/                         # C++ 和 CUDA 源代码
 20  │   ├── custom_op.h             # CustomOpBase 定义
 21  │   ├── custom_op.cc            # C++ 胶水层
 22  │   ├── cuda_kernels.h          # CUDA kernel 声明
 23  │   └── cuda_kernels.cu         # CUDA kernel 实现
 24  ├── python/                      # Python 绑定（后续整合到 src/python/）
 25  │   └── custom_ops.py
 26  └── tests/                       # 测试代码
 27      ├── test_inference.py
 28      └── test_correctness.py
 29  ```
 30  
 31  ## 构建
 32  
 33  ```bash
 34  mkdir -p build && cd build
 35  cmake .. -DONNXRUNTIME_ROOT=/path/to/onnxruntime
 36  make -j$(nproc)
 37  ```
 38  
 39  或使用 Makefile:
 40  
 41  ```bash
 42  make build-ort
 43  ```
 44  
 45  ## 运行测试
 46  
 47  ```bash
 48  # 安装依赖
 49  pip install -r 02_ORT_Custom_CUDA_Op/requirements.txt
 50  
 51  # 运行测试
 52  python -m pytest 02_ORT_Custom_CUDA_Op/tests/ -v
 53  ```
 54  
 55  或使用 Makefile:
 56  
 57  ```bash
 58  make test-ort
 59  ```
 60  
 61  ## GELU 激活函数实现
 62  
 63  ### 数学定义
 64  
 65  GELU (Gaussian Error Linear Unit):
 66  ```
 67  GELU(x) = x * Φ(x)
 68  ```
 69  
 70  其中 `Φ(x)` 是标准正态分布的累积分布函数。
 71  
 72  近似实现:
 73  ```
 74  GELU(x) ≈ 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³)))
 75  ```
 76  
 77  ### CUDA Kernel 接口
 78  
 79  ```cpp
 80  cudaError_t LaunchGeluKernel(
 81      const float* input,
 82      float* output,
 83      int64_t count,
 84      cudaStream_t stream
 85  );
 86  ```
 87  
 88  ### 自定义算子注册
 89  
 90  ```cpp
 91  struct GeluCustomOp : Ort::CustomOpBase<GeluCustomOp, GeluKernel> {
 92      const char* GetName() const { return "CustomGelu"; }
 93      const char* GetExecutionProviderType() const { return "CUDAExecutionProvider"; }
 94  };
 95  ```
 96  
 97  ### Python 使用
 98  
 99  ```python
100  import onnxruntime as ort
101  from python.custom_ops import create_session, run_inference
102  
103  session, library_path = create_session("model.onnx", use_cuda=True)
104  output = run_inference("model.onnx", input_data, use_cuda=True)
105  ```
106  
107  ## 性能基准 (1M 元素, A100)
108  
109  | 实现 | 延迟 (μs) | 带宽 (GB/s) |
110  |------|-----------|-------------|
111  | PyTorch GELU | ~50 | ~160 |
112  | Custom CUDA | **~35** | **~230** |
113  
114  ## 注意事项
115  
116  - `CustomGelu` 声明为 `CUDAExecutionProvider` 算子，**无 CPU fallback**
117  - 运行集成测试前需要先成功构建 native 动态库
118  - 若 ONNX Runtime 安装不包含 CUDA provider 或头文件，构建将失败
119  
120  ## 参考资源
121  
122  - [ONNX Runtime Custom Ops](https://onnxruntime.ai/docs/reference/operators/add-custom-op.html)
123  - [CUDA Programming Guide](https://docs.nvidia.com/cuda/cuda-c-programming-guide/)
124  - [GELU Paper](https://arxiv.org/abs/1606.08415)