external-cannbot-ops-triton-task-extractor

Operator Task Extraction Skill

Mode Determination

Determine which mode the input belongs to according to the following priority (a match on any one fixes the mode):

1. Source .py already defines the get_input_groups() function → multi case mode
2. Source .py same directoryexistssame name .json file → multi case mode
3. Otherwise → single case mode

Downstream verify.py / benchmark.py already has built-in determination (prefer get_input_groups, fall back to get_inputs）， Forbiddento downgrade a multi case source to a single case task file.

Target Format

Single case mode

The finally generated file must be a single self-contained Python file，containing only the following 4 parts：

1. import section: only torch / torch.nn / standard library are allowed
2. class Model(nn.Module): wraps the operator logic to be optimized (including __init__ and forward）
3. def get_inputs(): returns the input parameter list of forward() input parameter list of
4. def get_init_inputs(): returns the initialization parameter list of __init__() initialization parameter list of

Multi case mode

Output .py + .json a pair of files, both must be copied to the working directory at the same time, keeping the same-name same-directory relationship:

• {op_name}.py: contains Model + get_input_groups() + get_init_inputs(), where get_input_groups() reads, via os.path.dirname(__file__) , the {op_name}.json
• {op_name}.jsonin the same directory: JSONL format, one input specification per line

For detailed format specifications see @references/task-format.md

---

Extraction Flow

Step 1: Code Analysis and Mode Determination

• Read the source code file provided by the user
• Read the arch configuration (framework=torch、backend=ascend、dsl=triton_ascend are fixed values)
• Execute mode determination(see the "Mode Determination" section above)
• Record the determination result (mode = "single_case" | "multi_case"); Step 3 subsequently takes the corresponding branch

Step 2: Dependency Tracking

• Analyze the dependencies of the target code segment (at the AST level)
• Trace all called custom functions/classes
• Determine the external dependencies that need to be inlined
• Identify the import dependency chain, distinguishing standard library/PyTorch from custom modules

Step 3a: Build Task File (Single case mode)

• Wrap the target operator logic into Model.forward() in
• If the operator has initialization state (such as weights, parameters), put it into Model.__init__()
• Inline all dependent custom functions into the file (importing external modules is forbidden)
• Build get_inputs() based on shape/dtype information get_init_inputs()
• If the user does not provide shape/dtype, infer reasonable defaults from the code context
• Output:{工作目录}/{op_name}.py

Step 3b: Build Task File (Multi case mode)

Core principle: pass through as-is, do not rewrite the source code

1. Copy the source .py → {工作目录}/{op_name}.py
   • Do not modify any logic of Model,get_input_groups、get_init_inputs any logic in
   • If the source .py is missing get_init_inputs(required item), you may add one that returns [] implementation of
2. Copy .json → {工作目录}/{op_name}.json
   • automatically from the source .py discover the same-named .json
   • must be in the same directory with the same name as .py with the same name and in the same directory (do not rename, do not nest in a subdirectory), ensuring os.path.dirname(__file__) can still locate it
3. do not append get_inputs() compatibility layer: downstream verify.py / benchmark.py already calls it with priority get_input_groups(), no further injection is needed

Step 4: Verification (must be performed, skipping is prohibited)

Use the verification script bundled with this skillto perform static checks and runtime checks.

Verification command(execute using the bash tool):

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

where <skill-path> is the absolute path of the directory containing this skill.

Verification script check items：

1. Static check：class Model(nn.Module)、forward、get_init_inputs is required; input provider function get_inputs or get_input_groups at least one must exist
2. Runtime check：
   • exec → Model(*get_init_inputs()) → take the input provider
   • If get_input_groups() → traverse all groups, for each group execute forward + NaN/Inf + consistency checks in sequence
   • If only get_inputs() → single case execution
   • JSON output contains cases_tested / cases_passed

Result handling：

• exit code 0 and output [VALID] → verification passed, proceed to Step 5
• exit code non-0 and output [INVALID] → fix according to the error message (multi-case mode is prohibited from being changed to single case to "bypass"); retry at most 2 times
• still fails after 2 retries → report the error to the user and request assistance

Notes on multi-case mode：

• The verification script uses the real path of the source file as __file__, therefore .json must already be in the target location
• If cases_passed < cases_tested, it means some cases fail to run even under the reference implementation, this should be fed back to the user rather than trimming the cases

⛔ Prohibitions：

• It is prohibited to skip verification and proceed directly to Step 5
• It is prohibited to replace the verification script with a self-invented method
• It is prohibited to consider verification passed by only performing static checks (runtime checks must also pass)
• It is prohibited to downgrade a multi-case source to single-case output to pass verification

Step 5: User confirmation (must be performed, skipping is prohibited)

After verification passes,you must use the question toolto display the complete task file to the user and request confirmation.

1. Display the task code content (in multi-case mode also attach .json line count + first 2 lines as an example)
2. Ask the user

Operator generation is complete, please review the generated code:

Please choose:

1. Accept
2. <have the user enter modification requirements>

Specific requirements：

• You must call the question tool, you cannot just print text and let the user choose
• In question 's description, display .py Full content; in multi-case mode, append .json summary
• Provide two options: "Confirm" and "Needs modification"

Handle response：

• User confirms → operator extraction task complete
• User requests modification → return to Step 3a/3b to regenerate based on user feedback

⛔ Prohibitions：

• Do not end this skill without user confirmation
• Do not replace question tool call with a plain text message

---

Key constraints

Constraints	Description
Self-contained	Single case: all dependent functions must be inlined; multi-case: only dependencies with the same name are allowed .json
Executable	Single case:Model(*get_init_inputs())(*get_inputs()) must run directly; multi-case: must be runnable for all groups
Determinism	Given the same input, the output must be consistent
No NaN/Inf	forward output must not contain NaN or Inf
No rewriting	If the original function runs, reuse it directly without changing a single line; in multi-case mode, the source .py is passed through entirely as-is
Consistent return	Return type/shape must be consistent with the original implementation
Reasonable input	get_inputs should provide input of reasonable size (neither too small nor too large)
Mode cannot be downgraded	Source containing get_input_groups or same-named .json source ofmustgo through the multi-case flow

---

Example

Input (single case)

User says: "/path/to/model.py the matmul_with_bias function has room for optimization, shape information is in /path/to/config.py"

Output task_desc.py (single 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]

Input (multi-case)

User says: "/path/to/benchmarks/level430/1_LogicalAnd.py, multi-shape evaluation"

extractor automatically discovers same-directory 1_LogicalAnd.json(52-line JSONL).

Output (multi-case)

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

Validation script output:cases_tested: 52, cases_passed: 52。