# ms_ops

**Repository Path**: liangchenghui/ms_ops

## Basic Information

- **Project Name**: ms_ops
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 6
- **Created**: 2026-05-21
- **Last Updated**: 2026-05-21

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# ms_ops — MindSpore 自定义 ACLNN 算子库

ms_ops 是一个基于 MindSpore 的自定义算子库，采用"ACLNN 层 + Function 层"两层分离架构。

## 架构概览

```
Python 调用层
  └─ Function 层 (正向 + 反向 + autograd)
       └─ ACLNN 层 (ACLNN API 调用或复用 PyBoost)
```

## 安装

支持 **Dev** 和 **Release** 两种模式。

### 环境要求

- Python >= 3.10
- MindSpore (Ascend 后端，已安装并可用)
- Ascend CANN 工具包（包含 ACLNN 库）
- 环境变量 `ASCEND_OPP_PATH` 已设置

### Dev 模式（默认，JIT 编译）

Dev 模式使用 MindSpore 的 `CustomOpBuilder` 在导入时自动 JIT 编译，无需提前构建：

```bash
# 直接使用，自动编译
python -c "from ms_ops import ops; print(ops.relu(...))"
```

或使用统一编译脚本：

```bash
python compile.py                              # 编译全部
```

### Release 模式（预编译 .so）

适用于部署发布，无需在线编译：

```bash
# 构建 .so 到包目录
python compile.py --release --inplace

# 或通过 pip 安装
pip install .
```

Release 模式需设置环境变量后使用：

```bash
export MS_OPS_MODE=release
python -c "from ms_ops import ops; print(ops.relu(...))"
```

## 快速使用

```python
from ms_ops import ops

# 正向计算
y = ops.fast_gelu(x)

# 梯度计算（自动支持反向）
from mindspore import ops as ms_ops
grad_fn = ms_ops.value_and_grad(lambda x: ops.fast_gelu(x).sum())
y, grads = grad_fn(x)
```

已实现的算子参考[算子清单](docs/menu_ms_ops.md)。

## 环境变量

| 变量 | 默认值 | 说明 |
|------|--------|------|
| `MS_OPS_MODE` | (auto) | 加载模式（不设置时自动：先试 release 再 fallback dev）。`dev` 强制 JIT 编译，`release` 仅加载预编译 .so |
| `MS_OPS_DEV_VERBOSE` | `0` | Dev 模式下打印详细编译日志 |
| `ASCEND_OPP_PATH` | (必填) | Ascend CANN OPP 路径，需手动设置 |

## 编译模式说明

### Auto 模式（默认，不设置 MS_OPS_MODE）

- 先尝试加载 `ms_ops/_C/ops.so`（release 编译产物）
- 如果 .so 不存在，自动 fallback 到 Dev 模式 JIT 编译
- 适合开发阶段：第一次导入时 JIT 编译，后续可 `--release --inplace` 后自动切换到 release

### Dev 模式（MS_OPS_MODE=dev）

- 每次 `from ms_ops import ops` 时强制 JIT 编译所有 C++ 源码
- 适合开发和调试阶段
- 可通过 `MS_OPS_DEV_VERBOSE=1` 查看详细编译日志

### Release 模式（MS_OPS_MODE=release）

- 仅加载预编译的 `.so`，不存在时报错
- 编译一次，多次加载
- `setup.py` 将所有源文件编译为 `ms_ops/_C/ops.so`

## 配置格式

算子定义采用 YAML 配置，主要有三类配置文件：

| 文件 | 用途 |
|------|------|
| `config/pyboost_ops.yaml` | 复用已有 PyBoost 算子，只需指定 `exec` |
| `config/aclnn_ops.yaml` | 自定义新增 ACLNN 算子，需定义 `gen_opapi` |
| `config/function_definitions.yaml` | Function 层定义（正向 + saved + 反向公式） |

