已合并
fix: torch ops cpp sources #8489
丛吉钰创建于 17 天前
fix: torch ops cpp sources #8489
已合并
共 15 个文件变更+114-87
| @@ -62,7 +62,7 @@ class SwigluGroupOpBuilder(OpBuilder): | |||
| 62 | super().__init__("swiglu_group") | 62 | super().__init__("swiglu_group") |
| 63 | 63 | ||
| 64 | def sources(self): | 64 | def sources(self): |
| 65 | - return ["csrc/activation/swiglu_group.cpp"] | 65 | + return [self.resolve_source("swiglu_group.cpp")] |
| 66 | 66 | ||
| 67 | def schema(self): | 67 | def schema(self): |
| 68 | return ( | 68 | return ( |
| @@ -143,7 +143,7 @@ class SwigluGroupQuantOpBuilder(OpBuilder): | |||
| 143 | super().__init__("swiglu_group_quant") | 143 | super().__init__("swiglu_group_quant") |
| 144 | 144 | ||
| 145 | def sources(self): | 145 | def sources(self): |
| 146 | - return ["csrc/activation/swiglu_group_quant.cpp"] | 146 | + return [self.resolve_source("swiglu_group_quant.cpp")] |
| 147 | 147 | ||
| 148 | def schema(self): | 148 | def schema(self): |
| 149 | return ( | 149 | return ( |
| @@ -19,7 +19,7 @@ class SwigluGroupQuantBackwardOpBuilder(OpBuilder): | |||
| 19 | super().__init__("swiglu_group_quant_backward") | 19 | super().__init__("swiglu_group_quant_backward") |
| 20 | 20 | ||
| 21 | def sources(self): | 21 | def sources(self): |
| 22 | - return ["csrc/activation/swiglu_group_quant_grad.cpp"] | 22 | + return [self.resolve_source("swiglu_group_quant_grad.cpp")] |
| 23 | 23 | ||
| 24 | def schema(self): | 24 | def schema(self): |
| 25 | return ( | 25 | return ( |
| @@ -9,7 +9,7 @@ | |||
| 9 | # ----------------------------------------------------------------------------------------------------------- | 9 | # ----------------------------------------------------------------------------------------------------------- |
| 10 | 10 | ||
| 11 | # SwigluClamp PTA Python 前端 | 11 | # SwigluClamp PTA Python 前端 |
| 12 | -# 管理 JIT 编译(csrc/activation/swiglu_clamp.cpp)并把算子注册到 PyTorch Dispatcher。 | 12 | +# 管理 JIT 编译(csrc/swiglu_clamp.cpp)并把算子注册到 PyTorch Dispatcher。 |
| 13 | # 注册后可通过 torch.ops.cann_ops_nn.swiglu_clamp(x, limit) 调用。 | 13 | # 注册后可通过 torch.ops.cann_ops_nn.swiglu_clamp(x, limit) 调用。 |
| 14 | # 参考: torch_extension/README.md「新增算子」+ PR !5910 (commit d8af5e890)。 | 14 | # 参考: torch_extension/README.md「新增算子」+ PR !5910 (commit d8af5e890)。 |
| 15 | 15 | ||
| @@ -24,7 +24,7 @@ class SwigluClampOpBuilder(OpBuilder): | |||
| 24 | 24 | ||
| 25 | def sources(self): | 25 | def sources(self): |
| 26 | """C++ 源码路径(相对 cann_ops_nn 包根)。""" | 26 | """C++ 源码路径(相对 cann_ops_nn 包根)。""" |
| 27 | - return ["csrc/activation/swiglu_clamp.cpp"] | 27 | + return [self.resolve_source("swiglu_clamp.cpp")] |
| 28 | 28 | ||
| 29 | def schema(self) -> str: | 29 | def schema(self) -> str: |
| 30 | """PyTorch 算子签名。limit 默认 7.0(Step-3.7 专家路)。""" | 30 | """PyTorch 算子签名。limit 默认 7.0(Step-3.7 专家路)。""" |
| @@ -51,5 +51,5 @@ builder._ensure_initialized() | |||
| 51 | 51 | ||
| 52 | def swiglu_clamp(x: torch.Tensor, limit: float = 7.0): | 52 | def swiglu_clamp(x: torch.Tensor, limit: float = 7.0): |
| 53 | """Dispatcher 的 NPU 实现。PrivateUse1 是 NPU 后端分发键。""" | 53 | """Dispatcher 的 NPU 实现。PrivateUse1 是 NPU 后端分发键。""" |
| 54 | - op_module = builder.load() # JIT 编译/加载 csrc/activation/swiglu_clamp.cpp | 54 | + op_module = builder.load() # JIT 编译/加载 csrc/swiglu_clamp.cpp |
| 55 | return op_module.swiglu_clamp(x, limit) | 55 | return op_module.swiglu_clamp(x, limit) |
| @@ -21,7 +21,7 @@ class MatmulEmuSplitWeightOpBuilder(OpBuilder): | |||
| 21 | """ | 21 | """ |
| 22 | 22 | ||
| 23 | def __init__(self): | 23 | def __init__(self): |
| 24 | - super().__init__("matmul_emu_split_weight", category="matmul") | 24 | + super().__init__("matmul_emu_split_weight") |
| 25 | 25 | ||
| 26 | def sources(self) -> list: | 26 | def sources(self) -> list: |
| 27 | return [self.resolve_source("matmul_emu_split_weight.cpp")] | 27 | return [self.resolve_source("matmul_emu_split_weight.cpp")] |
| @@ -24,7 +24,7 @@ class QuantMatmulActivationQuantOpBuilder(OpBuilder): | |||
| 24 | super().__init__("quant_matmul_activation_quant") | 24 | super().__init__("quant_matmul_activation_quant") |
| 25 | 25 | ||
| 26 | def sources(self) -> list: | 26 | def sources(self) -> list: |
| 27 | - return ["csrc/matmul/quant_matmul_activation_quant.cpp"] | 27 | + return [self.resolve_source("quant_matmul_activation_quant.cpp")] |
| 28 | 28 | ||
| 29 | def schema(self) -> str: | 29 | def schema(self) -> str: |
| 30 | return ( | 30 | return ( |
| @@ -27,7 +27,7 @@ class RmsNormDynamicQuantOpBuilder(OpBuilder): | |||
| 27 | super().__init__("rms_norm_dynamic_quant") | 27 | super().__init__("rms_norm_dynamic_quant") |
| 28 | 28 | ||
| 29 | def sources(self) -> list: | 29 | def sources(self) -> list: |
| 30 | - return ["csrc/norm/rms_norm_dynamic_quant.cpp"] | 30 | + return [self.resolve_source("rms_norm_dynamic_quant.cpp")] |
| 31 | 31 | ||
| 32 | def schema(self) -> str: | 32 | def schema(self) -> str: |
| 33 | return ( | 33 | return ( |
| @@ -24,10 +24,10 @@ class FlatQuantOpBuilder(OpBuilder): | |||
| 24 | """ | 24 | """ |
| 25 | 25 | ||
| 26 | def __init__(self): | 26 | def __init__(self): |
| 27 | - super().__init__("flat_quant", "quant") | 27 | + super().__init__("flat_quant") |
| 28 | 28 | ||
| 29 | def sources(self) -> list: | 29 | def sources(self) -> list: |
| 30 | - return ["csrc/quant/flat_quant.cpp"] | 30 | + return [self.resolve_source("flat_quant.cpp")] |
| 31 | 31 | ||
| 32 | def schema(self) -> str: | 32 | def schema(self) -> str: |
| 33 | return ( | 33 | return ( |
Mquant/grouped_dynamic_mx_quant_with_dual_axis/torch_extension/grouped_dynamic_mx_quant_with_dual_axis.py+1-1
| @@ -21,7 +21,7 @@ class GroupedDynamicMxQuantWithDualAxisOpBuilder(OpBuilder): | |||
| 21 | """ | 21 | """ |
| 22 | 22 | ||
| 23 | def __init__(self): | 23 | def __init__(self): |
| 24 | - super().__init__("grouped_dynamic_mx_quant_with_dual_axis", category="quant") | 24 | + super().__init__("grouped_dynamic_mx_quant_with_dual_axis") |
| 25 | 25 | ||
| 26 | def sources(self) -> list: | 26 | def sources(self) -> list: |
| 27 | return [self.resolve_source("grouped_dynamic_mx_quant_with_dual_axis.cpp")] | 27 | return [self.resolve_source("grouped_dynamic_mx_quant_with_dual_axis.cpp")] |
| @@ -50,7 +50,7 @@ class MxToBlockMxQuantOpBuilder(OpBuilder): | |||
| 50 | super().__init__("mx_to_block_mx_quant") | 50 | super().__init__("mx_to_block_mx_quant") |
| 51 | 51 | ||
| 52 | def sources(self) -> list: | 52 | def sources(self) -> list: |
| 53 | - return ["csrc/quant/mx_to_block_mx_quant.cpp"] | 53 | + return [self.resolve_source("mx_to_block_mx_quant.cpp")] |
| 54 | 54 | ||
| 55 | def schema(self) -> str: | 55 | def schema(self) -> str: |
| 56 | return ( | 56 | return ( |
| @@ -92,7 +92,7 @@ def _check_swiglu_group_backward_inputs( | |||
| 92 | 92 | ||
| 93 | class SwigluGroupBackwardOpBuilder(OpBuilder): | 93 | class SwigluGroupBackwardOpBuilder(OpBuilder): |
| 94 | def __init__(self): | 94 | def __init__(self): |
| 95 | - super().__init__("swiglu_group_backward", category="quant") | 95 | + super().__init__("swiglu_group_backward") |
| 96 | 96 | ||
| 97 | def sources(self): | 97 | def sources(self): |
| 98 | return [self.resolve_source("swiglu_group_grad.cpp")] | 98 | return [self.resolve_source("swiglu_group_grad.cpp")] |
| @@ -69,10 +69,10 @@ TORCH_EXTENSION_OPS=swiglu_group TORCH_EXTENSION_VENDOR=custom python3 -m build | |||
| 69 | 69 | ||
| 70 | ```sh | 70 | ```sh |
| 71 | # 安装整包 | 71 | # 安装整包 |
| 72 | -python3 -m pip install dist/cann_ops_nn-*.whl --no-deps | 72 | +python3 -m pip install dist/cann_ops_nn-*.whl |
| 73 | 73 | ||
| 74 | # 安装单算子包 | 74 | # 安装单算子包 |
| 75 | -python3 -m pip install dist/cann_ops_nn_custom-*.whl --no-deps | 75 | +python3 -m pip install dist/cann_ops_nn_custom-*.whl |
| 76 | ``` | 76 | ``` |
| 77 | 77 | ||
| 78 | ### 整包与子包共存机制 | 78 | ### 整包与子包共存机制 |
| @@ -86,10 +86,10 @@ python3 -m pip install dist/cann_ops_nn_custom-*.whl --no-deps | |||
| 86 | 86 | ||
| 87 | ```sh | 87 | ```sh |
| 88 | # 安装整包 | 88 | # 安装整包 |
| 89 | -pip install cann_ops_nn-*.whl --no-deps | 89 | +pip install cann_ops_nn-*.whl |
| 90 | 90 | ||
| 91 | # 安装单算子包（覆盖整包中的同名算子） | 91 | # 安装单算子包（覆盖整包中的同名算子） |
| 92 | -pip install cann_ops_nn_custom-*.whl --no-deps | 92 | +pip install cann_ops_nn_custom-*.whl |
| 93 | 93 | ||
| 94 | # 卸载单算子包（整包算子自动恢复） | 94 | # 卸载单算子包（整包算子自动恢复） |
| 95 | pip uninstall cann-ops-nn-custom | 95 | pip uninstall cann-ops-nn-custom |
| @@ -100,29 +100,30 @@ pip uninstall cann-ops-nn-custom | |||
| 100 | ``` | 100 | ``` |
| 101 | ├── torch_extension | 101 | ├── torch_extension |
| 102 | │ ├── build.sh # 构建脚本（支持 --torch_extension 参数） | 102 | │ ├── build.sh # 构建脚本（支持 --torch_extension 参数） |
| 103 | -│ ├── setup.py # wheel 打包配置（支持单算子编包） | 103 | +│ ├── setup.py # wheel 打包配置（自动收集分布式算子） |
| 104 | │ ├── requirements.txt | 104 | │ ├── requirements.txt |
| 105 | -│ ├── cann_ops_nn # 整包安装目录 | 105 | +│ └── cann_ops_nn # 整包安装目录（框架公共部分） |
| 106 | -│ │ ├── __init__.py # 包入口，from .ops import * | 106 | +│ ├── __init__.py # 包入口，from .ops import * |
| 107 | -│ │ ├── op_builder/ | 107 | +│ ├── op_builder/ |
| 108 | -│ │ │ └── builder.py # OpBuilder 基类，JIT 编译管理 | 108 | +│ │ └── builder.py # OpBuilder 基类，JIT 编译管理 |
| 109 | -│ │ ├── common/ | 109 | +│ ├── common/ |
| 110 | -│ │ │ └── aclnn_common.h # ACLNN_CMD 宏等公共能力 | 110 | +│ │ └── aclnn_common.h # ACLNN_CMD 宏等公共能力 |
| 111 | -│ │ ├── csrc/ | 111 | +│ ├── csrc/ |
| 112 | -│ │ │ └── <category>/<op>.cpp # C++ kernel wrapper | 112 | +│ │ └── extension.cpp # 占位模块（满足 PyTorch 扩展要求） |
| 113 | -│ │ ├── ops/ | 113 | +│ ├── ops/ |
| 114 | -│ │ │ ├── __init__.py # 自动发现算子 + entry point 加载 | 114 | +│ │ └── __init__.py # 自动发现算子 + entry point 加载（构建时填充） |
| 115 | -│ │ │ └── <category>/<op>/<op>.py # Python 前端 | 115 | +│ └── docs/ |
| 116 | -│ │ └── docs/ | 116 | +│ └── torch_extension_guidelines.md # 开发规范 |
| 117 | -│ │ └── torch_extension_guidelines.md # 开发规范 | 117 | +├── <category>/<op>/torch_extension/ # 算子开发目录（仓库根目录下） |
| 118 | -│ └── dist/ # 构建产物输出目录 | 118 | +│ ├── <op>.py # Python 前端（OpBuilder + schema + meta） |
| 119 | -├── <category>/<op>/torch_extension/ # 分布式算子（仓库根目录下） | ||
| 120 | -│ ├── <op>.py | ||
| 121 | │ ├── __init__.py | 119 | │ ├── __init__.py |
| 122 | -│ └── csrc/<op>.cpp | 120 | +│ ├── graph_convert_<op>.py # （可选）图回退转换 |
| 121 | +│ └── csrc/<op>.cpp # C++ kernel wrapper | ||
| 123 | └── build_out/ # build.sh 输出目录 | 122 | └── build_out/ # build.sh 输出目录 |
| 124 | ``` | 123 | ``` |
| 125 | 124 | ||
| 125 | +> **构建时收集：** `setup.py` 在打包阶段自动遍历仓库中所有 `<category>/<op>/torch_extension/`，把 `<op>.py`、`csrc/<op>.cpp` 汇集到 wheel 包目录 `cann_ops_nn/ops/<category>/<op>/` 与 `cann_ops_nn/csrc/` 下。开发者无需手动在 `cann_ops_nn/` 下创建算子文件。 | ||
| 126 | + | ||
| 126 | ## 快速入门 | 127 | ## 快速入门 |
| 127 | 128 | ||
| 128 | ```python | 129 | ```python |
| @@ -144,19 +145,19 @@ print(result.shape) # torch.Size([16, 16]) | |||
| 144 | 145 | ||
| 145 | ## 开发者指南：新增算子 | 146 | ## 开发者指南：新增算子 |
| 146 | 147 | ||
| 147 | -以添加新算子 `new_operator` 为例，需提供 C++ 绑定和 Python 构建器。 | 148 | +以添加新算子 `new_operator`（类别 `my_cat`）为例，在算子目录 `<my_cat>/new_operator/` 下创建 `torch_extension/` 子目录，包含 C++ 绑定和 Python 构建器。 |
| 148 | 149 | ||
| 149 | -### 1. C++ 后端 (`cann_ops_nn/csrc/<op_category>/new_operator.cpp`) | 150 | +### 1. C++ 后端 (`<my_cat>/new_operator/torch_extension/csrc/new_operator.cpp`) |
| 150 | 151 | ||
| 151 | 该文件将 PyTorch 张量桥接到 ACLNN C-API。 | 152 | 该文件将 PyTorch 张量桥接到 ACLNN C-API。 |
| 152 | 153 | ||
| 153 | ```cpp | 154 | ```cpp |
| 154 | #include <torch/extension.h> | 155 | #include <torch/extension.h> |
| 155 | #include "aclnnop/aclnn_new_operator.h" | 156 | #include "aclnnop/aclnn_new_operator.h" |
| 156 | -#include "../common/aclnn_common.h" | 157 | +#include "aclnn_common.h" // 通过 -I cann_ops_nn/common 解析 |
| 157 | 158 | ||
| 158 | namespace cann_ops_nn { | 159 | namespace cann_ops_nn { |
| 159 | -namespace <op_category> { | 160 | +namespace my_cat { |
| 160 | 161 | ||
| 161 | at::Tensor new_operator( | 162 | at::Tensor new_operator( |
| 162 | const at::Tensor& input1, | 163 | const at::Tensor& input1, |
| @@ -178,16 +179,16 @@ at::Tensor new_operator( | |||
| 178 | return out; | 179 | return out; |
| 179 | } | 180 | } |
| 180 | 181 | ||
| 181 | -} // namespace <op_category> | 182 | +} // namespace my_cat |
| 182 | } // namespace cann_ops_nn | 183 | } // namespace cann_ops_nn |
| 183 | 184 | ||
| 184 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { | 185 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { |
| 185 | - m.def("new_operator", &cann_ops_nn::<op_category>::new_operator, | 186 | + m.def("new_operator", &cann_ops_nn::my_cat::new_operator, |
| 186 | "NewOperator on NPU"); | 187 | "NewOperator on NPU"); |
| 187 | } | 188 | } |
| 188 | ``` | 189 | ``` |
| 189 | 190 | ||
| 190 | -### 2. Python 前端 (`cann_ops_nn/ops/<op_category>/new_operator.py`) | 191 | +### 2. Python 前端 (`<my_cat>/new_operator/torch_extension/new_operator.py`) |
| 191 | 192 | ||
| 192 | 该文件管理 JIT 编译逻辑并将算子注册到 PyTorch Dispatcher。 | 193 | 该文件管理 JIT 编译逻辑并将算子注册到 PyTorch Dispatcher。 |
| 193 | 194 | ||
| @@ -203,7 +204,7 @@ class NewOperatorOpBuilder(OpBuilder): | |||
| 203 | 204 | ||
| 204 | def sources(self): | 205 | def sources(self): |
| 205 | """C++ 源码路径。""" | 206 | """C++ 源码路径。""" |
| 206 | - return ['csrc/<op_category>/new_operator.cpp'] | 207 | + return [self.resolve_source("new_operator.cpp")] |
| 207 | 208 | ||
| 208 | def schema(self) -> str: | 209 | def schema(self) -> str: |
| 209 | """PyTorch 算子签名。""" | 210 | """PyTorch 算子签名。""" |
| @@ -25,8 +25,7 @@ nn仓的算子采用**分布式目录结构**：每个算子的torch_extension | |||
| 25 | │ │ ├── common/ | 25 | │ │ ├── common/ |
| 26 | │ │ │ └── aclnn_common.h # ACLNN_CMD宏、类型转换等公共能力（一般无需修改） | 26 | │ │ │ └── aclnn_common.h # ACLNN_CMD宏、类型转换等公共能力（一般无需修改） |
| 27 | │ │ ├── csrc/ # C++源码staging目录（构建时自动收集） | 27 | │ │ ├── csrc/ # C++源码staging目录（构建时自动收集） |
| 28 | -│ │ │ └── ${category}/ | 28 | +│ │ │ └── ${op_api}.cpp |
| 29 | -│ │ │ └── ${op_api}.cpp | ||
| 30 | │ │ └── ops/ | 29 | │ │ └── ops/ |
| 31 | │ │ ├── __init__.py # 算子自动发现与加载 | 30 | │ │ ├── __init__.py # 算子自动发现与加载 |
| 32 | │ │ └── ${category}/ | 31 | │ │ └── ${category}/ |
| @@ -200,7 +199,7 @@ nn仓的算子采用**分布式目录结构**：每个算子的torch_extension | |||
| 200 | 负责JIT编译管理、schema/meta注册与对外接口封装： | 199 | 负责JIT编译管理、schema/meta注册与对外接口封装： |
| 201 | 200 | ||
| 202 | 1. **OpBuilder子类**：继承`OpBuilder`，在`__init__`中以`super().__init__("<schema算子名>")`传入算子名，并实现三个抽象方法： | 201 | 1. **OpBuilder子类**：继承`OpBuilder`，在`__init__`中以`super().__init__("<schema算子名>")`传入算子名，并实现三个抽象方法： |
| 203 | - - `sources()`：返回相对`cann_ops_nn`包根的C++源文件路径列表，如`["csrc/activation/swiglu_group.cpp"]`； | 202 | + - `sources()`：返回相对`cann_ops_nn`包根的C++源文件路径列表，统一通过`self.resolve_source("${op_api}.cpp")`获取（如`[self.resolve_source("swiglu_group.cpp")]`）； |
| 204 | - `schema()`：返回算子schema字符串（见[2.3 Schema标识符](#schema标识符算子签名)）； | 203 | - `schema()`：返回算子schema字符串（见[2.3 Schema标识符](#schema标识符算子签名)）； |
| 205 | - `register_meta()`：用`@impl(get_as_library(), self.name, "Meta")`注册Meta实现，仅做shape/dtype推导，不触碰真实NPU计算（FakeTensor/图模式必需）。 | 204 | - `register_meta()`：用`@impl(get_as_library(), self.name, "Meta")`注册Meta实现，仅做shape/dtype推导，不触碰真实NPU计算（FakeTensor/图模式必需）。 |
| 206 | 2. **实例化与初始化**：模块加载时实例化builder并调用`_ensure_initialized()`注册schema与meta： | 205 | 2. **实例化与初始化**：模块加载时实例化builder并调用`_ensure_initialized()`注册schema与meta： |
| @@ -225,8 +224,8 @@ nn仓的算子采用**分布式目录结构**：每个算子的torch_extension | |||
| 225 | def __init__(self): | 224 | def __init__(self): |
| 226 | super().__init__("swiglu_group") | 225 | super().__init__("swiglu_group") |
| 227 | 226 | ||
| 228 | - def sources(self): | 227 | + def sources(self): |
| 229 | - return ["csrc/activation/swiglu_group.cpp"] | 228 | + return [self.resolve_source("swiglu_group.cpp")] |
| 230 | 229 | ||
| 231 | def schema(self): | 230 | def schema(self): |
| 232 | return "swiglu_group(Tensor x, *, Tensor? weight=None, Tensor? group_index=None, float clamp_limit=-1.0) -> Tensor" | 231 | return "swiglu_group(Tensor x, *, Tensor? weight=None, Tensor? group_index=None, float clamp_limit=-1.0) -> Tensor" |
| @@ -9,6 +9,7 @@ | |||
| 9 | # ----------------------------------------------------------------------------------------------------------- | 9 | # ----------------------------------------------------------------------------------------------------------- |
| 10 | 10 | ||
| 11 | import os | 11 | import os |
| 12 | +import shutil | ||
| 12 | from abc import ABC, abstractmethod | 13 | from abc import ABC, abstractmethod |
| 13 | from typing import List, Union | 14 | from typing import List, Union |
| 14 | import torch | 15 | import torch |
| @@ -34,14 +35,12 @@ class OpBuilder(ABC): | |||
| 34 | 基于 aclnn 的算子构建基类 | 35 | 基于 aclnn 的算子构建基类 |
| 35 | 36 | ||
| 36 | :param name: 算子名称，如 'mat_mul_v3' | 37 | :param name: 算子名称，如 'mat_mul_v3' |
| 37 | - :param category: 算子类别目录名，如 'matmul'。用于解析源码路径：csrc/<category>/<name>.cpp | ||
| 38 | """ | 38 | """ |
| 39 | 39 | ||
| 40 | _loaded_ops = {} | 40 | _loaded_ops = {} |
| 41 | 41 | ||
| 42 | - def __init__(self, name, category=None): | 42 | + def __init__(self, name): |
| 43 | self.name = name | 43 | self.name = name |
| 44 | - self.category = category | ||
| 45 | self._initialized = False | 44 | self._initialized = False |
| 46 | 45 | ||
| 47 | def _ensure_initialized(self): | 46 | def _ensure_initialized(self): |
| @@ -79,10 +78,7 @@ class OpBuilder(ABC): | |||
| 79 | return [os.path.join(self._package_path, path) for path in paths] | 78 | return [os.path.join(self._package_path, path) for path in paths] |
| 80 | 79 | ||
| 81 | def resolve_source(self, cpp_filename): | 80 | def resolve_source(self, cpp_filename): |
| 82 | - self._ensure_initialized() | 81 | + return f"csrc/{cpp_filename}" |
 | |||
| 83 | - if self.category is None: | ||
| 84 | - return cpp_filename | ||
| 85 | - return f"csrc/{self.category}/{cpp_filename}" | ||
| 86 | 82 | ||
| 87 | def register_schema(self, op_schema: Union[str, List[str]]): | 83 | def register_schema(self, op_schema: Union[str, List[str]]): |
| 88 | if isinstance(op_schema, str): | 84 | if isinstance(op_schema, str): |
| @@ -175,14 +171,30 @@ class OpBuilder(ABC): | |||
| 175 | if self.name in OpBuilder._loaded_ops: | 171 | if self.name in OpBuilder._loaded_ops: |
| 176 | return OpBuilder._loaded_ops[self.name] | 172 | return OpBuilder._loaded_ops[self.name] |
| 177 | 173 | ||
| 178 | - op_module = load( | 174 | + if shutil.which("ninja") is None: |
| 179 | - name=self.name, | 175 | + raise RuntimeError( |
| 180 | - sources=self.get_absolute_paths(self.sources()), | 176 | + f"ninja is required to JIT compile operator '{self.name}', " |
| 181 | - extra_include_paths=self.get_absolute_paths(self.include_paths()), | 177 | + f"please install it via 'pip install ninja'" |
| 182 | - extra_cflags=self.cxx_args(), | 178 | + ) |
| 183 | - extra_ldflags=self.extra_ldflags(), | 179 | + |
| 184 | - verbose=verbose, | 180 | + try: |
| 185 | - ) | 181 | + op_module = load( |
| 182 | + name=self.name, | ||
| 183 | + sources=self.get_absolute_paths(self.sources()), | ||
| 184 | + extra_include_paths=self.get_absolute_paths(self.include_paths()), | ||
| 185 | + extra_cflags=self.cxx_args(), | ||
| 186 | + extra_ldflags=self.extra_ldflags(), | ||
| 187 | + verbose=verbose, | ||
| 188 | + ) | ||
| 189 | + except Exception as e: | ||
| 190 | + raise RuntimeError( | ||
| 191 | + f"Failed to JIT compile operator '{self.name}': {e}\n" | ||
| 192 | + f"Common causes:\n" | ||
| 193 | + f" 1. CANN toolkit not sourced: source <cann_path>/set_env.sh\n" | ||
| 194 | + f" 2. Missing compiler: ensure gcc/g++ in PATH\n" | ||
| 195 | + f" 3. Missing ninja: pip install ninja" | ||
| 196 | + ) from e | ||
| 197 | + | ||
| 186 | OpBuilder._loaded_ops[self.name] = op_module | 198 | OpBuilder._loaded_ops[self.name] = op_module |
| 187 | 199 | ||
| 188 | return op_module | 200 | return op_module |
| @@ -76,14 +76,14 @@ def _non_python_files(directory): | |||
| 76 | if _selected_ops is not None: | 76 | if _selected_ops is not None: |
| 77 | rel = os.path.relpath(path, directory) | 77 | rel = os.path.relpath(path, directory) |
| 78 | parts = rel.split(os.sep) | 78 | parts = rel.split(os.sep) |
| 79 | - if ( | 79 | + if len(parts) >= 2 and parts[0] == "ops" and parts[1] not in _selected_ops: |
| 80 | - len(parts) >= 2 | ||
| 81 | - and parts[0] in ("ops", "csrc") | ||
| 82 | - and parts[1] not in _selected_ops | ||
| 83 | - ): | ||
| 84 | continue | 80 | continue |
| 85 | for filename in filenames: | 81 | for filename in filenames: |
| 86 | if filename.endswith((".h", ".cpp")): | 82 | if filename.endswith((".h", ".cpp")): |
| 83 | + if _selected_ops is not None and filename.endswith(".cpp"): | ||
| 84 | + op_name = filename[:-4] | ||
| 85 | + if op_name not in _selected_ops: | ||
| 86 | + continue | ||
| 87 | paths.append(os.path.join(path, filename)) | 87 | paths.append(os.path.join(path, filename)) |
| 88 | return paths | 88 | return paths |
| 89 | 89 | ||
| @@ -103,6 +103,23 @@ _op_py_files = [] | |||
| 103 | _op_cpp_files = [] | 103 | _op_cpp_files = [] |
| 104 | 104 | ||
| 105 | 105 | ||
| 106 | +def _collect_cpp(csrc_dir): | ||
| 107 | + if not os.path.isdir(csrc_dir): | ||
| 108 | + return | ||
| 109 | + for cpp in os.listdir(csrc_dir): | ||
| 110 | + if not cpp.endswith(".cpp"): | ||
| 111 | + continue | ||
| 112 | + dst_rel = os.path.join("csrc", cpp) | ||
| 113 | + src_abs = os.path.join(csrc_dir, cpp) | ||
| 114 | + for existing_rel, existing_src in _op_cpp_files: | ||
| 115 | + if existing_rel == dst_rel: | ||
| 116 | + raise ValueError( | ||
| 117 | + "Duplicate csrc file '%s' from '%s' conflicts with " | ||
| 118 | + "already collected '%s'" % (dst_rel, src_abs, existing_src) | ||
| 119 | + ) | ||
| 120 | + _op_cpp_files.append((dst_rel, src_abs)) | ||
| 121 | + | ||
| 122 | + | ||
| 106 | def _collect_op(cat, name, torch_extension): | 123 | def _collect_op(cat, name, torch_extension): |
| 107 | # Collect one operator's torch_extension dir, staging its files under (cat, name). | 124 | # Collect one operator's torch_extension dir, staging its files under (cat, name). |
| 108 | _selected_op_categories.append((cat, name)) | 125 | _selected_op_categories.append((cat, name)) |
| @@ -120,13 +137,7 @@ def _collect_op(cat, name, torch_extension): | |||
| 120 | ) | 137 | ) |
| 121 | ) | 138 | ) |
| 122 | 139 | ||
| 123 | - csrc_dir = os.path.join(torch_extension, "csrc") | 140 | + _collect_cpp(os.path.join(torch_extension, "csrc")) |
| 124 | - if os.path.isdir(csrc_dir): | ||
| 125 | - for cpp in os.listdir(csrc_dir): | ||
| 126 | - if cpp.endswith(".cpp"): | ||
| 127 | - _op_cpp_files.append( | ||
| 128 | - (os.path.join("csrc", cat, cpp), os.path.join(csrc_dir, cpp)) | ||
| 129 | - ) | ||
| 130 | 141 | ||
| 131 | 142 | ||
| 132 | for cat in sorted(os.listdir(OPS_NN_ROOT)): | 143 | for cat in sorted(os.listdir(OPS_NN_ROOT)): |
| @@ -198,16 +209,7 @@ if os.path.isdir(_staged_ops_dir): | |||
| 198 | graph_src, | 209 | graph_src, |
| 199 | ) | 210 | ) |
| 200 | ) | 211 | ) |
| 201 | - csrc_dir = os.path.join(op_path, "csrc") | 212 | + _collect_cpp(os.path.join(op_path, "csrc")) |
| 202 | - if os.path.isdir(csrc_dir): | ||
| 203 | - for cpp in os.listdir(csrc_dir): | ||
| 204 | - if cpp.endswith(".cpp"): | ||
| 205 | - _op_cpp_files.append( | ||
| 206 | - ( | ||
| 207 | - os.path.join("csrc", cat, cpp), | ||
| 208 | - os.path.join(csrc_dir, cpp), | ||
| 209 | - ) | ||
| 210 | - ) | ||
| 211 | 213 | ||
| 212 | 214 | ||
| 213 | _sorted_selected_op_categories = sorted(set(_selected_op_categories)) | 215 | _sorted_selected_op_categories = sorted(set(_selected_op_categories)) |
| @@ -252,6 +254,15 @@ class BuildPyWithOps(_build_py): | |||
| 252 | base_dir = os.path.join(build_pkg, subdir) | 254 | base_dir = os.path.join(build_pkg, subdir) |
| 253 | if not os.path.isdir(base_dir): | 255 | if not os.path.isdir(base_dir): |
| 254 | continue | 256 | continue |
| 257 | + if subdir == "csrc": | ||
| 258 | + for name in os.listdir(base_dir): | ||
| 259 | + op_name = name[:-4] if name.endswith(".cpp") else name | ||
| 260 | + if op_name not in _selected_ops and name != "__pycache__": | ||
| 261 | + target = os.path.join(base_dir, name) | ||
| 262 | + if os.path.isfile(target): | ||
| 263 | + os.remove(target) | ||
| 264 | + logger.info("removing centralized %s/%s", subdir, name) | ||
| 265 | + continue | ||
| 255 | for cat in os.listdir(base_dir): | 266 | for cat in os.listdir(base_dir): |
| 256 | cat_dir = os.path.join(base_dir, cat) | 267 | cat_dir = os.path.join(base_dir, cat) |
| 257 | if not os.path.isdir(cat_dir): | 268 | if not os.path.isdir(cat_dir): |
| @@ -382,7 +393,11 @@ setup( | |||
| 382 | author="CANN", | 393 | author="CANN", |
| 383 | license="CANN Open Software License Agreement Version 2.0", | 394 | license="CANN Open Software License Agreement Version 2.0", |
| 384 | url="https://gitcode.com/cann/ops-nn/tree/master/torch_extension", | 395 | url="https://gitcode.com/cann/ops-nn/tree/master/torch_extension", |
| 385 | - install_requires=["torch>=2.6.0", "torch_npu"], | 396 | + install_requires=[ |
| 397 | + "ninja", | ||
| 398 | + "torch", | ||
| 399 | + "torch_npu", | ||
| 400 | + ], | ||
| 386 | packages=_all_packages, | 401 | packages=_all_packages, |
| 387 | package_data={PACKAGE_NAME: _non_python_files(_src_path)}, | 402 | package_data={PACKAGE_NAME: _non_python_files(_src_path)}, |
| 388 | entry_points={"cann_ops_nn.ops": _entry_points} if _entry_points else {}, | 403 | entry_points={"cann_ops_nn.ops": _entry_points} if _entry_points else {}, |
请尽量保证包里的目录结构层次和仓里的一致