算子任务提取 Skill

<role> 你是一个算子任务提取专家。你的任务是从用户提供的代码中提取出可优化的算子部分，并将其构建为算子任务格式的任务文件。 </role>

模式判定

按以下优先级判定输入属于哪种模式（任一命中即定型）：

源 .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 模式）

将目标算子逻辑包装到 Model.forward() 中
如果算子有初始化状态（如权重、参数），放入 Model.__init__()
将所有依赖的自定义函数内联到文件中（禁止 import 外部模块）
根据 shape/dtype 信息构建 get_inputs() 和 get_init_inputs()
如果用户未提供 shape/dtype，从代码上下文推断合理默认值
输出：{工作目录}/{op_name}.py

Step 3b: 构建任务文件（多 case 模式）

核心原则：原样透传，不改写源码

复制源 .py → {工作目录}/{op_name}.py
- 不修改 Model、get_input_groups、get_init_inputs 的任何逻辑
- 若源 .py 缺少 get_init_inputs（必需项），可补一个返回 [] 的实现
复制 .json → {工作目录}/{op_name}.json
- 自动从源 .py 同目录发现同名 .json
- 必须与 .py 同名同目录（不要改名、不要嵌套子目录），保证 os.path.dirname(__file__) 仍能定位
不追加 get_inputs() 兼容层：下游 verify.py / benchmark.py 已优先调用 get_input_groups()，无需再注入

Step 4: 验证（必须执行，禁止跳过）

使用本 skill 自带的验证脚本进行静态检查和运行时检查。

验证命令（使用 bash 工具执行）：

python3 <skill-path>/scripts/validate_task.py /abs/path/{op_name}.py --json

其中 <skill-path> 为本 skill 所在目录的绝对路径。

验证脚本检查项：

静态检查：class Model(nn.Module)、forward、get_init_inputs 必需；输入提供函数 get_inputs 或 get_input_groups 至少存在其一
运行时检查：
- exec → Model(*get_init_inputs()) → 取输入提供方
- 若提供 get_input_groups() → 遍历全部 groups，每组依次执行 forward + NaN/Inf + 一致性检查
- 若仅提供 get_inputs() → 单 case 执行
- JSON 输出含 cases_tested / cases_passed

结果处理：

退出码 0 且输出 [VALID] → 验证通过，进入 Step 5
退出码非 0 且输出 [INVALID] → 根据错误信息修复（多 case 模式禁止改成单 case 来"绕过"）；最多重试 2 次
重试 2 次仍失败 → 向用户报告错误，请求协助

多 case 模式注意：

验证脚本会用源文件的真实路径作为 __file__，因此 .json 必须已在目标位置
若 cases_passed < cases_tested，说明部分 case 在参考实现下就跑不通，应反馈给用户而非裁剪 cases

⛔ 禁止事项：

禁止跳过验证直接进入 Step 5
禁止用自创的方法替代验证脚本
禁止仅做静态检查就认为验证通过（必须同时通过运行时检查）
禁止将多 case 源降级为单 case 输出来通过验证
禁止对源 benchmark 路径（npu_benchmark/、ascendc-kernelgen-data*/）使用 Edit/Write——已由 PreToolUse hook 强制拦截，工具调用会被 deny 并返回原因。源文件有 bug 也不允许改，需在 report.md 标注 baseline_buggy: true 后失败退出

Step 4.5: 基线冻结（必须执行，禁止跳过）

验证通过后、用户确认前，立即调用基线冻结脚本，把工作目录里的 {op_name}.py 哈希锚定下来。

调用命令（用 bash 工具）：

python3 <triton-op-verifier-skill-path>/scripts/freeze_baseline.py \
    --op_name {op_name} \
    --work_dir {工作目录} \
    --mode user \
    --source_path <用户源 .py 绝对路径>

--mode 取值：

user：用户提供 benchmark（模式 A 标杆文件，绝大多数场景）—— 必传 --source_path
auto：Agent 自动生成 benchmark（模式 B 兜底翻译场景）—— 不传 --source_path

--source_path 传值（仅 mode=user 时）：

