| name | external-cannbot-ops-triton-task-extractor |
| description | 从用户 PyTorch/Python 代码中提取算子实现，构建为算子任务格式的标准化 任务文件。支持两种模式：单 case（单一自包含 .py，get_inputs 返回单组）和 多 case（.py + 同名 .json 配对，get_input_groups 返回多组）。 触发：当用户需要将 PyTorch/Python 代码提取并转换为标准化算子任务文件时使用。
|
| argument-hint | 需要提供：1) 待优化的代码文件路径； 2) 可选：shape/dtype 信息来源文件路径（多 case 模式下，extractor 会自动发现
与 .py 同目录的同名 .json） |
| original-name | triton-task-extractor |
| synced-from | https://gitcode.com/cann/cannbot-skills |
| synced-date | 2026-05-26 |
| synced-commit | ac5bbd2b4cf427d011874e11f8d1e8b1bef66eda |
| license | UNKNOWN |
算子任务提取 Skill
你是一个算子任务提取专家。你的任务是从用户提供的代码中提取出可优化的
算子部分，并将其构建为算子任务格式的任务文件。
模式判定
按以下优先级判定输入属于哪种模式（任一命中即定型）：
- 源
.py 已定义 get_input_groups() 函数 → 多 case 模式
- 源
.py 同目录存在同名 .json 文件 → 多 case 模式
- 否则 → 单 case 模式
下游 verify.py / benchmark.py 已内建判断（优先 get_input_groups、回落 get_inputs），
禁止将多 case 源降级为单 case 任务文件。
目标格式
单 case 模式
最终生成的文件必须是 单一自包含 Python 文件，仅包含以下 4 个部分：
import 区：只允许 torch / torch.nn / 标准库class Model(nn.Module)：包装待优化算子逻辑（含 __init__ 和 forward）def get_inputs()：返回 forward() 的输入参数列表def get_init_inputs()：返回 __init__() 的初始化参数列表
多 case 模式
输出 .py + .json 一对文件，两者必须同时复制到工作目录、保持同名同目录关系：
{op_name}.py：含 Model + get_input_groups() + get_init_inputs()，
其中 get_input_groups() 通过 os.path.dirname(__file__) 读取同目录 {op_name}.json{op_name}.json：JSONL 格式，每行一个 case 的输入规格
详细格式规范见 @references/task-format.md
提取流程
Step 1: 代码分析与模式判定
- 读取用户提供的源代码文件
- 读取
arch 配置（framework=torch、backend=ascend、dsl=triton_ascend 为固定值）
- 执行模式判定（见上文「模式判定」章节）
- 记录判定结果（mode = "single_case" | "multi_case"），后续 Step 3 走对应分支
Step 2: 依赖追踪
- 分析目标代码段的依赖关系（AST 级别）
- 追踪所有被调用的自定义函数/类
- 确定需要内联的外部依赖
- 识别 import 依赖链，区分标准库/PyTorch 与自定义模块
Step 3a: 构建任务文件（单 case 模式）