配置格式详见 [op_definitions_reference.md](docs/config/op_definitions_reference.md) 和 [function_definitions_reference.md](docs/config/function_definitions_reference.md)。

## 新增算子（流程 A/B）

算子的接入分为两种流程：

- **流程 A（全自动）**：当算子在 `reference/op-plugin` 已有 pytorch_npu 实现时，从 `derivatives.yaml` 和 `op_plugin_functions.yaml` 自动提取配置并生成代码。
- **流程 B（手动配置）**：对于新算子，手动编写 YAML 配置 + infer 函数，再由 codegen 生成代码。

详细接入流程见[开发指南](docs/ms_ops_development_guide.md#新增算子流程)。

## 代码生成

```bash
# 生成所有算子的代码
python codegen/regenerate_all.py

# 生成指定 ACLNN 算子
python codegen/gen_ops.py --op npu_fast_gelu

# 生成指定 Function 算子
python codegen/gen_functions.py --func fast_gelu
```

代码生成器架构见 [PLAN.md](PLAN.md#六代码生成器设计)。

## 测试

```bash
# 运行所有测试（Dev 模式默认）
pytest tests/

# Dev/Release 模式构建测试
pytest tests/test_build_modes.py -v

# 运行指定算子的测试
pytest tests/test_functions/test_fast_gelu.py -v
```

测试覆盖正向/反向/大张量/零输入/精度对比等场景。

测试文件使用 `from ms_ops import ops` 导入，不依赖 `compile.py` 的返回值。

## Sub-Agent 使用

本项目在 Claude Code 环境下提供一系列 sub-agent 来帮助完成算子接入的完整流程。使用 `Agent` 工具并指定 `subagent_type` 即可调用。

### 对接流程示意

```
您（需求） → 协调 Agent → 算子接入 Agent → Code Gen Agent
                    → 测试 Agent → Review Agent → Commit
```

### Sub-Agent 列表

| Agent | subagent_type | 职责 |
|-------|--------------|------|
| **Coordinator**（协调） | `coordinator` | 管理算子接入全流程，协调各 sub-agent 完成任务 |
| **Op Integration**（算子接入） | `op_integration` | 具体算子的接入执行者。负责 YAML 配置、运行代码生成、生成算子文档 |
| **Code Gen**（代码生成） | `code_gen` | 维护和优化 `codegen/` 下的代码生成逻辑、模板和 YAML 配置文件 |
| **Test**（测试验证） | `test` | 负责测试用例 review 和功能验证 |
| **Review**（代码审查） | `review` | 对照 checklist 检查算子交付件，无问题后提交 commit |

更多 sub-agent 配置细节见项目 memory 文件。

## 项目结构

```
ms_ops/
├── config/                # YAML 配置文件（算子定义）
├── src/
│   ├── common/            # 公共代码（ACLNN 执行器、infer 函数等）
│   ├── ops/               # ACLNN 层生成代码
│   └── functions/         # Function 层生成代码（正向 + 反向）
├── codegen/               # 代码生成器（Python + Jinja2 模板）
│   └── templates/         # Jinja2 模板（test.py.j2, test_aclnn.py.j2 等）
├── tests/                 # 测试用例
│   ├── test_functions/    # Function 层测试
│   ├── test_ops/          # ACLNN 层测试
│   └── test_build_modes.py  # Dev/Release 模式测试
├── docs/                  # 文档
├── ms_ops/                # Python 包
│   ├── __init__.py        # 包入口
│   ├── ops.py             # 算子接口（Dev/Release 模式分发）
│   ├── dev_loader.py      # Dev 模式 JIT 加载器
│   └── _C/                # Release 模式 .so 存放目录
├── reference/             # op-plugin 参考实现
├── compile.py             # 统一编译脚本
└── setup.py               # Release 模式打包（setuptools）
```

详细目录结构见 [directory_structure.md](docs/directory_structure.md)。