模式 A 单 case：用户最初提供的源 .py 绝对路径
模式 A 多 case：同上
模式 A 优先级 1（torch 标杆 + GPU kernel）：用户提供的 torch 标杆源 .py 路径

退出码处理：

0 → 锚文件写入成功，进入 Step 5
1 → 锚文件已存在（二次 freeze）。审计为何二次调用；若确需重 freeze，提示用户人工删除 {工作目录}/output/.baseline_anchor.json 后重跑
2 → {op_name}.py 不存在，或 --source_path 指向的源文件不存在。本步骤前置条件异常，回退修复
3 → mode=user 但未传 --source_path。补全参数后重跑
5 → 工作目录副本 sha256 ≠ 源 sha256。这意味着 Step 3a/3b 不是字节级 cp 而是改写过（如重写 Model 类、改 forward 签名、补 get_init_inputs）。严禁绕过：正确做法是在 report.md 标注 baseline_buggy: true 失败退出，告诉用户源 benchmark 自身有 bug 需要由人工修复源文件后重跑

关键约束：

Step 3a/3b 必须用字节级 cp（cp 命令或等价的二进制复制），禁止用 Read+Write 重建、禁止任何 Edit 修改。源文件可能有 bug（如 __init__ 参数与 get_init_inputs 不匹配），但那是源的问题，不归 Agent 修
{工作目录}/{op_name}.py 之后任何 Edit/Write 都会被 verify.py / benchmark.py 启动时的基线闸门检测到（exit 4，C 类终止）
锚文件 {工作目录}/output/.baseline_anchor.json 一旦写入禁止重写

Step 5: 用户确认（必须执行，禁止跳过）

验证通过后，必须使用 question 工具将完整的任务文件展示给用户，请求确认。

展示任务代码内容（多 case 模式同时附 .json 行数 + 前 2 行示例）
询问用户
算子生成完成，请查看生成代码：

请选择：
1. 接受
2. <让用户输入修改要求>

具体要求：

必须调用 question 工具，不能只打印文本让用户选择
在 question 的描述中展示 .py 完整内容；多 case 模式下附 .json 摘要
提供"确认"和"需要修改"两个选项

处理回复：

用户确认 → 算子提取任务完成
用户要求修改 → 结合用户反馈返回 Step 3a/3b 重新生成

⛔ 禁止事项：

禁止不经用户确认就结束此 skill
禁止用普通文本消息替代 question 工具调用

关键约束

约束	说明
自包含	单 case：所有依赖函数必须内联；多 case：仅允许依赖同名 `.json`
可执行	单 case：`Model(get_init_inputs())(get_inputs())` 必须直接运行；多 case：对所有 groups 都必须可运行
确定性	给定相同输入，输出必须一致
无 NaN/+Inf	forward 输出不能包含 NaN 或 +Inf；-Inf 允许（top-k/top-p 掩码取值）
禁止重写	原始函数可运行就直接复用，一行都不改；多 case 模式下源 `.py` 整体原样透传
返回一致	返回类型/形状必须与原始实现一致
合理输入	get_inputs 应提供合理大小的输入（不能过小或过大）
模式不可降级	含 `get_input_groups` 或同名 `.json` 的源必须走多 case 流程

示例

输入（单 case）

用户说："/path/to/model.py 的 matmul_with_bias 函数有优化空间，shape 信息在 /path/to/config.py"

输出 task_desc.py（单 case）

import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super(Model, self).__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.randn(out_features))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.matmul(x, self.weight.t()) + self.bias


def get_inputs():
    batch_size = 32
    in_features = 1024
    return [torch.randn(batch_size, in_features)]


def get_init_inputs():
    in_features = 1024
    out_features = 512
    return [in_features, out_features]

输入（多 case）

用户说："/path/to/benchmarks/level430/1_LogicalAnd.py，多 shape 评测"

extractor 自动发现同目录 1_LogicalAnd.json（52 行 JSONL）。

输出（多 case）

{工作目录}/1_LogicalAnd.py    # 复制源文件，不改一行
{工作目录}/1_LogicalAnd.json  # 复制源 JSON，不改一行

验证脚本输出：cases_tested: 52, cases_passed: 52。

triton-task-extractor

Files5

SKILL.mdops/triton-task-extractor/SKILL.md