SuperAPI AI 工具集

external-cannbot-ops-triton-task-extractor

编程与开发

查看项目

更新于 10-15 17:19

技能介绍

用途:从用户提供的 PyTorch/Python 代码中提取目标算子逻辑,整理为可供下游评测使用的标准化算子任务文件。

模式判定:源 .py 已定义 get_input_groups() 或同目录存在同名 .json 时判定为多 case 模式,否则为单 case。规则要求不得将多 case 源降级为单 case。

主要能力:单 case 输出单一自包含 Python 文件,含 import 区(仅 torch / torch.nn / 标准库)、class Model(nn.Module)、get_inputs()、get_init_inputs();多 case 输出 .py 与 .json 一对文件,二者同名同目录,.py 含 Model + get_input_groups() + get_init_inputs(),json 为 JSONL,每行一个 case。

使用方式:先做代码分析与模式判定,再进行 AST 级依赖追踪;单 case 需内联所有自定义依赖并依据 shape/dtype 构建输入函数,多 case 则原样复制源 .py 和 .json,不修改逻辑,仅在缺少 get_init_inputs 时补返回 [] 的实现。随后必须执行验证脚本:python3 <skill-path>/scripts/validate_task.py /abs/path/{op_name}.py --json。脚本先做静态检查(Model/forward/get_init_inputs 必需,get_inputs 或 get_input_groups 至少其一),再做运行时检查并输出 cases_tested / cases_passed。退出码 0 且 [VALID] 通过;[INVALID] 可修复,最多重试 2 次,仍失败需报告用户。验证通过后必须用 question 工具展示完整任务文件并请求确认,用户接受或提出修改,不能仅打印文本。

适用场景:为算子生成/优化评测准备标准化任务描述,尤其多 shape 批量评测场景。

限制:单 case 必须自包含且可直接运行、结果确定、无 NaN/Inf、返回一致;多 case 不得改写源码或加 get_inputs 兼容层;验证和用户确认均为必须步骤,禁止跳过。部分内部细节(如 arch 配置取值)以资料所述 framework=torch、backend=ascend、dsl=triton_ascend 为准,其余暂无说明。