From a7ec96a23bd173602cc050c9e230cc1f6bccba21 Mon Sep 17 00:00:00 2001 From: zhumingming Date: Sat, 30 May 2026 17:18:38 +0800 Subject: [PATCH] =?UTF-8?q?triton=5Foptimization=E6=94=B9=E5=8A=A8?= =?UTF-8?q?=E5=B9=B6=E5=85=A5main=E5=88=86=E6=94=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agents/triton-ascend-coder.md | 160 +++- memory/MEMORY.md | 11 + memory/archive/README.md | 35 + memory/archive/pad/pad_v1_20260522.py | 744 ++++++++++++++++++ memory/archive/pad/pad_v1_20260522_report.md | 78 ++ .../archive/pad/pad_v1_20260522_summary.json | 328 ++++++++ memory/archive/repeat/repeat_v2_20260526.py | 156 ++++ .../repeat/repeat_v2_20260526_report.md | 101 +++ .../repeat/repeat_v2_20260526_summary.json | 71 ++ memory/kernel-opt-framework.md | 120 +++ memory/kernel-opt-pad.md | 228 ++++++ memory/kernel-opt-repeat.md | 251 ++++++ skills/triton/kernel-designer/SKILL.md | 49 ++ skills/triton/kernel-generator/SKILL.md | 19 +- .../references/triton-ascend-fundamentals.md | 51 +- .../references/triton-ascend-reduce.md | 138 ++++ .../scripts/validate_triton_impl.py | 157 +++- skills/triton/latency-optimizer/SKILL.md | 160 +++- .../latency-optimizer/references/autotune.md | 10 + .../latency-optimizer/references/checklist.md | 14 +- .../references/constexpr_parameters.md | 58 +- .../grid-dispatch-specialization.md | 191 +++++ .../references/mixed_strategy.md | 196 +++++ .../references/scalar_to_vector.md | 32 +- .../references/tiling_optimization.md | 4 +- .../references/vector_core_partition.md | 16 + utils/exp-archive.py | 306 +++++++ utils/exp-check.py | 260 ++++++ utils/exp-edit.py | 509 ++++++++++++ utils/exp-init.py | 193 +++++ utils/exp_design.md | 167 ++++ 31 files changed, 4739 insertions(+), 74 deletions(-) create mode 100644 memory/MEMORY.md create mode 100644 memory/archive/README.md create mode 100644 memory/archive/pad/pad_v1_20260522.py create mode 100644 memory/archive/pad/pad_v1_20260522_report.md create mode 100644 memory/archive/pad/pad_v1_20260522_summary.json create mode 100644 memory/archive/repeat/repeat_v2_20260526.py create mode 100644 memory/archive/repeat/repeat_v2_20260526_report.md create mode 100644 memory/archive/repeat/repeat_v2_20260526_summary.json create mode 100644 memory/kernel-opt-framework.md create mode 100644 memory/kernel-opt-pad.md create mode 100644 memory/kernel-opt-repeat.md create mode 100644 skills/triton/latency-optimizer/references/grid-dispatch-specialization.md create mode 100644 skills/triton/latency-optimizer/references/mixed_strategy.md create mode 100644 utils/exp-archive.py create mode 100644 utils/exp-check.py create mode 100644 utils/exp-edit.py create mode 100644 utils/exp-init.py create mode 100644 utils/exp_design.md diff --git a/agents/triton-ascend-coder.md b/agents/triton-ascend-coder.md index 0bb0dc53..db5d2ee8 100644 --- a/agents/triton-ascend-coder.md +++ b/agents/triton-ascend-coder.md @@ -138,10 +138,22 @@ python3 -c "import datetime,random; ts=datetime.datetime.now().strftime('%Y%m%d_ 调用 `kernel-designer` skill，设计算法草图。 +**前置检查**： +1. 检查 `.claude/memory/kernel-opt-{category}.md` 是否存在。若存在，skill 调用方必须确保该文件被 skill 加载（通过显式传入路径或 skill 自动发现）。 +2. 若该文件存在，其 Layer 1 约束视为本次草图设计的**硬性边界**。 + **传入**：`op_name`、`task_desc`（任务文件完整内容）、`arch`、`user_requirements`（如有）。 **产出**：`{工作目录}/sketch.txt`。 +**Layer 1 合规检查门（强制）**： +- sketch 产出后，Agent 必须读取 `kernel-opt-{category}.md` 的 Layer 1 约束，逐条核对 `sketch.txt` 是否兼容。 +- 若发现冲突（如 Layer 1 禁止单 kernel 展平但草图设计为 flat-kernel；Layer 1 要求逐维度处理但草图无维度循环等），**视为 A 类错误**，必须： + 1. 不进入 Phase 3 + 2. 将冲突点作为 `conductor_suggestion` 反馈给 `kernel-designer` + 3. 重新执行 Phase 2，直到草图与 Layer 1 兼容 +- 该检查门最多重试 2 次，若仍无法通过，终止任务并报告"草图架构与历史 Layer 1 约束持续冲突"。 + 仅执行一次，后续 Phase 3 迭代不再重新设计草图。 --- @@ -331,11 +343,20 @@ while iteration < max_iterations: ``` opt_iteration = 0 +# max_opt_iterations 动态计算指令： +# Agent 必须在 Phase 4 开始时执行以下步骤： +# 1. 使用 Read 工具读取 .claude/skills/latency-optimizer/SKILL.md +# 2. 统计文本中 "### 优化点" 出现的次数（即为优化点个数） +# 3. 计算 max_opt_iterations = 优化点个数 + 1 +# 4. 若读取失败、文件不存在或统计失败，使用默认值 max_opt_iterations = 20 +max_opt_iterations = <由 Agent 按上述指令运行时计算> +target_speedup = 0.8 # 目标几何平均加速比 best_code = "" best_speedup = 0.0 baseline_code = Phase 3 产出的 generated_code.py phase3_last_iter = Phase 3 最后一次验证通过的 iter 编号 # 见 3.3 的记录 improvement_made = false +target_reached = false # 是否达到目标加速比 ``` ### Phase 4 入口硬断言（强制） @@ -360,7 +381,7 @@ improvement_made = false ### 迭代循环 ``` -while True: +while opt_iteration < max_opt_iterations: ── 4.1 代码分析 + 优化策略 + 代码重写 ──────────── 调用 latency-optimizer skill @@ -453,14 +474,21 @@ while True: - 若 `baseline_speedup` 或 `optimized_speedup` 任一为 `null`（全部 shape 异常， 无几何平均可算），直接判定为优化失败（拒绝优化），跳到 4.5 A 类分析。 + optimized_speedup >= target_speedup: + → 达到目标加速比，优化成功 + → 更新 best_code / best_speedup + → improvement_made = true + → target_reached = true + → break（直接退出循环） + optimized_speedup > baseline_speedup: - → 优化成功（几何平均加速比有提升） + → 有提升但未达目标 → 更新 best_code / best_speedup → improvement_made = true → opt_iteration++，continue 否则（含相等）: - → 视为无提升，opt_iteration++，continue + → 无提升，opt_iteration++，continue ── 4.5 分析决策 (验证失败时) ───────────────────── A 类 (优化引入逻辑错误) → 回退，调整策略，continue @@ -471,13 +499,16 @@ while True: continue ── 4.6 终局判定 ────────────────────────────────── - 无优化点时退出判定： + 循环退出后的终局判定： + + target_reached == true: + → 达到目标加速比，优化成功，进入 Phase 5 improvement_made == true: - → 优化成功，break，进入 Phase 5 + → 有提升但未达目标（迭代耗尽），进入 Phase 5 - improvement_made == false: - → 优化失败（做完所有尝试后没有效果），break，进入 Phase 5 + improvement_made == false 且 opt_iteration >= max_opt_iterations: + → 优化失败（无提升，且迭代次数耗尽），进入 Phase 5 ``` ### Phase 4 终局处理 @@ -584,6 +615,8 @@ while True: **写入 `{工作目录}/report.md`**： - 基本信息：arch、工作目录 - 生成结果：迭代次数、最终版本来源 +- **目标加速比**：target_speedup = 0.8，是否达到（target_reached） +- **实际最佳加速比**：best_speedup（保留 4 位小数） - **Shape 通过率（以 verify 为准）**：`passed_cases / total_cases` 必须从 `output/iter_{phase3_last_iter}/verify/verify_result.json` 读取。 ⚠️ **禁止**从 `perf_result.json` 取 passed_cases —— 后者是"benchmark exec 成功数" @@ -623,6 +656,9 @@ while True: "gen_iterations": 2, "opt_iterations": 1, "optimized": true, + "target_speedup": 0.8, + "target_reached": true, + "best_speedup": 0.85, "perf_method": "profiler", "skill_path": ".claude/skills/kernel-verifier", "perf_data": { @@ -647,6 +683,9 @@ while True: ``` **字段说明**： +- `target_speedup`: 目标几何平均加速比，固定为 0.8 +- `target_reached`: 是否达到目标加速比（optimized_speedup >= target_speedup） +- `best_speedup`: Phase 4 历史最佳几何平均加速比 - `speedup_vs_torch`: **几何平均**聚合 = `(∏ s_i)^(1/n)`（仅对通过且 `s_i` 为有限正数的 shape）；全部异常时为 `null` - `speedup_vs_baseline`: Phase 4 时 = `optimized.speedup_vs_torch / baseline.speedup_vs_torch`（两个几何平均之比） - `passed_cases` / `failed_cases`: 多 shape 时的通过 / 失败计数（策略 A 成功时应为 total / 0） @@ -709,13 +748,35 @@ Phase 4 入口断言失败（Phase 3 闸门被违反）： } ``` -Phase 4 失败时（Phase 3 成功，优化未成功）： +Phase 4 有提升但未达目标时： +```json +{ + "success": true, + "gen_iterations": 2, + "opt_iterations": 10, + "optimized": true, + "target_speedup": 0.8, + "target_reached": false, + "best_speedup": 0.65, + "perf_method": "profiler", + "skill_path": ".claude/skills/kernel-verifier", + "perf_data": { + "avg_latency_ms": 0.8000, + "speedup_vs_torch": 1.5000 + } +} +``` + +Phase 4 失败时（Phase 3 成功，优化无提升）： ```json { "success": true, "gen_iterations": 2, - "opt_iterations": 3, + "opt_iterations": 10, "optimized": false, + "target_speedup": 0.8, + "target_reached": false, + "best_speedup": 0.0, "perf_data": { "avg_latency_ms": 0.8000, "speedup_vs_torch": 1.5000 @@ -839,7 +900,7 @@ agent 收到 exit 2 时，必须按下表把它**等价映射**到对应 verify | GPU Kernel 模式 | `.pt` 必须与 `.py` 同名同目录；`vllm_gpu_perf.csv` 向上查找最多 3 级 | | Phase 3 单一 Kernel | Phase 3 必须且只能生成一个泛用 Kernel，禁止生成多个 Kernel 或调度器 | | Phase 3 最大迭代 | 5 次，禁止超出 | -| Phase 4 迭代策略 | 不做最大迭代次数限制，直到 latency-optimizer 报告无更多优化点则退出 | +| Phase 4 迭代策略 | max_opt_iterations = latency-optimizer 优化点个数 + 1，达到上限后，或者直到 latency-optimizer 报告无更多优化点则退出 | | Phase 4 成功底线 | 性能不劣化（speedup_vs_baseline ≥ 1.0） | | Phase 4 退出判定 | 有效果（speedup_vs_baseline ≥ 1.0）则成功；做完所有尝试后无效果则失败 | | Phase 4 基线复用 | 4.2/4.3 的基线侧 verify_result_baseline.json 和 baseline_perf_result.json 必须从 Phase 3 iter_{phase3_last_iter} 复制，禁止对基线代码重跑 verify.py 或 benchmark.py（基线代码与 Phase 3 generated_code.py 完全一致，重复执行只浪费时间） | @@ -861,3 +922,82 @@ agent 收到 exit 2 时，必须按下表把它**等价映射**到对应 verify - 专业、技术、简洁 - 每完成一个 Phase 提供一行状态更新 - 错误时清晰描述 + 建议操作 + +--- + +## Phase 7: 经验提炼与归档（算子探索成功后强制执行） + +⚠️ **本阶段为跨会话复用保障的关键闭环**。算子任务完成后，必须将验证过的设计决策和性能数据沉淀到项目级 memory，供后续同类算子复用。 + +### 触发条件 + +必须同时满足： +1. `summary.json` 中 `"success": true` +2. `passed_cases == total_cases > 0` +3. `speedup_vs_torch` 为有限正数（几何平均有效） + +### 执行步骤 + +**Step 1: 人工提炼 Layer 1-3（Agent 必须完成）** + +从本次探索中提取可复用经验，按**四层隔离模型**写入对应类别文件： + +- **Layer 1（设计约束）**：硬性必须遵守的规则（如 "constant 模式必须拆分为 fill + copy"） +- **Layer 2（算法骨架）**：核心并行策略的抽象描述（如 grid 分配模式、分支决策树） +- **Layer 3（关键技巧）**：5-15 行已验证有效的代码片段，标注"可替代方向" + +目标文件：`.claude/memory/kernel-opt-{category}.md` + +若该算子类别**首次归档**，先初始化经验文件模板： +```bash +python3 utils/exp-init.py {category} --op-name {op_name} +``` + +**Step 2: 物理归档 Layer 4（自动工具）** + +运行归档命令： +```bash +python3 utils/exp-archive.py {work_dir} --create-experience +``` + +该命令自动完成： +- 校验归档条件（success、精度全过、加速比有效） +- 复制 `{op_name}_generated.py` → `archive/{category}/{category}_v{N}_{date}.py` +- 复制 `report.md` → `archive/{category}/{category}_v{N}_{date}_report.md` +- 复制 `summary.json` → `archive/{category}/{category}_v{N}_{date}_summary.json` +- 版本号 N 按 archive 目录已有版本自动递增 +- 更新 `MEMORY.md` 索引 +- `--create-experience` 若该类别尚无经验文件，自动基于模板创建 + +**Step 3: 规范验证（强制）** + +运行检查命令： +```bash +python3 utils/exp-check.py +``` + +要求：**0 失败、0 警告**。任何失败项必须在结束会话前修复。 + +### 四层隔离复用规则（跨会话） + +| 层级 | 内容 | 受众 | 访问规则 | +|------|------|------|---------| +| Layer 1 | 设计约束、禁止事项 | `kernel-designer` | 必须作为 negative_prompt 遵守 | +| Layer 2 | 算法骨架、并行策略 | `kernel-designer` | 仅作参考方向，输出必须是全新草图 | +| Layer 3 | 关键代码片段 | `kernel-generator` / `latency-optimizer` | 技巧可参考但不可复制，变量名/结构必须重新设计 | +| Layer 4 | 完整历史代码路径 | **默认对 Agent 不可见** | 仅在用户明确指令对比时才可读取 | + +### 关键保障机制 + +1. **统一存储**：所有经验文件位于项目根目录 `.claude/memory/` 下，**所有会话共享同一套 memory** +2. **自动发现**：`kernel-designer` skill 在 Phase 2 必须查询并读取对应类别的 `kernel-opt-{category}.md`（仅 Layer 1-3） +3. **防复制**：Prompt 中必须包含"历史经验仅供启发，禁止直接复制代码结构" +4. **多样性保护**：若新实现采用与历史完全不同的思路且通过验证，将该思路**并列记录**到经验文件，而非覆盖旧经验 + +### 失败处理 + +| 场景 | 处理 | +|------|------| +| summary.json 不满足归档条件 | 禁止归档，在 report.md 中标注"未达归档标准" | +| exp-check.py 报告失败 | 必须修复失败项后方可结束会话 | +| 经验文件已存在 | `exp-archive.py` 仅更新 MEMORY.md 和 archive；Layer 1-3 由 Agent 手动追加到已有经验文件 | diff --git a/memory/MEMORY.md b/memory/MEMORY.md new file mode 100644 index 00000000..f2bdd2e6 --- /dev/null +++ b/memory/MEMORY.md @@ -0,0 +1,11 @@ +# Memory Index + +## 算子优化经验 +- [历史探索经验积累方案](kernel-opt-framework.md) — 算子分类体系、四层隔离存储模型、复用机制与防依赖策略 +- [Pad 算子优化经验](kernel-opt-pad.md) — 多 kernel 分支、维度压缩、constant 模式特化、边界映射模板 + +- [Repeat 算子优化经验](kernel-opt-repeat.md) — transformation-memory 类算子、逐维度串行处理、constexpr 循环展开、多核分区策略 +## 完整代码归档（Layer 4，Agent 默认不可读） +- [归档目录说明](archive/README.md) — 读取约束、归档规则、目录结构 +- [Pad](archive/pad/pad_v1_20260522.py) — 1.68x，51/51；[R](archive/pad/pad_v1_20260522_report.md)/[S](archive/pad/pad_v1_20260522_summary.json) +- [Repeat](archive/repeat/repeat_v2_20260526.py) — 0.88x，49/49 pass diff --git a/memory/archive/README.md b/memory/archive/README.md new file mode 100644 index 00000000..f52151b8 --- /dev/null +++ b/memory/archive/README.md @@ -0,0 +1,35 @@ +# Archive 目录说明 + +## 用途 + +本目录用于归档历史探索的最佳算子实现代码，作为**完整归档层 (Layer 4)** 的物理存储位置。 + +## 读取约束（Agent 必须遵守） + +- **默认不可见**：Agent 在 Phase 2/3/4 中**禁止**直接读取本目录下的任何 `.py` 文件 +- **禁止复制**：Agent **禁止**复制本目录中代码的变量命名、kernel 组织方式、函数签名或控制流结构 +- **仅用于人工复盘**：本目录的代码仅供人类开发者复盘历史探索路径时使用 +- **例外情况**：仅在用户明确指令"与历史实现对比"或"分析历史代码"时，Agent 才可读取指定文件 + +## 目录结构 + +``` +archive/ +├── README.md # 本文件 +├── pad/ +│ ├── pad_v1_20260522.py # Pad 算子最佳实现代码 (v1) +│ ├── pad_v1_20260522_report.md # 对应生成报告 (report.md) +│ └── pad_v1_20260522_summary.json # 对应性能摘要 (summary.json) +└── {op_name}/ # 其他算子归档目录 + └── {op_name}_v{N}_{date}.py +``` + +## 归档规则 + +1. 每个算子类别独立目录 +2. 代码文件命名格式：`{op_name}_v{version}_{YYYYMMDD}.py` +3. 配套报告命名格式：`{op_name}_v{version}_{YYYYMMDD}_report.md` +4. 配套摘要命名格式：`{op_name}_v{version}_{YYYYMMDD}_summary.json` +5. 仅归档**通过全部验证**且**几何平均加速比 > 1.0x** 的实现 +6. **每个算子类别仅保留历史最优版本**；新版本归档时，若几何平均加速比不优于当前最优版本，则终止归档（除非使用 `--force` 强制覆盖） +7. **配套报告与摘要**随代码一并归档，便于人工复盘时快速查阅历史性能数据，无需回溯原始工作目录 diff --git a/memory/archive/pad/pad_v1_20260522.py b/memory/archive/pad/pad_v1_20260522.py new file mode 100644 index 00000000..524fc76b --- /dev/null +++ b/memory/archive/pad/pad_v1_20260522.py @@ -0,0 +1,744 @@ +import torch +import torch.nn as nn +import triton +import triton.language as tl +import torch_npu + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 512}), + triton.Config({'BLOCK_SIZE': 1024}), + triton.Config({'BLOCK_SIZE': 2048}), + triton.Config({'BLOCK_SIZE': 4096}), + ], + key=[], +) +@triton.jit +def pad_kernel( + in_ptr, out_ptr, + in_d0: tl.constexpr, in_d1: tl.constexpr, in_d2: tl.constexpr, in_d3: tl.constexpr, + out_d0: tl.constexpr, out_d1: tl.constexpr, out_d2: tl.constexpr, out_d3: tl.constexpr, + pad_l0: tl.constexpr, pad_r0: tl.constexpr, pad_l1: tl.constexpr, pad_r1: tl.constexpr, + pad_l2: tl.constexpr, pad_r2: tl.constexpr, pad_l3: tl.constexpr, pad_r3: tl.constexpr, + in_s0: tl.constexpr, in_s1: tl.constexpr, in_s2: tl.constexpr, in_s3: tl.constexpr, + out_s0: tl.constexpr, out_s1: tl.constexpr, out_s2: tl.constexpr, out_s3: tl.constexpr, + mode: tl.constexpr, + fill_value: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + out_numel = out_d0 * out_d1 * out_d2 * out_d3 + elements_per_core = tl.cdiv(out_numel, num_cores) + core_start = pid * elements_per_core + core_end = core_start + elements_per_core + core_end = tl.minimum(core_end, out_numel) + num_blocks_per_core = tl.cdiv(core_end - core_start, BLOCK_SIZE) + + for block_idx in range(num_blocks_per_core): + block_start = core_start + block_idx * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < core_end + + c0 = offsets // out_s0 + rem0 = offsets - c0 * out_s0 + c1 = rem0 // out_s1 + rem1 = rem0 - c1 * out_s1 + c2 = rem1 // out_s2 + rem2 = rem1 - c2 * out_s2 + c3 = rem2 // out_s3 + + in_c0 = c0 - pad_l0 + in_c1 = c1 - pad_l1 + in_c2 = c2 - pad_l2 + in_c3 = c3 - pad_l3 + + if mode == 0: + c0_f = c0.to(tl.float32) + c1_f = c1.to(tl.float32) + c2_f = c2.to(tl.float32) + c3_f = c3.to(tl.float32) + + valid0 = (c0_f >= pad_l0) & (c0_f < pad_l0 + in_d0) + valid1 = (c1_f >= pad_l1) & (c1_f < pad_l1 + in_d1) + valid2 = (c2_f >= pad_l2) & (c2_f < pad_l2 + in_d2) + valid3 = (c3_f >= pad_l3) & (c3_f < pad_l3 + in_d3) + valid = valid0 & valid1 & valid2 & valid3 + + safe_c0 = tl.where(valid0, in_c0, 0) + safe_c1 = tl.where(valid1, in_c1, 0) + safe_c2 = tl.where(valid2, in_c2, 0) + safe_c3 = tl.where(valid3, in_c3, 0) + + in_idx = safe_c0 * in_s0 + safe_c1 * in_s1 + safe_c2 * in_s2 + safe_c3 * in_s3 + data = tl.load(in_ptr + in_idx, mask=valid & mask, other=fill_value) + elif mode == 1: + in_c0_f = in_c0.to(tl.float32) + in_c1_f = in_c1.to(tl.float32) + in_c2_f = in_c2.to(tl.float32) + in_c3_f = in_c3.to(tl.float32) + + in_c0 = tl.where(in_c0_f < 0.0, -in_c0, in_c0) + in_c0 = tl.where(in_c0_f >= in_d0, 2 * (in_d0 - 1) - in_c0, in_c0) + in_c1 = tl.where(in_c1_f < 0.0, -in_c1, in_c1) + in_c1 = tl.where(in_c1_f >= in_d1, 2 * (in_d1 - 1) - in_c1, in_c1) + in_c2 = tl.where(in_c2_f < 0.0, -in_c2, in_c2) + in_c2 = tl.where(in_c2_f >= in_d2, 2 * (in_d2 - 1) - in_c2, in_c2) + in_c3 = tl.where(in_c3_f < 0.0, -in_c3, in_c3) + in_c3 = tl.where(in_c3_f >= in_d3, 2 * (in_d3 - 1) - in_c3, in_c3) + + in_idx = in_c0 * in_s0 + in_c1 * in_s1 + in_c2 * in_s2 + in_c3 * in_s3 + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + elif mode == 2: + in_c0_f = in_c0.to(tl.float32) + in_c1_f = in_c1.to(tl.float32) + in_c2_f = in_c2.to(tl.float32) + in_c3_f = in_c3.to(tl.float32) + + in_c0 = tl.where(in_c0_f < 0.0, 0, in_c0) + in_c0 = tl.where(in_c0_f >= in_d0, in_d0 - 1, in_c0) + in_c1 = tl.where(in_c1_f < 0.0, 0, in_c1) + in_c1 = tl.where(in_c1_f >= in_d1, in_d1 - 1, in_c1) + in_c2 = tl.where(in_c2_f < 0.0, 0, in_c2) + in_c2 = tl.where(in_c2_f >= in_d2, in_d2 - 1, in_c2) + in_c3 = tl.where(in_c3_f < 0.0, 0, in_c3) + in_c3 = tl.where(in_c3_f >= in_d3, in_d3 - 1, in_c3) + + in_idx = in_c0 * in_s0 + in_c1 * in_s1 + in_c2 * in_s2 + in_c3 * in_s3 + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + else: + in_c0_f = in_c0.to(tl.float32) + in_c1_f = in_c1.to(tl.float32) + in_c2_f = in_c2.to(tl.float32) + in_c3_f = in_c3.to(tl.float32) + + in_c0 = tl.where(in_c0_f < 0.0, in_c0 + in_d0, in_c0) + in_c0 = tl.where(in_c0_f >= in_d0, in_c0 - in_d0, in_c0) + in_c1 = tl.where(in_c1_f < 0.0, in_c1 + in_d1, in_c1) + in_c1 = tl.where(in_c1_f >= in_d1, in_c1 - in_d1, in_c1) + in_c2 = tl.where(in_c2_f < 0.0, in_c2 + in_d2, in_c2) + in_c2 = tl.where(in_c2_f >= in_d2, in_c2 - in_d2, in_c2) + in_c3 = tl.where(in_c3_f < 0.0, in_c3 + in_d3, in_c3) + in_c3 = tl.where(in_c3_f >= in_d3, in_c3 - in_d3, in_c3) + + in_idx = in_c0 * in_s0 + in_c1 * in_s1 + in_c2 * in_s2 + in_c3 * in_s3 + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + + tl.store(out_ptr + offsets, data, mask=mask) + + +@triton.jit +def pad_kernel_2d( + in_ptr, out_ptr, + H: tl.constexpr, W: tl.constexpr, + H_out: tl.constexpr, W_out: tl.constexpr, + pad_t: tl.constexpr, pad_b: tl.constexpr, + pad_l: tl.constexpr, pad_r: tl.constexpr, + in_stride_h: tl.constexpr, + out_stride_h: tl.constexpr, + mode: tl.constexpr, + fill_value: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + rows_per_core = tl.cdiv(H_out, num_cores) + row_start = pid * rows_per_core + row_end = tl.minimum(row_start + rows_per_core, H_out) + num_rows = row_end - row_start + + for row_idx in range(num_rows): + row = row_start + row_idx + in_row = row - pad_t + + num_blocks = tl.cdiv(W_out, BLOCK_SIZE) + for block_idx in range(num_blocks): + col_start = block_idx * BLOCK_SIZE + cols = col_start + tl.arange(0, BLOCK_SIZE) + mask = cols < W_out + in_col = cols - pad_l + + if mode == 0: + c2_f = row.to(tl.float32) + c3_f = cols.to(tl.float32) + valid2 = (c2_f >= pad_t) & (c2_f < pad_t + H) + valid3 = (c3_f >= pad_l) & (c3_f < pad_l + W) + valid = valid2 & valid3 + safe_row = tl.where(valid2, in_row, 0) + safe_col = tl.where(valid3, in_col, 0) + in_idx = safe_row * in_stride_h + safe_col + data = tl.load(in_ptr + in_idx, mask=valid & mask, other=fill_value) + elif mode == 1: + in_row_f = in_row.to(tl.float32) + in_col_f = in_col.to(tl.float32) + safe_row = tl.where(in_row_f < 0.0, -in_row, in_row) + safe_row = tl.where(in_row_f >= H, 2 * (H - 1) - in_row, safe_row) + safe_col = tl.where(in_col_f < 0.0, -in_col, in_col) + safe_col = tl.where(in_col_f >= W, 2 * (W - 1) - in_col, safe_col) + in_idx = safe_row * in_stride_h + safe_col + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + elif mode == 2: + in_row_f = in_row.to(tl.float32) + in_col_f = in_col.to(tl.float32) + safe_row = tl.where(in_row_f < 0.0, 0, in_row) + safe_row = tl.where(in_row_f >= H, H - 1, safe_row) + safe_col = tl.where(in_col_f < 0.0, 0, in_col) + safe_col = tl.where(in_col_f >= W, W - 1, safe_col) + in_idx = safe_row * in_stride_h + safe_col + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + else: + in_row_f = in_row.to(tl.float32) + in_col_f = in_col.to(tl.float32) + safe_row = tl.where(in_row_f < 0.0, in_row + H, in_row) + safe_row = tl.where(in_row_f >= H, in_row - H, safe_row) + safe_col = tl.where(in_col_f < 0.0, in_col + W, in_col) + safe_col = tl.where(in_col_f >= W, in_col - W, safe_col) + in_idx = safe_row * in_stride_h + safe_col + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + + tl.store(out_ptr + row * out_stride_h + cols, data, mask=mask) + + +@triton.jit +def copy_kernel_2d( + in_ptr, out_ptr, + H: tl.constexpr, W: tl.constexpr, + H_out: tl.constexpr, W_out: tl.constexpr, + pad_t: tl.constexpr, pad_l: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + rows_per_core = tl.cdiv(H, num_cores) + row_start = pid * rows_per_core + row_end = tl.minimum(row_start + rows_per_core, H) + num_rows = row_end - row_start + + for row_idx in range(num_rows): + in_row = row_start + row_idx + out_row = in_row + pad_t + base_in = in_row * W + base_out = out_row * W_out + pad_l + + num_blocks = tl.cdiv(W, BLOCK_SIZE) + for block_idx in range(num_blocks): + col_start = block_idx * BLOCK_SIZE + cols = col_start + tl.arange(0, BLOCK_SIZE) + mask = cols < W + data = tl.load(in_ptr + base_in + cols, mask=mask) + tl.store(out_ptr + base_out + cols, data, mask=mask) + + +@triton.jit +def copy_kernel_3d( + in_ptr, out_ptr, + D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, + D0_out: tl.constexpr, D1_out: tl.constexpr, D2_out: tl.constexpr, + pad_d0: tl.constexpr, pad_d1: tl.constexpr, pad_d2: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + total_data_rows = D0 * D1 + rows_per_core = tl.cdiv(total_data_rows, num_cores) + row_start = pid * rows_per_core + row_end = tl.minimum(row_start + rows_per_core, total_data_rows) + num_rows = row_end - row_start + + in_plane = row_start // D1 + in_row = row_start - in_plane * D1 + + out_plane = in_plane + pad_d0 + out_row = in_row + pad_d1 + + base_in = row_start * D2 + base_out = (out_plane * D1_out + out_row) * D2_out + pad_d2 + + for row_idx in range(num_rows): + num_blocks = tl.cdiv(D2, BLOCK_SIZE) + for block_idx in range(num_blocks): + col_start = block_idx * BLOCK_SIZE + cols = col_start + tl.arange(0, BLOCK_SIZE) + mask = cols < D2 + data = tl.load(in_ptr + base_in + cols, mask=mask) + tl.store(out_ptr + base_out + cols, data, mask=mask) + + base_in += D2 + out_row += 1 + if out_row == D1 + pad_d1: + out_plane += 1 + out_row = pad_d1 + base_out = (out_plane * D1_out + out_row) * D2_out + pad_d2 + else: + base_out += D2_out + + +@triton.jit +def pad_kernel_2d_batched( + in_ptr, out_ptr, + Batch: tl.constexpr, H: tl.constexpr, W: tl.constexpr, + H_out: tl.constexpr, W_out: tl.constexpr, + pad_t: tl.constexpr, pad_b: tl.constexpr, + pad_l: tl.constexpr, pad_r: tl.constexpr, + mode: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + total_out_rows = Batch * H_out + if pid >= total_out_rows: + return + + batch_idx = pid // H_out + row = pid - batch_idx * H_out + in_row = row - pad_t + base_out = batch_idx * H_out * W_out + row * W_out + + num_blocks = tl.cdiv(W_out, BLOCK_SIZE) + for block_idx in range(num_blocks): + col_start = block_idx * BLOCK_SIZE + cols = col_start + tl.arange(0, BLOCK_SIZE) + mask = cols < W_out + in_col = cols - pad_l + + if mode == 1: + in_row_f = in_row.to(tl.float32) + in_col_f = in_col.to(tl.float32) + safe_row = tl.where(in_row_f < 0.0, -in_row, in_row) + safe_row = tl.where(in_row_f >= H, 2 * (H - 1) - in_row, safe_row) + safe_col = tl.where(in_col_f < 0.0, -in_col, in_col) + safe_col = tl.where(in_col_f >= W, 2 * (W - 1) - in_col, safe_col) + in_idx = batch_idx * H * W + safe_row * W + safe_col + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + elif mode == 2: + in_row_f = in_row.to(tl.float32) + in_col_f = in_col.to(tl.float32) + safe_row = tl.where(in_row_f < 0.0, 0, in_row) + safe_row = tl.where(in_row_f >= H, H - 1, safe_row) + safe_col = tl.where(in_col_f < 0.0, 0, in_col) + safe_col = tl.where(in_col_f >= W, W - 1, safe_col) + in_idx = batch_idx * H * W + safe_row * W + safe_col + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + else: + in_row_f = in_row.to(tl.float32) + in_col_f = in_col.to(tl.float32) + safe_row = tl.where(in_row_f < 0.0, in_row + H, in_row) + safe_row = tl.where(in_row_f >= H, in_row - H, safe_row) + safe_col = tl.where(in_col_f < 0.0, in_col + W, in_col) + safe_col = tl.where(in_col_f >= W, in_col - W, safe_col) + in_idx = batch_idx * H * W + safe_row * W + safe_col + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + + tl.store(out_ptr + base_out + cols, data, mask=mask) + + +@triton.jit +def pad_kernel_3d_constant_2d( + in_ptr, out_ptr, + D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, + D0_out: tl.constexpr, D1_out: tl.constexpr, D2_out: tl.constexpr, + pad_d0: tl.constexpr, pad_d1: tl.constexpr, pad_d2: tl.constexpr, + fill_value: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + d0 = tl.program_id(0) + d1 = tl.program_id(1) + + if d0 >= D0_out or d1 >= D1_out: + return + + in_d0 = d0 - pad_d0 + in_d1 = d1 - pad_d1 + valid0 = (in_d0 >= 0) & (in_d0 < D0) + valid1 = (in_d1 >= 0) & (in_d1 < D1) + + safe_d0 = tl.where(valid0, in_d0, 0) + safe_d1 = tl.where(valid1, in_d1, 0) + base_in_plane = safe_d0 * D1 * D2 + safe_d1 * D2 + base_out = d0 * D1_out * D2_out + d1 * D2_out + + num_blocks = tl.cdiv(D2_out, BLOCK_SIZE) + for block_idx in range(num_blocks): + col_start = block_idx * BLOCK_SIZE + cols = col_start + tl.arange(0, BLOCK_SIZE) + mask = cols < D2_out + in_d2 = cols - pad_d2 + in_d2_f = in_d2.to(tl.float32) + valid2 = (in_d2_f >= 0.0) & (in_d2_f < D2) + safe_d2 = tl.where(valid2, in_d2, 0) + valid = valid0 & valid1 & valid2 + in_idx = base_in_plane + safe_d2 + data = tl.load(in_ptr + in_idx, mask=valid & mask, other=fill_value) + tl.store(out_ptr + base_out + cols, data, mask=mask) + + +@triton.jit +def pad_kernel_3d_nonconstant_2d( + in_ptr, out_ptr, + D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, + D0_out: tl.constexpr, D1_out: tl.constexpr, D2_out: tl.constexpr, + pad_d0: tl.constexpr, pad_d1: tl.constexpr, pad_d2: tl.constexpr, + mode: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + d0 = tl.program_id(0) + d1 = tl.program_id(1) + + if d0 >= D0_out or d1 >= D1_out: + return + + in_d0 = d0 - pad_d0 + in_d1 = d1 - pad_d1 + in_d0_f = in_d0.to(tl.float32) + in_d1_f = in_d1.to(tl.float32) + + if mode == 1: + safe_d0 = tl.where(in_d0_f < 0.0, -in_d0, in_d0) + safe_d0 = tl.where(in_d0_f >= D0, 2 * (D0 - 1) - in_d0, safe_d0) + safe_d1 = tl.where(in_d1_f < 0.0, -in_d1, in_d1) + safe_d1 = tl.where(in_d1_f >= D1, 2 * (D1 - 1) - in_d1, safe_d1) + elif mode == 2: + safe_d0 = tl.maximum(0, tl.minimum(in_d0, D0 - 1)) + safe_d1 = tl.maximum(0, tl.minimum(in_d1, D1 - 1)) + else: + safe_d0 = tl.where(in_d0_f < 0.0, in_d0 + D0, in_d0) + safe_d0 = tl.where(in_d0_f >= D0, in_d0 - D0, safe_d0) + safe_d1 = tl.where(in_d1_f < 0.0, in_d1 + D1, in_d1) + safe_d1 = tl.where(in_d1_f >= D1, in_d1 - D1, safe_d1) + + base_in_plane = safe_d0 * D1 * D2 + safe_d1 * D2 + base_out = d0 * D1_out * D2_out + d1 * D2_out + + num_blocks = tl.cdiv(D2_out, BLOCK_SIZE) + for block_idx in range(num_blocks): + col_start = block_idx * BLOCK_SIZE + cols = col_start + tl.arange(0, BLOCK_SIZE) + mask = cols < D2_out + in_d2 = cols - pad_d2 + + if mode == 1: + in_d2_f = in_d2.to(tl.float32) + safe_d2 = tl.where(in_d2_f < 0.0, -in_d2, in_d2) + safe_d2 = tl.where(in_d2_f >= D2, 2 * (D2 - 1) - in_d2, safe_d2) + elif mode == 2: + safe_d2 = tl.maximum(0, tl.minimum(in_d2, D2 - 1)) + else: + in_d2_f = in_d2.to(tl.float32) + safe_d2 = tl.where(in_d2_f < 0.0, in_d2 + D2, in_d2) + safe_d2 = tl.where(in_d2_f >= D2, in_d2 - D2, safe_d2) + + in_idx = base_in_plane + safe_d2 + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + tl.store(out_ptr + base_out + cols, data, mask=mask) + + +@triton.jit +def pad_kernel_3d_constant_v2( + in_ptr, out_ptr, + D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, + D0_out: tl.constexpr, D1_out: tl.constexpr, D2_out: tl.constexpr, + pad_d0: tl.constexpr, pad_d1: tl.constexpr, pad_d2: tl.constexpr, + fill_value: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + total_rows = D0_out * D1_out + rows_per_core = tl.cdiv(total_rows, num_cores) + row_start = pid * rows_per_core + row_end = tl.minimum(row_start + rows_per_core, total_rows) + + d0 = row_start // D1_out + d1 = row_start - d0 * D1_out + + num_rows = row_end - row_start + for _ in range(num_rows): + base_out = (d0 * D1_out + d1) * D2_out + in_d0 = d0 - pad_d0 + in_d1 = d1 - pad_d1 + valid01 = (in_d0 >= 0) & (in_d0 < D0) & (in_d1 >= 0) & (in_d1 < D1) + safe_d0 = tl.where(valid01, in_d0, 0) + safe_d1 = tl.where(valid01, in_d1, 0) + base_in = safe_d0 * D1 * D2 + safe_d1 * D2 + + num_blocks = tl.cdiv(D2_out, BLOCK_SIZE) + for block_idx in range(num_blocks): + col_start = block_idx * BLOCK_SIZE + cols = col_start + tl.arange(0, BLOCK_SIZE) + mask = cols < D2_out + in_d2 = cols - pad_d2 + valid2 = (in_d2 >= 0) & (in_d2 < D2) + valid = valid01 & valid2 + safe_d2 = tl.where(valid2, in_d2, 0) + in_idx = base_in + safe_d2 + data = tl.load(in_ptr + in_idx, mask=valid & mask, other=fill_value) + tl.store(out_ptr + base_out + cols, data, mask=mask) + + d1 += 1 + if d1 == D1_out: + d1 = 0 + d0 += 1 + + +@triton.jit +def pad_kernel_3d_nonconstant_v2( + in_ptr, out_ptr, + D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, + D0_out: tl.constexpr, D1_out: tl.constexpr, D2_out: tl.constexpr, + pad_d0: tl.constexpr, pad_d1: tl.constexpr, pad_d2: tl.constexpr, + mode: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + total_rows = D0 * D1_out + rows_per_core = tl.cdiv(total_rows, num_cores) + row_start = pid * rows_per_core + row_end = tl.minimum(row_start + rows_per_core, total_rows) + + d0 = row_start // D1_out + d1 = row_start - d0 * D1_out + + num_rows = row_end - row_start + for _ in range(num_rows): + base_out = (d0 * D1_out + d1) * D2_out + base_in_plane = d0 * D1 * D2 + in_d1 = d1 - pad_d1 + + if mode == 1: + in_d1_f = in_d1.to(tl.float32) + safe_d1 = tl.where(in_d1_f < 0.0, -in_d1, in_d1) + safe_d1 = tl.where(in_d1_f >= D1, 2 * (D1 - 1) - in_d1, safe_d1) + elif mode == 2: + safe_d1 = tl.maximum(0, tl.minimum(in_d1, D1 - 1)) + else: + in_d1_f = in_d1.to(tl.float32) + safe_d1 = tl.where(in_d1_f < 0.0, in_d1 + D1, in_d1) + safe_d1 = tl.where(in_d1_f >= D1, in_d1 - D1, safe_d1) + + num_blocks = tl.cdiv(D2_out, BLOCK_SIZE) + for block_idx in range(num_blocks): + col_start = block_idx * BLOCK_SIZE + cols = col_start + tl.arange(0, BLOCK_SIZE) + mask = cols < D2_out + in_d2 = cols - pad_d2 + + if mode == 1: + in_d2_f = in_d2.to(tl.float32) + safe_d2 = tl.where(in_d2_f < 0.0, -in_d2, in_d2) + safe_d2 = tl.where(in_d2_f >= D2, 2 * (D2 - 1) - in_d2, safe_d2) + elif mode == 2: + safe_d2 = tl.maximum(0, tl.minimum(in_d2, D2 - 1)) + else: + in_d2_f = in_d2.to(tl.float32) + safe_d2 = tl.where(in_d2_f < 0.0, in_d2 + D2, in_d2) + safe_d2 = tl.where(in_d2_f >= D2, in_d2 - D2, safe_d2) + + in_idx = base_in_plane + safe_d1 * D2 + safe_d2 + data = tl.load(in_ptr + in_idx, mask=mask, other=0.0) + tl.store(out_ptr + base_out + cols, data, mask=mask) + + d1 += 1 + if d1 == D1_out: + d1 = 0 + d0 += 1 + + +def select_block_size(width): + if width <= 64: + return 64 + elif width <= 128: + return 128 + elif width <= 256: + return 256 + elif width <= 512: + return 512 + elif width <= 1024: + return 1024 + elif width <= 2048: + return 2048 + else: + return 4096 + + +class ModelNew(nn.Module): + def __init__(self): + super().__init__() + try: + self.VEC_CORE_NUM = torch_npu.npu.npu_config.get_device_limit(0).get("vector_core_num", 40) + except Exception: + self.VEC_CORE_NUM = 40 + + def forward(self, x, pad, mode='constant', value=None): + if value is None: + value = 0.0 + + pad_list = list(pad) + ndim_orig = x.ndim + + dim_deltas = {} + num_pad_dims = len(pad_list) // 2 + for i in range(num_pad_dims): + dim_idx = ndim_orig - 1 - i + dim_deltas[dim_idx] = pad_list[2 * i] + pad_list[2 * i + 1] + out_shape_orig = [x.shape[i] + dim_deltas.get(i, 0) for i in range(ndim_orig)] + + output = torch.empty(out_shape_orig, device=x.device, dtype=x.dtype) + + if not x.is_contiguous(): + x = x.contiguous() + + squeeze_count = 0 + while squeeze_count < ndim_orig - 1 and x.shape[squeeze_count] == 1: + dim_idx_from_right = ndim_orig - 1 - squeeze_count + pad_idx = 2 * dim_idx_from_right + if pad_idx < len(pad_list): + if pad_list[pad_idx] == 0 and pad_list[pad_idx + 1] == 0: + squeeze_count += 1 + else: + break + else: + squeeze_count += 1 + + x_kernel = x + output_kernel = output + pad_list_kernel = pad_list + ndim = ndim_orig + + if squeeze_count > 0: + squeeze_dims = tuple(range(squeeze_count)) + x_kernel = x.squeeze(squeeze_dims) + out_shape_squeezed = out_shape_orig[squeeze_count:] + output_kernel = output.view(out_shape_squeezed) + num_pad_pairs = len(pad_list) // 2 + implicit_pad_dims = ndim_orig - num_pad_pairs + if squeeze_count > implicit_pad_dims: + entries_to_remove = 2 * (squeeze_count - implicit_pad_dims) + pad_list_kernel = pad_list[:-entries_to_remove] + else: + pad_list_kernel = pad_list + ndim = x_kernel.ndim + + mode_map = {'constant': 0, 'reflect': 1, 'replicate': 2, 'circular': 3} + mode_val = mode_map.get(mode, 0) + + if ndim == 2: + H, W = x_kernel.shape[0], x_kernel.shape[1] + H_out, W_out = out_shape_squeezed[0] if squeeze_count > 0 else out_shape_orig[0], out_shape_squeezed[1] if squeeze_count > 0 else out_shape_orig[1] + pad_l = pad_list_kernel[0] if len(pad_list_kernel) > 0 else 0 + pad_r = pad_list_kernel[1] if len(pad_list_kernel) > 1 else 0 + pad_t = pad_list_kernel[2] if len(pad_list_kernel) > 2 else 0 + pad_b = pad_list_kernel[3] if len(pad_list_kernel) > 3 else 0 + block_size = select_block_size(W_out if mode != 'constant' else W) + + if mode == 'constant': + output_kernel.fill_(float(value)) + copy_kernel_2d[(self.VEC_CORE_NUM,)]( + x_kernel, output_kernel, + H=H, W=W, + H_out=H_out, W_out=W_out, + pad_t=pad_t, pad_l=pad_l, + num_cores=self.VEC_CORE_NUM, + BLOCK_SIZE=block_size, + ) + else: + pad_kernel_2d[(self.VEC_CORE_NUM,)]( + x_kernel, output_kernel, + H=H, W=W, + H_out=H_out, W_out=W_out, + pad_t=pad_t, pad_b=pad_b, pad_l=pad_l, pad_r=pad_r, + in_stride_h=W, + out_stride_h=W_out, + mode=mode_val, + fill_value=float(value), + num_cores=self.VEC_CORE_NUM, + BLOCK_SIZE=block_size, + ) + elif ndim == 3 and mode == 'constant': + D0, D1, D2 = x_kernel.shape[0], x_kernel.shape[1], x_kernel.shape[2] + D0_out, D1_out, D2_out = out_shape_squeezed[0] if squeeze_count > 0 else out_shape_orig[0], out_shape_squeezed[1] if squeeze_count > 0 else out_shape_orig[1], out_shape_squeezed[2] if squeeze_count > 0 else out_shape_orig[2] + pad_l2 = pad_list_kernel[0] if len(pad_list_kernel) > 0 else 0 + pad_l1 = pad_list_kernel[2] if len(pad_list_kernel) > 2 else 0 + pad_l0 = pad_list_kernel[4] if len(pad_list_kernel) > 4 else 0 + block_size = select_block_size(D2) + + output_kernel.fill_(float(value)) + copy_kernel_3d[(self.VEC_CORE_NUM,)]( + x_kernel, output_kernel, + D0=D0, D1=D1, D2=D2, + D0_out=D0_out, D1_out=D1_out, D2_out=D2_out, + pad_d0=pad_l0, pad_d1=pad_l1, pad_d2=pad_l2, + num_cores=self.VEC_CORE_NUM, + BLOCK_SIZE=block_size, + ) + elif ndim == 3 and mode != 'constant': + D0, D1, D2 = x_kernel.shape[0], x_kernel.shape[1], x_kernel.shape[2] + D0_out, D1_out, D2_out = out_shape_squeezed[0] if squeeze_count > 0 else out_shape_orig[0], out_shape_squeezed[1] if squeeze_count > 0 else out_shape_orig[1], out_shape_squeezed[2] if squeeze_count > 0 else out_shape_orig[2] + pad_l2 = pad_list_kernel[0] if len(pad_list_kernel) > 0 else 0 + pad_r2 = pad_list_kernel[1] if len(pad_list_kernel) > 1 else 0 + pad_l1 = pad_list_kernel[2] if len(pad_list_kernel) > 2 else 0 + pad_r1 = pad_list_kernel[3] if len(pad_list_kernel) > 3 else 0 + pad_l0 = pad_list_kernel[4] if len(pad_list_kernel) > 4 else 0 + block_size = select_block_size(D2_out) + + total_grid_rows = D0 * D1_out + if total_grid_rows > 3000: + pad_kernel_3d_nonconstant_v2[(self.VEC_CORE_NUM,)]( + x_kernel, output_kernel, + D0=D0, D1=D1, D2=D2, + D0_out=D0_out, D1_out=D1_out, D2_out=D2_out, + pad_d0=pad_l0, pad_d1=pad_l1, pad_d2=pad_l2, + mode=mode_val, + num_cores=self.VEC_CORE_NUM, + BLOCK_SIZE=block_size, + ) + else: + pad_kernel_3d_nonconstant_2d[(D0_out, D1_out)]( + x_kernel, output_kernel, + D0=D0, D1=D1, D2=D2, + D0_out=D0_out, D1_out=D1_out, D2_out=D2_out, + pad_d0=pad_l0, pad_d1=pad_l1, pad_d2=pad_l2, + mode=mode_val, + num_cores=self.VEC_CORE_NUM, + BLOCK_SIZE=block_size, + ) + else: + in_shape_padded = [1] * (4 - ndim) + list(x_kernel.shape) + out_shape_padded = [1] * (4 - ndim) + list(output_kernel.shape) + + pad_entries = list(zip(pad_list_kernel[::2], pad_list_kernel[1::2])) + pad_l_dict = {3 - i: e[0] for i, e in enumerate(pad_entries)} + pad_r_dict = {3 - i: e[1] for i, e in enumerate(pad_entries)} + pad_l = [pad_l_dict.get(i, 0) for i in range(4)] + pad_r = [pad_r_dict.get(i, 0) for i in range(4)] + + in_strides = [ + in_shape_padded[1] * in_shape_padded[2] * in_shape_padded[3], + in_shape_padded[2] * in_shape_padded[3], + in_shape_padded[3], + 1, + ] + out_strides = [ + out_shape_padded[1] * out_shape_padded[2] * out_shape_padded[3], + out_shape_padded[2] * out_shape_padded[3], + out_shape_padded[3], + 1, + ] + + pad_kernel[(self.VEC_CORE_NUM,)]( + x_kernel, output_kernel, + in_d0=in_shape_padded[0], in_d1=in_shape_padded[1], in_d2=in_shape_padded[2], in_d3=in_shape_padded[3], + out_d0=out_shape_padded[0], out_d1=out_shape_padded[1], out_d2=out_shape_padded[2], out_d3=out_shape_padded[3], + pad_l0=pad_l[0], pad_r0=pad_r[0], pad_l1=pad_l[1], pad_r1=pad_r[1], + pad_l2=pad_l[2], pad_r2=pad_r[2], pad_l3=pad_l[3], pad_r3=pad_r[3], + in_s0=in_strides[0], in_s1=in_strides[1], in_s2=in_strides[2], in_s3=in_strides[3], + out_s0=out_strides[0], out_s1=out_strides[1], out_s2=out_strides[2], out_s3=out_strides[3], + mode=mode_val, + fill_value=float(value), + num_cores=self.VEC_CORE_NUM, + ) + + return output diff --git a/memory/archive/pad/pad_v1_20260522_report.md b/memory/archive/pad/pad_v1_20260522_report.md new file mode 100644 index 00000000..643ec587 --- /dev/null +++ b/memory/archive/pad/pad_v1_20260522_report.md @@ -0,0 +1,78 @@ +# Pad 算子生成报告 + +## 基本信息 +- **算子名称**: 15_Pad +- **硬件架构**: ascend910 +- **工作目录**: /home/zmm/OpAgent-Pad/triton_ascend_output/op_0_15_Pad_20260522_0423_8652 + +## 生成结果 +- **迭代次数**: 0（直接复用参考实现） +- **最终版本来源**: 参考实现 +- **优化迭代**: 0（用户要求参考实现满足标准后直接输出，无需 Phase 4 优化） + +## 精度验证 +- **Shape 通过率**: 51/51（全部通过） + +## 性能数据 +- **几何平均加速比 (speedup_vs_torch)**: 1.6795 +- **框架平均延时 (avg_latency_ms)**: 0.1704 +- **实现平均延时 (avg_latency_ms)**: 0.0342 + +## 性能明细 + +| case_idx | shape_desc | status | framework(ms) | implementation(ms) | speedup | +|----------|-----------|--------|---------------|-------------------|---------| +| 1 | [128] | pass | 0.0048 | 0.0023 | 2.0870 | +| 2 | [256] | pass | 0.0052 | 0.0024 | 2.1667 | +| 3 | [512] | pass | 0.0047 | 0.0025 | 1.8800 | +| 4 | [1024] | pass | 0.0048 | 0.0026 | 1.8462 | +| 5 | [128, 256] | pass | 0.0093 | 0.0042 | 2.2143 | +| 6 | [256, 512] | pass | 0.0173 | 0.0050 | 3.4600 | +| 7 | [512, 1024] | pass | 0.0260 | 0.0058 | 4.4828 | +| 8 | [1024, 2048] | pass | 0.0132 | 0.0071 | 1.8592 | +| 9 | [64, 128, 256] | pass | 0.0128 | 0.0233 | 0.5494 | +| 10 | [32, 64, 128] | pass | 0.0097 | 0.0097 | 1.0000 | +| 11 | [16, 128, 256] | pass | 0.0092 | 0.0091 | 1.0110 | +| 12 | [8, 256, 512] | pass | 0.0130 | 0.0092 | 1.4130 | +| 13 | [1, 64, 128, 128] | pass | 0.0125 | 0.0226 | 0.5531 | +| 14 | [1, 128, 64, 64] | pass | 0.0099 | 0.1319 | 0.0751 | +| 15 | [1, 256, 32, 32] | pass | 0.0062 | 0.0356 | 0.1742 | +| 16 | [1, 512, 16, 16] | pass | 0.1454 | 0.0338 | 4.3018 | +| 17 | [1536] | pass | 0.0088 | 0.0028 | 3.1429 | +| 18 | [4096] | pass | 0.0073 | 0.0033 | 2.2121 | +| 19 | [8192] | pass | 0.0072 | 0.0043 | 1.6744 | +| 20 | [4096, 4096] | pass | 0.0498 | 0.0225 | 2.2133 | +| 21 | [4096, 11008] | pass | 0.8690 | 0.0476 | 18.2563 | +| 22 | [5120, 13824] | pass | 1.2730 | 0.1385 | 9.1913 | +| 23 | [3584, 18944] | pass | 1.0804 | 0.3057 | 3.5342 | +| 24 | [5120, 27648] | pass | 1.1979 | 0.3282 | 3.6499 | +| 25 | [1, 3, 224, 224] | pass | 0.0117 | 0.0058 | 2.0172 | +| 26 | [1, 3, 224, 224] | pass | 0.0100 | 0.0331 | 0.3021 | +| 27 | [1, 64, 56, 56] | pass | 0.0034 | 0.0247 | 0.1377 | +| 28 | [1, 128, 28, 28] | pass | 0.1308 | 0.0311 | 4.2058 | +| 29 | [1, 256, 14, 14] | pass | 0.0092 | 0.0096 | 0.9583 | +| 30 | [1, 512, 7, 7] | pass | 0.0083 | 0.0084 | 0.9881 | +| 31 | [100] | pass | 0.0049 | 0.0024 | 2.0417 | +| 32 | [200] | pass | 0.0049 | 0.0025 | 1.9600 | +| 33 | [34, 66] | pass | 0.0911 | 0.0022 | 41.4091 | +| 34 | [17, 33] | pass | 0.0071 | 0.0021 | 3.3810 | +| 35 | [65, 129] | pass | 0.0080 | 0.0023 | 3.4783 | +| 36 | [1, 16, 100, 100] | pass | 0.0085 | 0.0080 | 1.0625 | +| 37 | [1, 32, 50, 50] | pass | 0.0042 | 0.0314 | 0.1338 | +| 38 | [1, 64, 25, 25] | pass | 0.1106 | 0.0302 | 3.6623 | +| 39 | [1, 3, 112, 112] | pass | 0.0089 | 0.0045 | 1.9778 | +| 40 | [1, 64, 56, 56] | pass | 0.0036 | 0.0248 | 0.1452 | +| 41 | [128, 256] | pass | 0.0092 | 0.0037 | 2.4865 | +| 42 | [256, 512] | pass | 0.0064 | 0.0050 | 1.2800 | +| 43 | [32, 128, 128] | pass | 0.0092 | 0.0129 | 0.7132 | +| 44 | [16, 64, 64] | pass | 0.0111 | 0.0087 | 1.2759 | +| 45 | [1, 64, 64, 64] | pass | 0.0105 | 0.0131 | 0.8015 | +| 46 | [1, 128, 32, 32] | pass | 0.0066 | 0.0302 | 0.2185 | +| 47 | [1, 256, 16, 16] | pass | 0.1201 | 0.0163 | 7.3681 | +| 48 | [2048, 2048] | pass | 0.0194 | 0.0104 | 1.8654 | +| 49 | [3072, 3072] | pass | 0.0257 | 0.0126 | 2.0397 | +| 50 | [6144, 6144] | pass | 1.2020 | 0.1388 | 8.6599 | +| 51 | [8192, 8192] | pass | 2.0473 | 0.1099 | 18.6288 | + +## 代码路径 +- 最终代码: `/home/zmm/OpAgent-Pad/triton_ascend_output/op_0_15_Pad_20260522_0423_8652/15_Pad_generated.py` diff --git a/memory/archive/pad/pad_v1_20260522_summary.json b/memory/archive/pad/pad_v1_20260522_summary.json new file mode 100644 index 00000000..1ac163d9 --- /dev/null +++ b/memory/archive/pad/pad_v1_20260522_summary.json @@ -0,0 +1,328 @@ +{ + "success": true, + "gen_iterations": 0, + "opt_iterations": 0, + "optimized": false, + "perf_method": "profiler", + "skill_path": ".claude/skills/kernel-verifier", + "perf_data": { + "avg_latency_ms": 0.0342, + "speedup_vs_torch": 1.6795, + "total_cases": 51, + "passed_cases": 51, + "failed_cases": 0, + "nan_indices": [], + "inf_indices": [], + "zero_indices": [], + "negative_indices": [], + "none_indices": [], + "per_shape_results": [ + { + "case_idx": 1, + "status": "pass", + "shape_desc": "[128]", + "speedup_vs_torch": 2.087 + }, + { + "case_idx": 2, + "status": "pass", + "shape_desc": "[256]", + "speedup_vs_torch": 2.1667 + }, + { + "case_idx": 3, + "status": "pass", + "shape_desc": "[512]", + "speedup_vs_torch": 1.88 + }, + { + "case_idx": 4, + "status": "pass", + "shape_desc": "[1024]", + "speedup_vs_torch": 1.8462 + }, + { + "case_idx": 5, + "status": "pass", + "shape_desc": "[128, 256]", + "speedup_vs_torch": 2.2143 + }, + { + "case_idx": 6, + "status": "pass", + "shape_desc": "[256, 512]", + "speedup_vs_torch": 3.46 + }, + { + "case_idx": 7, + "status": "pass", + "shape_desc": "[512, 1024]", + "speedup_vs_torch": 4.4828 + }, + { + "case_idx": 8, + "status": "pass", + "shape_desc": "[1024, 2048]", + "speedup_vs_torch": 1.8592 + }, + { + "case_idx": 9, + "status": "pass", + "shape_desc": "[64, 128, 256]", + "speedup_vs_torch": 0.5494 + }, + { + "case_idx": 10, + "status": "pass", + "shape_desc": "[32, 64, 128]", + "speedup_vs_torch": 1.0 + }, + { + "case_idx": 11, + "status": "pass", + "shape_desc": "[16, 128, 256]", + "speedup_vs_torch": 1.011 + }, + { + "case_idx": 12, + "status": "pass", + "shape_desc": "[8, 256, 512]", + "speedup_vs_torch": 1.413 + }, + { + "case_idx": 13, + "status": "pass", + "shape_desc": "[1, 64, 128, 128]", + "speedup_vs_torch": 0.5531 + }, + { + "case_idx": 14, + "status": "pass", + "shape_desc": "[1, 128, 64, 64]", + "speedup_vs_torch": 0.0751 + }, + { + "case_idx": 15, + "status": "pass", + "shape_desc": "[1, 256, 32, 32]", + "speedup_vs_torch": 0.1742 + }, + { + "case_idx": 16, + "status": "pass", + "shape_desc": "[1, 512, 16, 16]", + "speedup_vs_torch": 4.3018 + }, + { + "case_idx": 17, + "status": "pass", + "shape_desc": "[1536]", + "speedup_vs_torch": 3.1429 + }, + { + "case_idx": 18, + "status": "pass", + "shape_desc": "[4096]", + "speedup_vs_torch": 2.2121 + }, + { + "case_idx": 19, + "status": "pass", + "shape_desc": "[8192]", + "speedup_vs_torch": 1.6744 + }, + { + "case_idx": 20, + "status": "pass", + "shape_desc": "[4096, 4096]", + "speedup_vs_torch": 2.2133 + }, + { + "case_idx": 21, + "status": "pass", + "shape_desc": "[4096, 11008]", + "speedup_vs_torch": 18.2563 + }, + { + "case_idx": 22, + "status": "pass", + "shape_desc": "[5120, 13824]", + "speedup_vs_torch": 9.1913 + }, + { + "case_idx": 23, + "status": "pass", + "shape_desc": "[3584, 18944]", + "speedup_vs_torch": 3.5342 + }, + { + "case_idx": 24, + "status": "pass", + "shape_desc": "[5120, 27648]", + "speedup_vs_torch": 3.6499 + }, + { + "case_idx": 25, + "status": "pass", + "shape_desc": "[1, 3, 224, 224]", + "speedup_vs_torch": 2.0172 + }, + { + "case_idx": 26, + "status": "pass", + "shape_desc": "[1, 3, 224, 224]", + "speedup_vs_torch": 0.3021 + }, + { + "case_idx": 27, + "status": "pass", + "shape_desc": "[1, 64, 56, 56]", + "speedup_vs_torch": 0.1377 + }, + { + "case_idx": 28, + "status": "pass", + "shape_desc": "[1, 128, 28, 28]", + "speedup_vs_torch": 4.2058 + }, + { + "case_idx": 29, + "status": "pass", + "shape_desc": "[1, 256, 14, 14]", + "speedup_vs_torch": 0.9583 + }, + { + "case_idx": 30, + "status": "pass", + "shape_desc": "[1, 512, 7, 7]", + "speedup_vs_torch": 0.9881 + }, + { + "case_idx": 31, + "status": "pass", + "shape_desc": "[100]", + "speedup_vs_torch": 2.0417 + }, + { + "case_idx": 32, + "status": "pass", + "shape_desc": "[200]", + "speedup_vs_torch": 1.96 + }, + { + "case_idx": 33, + "status": "pass", + "shape_desc": "[34, 66]", + "speedup_vs_torch": 41.4091 + }, + { + "case_idx": 34, + "status": "pass", + "shape_desc": "[17, 33]", + "speedup_vs_torch": 3.381 + }, + { + "case_idx": 35, + "status": "pass", + "shape_desc": "[65, 129]", + "speedup_vs_torch": 3.4783 + }, + { + "case_idx": 36, + "status": "pass", + "shape_desc": "[1, 16, 100, 100]", + "speedup_vs_torch": 1.0625 + }, + { + "case_idx": 37, + "status": "pass", + "shape_desc": "[1, 32, 50, 50]", + "speedup_vs_torch": 0.1338 + }, + { + "case_idx": 38, + "status": "pass", + "shape_desc": "[1, 64, 25, 25]", + "speedup_vs_torch": 3.6623 + }, + { + "case_idx": 39, + "status": "pass", + "shape_desc": "[1, 3, 112, 112]", + "speedup_vs_torch": 1.9778 + }, + { + "case_idx": 40, + "status": "pass", + "shape_desc": "[1, 64, 56, 56]", + "speedup_vs_torch": 0.1452 + }, + { + "case_idx": 41, + "status": "pass", + "shape_desc": "[128, 256]", + "speedup_vs_torch": 2.4865 + }, + { + "case_idx": 42, + "status": "pass", + "shape_desc": "[256, 512]", + "speedup_vs_torch": 1.28 + }, + { + "case_idx": 43, + "status": "pass", + "shape_desc": "[32, 128, 128]", + "speedup_vs_torch": 0.7132 + }, + { + "case_idx": 44, + "status": "pass", + "shape_desc": "[16, 64, 64]", + "speedup_vs_torch": 1.2759 + }, + { + "case_idx": 45, + "status": "pass", + "shape_desc": "[1, 64, 64, 64]", + "speedup_vs_torch": 0.8015 + }, + { + "case_idx": 46, + "status": "pass", + "shape_desc": "[1, 128, 32, 32]", + "speedup_vs_torch": 0.2185 + }, + { + "case_idx": 47, + "status": "pass", + "shape_desc": "[1, 256, 16, 16]", + "speedup_vs_torch": 7.3681 + }, + { + "case_idx": 48, + "status": "pass", + "shape_desc": "[2048, 2048]", + "speedup_vs_torch": 1.8654 + }, + { + "case_idx": 49, + "status": "pass", + "shape_desc": "[3072, 3072]", + "speedup_vs_torch": 2.0397 + }, + { + "case_idx": 50, + "status": "pass", + "shape_desc": "[6144, 6144]", + "speedup_vs_torch": 8.6599 + }, + { + "case_idx": 51, + "status": "pass", + "shape_desc": "[8192, 8192]", + "speedup_vs_torch": 18.6288 + } + ] + } +} \ No newline at end of file diff --git a/memory/archive/repeat/repeat_v2_20260526.py b/memory/archive/repeat/repeat_v2_20260526.py new file mode 100644 index 00000000..a91bde43 --- /dev/null +++ b/memory/archive/repeat/repeat_v2_20260526.py @@ -0,0 +1,156 @@ +import math +import torch +import torch.nn as nn +import triton +import triton.language as tl + + +@triton.jit +def repeat_small_kernel( + x_ptr, out_ptr, + inner_size, + r: tl.constexpr, + BLOCK: tl.constexpr, +): + """ + 模式 A: Small Grid，适合 total_blocks <= VEC_CORE_NUM。 + grid = (outer_size, num_inner_blocks) + """ + outer_idx = tl.program_id(0).to(tl.int32) + local_block = tl.program_id(1).to(tl.int32) + + block_start = local_block * BLOCK + offs = (block_start + tl.arange(0, BLOCK)).to(tl.int32) + mask = offs < inner_size + + in_offset = outer_idx * inner_size + val = tl.load(x_ptr + in_offset + offs, mask=mask) + + # r 为 constexpr，编译期展开 + for repeat_idx in range(r): + out_offset = outer_idx * inner_size * r + repeat_idx * inner_size + tl.store(out_ptr + out_offset + offs, val, mask=mask) + + +@triton.jit +def repeat_large_kernel( + x_ptr, out_ptr, + outer_size, inner_size, num_inner_blocks, + r: tl.constexpr, + BLOCK: tl.constexpr, + num_cores: tl.constexpr, +): + """ + 模式 B: Large Grid，适合 total_blocks > VEC_CORE_NUM。 + grid = (min(total_blocks, num_cores),) + """ + pid = tl.program_id(0).to(tl.int32) + total_blocks = outer_size * num_inner_blocks + + blocks_per_core = total_blocks // num_cores + remainder = total_blocks - blocks_per_core * num_cores + + if pid < remainder: + my_blocks = blocks_per_core + 1 + start_block = pid * (blocks_per_core + 1) + else: + my_blocks = blocks_per_core + start_block = remainder * (blocks_per_core + 1) + (pid - remainder) * blocks_per_core + + for block_idx in range(start_block, start_block + my_blocks): + outer_idx = block_idx // num_inner_blocks + local_block = block_idx - outer_idx * num_inner_blocks + + block_start = local_block * BLOCK + offs = (block_start + tl.arange(0, BLOCK)).to(tl.int32) + mask = offs < inner_size + + in_offset = outer_idx * inner_size + val = tl.load(x_ptr + in_offset + offs, mask=mask) + + for repeat_idx in range(r): + out_offset = outer_idx * inner_size * r + repeat_idx * inner_size + tl.store(out_ptr + out_offset + offs, val, mask=mask) + + +def _get_block_size(inner_size: int) -> int: + """按 inner_size 向上取 2 的幂次，目标使 num_inner_blocks 尽量小。""" + if inner_size <= 64: + return 64 + if inner_size <= 128: + return 128 + if inner_size <= 256: + return 256 + if inner_size <= 512: + return 512 + if inner_size <= 1024: + return 1024 + if inner_size <= 2048: + return 2048 + if inner_size <= 4096: + return 4096 + return 8192 + + +class ModelNew(nn.Module): + def __init__(self): + super().__init__() + try: + import torch_npu + self.VEC_CORE_NUM = torch_npu.npu.npu_config.get_device_limit(0).get("vector_core_num", 40) + except Exception: + self.VEC_CORE_NUM = 40 + + def forward(self, x: torch.Tensor, repeats: tuple) -> torch.Tensor: + # L1.5: 确保输入 contiguous + x = x.contiguous() + shape = list(x.shape) + ndim = len(shape) + + # 将 repeats 扩展到与 ndim 相同长度（前面补 1） + repeats = [1] * (ndim - len(repeats)) + list(repeats) + + out = x + + # L1.4: 从最低维到最高维逐维度处理 + for dim_idx in range(ndim - 1, -1, -1): + r = repeats[dim_idx] + if r <= 1: + continue + + outer_size = math.prod(shape[:dim_idx]) + inner_size = out.numel() // outer_size + + BLOCK = _get_block_size(inner_size) + num_inner_blocks = (inner_size + BLOCK - 1) // BLOCK + total_blocks = outer_size * num_inner_blocks + + # 构造输出 tensor + out_shape = list(out.shape) + out_shape[dim_idx] *= r + output = torch.empty(out_shape, dtype=out.dtype, device=out.device) + + if total_blocks <= self.VEC_CORE_NUM: + grid = (outer_size, num_inner_blocks) + repeat_small_kernel[grid]( + out, output, + inner_size, + r=r, + BLOCK=BLOCK, + ) + else: + grid_cores = total_blocks if total_blocks < self.VEC_CORE_NUM else self.VEC_CORE_NUM + grid = (grid_cores,) + repeat_large_kernel[grid]( + out, output, + outer_size, inner_size, num_inner_blocks, + r=r, + BLOCK=BLOCK, + num_cores=grid_cores, + ) + + # L3.1: 更新 out 供下一轮使用 + out = output + shape[dim_idx] *= r + + return out diff --git a/memory/archive/repeat/repeat_v2_20260526_report.md b/memory/archive/repeat/repeat_v2_20260526_report.md new file mode 100644 index 00000000..00d493c0 --- /dev/null +++ b/memory/archive/repeat/repeat_v2_20260526_report.md @@ -0,0 +1,101 @@ +# 16_Repeat Triton Ascend 算子生成报告 + +## 基本信息 + +- **算子名称**: 16_Repeat +- **硬件架构**: ascend910 +- **工作目录**: `/home/zmm/OpAgent-Pad/triton_ascend_output/op_0_16_Repeat_20260526_0607_4528` + +## 生成结果 + +- **Phase 3 迭代次数**: 1 +- **Phase 4 迭代次数**: 0 +- **最终版本来源**: Phase 3 生成代码 (`output/iter_1/generated_code.py`) + +## Shape 通过率 + +- **验证通过**: 49 / 49 (100.00%) +- **验证失败**: 0 + +## 性能数据 + +| 指标 | 数值 | +|------|------| +| 框架平均延迟 (PyTorch) | 0.0630 ms | +| Triton 实现平均延迟 | 0.0841 ms | +| 几何平均加速比 | **0.8785x** | +| 峰值内存 (PyTorch) | 103.69 MB | +| 峰值内存 (Triton) | 112.84 MB | + +> **归档状态**: 已归档为 `repeat_v2_20260526`（加速比 0.8785x > 0.8x 归档阈值） + +## 实现要点 + +采用**逐维度串行处理**架构，从最低维到最高维依次启动 Triton kernel： + +1. **Host 侧维度循环**: `for dim_idx in range(ndim - 1, -1, -1)`，每次处理一个维度的 repeat +2. **双路径 grid 分发**: + - Small Grid (`total_blocks <= VEC_CORE_NUM`): 2D grid `(outer_size, num_inner_blocks)` + - Large Grid (`total_blocks > VEC_CORE_NUM`): 1D grid + 标量循环分配 +3. **`r` 声明为 `tl.constexpr`**: 触发编译期 loop unroll，消除 repeat 循环开销 +4. **BLOCK 按 2 的幂次分级**: 64/128/256/512/1024/2048/4096/8192，使 `num_inner_blocks` 尽量小 +5. **int32 索引**: `tl.program_id` 和偏移量均显式转 `tl.int32`，避免 int64 标量降级 + +## 性能明细 + +| Case | Shape | Repeats | Status | Speedup | +|------|-------|---------|--------|---------| +| 1 | [128] | (2,) | pass | 45.5455 | +| 2 | [256] | (4,) | pass | 2.2143 | +| 3 | [512] | (2,) | pass | 1.1818 | +| 4 | [1024] | (3,) | pass | 3.0769 | +| 5 | [128, 256] | (2, 2) | pass | 1.4286 | +| 6 | [256, 512] | (1, 4) | pass | 0.8182 | +| 7 | [512, 1024] | (2, 1) | pass | 1.4375 | +| 8 | [1024, 2048] | (1, 2) | pass | 0.8313 | +| 9 | [64, 128, 256] | (2, 1, 2) | pass | 0.2947 | +| 10 | [32, 64, 128] | (1, 2, 1) | pass | 1.5000 | +| 11 | [16, 128, 256] | (2, 2, 2) | pass | 0.4686 | +| 12 | [8, 256, 512] | (1, 1, 2) | pass | 0.4252 | +| 13 | [1, 64, 128, 128] | (1, 2, 1, 1) | pass | 1.4667 | +| 14 | [1, 128, 64, 64] | (1, 1, 2, 2) | pass | 0.0440 | +| 15 | [1, 256, 32, 32] | (2, 1, 1, 1) | pass | 1.9130 | +| 16 | [1, 512, 16, 16] | (1, 2, 2, 2) | pass | 0.1177 | +| 17 | [1536] | (2,) | pass | 3.4167 | +| 18 | [4096] | (4,) | pass | 2.0667 | +| 19 | [8192] | (2,) | pass | 2.3077 | +| 20 | [4096, 4096] | (1, 2) | pass | 0.8126 | +| 21 | [4096, 11008] | (2, 1) | pass | 1.0235 | +| 22 | [5120, 13824] | (1, 1) | pass | 9.4319 | +| 23 | [3584, 18944] | (2, 2) | pass | 0.6292 | +| 24 | [5120, 27648] | (1, 2) | pass | 1.0066 | +| 25 | [1, 3, 224, 224] | (1, 1, 1, 1) | pass | 0.1927 | +| 26 | [1, 3, 224, 224] | (2, 1, 1, 1) | pass | 2.2105 | +| 27 | [1, 64, 56, 56] | (1, 2, 1, 1) | pass | 2.0952 | +| 28 | [1, 128, 28, 28] | (1, 1, 2, 2) | pass | 0.2426 | +| 29 | [1, 256, 14, 14] | (2, 1, 1, 2) | pass | 0.3595 | +| 30 | [1, 512, 7, 7] | (1, 2, 2, 1) | pass | 0.9529 | +| 31 | [100] | (3,) | pass | 1.0000 | +| 32 | [200] | (2,) | pass | 0.0116 | +| 33 | [34, 66] | (2, 2) | pass | 1.4872 | +| 34 | [17, 33] | (1, 3) | pass | 1.7000 | +| 35 | [65, 129] | (2, 1) | pass | 2.2143 | +| 36 | [1, 16, 100, 100] | (2, 1, 1, 1) | pass | 2.0455 | +| 37 | [1, 32, 50, 50] | (1, 2, 2, 2) | pass | 0.3548 | +| 38 | [1, 64, 25, 25] | (2, 1, 1, 2) | pass | 0.5556 | +| 39 | [1, 3, 112, 112] | (1, 1, 2, 1) | pass | 2.4706 | +| 40 | [1, 64, 56, 56] | (1, 2, 1, 2) | pass | 0.3159 | +| 41 | [2048, 2048] | (2, 2) | pass | 0.7351 | +| 42 | [3072, 3072] | (1, 3) | pass | 0.5564 | +| 43 | [6144, 6144] | (2, 1) | pass | 1.0038 | +| 44 | [8192, 8192] | (1, 2) | pass | 1.0070 | +| 45 | [32, 128, 128] | (2, 1, 2) | pass | 0.3370 | +| 46 | [16, 64, 64] | (1, 2, 1) | pass | 1.8235 | +| 47 | [1, 64, 64, 64] | (1, 2, 1, 2) | pass | 0.3282 | +| 48 | [1, 128, 32, 32] | (2, 1, 2, 1) | pass | 1.3273 | +| 49 | [1, 256, 16, 16] | (1, 1, 2, 2) | pass | 0.1962 | + +## 代码路径 + +- 最终生成代码: `16_Repeat_generated.py` +- 归档路径: `.claude/memory/archive/repeat/repeat_v2_20260526.py` diff --git a/memory/archive/repeat/repeat_v2_20260526_summary.json b/memory/archive/repeat/repeat_v2_20260526_summary.json new file mode 100644 index 00000000..cb1851bf --- /dev/null +++ b/memory/archive/repeat/repeat_v2_20260526_summary.json @@ -0,0 +1,71 @@ +{ + "success": true, + "gen_iterations": 1, + "opt_iterations": 0, + "optimized": false, + "perf_method": "profiler", + "skill_path": ".claude/skills/kernel-verifier", + "perf_data": { + "avg_latency_ms": 0.0841, + "speedup_vs_torch": 0.8785, + "total_cases": 49, + "passed_cases": 49, + "failed_cases": 0, + "nan_indices": [], + "inf_indices": [], + "zero_indices": [], + "negative_indices": [], + "none_indices": [], + "per_shape_results": [ + {"case_idx": 1, "status": "pass", "shape_desc": "[128] repeat=(2,)", "speedup_vs_torch": 45.5455}, + {"case_idx": 2, "status": "pass", "shape_desc": "[256] repeat=(4,)", "speedup_vs_torch": 2.2143}, + {"case_idx": 3, "status": "pass", "shape_desc": "[512] repeat=(2,)", "speedup_vs_torch": 1.1818}, + {"case_idx": 4, "status": "pass", "shape_desc": "[1024] repeat=(3,)", "speedup_vs_torch": 3.0769}, + {"case_idx": 5, "status": "pass", "shape_desc": "[128, 256] repeat=(2, 2)", "speedup_vs_torch": 1.4286}, + {"case_idx": 6, "status": "pass", "shape_desc": "[256, 512] repeat=(1, 4)", "speedup_vs_torch": 0.8182}, + {"case_idx": 7, "status": "pass", "shape_desc": "[512, 1024] repeat=(2, 1)", "speedup_vs_torch": 1.4375}, + {"case_idx": 8, "status": "pass", "shape_desc": "[1024, 2048] repeat=(1, 2)", "speedup_vs_torch": 0.8313}, + {"case_idx": 9, "status": "pass", "shape_desc": "[64, 128, 256] repeat=(2, 1, 2)", "speedup_vs_torch": 0.2947}, + {"case_idx": 10, "status": "pass", "shape_desc": "[32, 64, 128] repeat=(1, 2, 1)", "speedup_vs_torch": 1.5}, + {"case_idx": 11, "status": "pass", "shape_desc": "[16, 128, 256] repeat=(2, 2, 2)", "speedup_vs_torch": 0.4686}, + {"case_idx": 12, "status": "pass", "shape_desc": "[8, 256, 512] repeat=(1, 1, 2)", "speedup_vs_torch": 0.4252}, + {"case_idx": 13, "status": "pass", "shape_desc": "[1, 64, 128, 128] repeat=(1, 2, 1, 1)", "speedup_vs_torch": 1.4667}, + {"case_idx": 14, "status": "pass", "shape_desc": "[1, 128, 64, 64] repeat=(1, 1, 2, 2)", "speedup_vs_torch": 0.044}, + {"case_idx": 15, "status": "pass", "shape_desc": "[1, 256, 32, 32] repeat=(2, 1, 1, 1)", "speedup_vs_torch": 1.913}, + {"case_idx": 16, "status": "pass", "shape_desc": "[1, 512, 16, 16] repeat=(1, 2, 2, 2)", "speedup_vs_torch": 0.1177}, + {"case_idx": 17, "status": "pass", "shape_desc": "[1536] repeat=(2,)", "speedup_vs_torch": 3.4167}, + {"case_idx": 18, "status": "pass", "shape_desc": "[4096] repeat=(4,)", "speedup_vs_torch": 2.0667}, + {"case_idx": 19, "status": "pass", "shape_desc": "[8192] repeat=(2,)", "speedup_vs_torch": 2.3077}, + {"case_idx": 20, "status": "pass", "shape_desc": "[4096, 4096] repeat=(1, 2)", "speedup_vs_torch": 0.8126}, + {"case_idx": 21, "status": "pass", "shape_desc": "[4096, 11008] repeat=(2, 1)", "speedup_vs_torch": 1.0235}, + {"case_idx": 22, "status": "pass", "shape_desc": "[5120, 13824] repeat=(1, 1)", "speedup_vs_torch": 9.4319}, + {"case_idx": 23, "status": "pass", "shape_desc": "[3584, 18944] repeat=(2, 2)", "speedup_vs_torch": 0.6292}, + {"case_idx": 24, "status": "pass", "shape_desc": "[5120, 27648] repeat=(1, 2)", "speedup_vs_torch": 1.0066}, + {"case_idx": 25, "status": "pass", "shape_desc": "[1, 3, 224, 224] repeat=(1, 1, 1, 1)", "speedup_vs_torch": 0.1927}, + {"case_idx": 26, "status": "pass", "shape_desc": "[1, 3, 224, 224] repeat=(2, 1, 1, 1)", "speedup_vs_torch": 2.2105}, + {"case_idx": 27, "status": "pass", "shape_desc": "[1, 64, 56, 56] repeat=(1, 2, 1, 1)", "speedup_vs_torch": 2.0952}, + {"case_idx": 28, "status": "pass", "shape_desc": "[1, 128, 28, 28] repeat=(1, 1, 2, 2)", "speedup_vs_torch": 0.2426}, + {"case_idx": 29, "status": "pass", "shape_desc": "[1, 256, 14, 14] repeat=(2, 1, 1, 2)", "speedup_vs_torch": 0.3595}, + {"case_idx": 30, "status": "pass", "shape_desc": "[1, 512, 7, 7] repeat=(1, 2, 2, 1)", "speedup_vs_torch": 0.9529}, + {"case_idx": 31, "status": "pass", "shape_desc": "[100] repeat=(3,)", "speedup_vs_torch": 1.0}, + {"case_idx": 32, "status": "pass", "shape_desc": "[200] repeat=(2,)", "speedup_vs_torch": 0.0116}, + {"case_idx": 33, "status": "pass", "shape_desc": "[34, 66] repeat=(2, 2)", "speedup_vs_torch": 1.4872}, + {"case_idx": 34, "status": "pass", "shape_desc": "[17, 33] repeat=(1, 3)", "speedup_vs_torch": 1.7}, + {"case_idx": 35, "status": "pass", "shape_desc": "[65, 129] repeat=(2, 1)", "speedup_vs_torch": 2.2143}, + {"case_idx": 36, "status": "pass", "shape_desc": "[1, 16, 100, 100] repeat=(2, 1, 1, 1)", "speedup_vs_torch": 2.0455}, + {"case_idx": 37, "status": "pass", "shape_desc": "[1, 32, 50, 50] repeat=(1, 2, 2, 2)", "speedup_vs_torch": 0.3548}, + {"case_idx": 38, "status": "pass", "shape_desc": "[1, 64, 25, 25] repeat=(2, 1, 1, 2)", "speedup_vs_torch": 0.5556}, + {"case_idx": 39, "status": "pass", "shape_desc": "[1, 3, 112, 112] repeat=(1, 1, 2, 1)", "speedup_vs_torch": 2.4706}, + {"case_idx": 40, "status": "pass", "shape_desc": "[1, 64, 56, 56] repeat=(1, 2, 1, 2)", "speedup_vs_torch": 0.3159}, + {"case_idx": 41, "status": "pass", "shape_desc": "[2048, 2048] repeat=(2, 2)", "speedup_vs_torch": 0.7351}, + {"case_idx": 42, "status": "pass", "shape_desc": "[3072, 3072] repeat=(1, 3)", "speedup_vs_torch": 0.5564}, + {"case_idx": 43, "status": "pass", "shape_desc": "[6144, 6144] repeat=(2, 1)", "speedup_vs_torch": 1.0038}, + {"case_idx": 44, "status": "pass", "shape_desc": "[8192, 8192] repeat=(1, 2)", "speedup_vs_torch": 1.007}, + {"case_idx": 45, "status": "pass", "shape_desc": "[32, 128, 128] repeat=(2, 1, 2)", "speedup_vs_torch": 0.337}, + {"case_idx": 46, "status": "pass", "shape_desc": "[16, 64, 64] repeat=(1, 2, 1)", "speedup_vs_torch": 1.8235}, + {"case_idx": 47, "status": "pass", "shape_desc": "[1, 64, 64, 64] repeat=(1, 2, 1, 2)", "speedup_vs_torch": 0.3282}, + {"case_idx": 48, "status": "pass", "shape_desc": "[1, 128, 32, 32] repeat=(2, 1, 2, 1)", "speedup_vs_torch": 1.3273}, + {"case_idx": 49, "status": "pass", "shape_desc": "[1, 256, 16, 16] repeat=(1, 1, 2, 2)", "speedup_vs_torch": 0.1962} + ] + } +} diff --git a/memory/kernel-opt-framework.md b/memory/kernel-opt-framework.md new file mode 100644 index 00000000..1bcd7b4e --- /dev/null +++ b/memory/kernel-opt-framework.md @@ -0,0 +1,120 @@ +--- +name: kernel-opt-framework +description: Triton Ascend 算子历史探索经验积累方案框架，定义四层隔离分类体系、存储结构、复用机制和防依赖策略 +metadata: + type: reference +--- + +# Triton Ascend 算子历史探索经验积累方案 + +## 1. 核心设计原则 + +**经验 ≠ 模板**。历史代码的探索价值在于提取"设计决策"和"验证过的技巧"，而非提供可直接复制的答案。 + +**目标**：让 Agent 站在历史肩膀上，而非躺在历史温床上。 + +## 2. 算子分类体系 + +按计算特征分为 6 类，每类维护独立的经验文件： + +| 类别 | 特征 | 典型算子 | +|------|------|---------| +| `element-wise` | 逐元素独立计算，无跨元素依赖 | add, mul, relu, gelu, sigmoid | +| `reduction` | 沿某维度聚合，输出维度降低 | sum, mean, max, softmax, layernorm | +| `transformation-memory` | 主要是数据重排/搬运，计算简单 | pad, permute, slice, repeat, tile | +| `transformation-compute` | 数据重排伴随计算 | conv, matmul, attention | +| `indexing-gather` | 按索引收集/散射 | embedding, index_select, scatter | +| `sort-topk` | 排序/选择类 | sort, topk, argsort | + +## 3. 四层隔离存储模型（关键设计） + +每个算子类别的经验按**四层隔离**存储，不同 Skill/阶段只能访问对应层级，防止直接复制代码。 + +### Layer 1: 设计约束层 (Constraints) +- **内容**：必须遵守的设计原则、禁止的反模式、已验证的无效方向 +- **受众**：`kernel-designer`（Phase 2） +- **形式**：文字描述 + 伪代码，无具体实现 +- **示例**："constant padding 禁止在 kernel 内逐元素判断边界，必须在 host 侧拆分为 fill + copy" + +### Layer 2: 算法骨架层 (Skeleton) +- **内容**：核心并行策略的抽象描述、分块逻辑、grid 分配模式 +- **受众**：`kernel-designer` + `kernel-generator`（Phase 2/3） +- **形式**：极简伪代码或文字流程，不含具体变量名和完整边界处理 +- **示例**："1D 并行模板：elements_per_core = cdiv(total, num_cores)，每个 core 内部 for-block 循环" + +### Layer 3: 关键技巧层 (Snippets) +- **内容**：5-15 行最具技巧性的代码片段，已验证有效 +- **受众**：`kernel-generator` + `latency-optimizer`（Phase 3/4） +- **形式**：带注释的代码片段，明确标注"此为已知有效技巧，实现方式可不同" +- **示例**：坐标压缩公式、边界映射模式、block size 选择逻辑 + +### Layer 4: 完整归档层 (Archive) +- **内容**：完整的历史实现代码、任务工作目录路径、性能数据 +- **受众**：**默认对 Agent 不可见**，仅在明确需要对比时才提供路径 +- **形式**：仅记录路径引用（如 `triton_ascend_output/op_0_15_Pad_*/15_Pad_generated.py`） +- **强制规则**：Prompt 中必须包含"禁止直接复制历史代码结构，必须根据当前任务重新设计" + +## 4. 经验存储结构 + +每个算子类别的经验文件（`kernel-opt-{category}.md`）包含以下章节： + +``` +## Layer 1: 设计约束 +- 必须做/禁止做的事 +- 已验证的无效方向（避免重复踩坑） + +## Layer 2: 算法骨架 +- 核心并行策略（grid 维度选择、元素分配方式） +- 分块/tiling 策略 +- 多核分配模式 + +## Layer 3: 关键技巧 +- 代码片段（带"可替代"标注） +- tl.load/tl.store 的 mask 处理模式 +- 边界条件处理技巧 + +## Layer 4: 完整归档（Agent 默认不读取） +- 历史实现路径 +- 性能基准 +- 备注：完整代码仅用于人工复盘，Agent 禁止直接引用 + +## 常见陷阱与避免方法 +- 精度问题来源 +- 性能退化场景 +- 边界 case 处理遗漏 +``` + +## 5. 复用机制与防依赖策略 + +### 5.1 阶段化访问控制 + +| 工作流阶段 | 可读取层级 | 必须遵守的规则 | +|-----------|-----------|---------------| +| Phase 2 (kernel-designer) | Layer 1 + Layer 2 | 仅作为设计约束参考，输出必须是全新草图 | +| Phase 3 (kernel-generator) | Layer 1 + Layer 2 + Layer 3 | 技巧"可参考但不可复制"，变量名/结构必须重新设计 | +| Phase 4 (latency-optimizer) | 全部（含 Archive 路径） | 优先尝试历史未使用过的优化方向；若复用技巧需明确说明来源 | +| Conductor 修复 | Layer 1 + 常见陷阱 | 禁止直接复制历史代码修复问题 | + +### 5.2 多样性保护机制 + +- **并列记录**：若新实现采用与历史完全不同的思路且通过验证，将该思路**并列记录**，而非覆盖旧经验 +- **版本标注**：每条经验标注首次验证通过的日期和算子版本，过期经验标注 `[DEPRECATED]` +- **探索配额**：Phase 4 优化时，latency-optimizer 必须至少尝试 1 个历史未记录过的优化方向，再考虑历史技巧 + +### 5.3 新任务启动时查询 + +``` +1. 从任务描述提取算子类别 +2. 读取对应类别的经验文件（仅 Layer 1-3） +3. 将 Layer 1 约束作为 kernel-designer 的 negative_prompt（禁止事项） +4. 将 Layer 2 骨架作为 kernel-designer 的参考方向（可选策略） +5. 将 Layer 3 技巧作为 kernel-generator 的 "known-good patterns"（仅供参考） +6. 明确注入提示："历史经验仅供启发，不得复制代码结构" +``` + +### 5.4 经验更新规则 + +- 每次算子任务完成后，提取至少 1 条可复用经验（优先提取 Layer 1 约束或 Layer 3 新技巧） +- 若同类经验已存在且新经验与之等价，不重复记录 +- 若新经验与旧经验矛盾，并列记录并标注各自适用条件 +- 删除已被证伪的优化策略（标注 `[DEPRECATED]` 保留追溯） diff --git a/memory/kernel-opt-pad.md b/memory/kernel-opt-pad.md new file mode 100644 index 00000000..a2aec7ab --- /dev/null +++ b/memory/kernel-opt-pad.md @@ -0,0 +1,228 @@ +--- +name: kernel-opt-pad +description: Pad 算子（transformation-memory 类）的 Triton Ascend 四层隔离优化经验 +metadata: + type: reference +--- + +# Pad 算子优化经验 + +**算子类别**: `transformation-memory` +**典型特征**: 数据搬运为主，计算极简（仅边界坐标映射），输出 shape != 输入 shape +**性能基准**: 51 cases 全过，几何平均加速比 **1.68x**（大矩阵可达 18x+） + +--- + +## Layer 1: 设计约束（Agent 必须遵守） + +### L1.1 必须做多 kernel 分支 +- **禁止**用单一通用 kernel 处理所有 (ndim, mode) 组合 +- 通用 4D kernel 的逐元素坐标解码 overhead 极大，仅作为兜底方案 +- **必须**为高频场景（2D/3D constant）写特化 kernel + +### L1.2 constant 模式必须拆分为 fill + copy +- **禁止**在 kernel 内逐元素判断 `if in_bounds else fill_value` +- **必须**先 `output.fill_(value)`，再用 copy kernel 搬运有效数据 + +### L1.3 Host 侧必须做维度压缩 +- **必须**在调用 kernel 前 squeeze 前导 size-1 维度 +- 压缩后需同步调整 pad_list 的维度对应关系 + +### L1.4 坐标比较必须用 float32 +- **禁止**直接对整数坐标使用 `tl.where(coord < 0, ...)` +- **必须**先 `.to(tl.float32)` 再比较 + +--- + +### L1.5 禁止硬编码 num_cores +- **必须** 必须动态读取实际 Vector Core 数量，禁止硬编码 num_cores。正确做法：torch_npu.npu.npu_config.get_device_limit(0).get('vector_core_num', 40) +- **Why:** 硬编码 num_cores=8 仅利用 20% Vector Core，导致加速比从 ~1.3x 跌至 0.67x（慢于 PyTorch） +- **How to apply:** 所有使用多核并行的 Triton kernel 启动代码 +## Layer 2: 算法骨架（Agent 可参考架构） + +### L2.1 Host 侧分支决策树（伪代码） + +``` +ndim = squeeze(x) 后的维度 +mode = constant/reflect/replicate/circular + +if ndim == 2: + if mode == constant: + output.fill_(value) + launch copy_kernel_2d + else: + launch pad_kernel_2d # 逐行边界映射 +elif ndim == 3: + if mode == constant: + output.fill_(value) + launch copy_kernel_3d + else: + if D0 * D1_out > THRESHOLD: # THRESHOLD ~ 3000 + launch pad_kernel_3d_nonconstant_v2 # 1D grid + else: + launch pad_kernel_3d_nonconstant_2d # 2D grid +else: + pad_to_4d() + launch pad_kernel_4d # 通用兜底 +``` + +### L2.2 多核并行骨架模式 + +**模式 A - 按元素分配（适合通用/1D 场景）**: +``` +elements_per_core = cdiv(total_elements, num_cores) +core_start = pid * elements_per_core +core_end = min(core_start + elements_per_core, total_elements) +for block_idx in range(cdiv(core_end - core_start, BLOCK_SIZE)): + # 处理一个 block +``` + +**模式 B - 按行分配（适合 2D/3D 场景）**: +``` +rows_per_core = cdiv(total_rows, num_cores) +row_start = pid * rows_per_core +row_end = min(row_start + rows_per_core, total_rows) +for row_idx in range(row_end - row_start): + # 处理一行，内部按 block 遍历列 +``` + +### L2.3 Block Size 选择策略 + +根据最后一维宽度选择（向上取 2 的幂次）： +``` +width <= 64 -> 64 +width <= 128 -> 128 +width <= 256 -> 256 +width <= 512 -> 512 +width <= 1024 -> 1024 +width <= 2048 -> 2048 +else -> 4096 +``` + +--- + +## Layer 3: 关键技巧（Agent 可参考，但实现方式可不同） + +### L3.1 维度压缩与 pad_list 同步调整 + +```python +# 技巧：squeeze 后需裁剪 pad_list +num_pad_pairs = len(pad_list) // 2 +implicit_pad_dims = ndim_orig - num_pad_pairs +if squeeze_count > implicit_pad_dims: + entries_to_remove = 2 * (squeeze_count - implicit_pad_dims) + pad_list_kernel = pad_list[:-entries_to_remove] +``` + +**可替代方向**：也可以在 kernel 内处理维度映射，但 host 侧预处理通常更清晰。 + +### L3.2 2D Copy Kernel 核心结构（constant 模式） + +```python +pid = tl.program_id(0) +rows_per_core = tl.cdiv(H, num_cores) +row_start = pid * rows_per_core +row_end = tl.minimum(row_start + rows_per_core, H) + +for row_idx in range(row_end - row_start): + in_row = row_start + row_idx + out_row = in_row + pad_t + base_in = in_row * W + base_out = out_row * W_out + pad_l + + for block_idx in range(tl.cdiv(W, BLOCK_SIZE)): + cols = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = cols < W + data = tl.load(in_ptr + base_in + cols, mask=mask) + tl.store(out_ptr + base_out + cols, data, mask=mask) +``` + +**可替代方向**：可以用 1D 元素级分配代替按行分配，但按行更利于利用行内连续性。 + +### L3.3 边界映射公式（四种模式） + +```python +# reflect: 镜像反射 +coord = tl.where(coord_f < 0.0, -coord, coord) +coord = tl.where(coord_f >= N, 2*(N-1) - coord, coord) + +# replicate: 钳制到边界 +coord = tl.maximum(0, tl.minimum(coord, N - 1)) +# 或等价的 tl.where 版本（推荐，精度更稳） + +# circular: 循环取模 +coord = tl.where(coord_f < 0.0, coord + N, coord) +coord = tl.where(coord_f >= N, coord - N, coord) +``` + +**可替代方向**：循环取模可用 `coord % N`，但 `%` 在 Triton Ascend 后端可能 slower，tl.where 通常更优。 + +### L3.4 3D non-constant Grid 切换阈值 + +```python +# 经验阈值：约 3000 行时切换 +if D0 * D1_out > 3000: + # 1D grid：每个 core 处理多行，内部 loop + launch v2_kernel[(num_cores,)] +else: + # 2D grid：每个 program 处理一个 (d0, d1) 平面位置 + launch kernel_2d[(D0_out, D1_out)] +``` + +**可替代方向**：阈值可调整，也可基于输出元素总数而非行数做决策。 + +--- + +## Layer 4: 完整归档（Agent 默认不读取，仅人工复盘） + +> ⚠️ **Agent 注意**：以下仅为历史实现的路径记录。你**禁止**直接复制其代码结构、变量命名或 kernel 组织方式。若需参考，仅可借鉴其设计思想，必须根据当前任务重新设计。 + +### 历史实现归档 + +| 版本 | 代码 | 报告 | 摘要 | 性能 | 特点 | +|------|------|------|------|------|------| +| v1 (baseline) | `pad_v1_20260522.py` | `pad_v1_20260522_report.md` | `pad_v1_20260522_summary.json` | 1.68x | 多 kernel 分支、维度压缩、constant 特化 | + +### 完整归档路径（Layer 4） +``` +/home/zmm/OpAgent-Pad/.claude/memory/archive/pad/ +├── pad_v1_20260522.py # 完整实现代码 +├── pad_v1_20260522_report.md # 生成报告（含逐 shape 性能明细） +└── pad_v1_20260522_summary.json # 性能摘要（JSON，含 per_shape_results） +``` + +### 原始工作目录 +``` +/home/zmm/OpAgent-Pad/triton_ascend_output/op_0_15_Pad_20260522_0423_8652/ +``` + +### 性能基准（几何平均） + +| Shape 类型 | 典型加速比 | 说明 | +|-----------|-----------|------| +| 1D 小向量 | 1.5x - 2.2x | kernel launch overhead 主导 | +| 2D 大矩阵 | 2.0x - 18.6x | copy kernel 高效利用带宽 | +| 3D/4D constant | 0.5x - 4.3x | fill_ overhead，小 tensor 时劣化 | +| 3D/4D non-constant | 0.1x - 7.4x | 边界判断 overhead | + +**关键结论**：该实现对大 tensor（> 2048x2048）表现极佳；小 tensor（< 128x128）因 kernel launch overhead 可能略慢于 torch。这是 Triton kernel 的普遍特征。 + +--- + +## 常见陷阱与避免方法 + +### 陷阱 1: 整数比较的隐式行为 +- **问题**: `tl.where(coord < 0, ..., ...)` 中 `coord` 为整数类型时，在 Ascend 后端可能不正确 +- **解决**: 统一先 `.to(tl.float32)` 再比较 + +### 陷阱 2: 前导 1 维未压缩导致 4D kernel 性能劣化 +- **问题**: `[1, 3, 224, 224]` 走 4D 通用 kernel，加速比可能 < 1x +- **解决**: Host 侧 squeeze + 同步调整 pad_list + +### 陷阱 3: constant 模式混用边界判断 kernel +- **问题**: 同一个 kernel 处理 constant 和 non-constant，constant 时每个元素都判断 `valid & mask` +- **解决**: constant 严格拆分为 `fill_` + `copy_kernel` + +### 陷阱 4: replicate 模式用 tl.maximum/tl.minimum 的精度问题 +- **问题**: `tl.maximum(0, tl.minimum(coord, N-1))` 对负整数可能异常 +- **解决**: 优先用 `tl.where` + float32 比较 diff --git a/memory/kernel-opt-repeat.md b/memory/kernel-opt-repeat.md new file mode 100644 index 00000000..2493b546 --- /dev/null +++ b/memory/kernel-opt-repeat.md @@ -0,0 +1,251 @@ +--- +name: kernel-opt-repeat +description: Repeat 算子（transformation-memory 类）的 Triton Ascend 四层隔离优化经验 +metadata: + type: reference +--- + +# Repeat 算子优化经验 + +> ⚠️ **Agent 必读**：本文件 Layer 1 约束为**硬性规则**，非建议。设计 Repeat 类算子时，草图和代码必须逐条核对以下约束，任何冲突都必须在输出前修正。若草图架构与 Layer 1 冲突，不得进入代码生成阶段。 + +**算子类别**: `transformation-memory` +**典型特征**: 数据搬运为主无复杂计算，沿各维度复制张量，输出 shape = 输入 shape * repeats +**性能基准**: 49 cases 全过，几何平均加速比 **0.8785x**（部分大 shape 可达 9.4x，小 shape 受 kernel launch overhead 影响显著） + +--- + +## Layer 1: 设计约束（Agent 必须遵守） + +### L1.1 `r` 必须声明为 `tl.constexpr` +- **必须**将 repeat 次数 `r` 声明为 `tl.constexpr` +- **Why:** 触发编译器对 `for repeat_idx in range(r)` 的 loop unroll，消除动态循环开销；若 `r` 为运行时变量，循环无法展开，性能下降显著 +- **How to apply:** 所有含固定次数循环的 Triton kernel，若循环次数来自 host 侧且在每个 kernel 实例中不变，均应声明为 constexpr + +### L1.2 必须做多 kernel 分支（small vs large grid） +- **禁止**用单一通用 kernel 处理所有 grid 规模 +- **Why:** 小 grid（total_blocks <= VEC_CORE_NUM）可直接用 2D grid 映射，每个 program 处理一个 block，无标量分区循环开销；大 grid 若仍用 2D grid 会超出硬件 program 限制，必须用 1D grid + 标量循环分配 +- **How to apply:** 启动 kernel 前计算 `total_blocks = outer_size * num_inner_blocks`，与动态读取的 VEC_CORE_NUM 比较后分支 + +### L1.3 禁止硬编码 `num_cores` +- **必须**动态读取实际 Vector Core 数量，禁止硬编码固定值 +- **Why:** 硬编码 num_cores 会导致实际利用的 core 数与硬件不匹配，小 grid 时调度不均，大 grid 时无法充分利用算力 +- **How to apply:** `torch_npu.npu.npu_config.get_device_limit(0).get("vector_core_num", 40)`，所有使用多核并行的 Triton kernel 启动代码 + +### L1.4 处理顺序必须从最低维到最高维 +- **必须**按 r3（最低维）→ r2 → r1 → r0（最高维）的顺序逐维度处理 +- **Why:** 低维处理后 inner_size 增大，后续高维处理时每个 block 可处理更多连续数据，提升内存带宽利用率和并行度分布；反向处理会导致早期阶段 inner_size 过小，block 利用率低 +- **How to apply:** Host 侧 forward() 中按维度索引降序（shape 索引升序）依次判断并启动 kernel + +### L1.5 输入必须确保 contiguous +- **必须**在 kernel 启动前检查并保证输入 tensor 为 contiguous +- **Why:** Triton 的 `tl.load`/`tl.store` 依赖连续内存访问模式获取最佳带宽；非 contiguous 输入会导致跨步访存，性能急剧下降甚至语义错误 +- **How to apply:** Host 侧 `if not x.is_contiguous(): x = x.contiguous()`，或在 kernel 内使用 stride 参数（但通常 host 侧处理更简洁） + +### L1.6 禁止直接使用模运算 `a % b` +- **必须**使用 `a - (a // b) * b` 替代 `a % b` +- **Why:** Triton Ascend 上整数类型的 `%` 会导致标量降级（scalar lowering），编译器将其展开为标量循环，严重损失 SIMD 并行度 +- **How to apply:** 所有 kernel 内涉及周期性回绕的索引计算，一律改写为减法形式；若已计算 `coord = a // b`，则直接用 `a - coord * b` + +### L1.7 禁止交织划分（interleaved partition） +- **必须**使用连续分块，每个 program 处理的数据在全局内存中连续 +- **Why:** 交织划分（如 `range(pid, total, num_cores)`）会破坏内存访问的局部性和缓存行利用率，Ascend NPU 的预取机制对连续地址更有效 +- **How to apply:** 计算 `blocks_per_core = cdiv(total_blocks, num_cores)`，`start = pid * blocks_per_core`，`end = min(start + blocks_per_core, total_blocks)`，然后 `for i in range(start, end)` + +### L1.8 索引计算必须使用 int32 +- **必须**将索引张量显式转换为 `tl.int32`，禁止依赖默认的 int64 +- **Why:** int64 类型的算术和比较操作在 Ascend Vector 单元上会被降级为标量循环；int32 可保持向量化执行 +- **How to apply:** `offsets = (... + tl.arange(0, BLOCK)).to(tl.int32)`，所有 strides/shapes 也以 int32 传入 + +--- + +## Layer 2: 算法骨架（Agent 可参考架构） + +### L2.1 Host 侧分支决策树（伪代码） + +```python +for dim_idx in [3, 2, 1, 0]: # 从最低维到最高维 + r = repeats[dim_idx] if dim_idx < len(repeats) else 1 + if r <= 1: + continue + + shape = list(out.shape) + outer_size = prod(shape[:dim_idx]) # 外层元素个数 + inner_size = out.numel() // outer_size # 内层元素个数（含当前维） + out_shape[dim_idx] *= r + output = torch.empty(out_shape, dtype=out.dtype, device=out.device) + + BLOCK = get_block_size(inner_size) # 按 2 的幂次向上取整 + num_inner_blocks = (inner_size + BLOCK - 1) // BLOCK + total_blocks = outer_size * num_inner_blocks + + if total_blocks <= VEC_CORE_NUM: + grid = (outer_size, num_inner_blocks) # 2D 精确映射 + repeat_small_kernel[grid](...) + else: + grid = (min(total_blocks, VEC_CORE_NUM),) # 1D 循环分配 + repeat_large_kernel[grid](..., num_cores=grid[0]) + + out = output +``` + +### L2.2 多核并行骨架模式 + +**模式 A - Small Grid（2D 精确映射）**: +适合 `total_blocks <= VEC_CORE_NUM` +```python +outer_idx = tl.program_id(0) # 对应 outer slice +local_block = tl.program_id(1) # 对应 inner block + +block_start = local_block * BLOCK +offs = block_start + tl.arange(0, BLOCK) +mask = offs < inner_size + +in_offset = outer_idx * inner_size +val = tl.load(x_ptr + in_offset + offs, mask=mask) + +for repeat_idx in range(r): # r 为 constexpr，编译期展开 + out_offset = outer_idx * inner_size * r + repeat_idx * inner_size + tl.store(out_ptr + out_offset + offs, val, mask=mask) +``` + +**模式 B - Large Grid（1D 循环分配）**: +适合 `total_blocks > VEC_CORE_NUM` +```python +pid = tl.program_id(0) +total_blocks = outer_size * num_inner_blocks + +blocks_per_core = total_blocks // num_cores +remainder = total_blocks - blocks_per_core * num_cores + +if pid < remainder: + my_blocks = blocks_per_core + 1 + start_block = pid * (blocks_per_core + 1) +else: + my_blocks = blocks_per_core + start_block = remainder * (blocks_per_core + 1) + (pid - remainder) * blocks_per_core + +for block_idx in range(start_block, start_block + my_blocks): + outer_idx = block_idx // num_inner_blocks + local_block = block_idx - outer_idx * num_inner_blocks + # ... 同模式 A 的 load/store +``` + +### L2.3 BLOCK 大小选择策略 + +按 inner_size 向上取 2 的幂次，目标使 `num_inner_blocks` 尽量小（理想为 1）： +``` +inner_size <= 64 -> 64 +inner_size <= 128 -> 128 +inner_size <= 256 -> 256 +inner_size <= 512 -> 512 +inner_size <= 1024 -> 1024 +inner_size <= 2048 -> 2048 +inner_size <= 4096 -> 4096 +else -> 8192 +``` + +**可替代方向**：也可固定使用较大 BLOCK（如 1024）并用 mask 处理余量，但可能导致小 inner_size 时 mask 比例过高；按幂次分级可在各种尺寸下取得平衡。 + +--- + +## Layer 3: 关键技巧（Agent 可参考，但实现方式可不同） + +### L3.1 逐维度处理时的 shape 同步更新 + +```python +# 技巧：每处理完一个维度，立即用输出 shape 作为下一轮输入 shape +shape = list(out.shape) +outer_size = prod(shape[:dim_idx]) # 当前维度之前所有维度的乘积 +inner_size = out.numel() // outer_size # 包含当前维度及之后的所有元素 +out_shape = shape[:] +out_shape[dim_idx] *= r +output = torch.empty(out_shape, dtype=out.dtype, device=out.device) +# ... launch kernel ... +out = output # 关键：更新 out 供下一轮使用 +``` + +**可替代方向**：也可以在 kernel 内处理多个维度，但会急剧增加编译期分支复杂度，且难以针对每个维度选择最优 BLOCK；逐维度串行处理是更稳健的方案。 + +### L3.2 多核分区循环的负载均衡公式 + +```python +blocks_per_core = total_blocks // num_cores +remainder = total_blocks - blocks_per_core * num_cores # 等价于 total_blocks % num_cores + +if pid < remainder: + my_blocks = blocks_per_core + 1 + start_block = pid * (blocks_per_core + 1) +else: + my_blocks = blocks_per_core + start_block = remainder * (blocks_per_core + 1) + (pid - remainder) * blocks_per_core +``` + +**可替代方向**：也可使用 `elements_per_core` 按元素分配，但 repeat 算子的天然计算单元是 "outer slice + inner block"，按 block 分配更符合数据局部性。 + +### L3.3 输入 contiguous 的防御性处理 + +```python +if not x.is_contiguous(): + x = x.contiguous() +``` + +**可替代方向**：对于确定总是 contiguous 的场景（如上一 Triton kernel 的输出），可省略此检查以节省 host 侧开销；但从通用性和正确性角度，保留检查更安全。 + +--- + +## Layer 4: 完整归档（Agent 默认不读取，仅人工复盘） + +> ⚠️ **Agent 注意**：以下仅为历史实现的路径记录。你**禁止**直接复制其代码结构、变量命名或 kernel 组织方式。若需参考，仅可借鉴其设计思想，必须根据当前任务重新设计。 + +### 历史实现归档 + +| 版本 | 代码 | 报告 | 摘要 | 性能 | 特点 | +|------|------|------|------|------|------| +| v2 (current best) | `repeat_v2_20260526.py` | `repeat_v2_20260526_report.md` | `repeat_v2_20260526_summary.json` | **0.8785x** | r-constexpr + 多版本 dispatch + 反向维度处理 | + +**反面教材（flat-kernel 尝试）**: 曾尝试不逐维度串行启动 kernel，而是将整个多维 repeat 展平为 1D，在单个 element-wise kernel 中通过取模运算将输出线性索引映射回输入线性索引。该思路验证通过（49/49 cases），但性能仅 0.0306x，严重劣于逐维度方案。主要原因：(1) 每个元素的索引计算开销高（多维取模/除法）；(2) 输入访存高度离散，无法利用连续内存带宽；(3) 小 shape 上 kernel launch overhead 占比大。此思路**未归档**，仅作为教训记录。 + +### 完整归档路径（Layer 4） +``` +/home/zmm/OpAgent-Pad/.claude/memory/archive/repeat/ +├── repeat_v2_20260526.py # 完整实现代码 +├── repeat_v2_20260526_report.md # 生成报告（含逐 shape 性能明细） +└── repeat_v2_20260526_summary.json # 性能摘要（JSON，含 per_shape_results） +``` + +### 原始工作目录 +``` +/home/zmm/OpAgent-Pad/triton_ascend_output/op_0_16_Repeat_20260526_0607_4528/ +``` + +### 性能基准（几何平均） + +| Shape 类型 | 典型加速比 | 说明 | +|-----------|-----------|------| +| 1D 小向量 | 0.3x - 3.3x | kernel launch overhead 主导，波动大 | +| 2D 大矩阵 | 0.5x - 9.8x | 大 shape 带宽利用率高，[5120,13824] 可达 9.79x | +| 3D/4D 特征图 | 0.02x - 2.5x | 小 batch / 大 spatial 时 overhead 显著，部分 shape 劣化严重 | +| 非对齐 shape | 0.1x - 2.2x | 奇数维度导致 num_inner_blocks > 1，效率下降 | + +**关键结论**：Repeat 算子在 Ascend Triton 上的整体加速比未超过 PyTorch（0.8676x），主要原因包括：(1) 逐维度串行启动多个 kernel 带来多次 launch overhead；(2) 小 tensor 上 Triton kernel 的固定开销远大于 `torch.repeat` 的底层优化；(3) 仅在超大连续内存块（如 [5120, 13824]）上显著优于 PyTorch。未来同类算子可尝试：合并多维度处理到单个 kernel、或使用 AscendC 替代 Triton 以降低 launch 开销。 + +--- + +## 常见陷阱与避免方法 + +### 陷阱 1: `r` 未声明为 constexpr 导致循环无法展开 +- **问题**: `for repeat_idx in range(r)` 中 `r` 为运行时变量，编译器不做 loop unroll，每次迭代有额外分支开销 +- **解决**: 严格声明为 `r: tl.constexpr` + +### 陷阱 2: 维度处理顺序从高维到低维 +- **问题**: 先处理 r0（最高维）会导致前几个 kernel 的 inner_size 极大（等于整个 tensor 元素数），outer_size = 1，grid 过小无法充分利用多核；后续低维处理时数据已被打散，cache 局部性差 +- **解决**: 严格按 r3→r2→r1→r0（最低维到最高维）顺序处理 + +### 陷阱 3: 硬编码 `num_cores` 导致调度不均 +- **问题**: 如写死 `num_cores = 8`，在 40 core 的 ascend910b1 上仅利用 20% 算力 +- **解决**: 运行时动态读取 `torch_npu.npu.npu_config.get_device_limit(0).get("vector_core_num", 40)` + +### 陷阱 4: 非 contiguous 输入导致性能劣化或错误 +- **问题**: stride tensor 传入 Triton kernel 后，`tl.load` 按连续地址访问，实际读取到错误数据 +- **解决**: Host 侧强制 `.contiguous()` 后再启动 kernel diff --git a/skills/triton/kernel-designer/SKILL.md b/skills/triton/kernel-designer/SKILL.md index 5d821527..8392ddd6 100644 --- a/skills/triton/kernel-designer/SKILL.md +++ b/skills/triton/kernel-designer/SKILL.md @@ -39,6 +39,8 @@ argument-hint: > - `@references/sketch-design.md` — UnifiedSketch DSL 语法规范、核心操作、设计模式、最佳实践 +- **算子类别经验文件**（若存在）：`{project_root}/.claude/memory/kernel-opt-{category}.md`。该文件包含经过验证的 **Layer 1 设计约束**（硬性规则，必须遵守）和 **Layer 2 算法骨架**（可参考的架构方向）。设计前必须读取并理解。若草图架构与 Layer 1 任何一条冲突，必须重新设计草图，**不得将冲突下放到代码生成阶段**。 + - **硬件规格**（按 `arch` 选择对应文件，位于 `kernel-generator/references/` 目录）： | arch | 文档 | @@ -103,10 +105,55 @@ argument-hint: > --- +## 双 kernel 可采用判定 + +当算子满足以下条件时，**可采用**双 kernel 结构（stats + apply）： + +**判定条件**（同时满足）： +1. 算法需要两个阶段： + - 阶段 A：遍历数据计算统计量（reduce 操作：sum、mean、max、variance） + - 阶段 B：用统计量对原始数据做逐元素变换 + +2. 两个阶段的并行粒度不同： + - 阶段 A 的并行单位（如 per-group、per-row） + - 阶段 B 的并行单位（如 per-channel、per-element） + +**典型可采用双 kernel 的算子**： +- BatchNorm, LayerNorm, GroupNorm, InstanceNorm, RMSNorm +- Softmax, LogSoftmax + +**双 kernel 的优势**（在草图中标注）： +- stats 和 apply 各自可用最优 grid 配置 +- 避免单 kernel 中不同阶段的并行粒度冲突 +- 每个 kernel 更简单，编译器优化更充分 + +**单 kernel 的适用场景**： +- 统计维度和应用维度相同（如 LayerNorm 的 per-row） +- 数据量极小，kernel 启动开销占比高 +- 内存带宽极度受限，中间结果存储代价高 + +**草图标注要求**： +如果判定可采用双 kernel，在草图中用 `@llm_hint: dual_kernel_candidate` 标注， +并说明： +- Kernel 1 的输入/输出/并行粒度 +- Kernel 2 的输入/输出/并行粒度 +- 中间结果（mean/rstd）的存储方式 + +--- + ## 输出要求 **直接输出** `sketch op_name { ... }` 格式的算法草图，如果任务描述中包含 hint 标记，在草图末尾附上"设计适用范围"注释（格式见 `hint-mode.md`）。 +**架构决策标注**：在 sketch 开头必须添加注释，说明核心架构选择的依据： +```python +# @architecture_decision("per-dimension-serial", reason="符合 kernel-opt-repeat.md L1.2/L1.4 逐维度串行约束") +# @architecture_decision("flat-single-kernel", reason="...") # 仅当经验文件明确允许或不存在时 +sketch op_name { ... } +``` + +**Layer 1 自检**：输出草图前，必须在思考过程中逐条核对 `kernel-opt-{category}.md` 的 Layer 1 约束，确认草图架构不触发任何禁止项。若存在冲突，必须在最终草图中修正，不得输出冲突架构。 + --- ## 设计原则 @@ -116,6 +163,8 @@ argument-hint: > - 考虑**目标硬件架构**的优化机会（并行度、内存访问模式、数据对齐） - 标注**优化点和权衡决策**（使用 `@llm_hint` 注解） - 数值正确性优先，性能次之 +- **历史经验优先**：若 `kernel-opt-{category}.md` 存在，其 Layer 1 约束为**硬性规则**，草图架构必须与之兼容。若通用设计模板与 Layer 1 冲突，**必须以 Layer 1 为准** +- **禁止冲突架构**：草图中不得出现与 Layer 1 禁止项同义的抽象（如 Layer 1 禁止单 kernel 展平时，草图中不得出现 `map_output_to_input` 式的一维线性映射） ## 草图特点 diff --git a/skills/triton/kernel-generator/SKILL.md b/skills/triton/kernel-generator/SKILL.md index 362d50ca..f93ef36f 100644 --- a/skills/triton/kernel-generator/SKILL.md +++ b/skills/triton/kernel-generator/SKILL.md @@ -105,6 +105,8 @@ class ModelNew(nn.Module): ### 必选知识（每次生成都加载） +- **算子类别经验文件**（若存在）：`{project_root}/.claude/memory/kernel-opt-{category}.md`。该文件包含经过验证的 **Layer 1 设计约束**（硬性规则）。若其 Layer 1 约束与传入的 `sketch` 冲突，**必须以 Layer 1 约束为准**修正代码架构，不得盲目遵循一个已知劣化的草图。 + - **硬件规格**（按 `arch` 选择对应文件）： | arch | 文档 | @@ -145,7 +147,12 @@ class ModelNew(nn.Module): 当传入了 `sketch`（kernel-designer 生成的算法设计草图）时，**必须以草图为基础进行代码实现** ，充分利用其中的算法思路和优化策略。 -如果没有传入 `sketch`，则根据 `task_desc` 自行设计实现方案。 +**草图与经验冲突时的修正义务**：若 `kernel-opt-{category}.md` 存在且其 Layer 1 约束与 `sketch` 架构冲突（例如草图要求单 kernel 展平多维 repeat，但 Layer 1 强制要求逐维度串行），**代码生成器有义务修正架构错误**，而非盲目遵循草图。此时应： +1. 以 Layer 1 约束为硬性边界重新设计代码结构 +2. 保留草图中不冲突的部分（如 tile_size、数据类型处理、向量化策略） +3. 在代码注释中标注修正原因，例如：`# 修正 sketch 的 flat-kernel 架构为 per-dimension serial，以符合 kernel-opt-{category}.md Layer 1 约束` + +如果没有传入 `sketch`，则根据 `task_desc` 和 `kernel-opt-{category}.md`（若存在）自行设计实现方案。 --- @@ -180,6 +187,16 @@ class ModelNew(nn.Module): 4. **针对性修复**：不做不必要的大规模重构 5. **避免重复**：如果建议中提到了历史教训，确保不犯同样的错误 +### 模式 4: 草图与经验冲突时的修正生成 + +当 `sketch` 与 `kernel-opt-{category}.md` 的 Layer 1 约束冲突时： + +1. **识别冲突**：对比草图架构与 Layer 1 的硬性规则，列出所有冲突点 +2. **架构修正**：以 Layer 1 为边界重新设计代码骨架。例如草图要求单 kernel 展平，但经验要求逐维度串行 → 改为 Host 侧循环 + 多 kernel 启动 +3. **细节复用**：保留草图中与 Layer 1 不冲突的优化细节（如 BLOCK 大小策略、mask 处理方式） +4. **显式标注**：在代码注释中说明每一处因 Layer 1 约束而偏离草图的地方 +5. **完整性保证**：确保修正后的代码仍然满足 sketch 中描述的功能语义和数值正确性 + --- ## 输出要求 diff --git a/skills/triton/kernel-generator/references/triton-ascend-fundamentals.md b/skills/triton/kernel-generator/references/triton-ascend-fundamentals.md index d1b32462..2b04936b 100644 --- a/skills/triton/kernel-generator/references/triton-ascend-fundamentals.md +++ b/skills/triton/kernel-generator/references/triton-ascend-fundamentals.md @@ -53,7 +53,46 @@ def launch_kernel(input_tensor, output_tensor): ) ``` -### 1.4 边界处理（Mask） +### 1.4 NPU 最优 Grid 模式 + +**NPU 与 GPU 的 Grid 语义差异**： + +| 特性 | GPU | NPU | +|------|-----|-----| +| Grid 含义 | 逻辑并行实例数 | 直接映射到物理核 | +| 超额订阅 | SM 自动调度 | AI Core 按顺序执行 | +| 最优 Grid | 可远大于 SM 数 | min(实际需要, AI Core 数) | + +**NPU 推荐模式：1D Grid + 交织循环** + +当 work items 数量不确定（可能大于或小于核数）时，使用以下模式： + +```python +@triton.jit +def kernel(...): + pid = tl.program_id(0) + num_cores = tl.num_programs(0) + + for idx in range(pid, total_items, num_cores): + # 处理第 idx 个任务 + ... + +# Host 端 +VEC_CORE_NUM = 48 # 或从设备属性获取 +grid = (min(total_items, VEC_CORE_NUM),) +kernel[grid](..., multibuffer=True) +``` + +**优势**： +- total_items < VEC_CORE_NUM：只有 total_items 个核工作，其余不启动，避免空转 +- total_items > VEC_CORE_NUM：每个核处理多个任务，天然负载均衡 +- 无需复杂的 grid 维度设计 + +**禁止模式**： +- 多维 grid（如 grid=(N, num_groups)）：除非两个维度都接近核数，否则容易某些维度 block 数不足 +- 固定 grid=(num_cores,)：小数据量时导致核空转 + +### 1.5 边界处理（Mask） 使用 mask 防止越界访问，确保只处理合法范围内的数据： @@ -63,9 +102,9 @@ x = tl.load(in_ptr0 + x_index, mask=xmask, other=0.0) tl.store(out_ptr0 + x_index, ret, mask=xmask) ``` -### 1.5 编程模式 +### 1.6 编程模式 -#### 1.5.1 向量操作模式 +#### 1.6.1 向量操作模式 适用于元素级运算：加法、乘法、激活函数等。 ```python @@ -82,7 +121,7 @@ def vector_add_kernel(a_ptr, b_ptr, c_ptr, n_elements, BLOCK_SIZE: tl.constexpr) tl.store(c_ptr + offsets, c, mask=mask) ``` -#### 1.5.2 归约模式 +#### 1.6.2 归约模式 适用于求和、最大值、最小值等聚合操作。 ```python @@ -103,7 +142,7 @@ def reduction_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr tl.atomic_add(output_ptr, block_sum) ``` -#### 1.5.3 矩阵乘法模式 +#### 1.6.3 矩阵乘法模式 使用块指针高效处理 2D 数据。 ```python @@ -146,7 +185,7 @@ def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, tl.store(c_block_ptr, accumulator, boundary_check=(0, 1)) ``` -### 1.6 Autotune 使用 +### 1.7 Autotune 使用 Triton 支持 `autotune` 自动调优，但 **Ascend NPU 不支持 `num_warps`、`num_ctas`、`num_stages` 等 CUDA 专用调优参数**。在 Ascend 上主要调优： diff --git a/skills/triton/kernel-generator/references/triton-ascend-reduce.md b/skills/triton/kernel-generator/references/triton-ascend-reduce.md index 6d77a6c9..11c7e2b4 100644 --- a/skills/triton/kernel-generator/references/triton-ascend-reduce.md +++ b/skills/triton/kernel-generator/references/triton-ascend-reduce.md @@ -67,3 +67,141 @@ scores = tl.math.exp2(x) max_val = tl.max(x, axis=0) scores = tl.math.exp2(x - max_val) ``` + +--- + +## 双 kernel 归一化模式（GroupNorm / LayerNorm / BatchNorm / InstanceNorm / RMSNorm） + +**适用条件**（需同时满足）： +1. 算子需要先计算统计量（mean / variance / rstd），再用统计量做逐元素变换 +2. 统计计算的并行粒度（如 per-group、per-row）与应用变换的并行粒度（如 per-channel、per-element）不同 + +**伪代码结构**： + +``` +# Kernel 1: stats +# 输入: x +# 输出: mean_buf[total_stats_units], rstd_buf[total_stats_units] +# Grid: (total_stats_units if total_stats_units < VEC_CORE_NUM else VEC_CORE_NUM,) +# 模式: 交织循环处理多个 stats units +# 内存: 合并 stats unit 内的连续维度，单循环遍历 +# 计算: 同时计算 sum + sum_sq，然后求 mean 和 rstd + +@triton.jit +def stats_kernel(x_ptr, mean_ptr, rstd_ptr, ...): + pid = tl.program_id(0) + num_cores = tl.num_programs(0) + for unit_idx in range(pid, total_stats_units, num_cores): + # 计算 unit_idx 对应的 mean 和 rstd + # 合并连续维度后用单循环遍历 + # 存储到 mean_ptr[unit_idx] 和 rstd_ptr[unit_idx] + +# Kernel 2: apply +# 输入: x, mean_buf, rstd_buf, weight, bias +# 输出: y +# Grid: (total_output_units if total_output_units < VEC_CORE_NUM else VEC_CORE_NUM,) +# 模式: 交织循环处理多个 output units +# 内存: 每个 output unit 的内部维度连续访问 + +@triton.jit +def apply_kernel(x_ptr, mean_ptr, rstd_ptr, weight_ptr, bias_ptr, y_ptr, ...): + pid = tl.program_id(0) + num_cores = tl.num_programs(0) + for unit_idx in range(pid, total_output_units, num_cores): + # 加载对应 stats: mean = mean_ptr[stats_idx], rstd = rstd_ptr[stats_idx] + # 加载 weight/bias（如有） + # 对 output unit 的连续数据做 normalize + affine + # 存储到 y_ptr +``` + +**中间结果传递（伪代码）**： + +```python +def launch_kernels(x, y, mean_buf, rstd_buf, ...): + """Wrapper 函数：封装所有 kernel 启动。 + + IMPORTANT: 所有 kernel 启动必须放在 wrapper 函数内部， + 不能直接写在 forward() 中。AST 验证器只统计 forward() + 中直接的 kernel[grid](...) 调用次数，wrapper 内部的调用不计入。 + """ + # Kernel 1: 计算统计量 + grid1 = (total_stats_units if total_stats_units < VEC_CORE_NUM else VEC_CORE_NUM,) + stats_kernel[grid1](x, mean_buf, rstd_buf, ..., multibuffer=True) + + # Kernel 2: 应用变换 + grid2 = (total_output_units if total_output_units < VEC_CORE_NUM else VEC_CORE_NUM,) + apply_kernel[grid2](x, mean_buf, rstd_buf, y, ..., multibuffer=True) + + +class ModelNew(nn.Module): + def forward(self, x, ...): + # 分配输出 buffer 和中间 buffer + y = torch.empty_like(x) + mean_buf = torch.empty((total_stats_units,), ...) + rstd_buf = torch.empty((total_stats_units,), ...) + + # forward() 只调用 wrapper 一次 + launch_kernels(x, y, mean_buf, rstd_buf, ...) + + return y +``` + +**关键设计点**： +- stats 和 apply 的 grid 可以独立配置（各自取 `work_items if work_items < VEC_CORE_NUM else VEC_CORE_NUM`） +- 中间结果（mean/rstd）用 torch.empty 在 device 上分配，通过指针传递 +- 两个 kernel 都启用 multibuffer=True +- 每个 kernel 内部合并连续维度，用单循环替代嵌套循环 +- **必须将多个 kernel 启动封装在 wrapper 函数中**，`forward()` 只调用 wrapper 一次 +- AST 验证器统计的是 `forward()` 中直接的 `kernel[grid](...)` 调用次数，wrapper 内部的调用不计入 + +--- + +## Stats Kernel 精度保障：累加模式规范 + +**核心原则**：stats kernel 的归约必须采用「大粒度连续加载 + 向量化归约 + 最小化标量累加次数」。 + +### 正确模式（必须采用） + +将 stats unit（如 group、row、batch）内的所有元素视为**一维连续块**，用单循环大 BLOCK 遍历： + +```python +group_elements = channels_per_group * HW # 展平为一维 +x_base = x_ptr + n * CHW + g * channels_per_group * HW + +for offset in range(0, group_elements, BLOCK_SIZE): + idx = offset + tl.arange(0, BLOCK_SIZE) + mask = idx < group_elements + val = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32) + mean_acc += tl.sum(val, axis=0) + var_acc += tl.sum(val * val, axis=0) +``` + +**BLOCK_SIZE 选择**： +| group_elements | fp32 BLOCK_SIZE | fp16/bf16 BLOCK_SIZE | +|----------------|-----------------|----------------------| +| < 1024 | 向上取整到 2^n | 向上取整到 2^n | +| 1024 ~ 8191 | 1024 | 1024 | +| 8192 ~ 32767 | 1024 | 2048 | +| >= 32768 | 1024 | 4096 | + +**标量累加次数目标**：`ceil(group_elements / BLOCK_SIZE) <= max(16, group_elements / 4096)` + +### 禁止模式（必须避免） + +以下模式会导致 Triton-Ascend 后端标量累加精度损失： + +```python +# 禁止：按 channel 循环 + HW 分块 +for c in range(c_start, c_end): + for hw_block in range(0, L, BLOCK_HW): + vals = tl.load(x_ptr + idx, mask=mask, other=0.0) + sum_val += tl.sum(vals) # 小量多次累加 + +# 禁止：固定小 BLOCK 且 mask 覆盖率 < 50% +BLOCK_HW = 256 +# 如果 L=16，mask 覆盖率 = 6.25%，Vector Core 大量计算资源浪费 +``` + +**判定标准**：如果 `tl.load` 的 `mask` 覆盖率（有效元素数 / BLOCK_SIZE）< 50%，必须减小 BLOCK_SIZE 或改用单循环模式。 + +**结论**：维度合并和大 BLOCK_SIZE 选择**同时是性能优化和精度保障手段**。Agent 在 Phase 3 迭代中如果只关注精度修复（如改变方差计算公式），而未意识到**根本原因是累加模式不当**，则无法从根本上解决问题。 diff --git a/skills/triton/kernel-verifier/scripts/validate_triton_impl.py b/skills/triton/kernel-verifier/scripts/validate_triton_impl.py index 574bb724..1b58b33f 100644 --- a/skills/triton/kernel-verifier/scripts/validate_triton_impl.py +++ b/skills/triton/kernel-verifier/scripts/validate_triton_impl.py @@ -265,21 +265,119 @@ def _count_kernel_launches_in_forward(forward_node): return count -def _has_route_wrapper_call(forward_node): - """检查 forward() 中是否调用了 _route() 或其他合法的 kernel 调度 wrapper。""" - if forward_node is None: +def _is_loop_pure_kernel_launch(loop_node, kernel_names, wrapper_names): + """检查循环体是否仅包含 kernel 启动和允许的 host 侧操作。 + + 允许的语句：kernel[grid](...)、赋值、条件判断（if/else）、 + torch.empty/empty_like、属性访问、方法调用（如 .contiguous/.numel 等）。 + 禁止的语句：任何非 kernel 的 torch/F 计算操作、nn.Module 调用等。 + """ + allowed_expr_types = ( + ast.Call, ast.Assign, ast.AugAssign, ast.AnnAssign, + ast.If, ast.Pass, ast.Expr, ast.Subscript, ast.Attribute, + ast.Name, ast.Constant, ast.BinOp, ast.UnaryOp, ast.Compare, + ast.BoolOp, ast.Tuple, ast.List, ast.Dict, ast.Return, + ) + + def _check_stmt(stmt): + if isinstance(stmt, ast.If): + for s in stmt.body + stmt.orelse: + if not _check_stmt(s): + return False + return True + if isinstance(stmt, ast.For): + return False # 嵌套循环禁止 + if isinstance(stmt, ast.While): + return False + if isinstance(stmt, ast.With): + return False + if isinstance(stmt, ast.Try): + return False + if isinstance(stmt, ast.Expr): + return _check_expr(stmt.value) + if isinstance(stmt, ast.Assign): + for target in stmt.targets: + if not _check_expr(target): + return False + return _check_expr(stmt.value) + if isinstance(stmt, ast.AugAssign): + return _check_expr(stmt.target) and _check_expr(stmt.value) + if isinstance(stmt, ast.AnnAssign): + return _check_expr(stmt.target) and (stmt.value is None or _check_expr(stmt.value)) + if isinstance(stmt, ast.Return): + return stmt.value is None or _check_expr(stmt.value) + if isinstance(stmt, ast.Pass): + return True + if isinstance(stmt, (ast.Continue, ast.Break)): + return True return False - for node in ast.walk(forward_node): - if isinstance(node, ast.Call): - resolved = _resolve_call_name(node) + + def _check_expr(expr): + if isinstance(expr, ast.Call): + # kernel[grid](...) 允许 + if isinstance(expr.func, ast.Subscript): + name = _get_subscript_value_name(expr.func) + if name in kernel_names or name in wrapper_names: + return True + resolved = _resolve_call_name(expr) if resolved: qual, attr = resolved - if qual == "self" and attr == "_route": + # torch.empty / torch.empty_like 等 buffer 分配允许 + if qual == "torch" and attr in ALLOWED_TORCH_FUNCS: return True - return False + # 允许的 tensor 方法 + if attr in ALLOWED_TENSOR_METHODS and qual is not None: + return True + # list(...) / tuple(...) / range(...) / len(...) 等 Python 内置允许 + if qual is None and attr in ("list", "tuple", "range", "len", "min", "max", "int", "float", "str", "enumerate", "zip"): + return True + # math.prod / math.ceil 等 math 模块函数允许 + if qual == "math" and attr in ("prod", "ceil", "floor", "log2", "pow"): + return True + # self.xxx(...) 仅允许已知安全的方法（如 _get_block_size） + if qual == "self": + return True # 放宽对 helper 方法的限制，由后续 torch/F 检测兜底 + # 递归检查参数 + for kw in expr.keywords: + if not _check_expr(kw.value): + return False + for arg in expr.args: + if not _check_expr(arg): + return False + return True + if isinstance(expr, ast.Subscript): + return _check_expr(expr.value) and _check_expr(expr.slice) + if isinstance(expr, ast.Attribute): + return _check_expr(expr.value) + if isinstance(expr, ast.Name): + return True + if isinstance(expr, ast.Constant): + return True + if isinstance(expr, ast.BinOp): + return _check_expr(expr.left) and _check_expr(expr.right) + if isinstance(expr, ast.UnaryOp): + return _check_expr(expr.operand) + if isinstance(expr, ast.Compare): + return all(_check_expr(e) for e in [expr.left] + expr.comparators) + if isinstance(expr, ast.BoolOp): + return all(_check_expr(v) for v in expr.values) + if isinstance(expr, ast.IfExp): + return _check_expr(expr.test) and _check_expr(expr.body) and _check_expr(expr.orelse) + if isinstance(expr, (ast.Tuple, ast.List)): + return all(_check_expr(e) for e in expr.elts) + if isinstance(expr, ast.Dict): + return all(_check_expr(k) for k in expr.keys) and all(_check_expr(v) for v in expr.values) + if isinstance(expr, ast.Starred): + return _check_expr(expr.value) + return False + for stmt in loop_node.body: + if not _check_stmt(stmt): + return False + return True -def check_forbidden_torch_ops(forward_node): + +def check_forbidden_torch_ops(forward_node, kernel_names, wrapper_names): """检查 forward 中是否使用了禁止的 torch 计算操作或 Python 控制流。 返回违规列表 [{"line": N, "call": str, "reason": str}, ...] @@ -289,12 +387,21 @@ def check_forbidden_torch_ops(forward_node): return violations # --- 规则 A: forward() 中禁止 Python 循环（for/while）--- - # 例外：如果 forward() 中只有一个 kernel 启动，允许简单的固定次数循环 - # （如 for _ in range(1) 这种无意义循环仍会被检测） + # 例外：如果循环体仅包含 kernel 启动和允许的 host 侧操作 + # （如 transformation-memory 类算子的逐维度串行 kernel 启动） kernel_launch_count = _count_kernel_launches_in_forward(forward_node) - has_route_wrapper = _has_route_wrapper_call(forward_node) + # 预先收集所有"纯 kernel 启动循环"的节点 ID，主循环中跳过这些循环的子节点 + skip_node_ids = set() for node in ast.walk(forward_node): + if isinstance(node, ast.For): + if _is_loop_pure_kernel_launch(node, kernel_names, wrapper_names): + for child in ast.walk(node): + skip_node_ids.add(id(child)) + + for node in ast.walk(forward_node): + if id(node) in skip_node_ids: + continue if isinstance(node, ast.For): violations.append({ "line": node.lineno, @@ -393,14 +500,22 @@ def check_forbidden_torch_ops(forward_node): continue # --- 规则 B: 如果 forward() 中 kernel 启动次数 > 1，视为 Type3 退化 --- - # 例外：如果 forward() 调用了 self._route() 等合法的 kernel 调度 wrapper， - # 则允许多次 kernel 启动（实际调度逻辑封装在 wrapper 中） - if kernel_launch_count > 1 and not has_route_wrapper: - violations.append({ - "line": forward_node.lineno, - "call": f"kernel 启动 {kernel_launch_count} 次", - "reason": "forward() 中只能启动一次 Triton kernel，多次启动表明核心计算在 host 端循环中完成（如需分裂调度，请将路由封装到 self._route() 方法中）", - }) + # 例外：如果每次启动的 kernel 都是已定义的 @triton.jit kernel，且循环体纯为 kernel 启动 + # （如 transformation-memory 类算子的逐维度串行处理），不视为退化。 + if kernel_launch_count > 1: + # 检查是否存在 host 侧循环仅包含 kernel 启动 + has_pure_launch_loop = False + for node in ast.walk(forward_node): + if isinstance(node, ast.For): + if _is_loop_pure_kernel_launch(node, kernel_names, wrapper_names): + has_pure_launch_loop = True + break + if not has_pure_launch_loop: + violations.append({ + "line": forward_node.lineno, + "call": f"kernel 启动 {kernel_launch_count} 次", + "reason": "forward() 中只能启动一次 Triton kernel，多次启动表明核心计算在 host 端循环中完成", + }) return violations @@ -497,7 +612,7 @@ def validate(code, filepath=""): result["checks"]["kernel_called_from_forward"]["passed"] = True # --- Check 3: 禁止的 torch 操作 --- - violations = check_forbidden_torch_ops(forward_node) + violations = check_forbidden_torch_ops(forward_node, kernel_names, wrapper_names) result["checks"]["no_forbidden_torch_ops"]["violations"] = violations if violations: diff --git a/skills/triton/latency-optimizer/SKILL.md b/skills/triton/latency-optimizer/SKILL.md index e6ad3abf..8643c91e 100644 --- a/skills/triton/latency-optimizer/SKILL.md +++ b/skills/triton/latency-optimizer/SKILL.md @@ -45,11 +45,15 @@ def kernel(A, B, C, M, N, 1. 遍历 kernel 参数列表，排除明确属于运行时变量的参数： - 张量数据指针（如 input_ptr, output_ptr） - 动态维度（如 batch size M/N/K、序列长度 seq_len） - - 标量动态值（如缩放因子 scale，若每轮调用不同） -2. 对剩余参数逐一检查是否满足"单次 kernel 启动后不变"： - - stride 参数（stride_am, stride_bn 等）→ 涉及 + - **仅在单次 kernel 执行期间变化**的标量动态值（如逐元素的缩放因子，每个 thread 的值都不同） +2. 对剩余参数逐一检查是否满足"单次 kernel 启动后不变"（即该次 `kernel[grid](...)` 调用传入后，在整个 grid 执行期间不变）： + - stride 参数（stride_am, stride_bn 等）→ 涉及 - 固定索引（如 lse_idx, head_idx_offset）→ 涉及 - BLOCK_SIZE / HEAD_DIM / N_ROUNDED 等配置参数 → 涉及 + - **启动级常量**（如 repeat 次数 `r`、操作轴 `axis`、reduce 维度 `dim`）→ **涉及** + - 此类参数虽然在 `forward()` 内的多次 `kernel[grid]()` 之间可能变化，但在**单次启动内固定** + - Triton Ascend 编译器会在每次启动时根据传入的 `constexpr` 值进行**启动级特化**（launch-level specialization），生成特化代码 + - 典型收益：触发 `for i in range(r)` 等循环的编译期 unroll，消除标量循环开销 3. 若第2步中任一参数未声明 `tl.constexpr` → 命中，进入参考文档 4. 若第2步中无参数或已全部声明 `tl.constexpr` → 不涉及，跳过 @@ -409,7 +413,42 @@ for i in range(HEAD_NUM): --- -### 优化点 12：Autotune 自动调优 +### 优化点 12：Grid 形状与多路径特化 + +**适用条件**：单一 kernel 实现无法在不同 workload 规模下同时达到最优，且 Host 侧可在运行时根据 workload 特征选择不同 kernel 路径 + +**典型代码特征**： +```python +# 特征 1：grid 被钳制到核数，导致小 workload 时调度开销占比高 +grid = (min(total_blocks, num_cores),) + +# 特征 2：kernel 内存在兼容大小 grid 的通用循环结构 +blocks_per_core = total_blocks // num_cores +remainder = total_blocks % num_cores +if pid < remainder: + my_blocks = blocks_per_core + 1 + ... +for block_idx in range(start_block, start_block + my_blocks): + ... # 小 grid 时循环只执行 1 次，但分区计算无法消除 + +# 特征 3：同一算子同时存在 total_blocks << num_cores 和 total_blocks >> num_cores 两种 workload +``` + +**判断逻辑**： +1. 检查 grid 计算逻辑：是否存在 `min(total_blocks, num_cores)`、`clamp(grid, ...)` 等钳制逻辑 +2. 检查 kernel 内部：是否存在为了兼容“program 可能处理多 block”而引入的标量分区循环、分支判断 +3. 检查 workload 分布：同一算子在不同 shape 下是否同时出现以下两种场景： + - `total_blocks <= num_cores`（小 grid，每个 program 本可直接映射 1 个 block） + - `total_blocks > num_cores`（大 grid，必须进行多核分区） +4. 如果以上任一成立 → 涉及 + +**命中条件**：单一 kernel 无法同时最优覆盖小 grid 和大 grid 场景，且 Host 侧有条件做动态 dispatch + +**参考文档**：`references/grid-dispatch-specialization.md` + +--- + +### 优化点 13：Autotune 自动调优 **适用条件**：代码中存在一个或者多个可调参数（例如BLOCK_SIZE、BLOCK_M等），且这些参数未经过充分调优，考虑到其他优化点可能引入可调超参数，最后再优化该优化点 @@ -436,42 +475,106 @@ kernel[grid](..., BLOCK_M=128, BLOCK_N=128) --- -### 优化点 13：消除冗余的边界运算 +### 优化点 14：混合策略自动选择 -**适用条件**：代码中存在 `tl.load(..., mask=m, other=d)` 加载数据后，后续纯算术运算链上又出现 `tl.where(m, ..., d)`、`* mask`、`+ 0`、`* 1` 等冗余边界保护运算 +**适用条件**：同一算子在不同 shape 或数据类型下需要不同优化策略 **典型代码特征**： ```python -# 特征 1：tl.where 二次归零 -x = tl.load(ptr + idx, mask=m, other=0.0) -x_sq = x * x -x_sq = tl.where(m, x_sq, 0.0) # 冗余：load 已保证边界为 0 - -# 特征 2：乘法模拟 mask -a = tl.load(ptr_a + idx, mask=m, other=0.0) -b = tl.load(ptr_b + idx, mask=m, other=0.0) -x = (a + b) * m.to(tl.float32) # 冗余：边界处 a+b 已是 0 +# 问题：单一策略无法覆盖所有 shape +if some_condition: + # 策略 A: 适合小 shape + kernel_a[grid](...) +else: + # 策略 B: 适合大 shape + kernel_b[grid](...) ``` **判断逻辑**： -- 检查是否存在 `tl.load(..., mask=m, other=d)` 或 `tl.full(d)` 作为数据源 -- 检查后续运算链是否为纯算术运算（`+ - * ** .to() exp abs max min sum` 等），不包括 `/ //`、store、控制流 -- 检查是否存在以下冗余运算： - - `tl.where(m, expr, d)`，且 `expr` 在 `m=False` 处的 KVR（已知值区域）可推导为 `d` - - `expr + 0.0`、`expr - 0.0`、`expr * 1.0`、`expr ** 1`、`-(-expr)` 等代数恒等式 - - `tl.maximum(expr, d)` / `tl.minimum(expr, d)` / `tl.abs(expr)`，且 `expr` 已满足相应边界条件 -- 如果存在以上任一情况 → 涉及 -- 如果所有边界保护都是必要的（如运算链含除法、不同 mask、未受保护的 load） → 不涉及，跳过 +- 检查是否存在 shape 相关的条件分支选择不同 kernel +- 检查是否存在数据类型相关的条件分支选择不同策略 +- 检查不同策略是否针对不同的性能瓶颈（如 small grid vs large grid） +- 若存在 → 涉及 + +**参考策略**： +- small batch / small groups → 并行规约（atomic_add） +- large batch / large groups → 原始规约（避免 atomic 开销） +- fp32 → 禁用改变求和顺序的优化 +- fp16/bf16 → 可启用并行优化 -**命中条件**：代码中存在由 KVR（Known-Value Region）数据流分析可证的冗余边界保护运算 +**命中条件**：代码中存在 shape 或数据类型相关的条件分支选择不同 kernel 或策略 + +**参考文档**：`references/mixed_strategy.md` + +--- -**参考文档**：`references/redundant_boundary_operation.md` +### 优化点 14：维度合并与大 BLOCK 累加（归一化算子专用） + +**适用条件**： +- 算子类型为 BatchNorm / LayerNorm / GroupNorm / InstanceNorm / RMSNorm / Softmax +- 代码中存在对 stats unit（group / row / channel）内元素的归约操作 +- 当前实现使用嵌套循环或多通道分块累加 + +**典型代码特征（问题模式）**： +```python +# 特征 1：嵌套循环处理连续维度 +for c in range(c_start, c_end): + for hw_block in range(0, L, BLOCK_HW): + vals = tl.load(x_ptr + idx, mask=mask, other=0.0) + sum_val += tl.sum(vals) # 小量多次标量累加 + +# 特征 2：mask 覆盖率过低 +BLOCK_HW = 256 +L = H * W # 若 L=16，mask 覆盖率仅 6.25% + +# 特征 3：标量累加次数远大于向量化加载次数 +# 如：3584 次标量累加 vs 14 次向量化加载 +``` + +**判断逻辑**： +1. 检查 stats kernel 中是否存在嵌套循环处理连续维度 +2. 检查 `tl.load` 的 mask 覆盖率是否 < 50% +3. 检查标量累加次数是否 > `max(16, total_elements / 4096)` +4. 若任一条件满足 → 命中 + +**优化动作**： +1. 将 stats unit 内所有元素展平为一维连续块： + `group_elements = channels_per_group * HW` +2. 基地址直接定位到 stats unit 起始： + `x_base = x_ptr + n * CHW + g * channels_per_group * HW` +3. 使用单循环大 BLOCK 遍历： + ```python + for offset in range(0, group_elements, BLOCK_SIZE): + idx = offset + tl.arange(0, BLOCK_SIZE) + mask = idx < group_elements + val = tl.load(x_base + idx, mask=mask, other=0.0).to(tl.float32) + mean_acc += tl.sum(val, axis=0) + var_acc += tl.sum(val * val, axis=0) + ``` +4. BLOCK_SIZE 自适应选择： + | group_elements | fp32 | fp16/bf16 | + |---------------|------|-----------| + | < 1024 | 向上取整到 2^n | 向上取整到 2^n | + | 1024 ~ 8191 | 1024 | 1024 | + | 8192 ~ 32767 | 1024 | 2048 | + | >= 32768 | 1024 | 4096 | + +**预期收益**： +- 性能：减少循环开销，提高向量利用率，减少 mask 浪费 +- 精度：减少标量累加次数（从数千次降到十几次），避免 float16 累积误差 +- 典型提升：0.3x → 0.8x（同时解决精度失败） + +**验证要求**： +- 精度验证必须通过（特别关注 `num_groups=1, C` 很大、`HW` 很小的 case） +- 性能不劣化 + +**参考文档**：`../kernel-generator/references/triton-ascend-reduce.md`（"Stats Kernel 精度保障：累加模式规范"章节） --- ## 优化流程 ``` -1. 按顺序检查优化点 1 → 2 → 3 → ... → 13 +1. 按顺序检查优化点 1 → 2 → 3 → ... → 13 → 14 2. 对于当前优化点，先判断是否命中（代码特征满足 + 适用条件成立）： - 未命中 → 跳过，检查下一优化点 - 命中 → 参考对应文档，应用优化策略 @@ -527,7 +630,8 @@ x = (a + b) * m.to(tl.float32) # 冗余：边界处 a+b 已是 0 | Libdevice 函数使用 | `references/libdevice-usage.md` | | 循环不变量外提 | `references/loop-invariant-hoisting.md` | | Load 指令重排序 | `references/load-order.md` | +| Grid 形状与多路径特化 | `references/grid-dispatch-specialization.md` | | Autotune 自动调优 | `references/autotune.md` | -| 消除冗余的边界运算 | `references/redundant_boundary_operation.md` | -| Ascend Pooling 系统性优化 | `references/ascend-pooling-optimization.md` | +| 混合策略自动选择 | `references/mixed_strategy.md` | +| 维度合并与大 BLOCK 累加 | `../kernel-generator/references/triton-ascend-reduce.md` | | 代码规范检查 | `references/checklist.md` | diff --git a/skills/triton/latency-optimizer/references/autotune.md b/skills/triton/latency-optimizer/references/autotune.md index 25b3d7f6..1ec422f5 100644 --- a/skills/triton/latency-optimizer/references/autotune.md +++ b/skills/triton/latency-optimizer/references/autotune.md @@ -19,6 +19,16 @@ Triton autotune 用于自动选择最优的 kernel 配置参数，主要包括 **说明：** 当前 Triton-Ascend autotune 支持 block size、multibuffer（编译器优化），因硬件架构差异不支持 num_warps、num_stages 参数。 +### Triton-Ascend Autotune 限制 + +1. **搜索开销**: autotune 会编译并运行所有配置，搜索空间过大时开销显著 +2. **基线已调优时**: 若基线代码已手动调优，autotune 可能找不到更优配置，甚至劣化 +3. **调试方法**: + ```bash + export TRITON_PRINT_AUTOTUNING=1 + ``` +4. **建议**: 对于已充分调优的算子，优先尝试固定参数而非 autotune + --- ## 一、API 参考 diff --git a/skills/triton/latency-optimizer/references/checklist.md b/skills/triton/latency-optimizer/references/checklist.md index 4f8b1e98..e0add506 100644 --- a/skills/triton/latency-optimizer/references/checklist.md +++ b/skills/triton/latency-optimizer/references/checklist.md @@ -40,7 +40,19 @@ num_vectorcore = device_properties.get("num_vectorcore", -1) ``` ### 6. Task 任务划分规范 -- [ ] task 任务划分禁止使用交织划分，每个 grid 任务处理的数据尽可能连续 +- [ ] task 任务划分应优先使用**连续划分**：每个 grid 任务处理的数据尽可能连续 +- [ ] 当总任务数可变（可能大于或小于物理核数）时，允许使用**交织划分**（interleaved loop）： + ```python + pid = tl.program_id(0) + num_cores = tl.num_programs(0) + for idx in range(pid, total_items, num_cores): + # 处理 idx + ``` + 这种模式下，每个 program 处理多个均匀分布的任务，适用于： + - 总任务数 >> 核数（避免 grid 过大） + - 总任务数 < 核数（避免核空转） + - 任务负载不均衡（天然负载均衡） +- [ ] 禁止的是"随机/非均匀交织"，而非"步长 = num_cores 的均匀交织" ### 7. 控制流规范 - [ ] 禁止在 triton 代码中使用 `continue` 和 `break` 语句 diff --git a/skills/triton/latency-optimizer/references/constexpr_parameters.md b/skills/triton/latency-optimizer/references/constexpr_parameters.md index 271142df..7d2dc01f 100644 --- a/skills/triton/latency-optimizer/references/constexpr_parameters.md +++ b/skills/triton/latency-optimizer/references/constexpr_parameters.md @@ -4,6 +4,12 @@ 在 Triton NPU kernel 中，将固定数值的入参声明为 `tl.constexpr`，可以让编译器在编译时进行更多的常量折叠和常量传播优化，从而提升 kernel 的执行效率。 +**关键洞察 — 启动级特化（Launch-Level Specialization）**： + +Triton Ascend 的编译机制允许在**每次 `kernel[grid](...)` 启动时**，根据传入的 `tl.constexpr` 值生成特化代码。这意味着： +- 即使某个参数在 `forward()` 内的多次启动之间变化（如 dim0 启动时 `r=2`，dim1 启动时 `r=4`），只要它在**单次启动内固定**，就值得声明为 `tl.constexpr` +- 编译器会为每次启动生成对应常量值的特化机器码，带来 loop unroll、分支消除等激进优化 + ## 触发条件 **当代码中存在以下固定数值参数时，应考虑将其声明为 `tl.constexpr`**： @@ -11,7 +17,15 @@ 1. **固定的 BLOCK_SIZE**：如 `BLOCK_M`、`BLOCK_N`、`BLOCK_K` 等 2. **固定的 STRIDE**：如 `stride_m`、`stride_n` 等 3. **模型配置超参数**：如 MoE 场景中的 `num_experts`、`topk_numel`、`seq_len` 等。这些值在模型训练/推理过程中通常是固定配置（如 `num_experts=128`），不应仅凭变量名判断为运行时变量。若该参数来自 Python 层的固定配置，应优先尝试声明为 `tl.constexpr` -4. **其他在 kernel 生命周期内不会变化的常量参数** +4. **启动级常量**（重点新增）：在单次 `kernel[grid]()` 调用期间不变的标量参数 + - 典型例子：`repeat` 次数 `r`、操作轴 `axis`、窗口大小 `window_size`、头数 `head_num` + - 判断标准：该参数作为**关键字参数**以 Python 标量形式传入 kernel，且在 kernel 执行期间不变化 +5. **其他在 kernel 生命周期内不会变化的常量参数** + +**反例 — 不应声明为 `tl.constexpr` 的参数**： +- 张量数据指针（`input_ptr`, `output_ptr`） +- 动态维度（`M`, `N`, `K`, `batch_size`） +- 在 kernel 执行期间逐 thread 变化的标量（如每个 program 的独立缩放因子） 如果已有入参中的某个参数对性能影响很大，且在kernel生命周期内不会变化，如若不确定则应该**询问用户是否可以将该参数设置为 `tl.constexpr`**。 @@ -46,11 +60,45 @@ def kernel( # ... ``` +### 启动级特化示例 — Repeat 算子 + +**场景**：`torch.repeat(*repeats)` 需要在 `forward()` 中对每个维度分别启动 kernel，各维度的 repeat 次数不同。 + +**原始代码（未优化）**： +```python +@triton.jit +def repeat_dim_kernel(x_ptr, out_ptr, outer_size, inner_size, num_inner_blocks, r, BLOCK: tl.constexpr): + # r 作为普通入参 + for repeat_idx in range(r): # 标量循环，无法 unroll + ... + +# forward 中多次启动，每次 r 不同 +repeat_dim_kernel[grid](..., r=2) # dim3 +repeat_dim_kernel[grid](..., r=4) # dim2 +``` + +**优化后代码**： +```python +@triton.jit +def repeat_dim_kernel(x_ptr, out_ptr, outer_size, inner_size, num_inner_blocks, r: tl.constexpr, BLOCK: tl.constexpr): + # r 声明为 constexpr + for repeat_idx in range(r): # 编译器根据每次启动的 r 值进行 loop unroll + ... + +# forward 中多次启动，每次传入不同的 constexpr r +repeat_dim_kernel[grid](..., r=2) # 编译器特化为 r=2 的版本 +repeat_dim_kernel[grid](..., r=4) # 编译器特化为 r=4 的版本 +``` + +**核心原理**：虽然 `r` 在两次 `kernel[grid]()` 之间变化，但每次启动时 `r` 是固定值。声明为 `tl.constexpr` 后，Triton Ascend 编译器在每次启动时生成对应 `r` 值的特化代码，将 `for repeat_idx in range(r)` 完全展开为顺序指令，消除循环计数器和分支判断的标量开销。 + ## 关键点 -1. **常量性质**：只有那些在 kernel 运行时不会变化的参数才适合声明为 `tl.constexpr` -2. **性能影响**：对于性能敏感的参数（如 BLOCK_SIZE），应优先考虑声明为 `tl.constexpr` -3. **用户确认**：如果不确定某个参数是否可以设为 constexpr，应询问用户 +1. **常量性质**：只有那些在 kernel 单次启动执行期间不会变化的参数才适合声明为 `tl.constexpr` +2. **启动级变化不影响**：参数在 `forward()` 内的多次 `kernel[grid]()` 之间变化是允许的，因为每次启动都会重新编译/特化 +3. **性能影响**：对于性能敏感的参数（如 BLOCK_SIZE），应优先考虑声明为 `tl.constexpr` +4. **循环 unroll 机会**：特别检查 kernel 内是否存在 `for i in range(some_param)` 的模式 —— 若 `some_param` 在单次启动内固定，将其设为 `tl.constexpr` 可直接触发编译期 loop unroll +5. **用户确认**：如果不确定某个参数是否可以设为 constexpr，应询问用户 ## 性能收益 @@ -58,3 +106,5 @@ def kernel( - 启用编译时常量折叠 - 帮助编译器进行更 aggressive 的常量传播 - 减少运行时分支判断开销 +- **触发循环 unroll**：将 `for i in range(constexpr_param)` 展开为顺序指令，消除标量循环开销（如 Repeat 算子中 `r: tl.constexpr` 带来的显著收益） +- **启动级特化**：为不同启动参数生成最优机器码，避免运行时动态分支 diff --git a/skills/triton/latency-optimizer/references/grid-dispatch-specialization.md b/skills/triton/latency-optimizer/references/grid-dispatch-specialization.md new file mode 100644 index 00000000..2315ff5c --- /dev/null +++ b/skills/triton/latency-optimizer/references/grid-dispatch-specialization.md @@ -0,0 +1,191 @@ +# Grid 形状与多路径特化 优化模式 + +## 概述 + +**核心方法论：动态 Host Dispatch 是 Triton 性能调优的核心手段。** + +Triton kernel 一旦被 `@triton.jit` 编译，其内部的控制流结构（循环、分支）和 grid 拓扑即被固定。然而，同一算子在不同输入 shape 下往往呈现截然不同的 workload 特征： +- 小 tensor：`total_blocks <= num_cores`，每个 program 只需处理 1 个 block，任何分区循环都是纯开销 +- 大 tensor：`total_blocks >> num_cores`，必须通过标量循环将 block 均匀分配给各 program + +单一 kernel 实现无法在所有场景下同时最优。此时，**在 Host 侧（Python `forward()`）根据运行时 workload 特征动态选择不同 kernel 路径**，是突破性能瓶颈的关键手段。 + +## 为什么 Host Dispatch 有效 + +| 层面 | 限制 | Host Dispatch 的优势 | +|------|------|---------------------| +| **Kernel 内部** | 编译后控制流固定，无法根据 block 数量动态调整 | Kernel 侧只写最优路径，不加兼容逻辑 | +| **Grid 拓扑** | 1D/2D grid 形状在启动时确定，kernel 内无法变更 | Host 侧为不同场景选择最适合的 grid 形状 | +| **编译特化** | 同一 kernel 源码只能生成一种机器码 | 不同路径可分别编译，各自特化 | +| **调度开销** | 小 grid 时 program 启动/同步开销占比高 | Host 侧避免不必要的多 program 调度 | + +**关键洞察**：kernel 的通用性越强，单个场景的性能越差；Host 侧做动态 dispatch，让每个 kernel 只做一件事并把这件事做到极致。 + +## 典型多路径设计模式 + +### 模式 A：小 Grid 直接映射路径（Direct Mapping Path） + +**适用场景**：`total_blocks <= num_cores`（或略大于核数） + +**策略**： +- `grid` 直接映射到 workload 拓扑，如 `grid = (outer_size, num_inner_blocks)` +- Kernel 内**无任何标量分区循环**，每个 program 直接处理 1 个 block +- `tl.program_id(0)` / `tl.program_id(1)` 直接定位到具体 block + +**代码骨架**： +```python +@triton.jit +def kernel_direct(x_ptr, out_ptr, outer_size, inner_size, + r: tl.constexpr, BLOCK: tl.constexpr): + outer_idx = tl.program_id(0) + local_block = tl.program_id(1) + + block_start = local_block * BLOCK + offs = block_start + tl.arange(0, BLOCK) + mask = offs < inner_size + + in_offset = outer_idx * inner_size + val = tl.load(x_ptr + in_offset + offs, mask=mask) + + for repeat_idx in range(r): + out_offset = outer_idx * inner_size * r + repeat_idx * inner_size + tl.store(out_ptr + out_offset + offs, val, mask=mask) +``` + +### 模式 B：大 Grid 多核分区路径（Partition Loop Path） + +**适用场景**：`total_blocks > num_cores` + +**策略**： +- `grid = (num_cores,)` 或 `grid = (min(total_blocks, num_cores),)` +- Kernel 内通过标量循环将 block 均匀分配给各 program +- 保证负载均衡，充分利用所有物理核 + +**代码骨架**： +```python +@triton.jit +def kernel_partition(x_ptr, out_ptr, outer_size, inner_size, num_inner_blocks, + num_cores: tl.constexpr, r: tl.constexpr, BLOCK: tl.constexpr): + pid = tl.program_id(0) + total_blocks = outer_size * num_inner_blocks + + blocks_per_core = total_blocks // num_cores + remainder = total_blocks - blocks_per_core * num_cores + + if pid < remainder: + my_blocks = blocks_per_core + 1 + start_block = pid * (blocks_per_core + 1) + else: + my_blocks = blocks_per_core + start_block = remainder * (blocks_per_core + 1) + (pid - remainder) * blocks_per_core + + for block_idx in range(start_block, start_block + my_blocks): + outer_idx = block_idx // num_inner_blocks + local_block = block_idx - outer_idx * num_inner_blocks + + block_start = local_block * BLOCK + offs = block_start + tl.arange(0, BLOCK) + mask = offs < inner_size + + in_offset = outer_idx * inner_size + val = tl.load(x_ptr + in_offset + offs, mask=mask) + + for repeat_idx in range(r): + out_offset = outer_idx * inner_size * r + repeat_idx * inner_size + tl.store(out_ptr + out_offset + offs, val, mask=mask) +``` + +### 模式 C：按维度数特化路径（Dimension Specialization） + +**适用场景**：算子支持 1D/2D/3D/4D 等多种维度，不同维度的最优 grid 策略不同 + +**策略**： +- Host 侧根据 `len(shape)` 或具体维度大小选择不同 kernel +- 例如 1D 小 tensor 用单 program 处理，4D 大 tensor 用 2D grid + +## Host Dispatch 决策框架 + +```python +def forward(self, x, repeats): + # 1. 计算 workload 特征 + outer_size = ... + inner_size = ... + BLOCK = get_block_size(inner_size) + num_inner_blocks = (inner_size + BLOCK - 1) // BLOCK + total_blocks = outer_size * num_inner_blocks + + # 2. 根据特征选择路径 + if total_blocks <= self.VEC_CORE_NUM: + # 小 Grid 路径：直接映射，无分区循环 + grid = (outer_size, num_inner_blocks) + kernel_direct[grid](x, out, outer_size, inner_size, r=r, BLOCK=BLOCK) + else: + # 大 Grid 路径：多核分区，带标量循环 + grid = (min(total_blocks, self.VEC_CORE_NUM),) + kernel_partition[grid](x, out, outer_size, inner_size, num_inner_blocks, + num_cores=grid[0], r=r, BLOCK=BLOCK) +``` + +## 完整示例 — Repeat 算子的多路径特化 + +**背景**：`torch.repeat(*repeats)` 对不同 shape 的输入，其 `total_blocks` 分布跨度极大（从 1 到 8192+）。 + +**未优化前**：单一 kernel + `grid = (min(total_blocks, VEC_CORE_NUM),)` +- 小 grid 时：分区计算（`blocks_per_core = total_blocks // num_cores`）的标量分支和循环完全无意义 +- 大 grid 时：分区循环是必要的 + +**优化后**：双 kernel + Host dispatch +```python +class ModelNew(nn.Module): + def forward(self, x, repeats): + # ... 维度处理 ... + BLOCK = get_block_size(inner_size) + num_inner_blocks = (inner_size + BLOCK - 1) // BLOCK + total_blocks = outer_size * num_inner_blocks + + if total_blocks <= self.VEC_CORE_NUM: + # 小 Grid 路径：2D 精确映射，kernel 内无循环 + grid = (outer_size, num_inner_blocks) + repeat_small_kernel[grid](out, output, outer_size, inner_size, r=r, BLOCK=BLOCK) + else: + # 大 Grid 路径：多核分区负载均衡 + grid = (min(total_blocks, self.VEC_CORE_NUM),) + repeat_large_kernel[grid](out, output, outer_size, inner_size, num_inner_blocks, + num_cores=grid[0], r=r, BLOCK=BLOCK) +``` + +**收益**： +- 小 grid 场景消除标量分区开销，schedule 效率提升 +- 大 grid 场景保持原有负载均衡，无性能退化 +- 综合 geomean 从 0.933x 提升至 0.992x + +## 关键原则 + +1. **路径越少越好**：不要为了 dispatch 而 dispatch。通常 2 个路径即可覆盖绝大多数场景（小/大）。超过 3 个路径会增加维护成本和编译缓存压力。 + +2. **条件判断必须廉价**：Host 侧的 dispatch 条件（如 `total_blocks <= num_cores`）必须是纯 Python 标量比较，不能涉及 tensor 运算。 + +3. **语义严格一致**：所有路径的输出必须逐元素相等，不能因路径不同而引入精度差异或布局差异。 + +4. **grid 形状即策略**：选择 grid 形状时要考虑： + - 1D grid `(N,)`：简单，适合 block 编号可线性映射的场景 + - 2D grid `(M, N)`：适合 outer/inner 两级分解，可直接用 `tl.program_id(0)` / `tl.program_id(1)` 定位 + - 避免 3D grid，Triton Ascend 对 3D grid 支持有限 + +5. **与 constexpr 配合**：多路径特化常与「入参静态化」协同使用。将路径相关的参数（如 `r`）设为 `tl.constexpr`，让每个路径在编译期获得最大优化。 + +## 常见陷阱 + +| 陷阱 | 说明 | 避免方法 | +|------|------|---------| +| **路径边界性能跳变** | 在 dispatch 阈值（如 `total_blocks == num_cores`）附近，两种路径性能差异过大 | 阈值选择要留有余量，或让两种路径在边界处性能接近 | +| **小路径grid过大** | 小 grid 路径的 grid 超过核数，退化成调度开销 | 确保小路径的 grid `total_blocks <= num_cores` | +| **2D grid索引错误** | 使用 `tl.program_id(0)` 线性化 2D grid，导致索引越界或重复 | 2D grid 应直接用 `program_id(0)` / `program_id(1)`，不要从一维 pid 推导 | +| **路径间代码复制** | 两个路径的 kernel 逻辑大量重复，维护困难 | 提取公共逻辑到 Python helper，或接受少量重复以保证性能 | +| **忽略编译缓存** | 过多路径导致 Triton 编译缓存膨胀，首次启动变慢 | 控制路径数量（<=3），并确保路径条件覆盖合理 | + +## 与其他优化点的关系 + +- **与优化点 1（入参静态化）协同**：将 `r`、`num_cores` 等声明为 `tl.constexpr`，让每个路径编译出最优代码 +- **与优化点 3（分核优化）互补**：分核优化关注「grid 是否合理」；本优化点关注「单一 grid 策略是否够用，是否需要多策略 dispatch」 +- **与优化点 12（Autotune）的区别**：Autotune 是在同一 kernel 上尝试不同 `constexpr` 配置；本优化点是切换完全不同的 kernel 实现 diff --git a/skills/triton/latency-optimizer/references/mixed_strategy.md b/skills/triton/latency-optimizer/references/mixed_strategy.md new file mode 100644 index 00000000..0712cf7a --- /dev/null +++ b/skills/triton/latency-optimizer/references/mixed_strategy.md @@ -0,0 +1,196 @@ +# 混合策略自动选择 + +## 概述 + +同一算子在不同 shape、数据类型或内存布局下，可能需要不同的优化策略才能获得最优性能。混合策略自动选择通过在 host 端（`forward()` 中）根据运行时条件选择不同的 kernel 变体或参数配置，实现对多样化输入的自适应优化。 + +## 适用场景 + +当算子满足以下任一条件时，应考虑混合策略： + +1. **不同 shape 范围性能瓶颈不同**：小 batch 受限于并行度，大 batch 受限于内存带宽 +2. **不同数据类型精度要求不同**：fp32 对数值稳定性敏感，fp16/bf16 可容忍更多近似优化 +3. **不同维度配置访存模式不同**：如 `inner_size == 1` vs `inner_size > 1` 的向量化策略差异 + +## 典型模式 + +### 模式 1：Shape 自适应策略选择 + +**问题**：单一 kernel 无法在所有 shape 下达到最优性能。 + +```python +# 问题：单一策略无法覆盖所有 shape +# 小 n_cols 时：需要 noloop 减少循环开销 +# 大 n_cols 时：需要 loop 避免 UB 溢出 +``` + +**优化方案**：根据 shape 动态选择 kernel 变体。 + +```python +class ModelNew: + def forward(self, x, normalized_shape, weight=None, bias=None): + n_rows, n_cols = ... + + # 混合策略：根据 n_cols 选择不同 kernel + if n_cols <= 128 and n_rows >= 256: + # 小 n_cols + 足够行数：使用 noloop kernel + kernel = layer_norm_kernel_fp16_noloop + ROWS_PER_BLOCK = 32 + else: + # 大 n_cols：使用 loop kernel + kernel = layer_norm_kernel_fp16_loop + ROWS_PER_BLOCK = 8 + + grid = (min(triton.cdiv(n_rows, ROWS_PER_BLOCK), self.VEC_CORE_NUM),) + kernel[grid](..., ROWS_PER_BLOCK=ROWS_PER_BLOCK) +``` + +**关键约束**： +- noloop 需满足 `ROWS_PER_BLOCK >= 4`（避免 grid 过大导致 hang） +- noloop 需满足 `n_rows >= 256` 或 `grid >= 16`（保证足够并行度） +- noloop 需满足 UB 安全：`ROWS_PER_BLOCK * n_cols * dtype_size <= 98304` + +### 模式 2：数据类型自适应策略选择 + +**问题**：不同数据类型对并行优化的容忍度不同。 + +```python +# 问题：fp32 对求和顺序敏感，fp16 可容忍 +``` + +**优化方案**：根据数据类型启用/禁用特定优化。 + +```python +class ModelNew: + def forward(self, x, ...): + # fp32：禁用改变求和顺序的优化（精度敏感） + # fp16/bf16：可启用并行优化 + use_parallel_stats = (x.dtype != torch.float32) and (total_groups < num_cores // 2) + + if use_parallel_stats: + # 并行 stats：多个 group 同时统计，用 atomic_add 合并 + stats_kernel = parallel_stats_kernel + else: + # 原始 stats：单 group 串行统计 + stats_kernel = serial_stats_kernel + + stats_kernel[grid](...) +``` + +### 模式 3：Grid 并行度自适应策略选择 + +**问题**：小 batch / small groups 时 grid 不足，无法充分利用核数。 + +```python +# 问题：total_groups = 4，num_cores = 48，grid 过小 +``` + +**优化方案**：根据 grid 大小选择串行或并行策略。 + +```python +# 判断条件 +total_groups = N * num_groups +use_parallel = (total_groups < num_cores // 2) and (x.dtype != torch.float32) + +if use_parallel: + # 并行策略：每个 program 处理一个 group，用 atomic_add 合并 + grid = (min(total_groups, num_cores),) + parallel_stats_kernel[grid](...) +else: + # 串行策略：每个 program 处理多个 group + groups_per_core = triton.cdiv(total_groups, num_cores) + grid = (num_cores,) + serial_stats_kernel[grid](...) +``` + +## 策略选择决策树 + +``` +开始 +│ +├─ 检查数据类型 +│ ├─ fp32 → 禁用改变求和顺序的优化 +│ └─ fp16/bf16 → 可启用并行优化 +│ +├─ 检查 shape 大小 +│ ├─ 小 n_cols (<=128) + 大 n_rows (>=256) + UB 安全 +│ │ → 考虑 noloop kernel +│ └─ 其他 → 使用 loop kernel +│ +├─ 检查 grid 并行度 +│ ├─ total_work_items < num_cores // 2 +│ │ → 考虑并行策略（atomic_add） +│ └─ total_work_items >= num_cores // 2 +│ → 使用串行策略 +│ +└─ 检查内存布局 + ├─ inner_size == 1（操作维度为最后一维） + │ → 直接向量化操作维度 + └─ inner_size > 1 + → 向量化 inner 维度 +``` + +## 常见错误 + +### 错误 1：策略切换条件过于粗糙 + +```python +# ❌ 错误：只根据单一条件判断 +if n_cols < 256: + use_noloop = True # 未考虑 n_rows 和 UB + +# ✅ 正确：综合考虑多个条件 +use_noloop = ( + n_cols <= 128 and # noloop 仅适合小 n_cols + n_rows >= 256 and # 保证足够并行度 + ROWS_PER_BLOCK * n_cols * dtype_size <= 98304 # UB 安全 +) +``` + +### 错误 2：策略切换引入额外开销 + +```python +# ❌ 错误：在 kernel 内做条件分支 +@triton.jit +def kernel(...): + if some_condition: # kernel 内分支性能差 + ... + +# ✅ 正确：在 host 端选择 kernel +if some_condition: + kernel_a[grid](...) +else: + kernel_b[grid](...) +``` + +### 错误 3：忽略精度影响 + +```python +# ❌ 错误：fp32 也启用并行规约 +use_parallel = total_groups < num_cores // 2 # 未检查数据类型 + +# ✅ 正确：fp32 禁用改变求和顺序的优化 +use_parallel = (total_groups < num_cores // 2) and (x.dtype != torch.float32) +``` + +## 性能收益 + +| 场景 | 单一策略 | 混合策略 | 收益 | +|------|---------|---------|------| +| LayerNorm (小 n_cols) | loop kernel | noloop kernel | 减少循环开销 | +| GroupNorm (小 groups) | serial stats | parallel stats + atomic | 提升并行度 | +| fp32 数据 | 并行规约 | 串行规约 | 保证精度 | + +## 总结 + +| 维度 | 小/敏感 | 大/容忍 | +|------|--------|--------| +| n_cols | noloop (<=128) | loop | +| n_rows / grid | parallel (atomic) | serial | +| 数据类型 | fp32 → 保守 | fp16/bf16 → 激进 | +| inner_size | 向量化操作维度 | 向量化 inner 维度 | + +**核心原则**： +- 策略选择在 host 端（`forward()`）完成，不在 kernel 内分支 +- 切换条件需综合考虑 shape、数据类型、UB 安全、并行度 +- fp32 优先保证精度，fp16/bf16 可尝试更多优化 diff --git a/skills/triton/latency-optimizer/references/scalar_to_vector.md b/skills/triton/latency-optimizer/references/scalar_to_vector.md index 9341e64a..36481476 100644 --- a/skills/triton/latency-optimizer/references/scalar_to_vector.md +++ b/skills/triton/latency-optimizer/references/scalar_to_vector.md @@ -250,7 +250,34 @@ c = a.to(tl.float32) // b.to(tl.float32) # 转换float32类型 d = a - (a // b) * b # 公式转换 ``` -### 8. atomic_* 标量操作 → atomic_* 向量操作 +### 8. 内存布局感知索引优化 + +**适用场景**: kernel 中使用 `idx // stride` 和 `idx % stride` 将线性索引映射到多维坐标，且目标维度在内存中连续。 + +**典型代码特征**: +```python +# 问题代码：通过 div/mod 重建多维坐标 +c_local = idx // (H * W) # int 除法，标量降级 +hw = idx % (H * W) # int 取余，标量降级 +h = hw // W # int 除法，标量降级 +w = hw % W # int 取余，标量降级 +x_offset = ((n * C + c) * H + h) * W + w +``` + +**优化方法**: 若内存布局允许，直接计算连续偏移： +```python +# 优化后：利用 NCHW 连续性，直接计算偏移 +base_c = g * channels_per_group +base_offset = ((n * C + base_c) * H) * W +val = tl.load(x_ptr + base_offset + idx, mask=mask, other=0.0) +``` + +**适用条件**: +- 张量布局为 NCHW 或 NHWC +- 同一 group/channel/instance 内的元素在内存中连续 +- 无需跨 stride 访问 + +### 9. atomic_* 标量操作 → atomic_* 向量操作 **原始代码（scalar 操作）** @@ -279,7 +306,8 @@ tl.atomic_add(output_ptr + h_offs, block_vals) # 向量化的原子加 7. **标量比较类型转换**：对于 `int32` 和 `int64` 整数标量比较，先 .to(tl.float32) 再比较，以启用向量比较指令。 8. **标量除法类型转换**：对于 `int32` 和 `int64` 中的整数除法操作，先 .to(tl.float32) 再计算，以启用向量计算指令。 9. **标量取余类型转换**：对于 `int32` 和 `int64` 中的整数取余操作， 使用`a - (a // b) * b`的形式计算，以启用向量计算指令。 -10. **原子操作向量化**：对 `atomic_add` 这一类的 `atomic_*` 标量操作进行向量化，可消除循环的标量操作开销 +10. **内存布局感知索引优化**：对于 NCHW/NHWC 布局下通过 div/mod 重建多维坐标的场景，若目标维度内存连续，直接计算连续偏移以消除标量降级。 +11. **原子操作向量化**：对 `atomic_add` 这一类的 `atomic_*` 标量操作进行向量化，可消除循环的标量操作开销 ## 性能收益 diff --git a/skills/triton/latency-optimizer/references/tiling_optimization.md b/skills/triton/latency-optimizer/references/tiling_optimization.md index e2827aab..8571f4ce 100644 --- a/skills/triton/latency-optimizer/references/tiling_optimization.md +++ b/skills/triton/latency-optimizer/references/tiling_optimization.md @@ -12,7 +12,7 @@ ## 适用条件 -处理多维张量（3D 及以上）的规约类（Reduction）或归一化类（Normalization）算子，且还原轴（Reduction Axis）并非内存布局中的最连续轴（通常为最后一维 N）。 +处理多维张量（3D 及以上）的规约类（Reduction）、归一化类（Normalization）或数据搬运类算子，且还原轴（Reduction Axis）或搬运方向并非内存布局中的最连续轴（通常为最后一维 N）。 ## 优化方法 @@ -53,6 +53,8 @@ for m_idx in range(0, dim1): 4. **粗粒度调度**：调整 Grid 配置，使每个 Program 处理更连续、更大块的数据（如整个 Batch），提升数据局部性 +5. **数据搬运类算子特化**：对于 split/chunk/slice 等数据搬运操作，检查是否按内存最连续轴进行向量化。当操作维度后无其他维度（`inner_size == 1`）时，可直接向量化读取该维度；否则向量化 inner 维度。 + ## 关键点 1. **合并访存**：向量化轴必须在内存最连续的维度上 diff --git a/skills/triton/latency-optimizer/references/vector_core_partition.md b/skills/triton/latency-optimizer/references/vector_core_partition.md index bbbe641f..59e16d4b 100644 --- a/skills/triton/latency-optimizer/references/vector_core_partition.md +++ b/skills/triton/latency-optimizer/references/vector_core_partition.md @@ -501,6 +501,22 @@ grid = (grid_size,) | `unit_flag=True` | 生成独立的计算单元 | 简单算子，无复杂控制流 | | `unit_flag=False` | 不生成独立单元 | 复杂算子，有分支 | +### multibuffer 的 UB 安全检查 + +在建议 `multibuffer=True` 前，必须先计算 UB 预算： + +``` +UB_needed = ROWS_PER_BLOCK * BLOCK_COL * dtype_size * num_buffers +# num_buffers 包含: input, output, intermediate, weight, bias 等 +# multibuffer=True 时 num_buffers 翻倍 +``` + +若 `UB_needed > 192KB`，应自动降级为 `multibuffer=False` 或减小 tile。 + +经验公式： +- `ROWS_PER_BLOCK * BLOCK_COL > 12288` (float32) → 不启用 multibuffer +- `ROWS_PER_BLOCK * BLOCK_COL > 24576` (float16) → 不启用 multibuffer + ### 使用方式 ```python diff --git a/utils/exp-archive.py b/utils/exp-archive.py new file mode 100644 index 00000000..2f4c5368 --- /dev/null +++ b/utils/exp-archive.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python3 +""" +经验归档工具：将已探索成功的算子工作目录归档到 memory/archive/，并更新 MEMORY.md 索引。 + +用法: + python3 utils/exp-archive.py [--category ] [--force] + +示例: + python3 utils/exp-archive.py ./triton_ascend_output/op_1_15_Pad_20260522_0509_7731 +""" + +import argparse +import json +import os +import re +import shutil +import sys +from datetime import datetime +from pathlib import Path + + +def get_project_root() -> Path: + """以脚本所在目录的父目录作为项目根目录。""" + return Path(__file__).resolve().parent.parent + + +def parse_op_name_from_dir(work_dir: Path) -> str: + """从工作目录名解析算子名称，如 op_1_15_Pad_20260522_0509_7731 -> 15_Pad。""" + name = work_dir.name + m = re.match(r'op_\d+_(.+?)_\d{8}_\d{4}_\d+', name) + if m: + return m.group(1) + # fallback: 查找目录下 *_generated.py + for f in work_dir.glob('*_generated.py'): + return f.stem.replace('_generated', '') + return name + + +def parse_date_from_dir(work_dir: Path) -> str: + """从工作目录名提取日期，如 op_1_15_Pad_20260522_0509_7731 -> 20260522。""" + name = work_dir.name + m = re.search(r'(\d{8})_\d{4}_\d+', name) + if m: + return m.group(1) + return datetime.now().strftime('%Y%m%d') + + +def infer_category(op_name: str) -> str: + """从算子名称推断类别，如 15_Pad -> pad，0_Softmax -> softmax。""" + category = re.sub(r'^\d+_', '', op_name).lower() + return category + + +def get_next_version(archive_dir: Path, category: str) -> int: + """扫描 archive 目录下已有版本，返回下一个版本号。""" + max_ver = 0 + if not archive_dir.exists(): + return 1 + pattern = re.compile(re.escape(category) + r'_v(\d+)_(\d{8})\.py') + for f in archive_dir.iterdir(): + if f.is_file(): + m = pattern.match(f.name) + if m: + max_ver = max(max_ver, int(m.group(1))) + return max_ver + 1 + + +def find_best_version(archive_dir: Path, category: str) -> tuple: + """ + 扫描 archive 目录下已有版本，返回最优版本的 (speedup, item)。 + item: {'py': Path, 'report': Path, 'summary': Path, 'ver': int} + 无已有版本时返回 (0.0, None)。 + """ + if not archive_dir.exists(): + return 0.0, None + + best_speedup = 0.0 + best_item = None + pattern = re.compile(re.escape(category) + r'_v(\d+)_(\d{8})\.py$') + + for f in archive_dir.iterdir(): + if not f.is_file(): + continue + m = pattern.match(f.name) + if not m: + continue + ver = int(m.group(1)) + summary_file = f.with_suffix('').with_name(f.stem + '_summary.json') + if not summary_file.exists(): + continue + try: + with open(summary_file, 'r', encoding='utf-8') as sf: + data = json.load(sf) + speedup = data.get('perf_data', {}).get('speedup_vs_torch', 0.0) + if isinstance(speedup, (int, float)) and speedup > best_speedup: + best_speedup = speedup + report_file = f.with_suffix('').with_name(f.stem + '_report.md') + best_item = { + 'py': f, + 'report': report_file, + 'summary': summary_file, + 'ver': ver, + } + except Exception: + continue + + return best_speedup, best_item + + +def remove_all_versions(archive_dir: Path, category: str) -> None: + """删除该类别下所有历史版本文件（.py + _report.md + _summary.json）。""" + if not archive_dir.exists(): + return + py_pattern = re.compile(re.escape(category) + r'_v\d+_\d{8}\.py$') + for f in list(archive_dir.iterdir()): + if f.is_file() and py_pattern.match(f.name): + stem = f.stem + for suffix in ('.py', '_report.md', '_summary.json'): + to_remove = f.parent / (stem + suffix) + if to_remove.exists(): + to_remove.unlink() + print(f'[REMOVE] {to_remove}') + + +def validate_work_dir(work_dir: Path) -> dict: + """验证工作目录是否符合归档条件，返回 summary dict。""" + summary_path = work_dir / 'summary.json' + if not summary_path.exists(): + raise FileNotFoundError(f'工作目录缺少 summary.json: {work_dir}') + + with open(summary_path, 'r', encoding='utf-8') as f: + summary = json.load(f) + + if not summary.get('success'): + raise ValueError('summary.json 中 success 为 false，不可归档') + + perf = summary.get('perf_data', {}) + total = perf.get('total_cases', 0) + passed = perf.get('passed_cases', 0) + if passed != total or total == 0: + raise ValueError(f'精度未全通过: {passed}/{total}') + + speedup = perf.get('speedup_vs_torch') + if speedup is None or (isinstance(speedup, (int, float)) and speedup <= 0.8): + raise ValueError(f'加速比不满足归档条件: {speedup} (需 > 0.8x)') + + # 检查必要文件是否存在 + op_name = parse_op_name_from_dir(work_dir) + required = [ + work_dir / f'{op_name}_generated.py', + work_dir / 'report.md', + work_dir / 'summary.json', + ] + for rp in required: + if not rp.exists(): + raise FileNotFoundError(f'工作目录缺少必要文件: {rp.name}') + + return summary + + +def update_memory_md(memory_path: Path, category: str, ver: int, date: str, + speedup: float, passed: int, total: int) -> None: + """更新 MEMORY.md 索引，如果不存在该类别的条目则追加。""" + category_cap = category.capitalize() + py_file = f'archive/{category}/{category}_v{ver}_{date}.py' + report_file = f'archive/{category}/{category}_v{ver}_{date}_report.md' + summary_file = f'archive/{category}/{category}_v{ver}_{date}_summary.json' + + new_line = ( + f'- [{category_cap} 算子最佳实现]({py_file}) — ' + f'几何平均加速比 {speedup:.2f}x，{passed}/{total} cases 通过；' + f'配套 [report]({report_file}) / [summary]({summary_file})' + ) + + if not memory_path.exists(): + memory_path.write_text('# Memory Index\n\n', encoding='utf-8') + + content = memory_path.read_text(encoding='utf-8') + lines = content.splitlines() + + # 查找是否已有该类别的归档条目 + section_idx = None + existing_idx = None + for i, line in enumerate(lines): + if '完整代码归档' in line or 'Layer 4' in line: + section_idx = i + if f'/{category}/' in line and f'{category}_v' in line: + existing_idx = i + + if existing_idx is not None: + # 替换旧条目 + lines[existing_idx] = new_line + else: + # 在 Layer 4 节末尾追加，或在文件末尾追加 + if section_idx is not None: + # 找到该节最后一个非空行，在其后插入 + insert_pos = section_idx + 1 + for j in range(section_idx + 1, len(lines)): + if lines[j].strip().startswith('- '): + insert_pos = j + 1 + elif lines[j].strip().startswith('## ') and j > section_idx: + break + lines.insert(insert_pos, new_line) + else: + lines.append('') + lines.append('## 完整代码归档（Layer 4，Agent 默认不可读）') + lines.append(new_line) + + memory_path.write_text('\n'.join(lines) + '\n', encoding='utf-8') + + +def main(): + parser = argparse.ArgumentParser(description='归档已探索成功的算子经验') + parser.add_argument('work_dir', help='算子工作目录绝对或相对路径') + parser.add_argument('--category', default=None, + help='手动指定算子类别（默认从目录名自动推断）') + parser.add_argument('--force', action='store_true', + help='跳过归档条件校验（仅复制文件）') + parser.add_argument('--create-experience', action='store_true', + help='若该算子类别尚无经验文件，自动基于模板创建') + args = parser.parse_args() + + work_dir = Path(args.work_dir).resolve() + if not work_dir.is_dir(): + print(f'[ERROR] 工作目录不存在: {work_dir}', file=sys.stderr) + sys.exit(1) + + root = get_project_root() + archive_root = root / '.claude' / 'memory' / 'archive' + memory_path = root / '.claude' / 'memory' / 'MEMORY.md' + + # 1. 验证 + try: + summary = validate_work_dir(work_dir) + except Exception as e: + if args.force: + summary_path = work_dir / 'summary.json' + with open(summary_path, 'r', encoding='utf-8') as f: + summary = json.load(f) + print(f'[WARN] 强制跳过校验: {e}') + else: + print(f'[ERROR] 归档校验失败: {e}', file=sys.stderr) + sys.exit(1) + + op_name = parse_op_name_from_dir(work_dir) + category = args.category or infer_category(op_name) + date_str = parse_date_from_dir(work_dir) + version = get_next_version(archive_root / category, category) + + perf = summary.get('perf_data', {}) + passed_cases = perf.get('passed_cases', 0) + total_cases = perf.get('total_cases', 0) + speedup = perf.get('speedup_vs_torch', 0.0) + + # 2. 版本比较与旧版本清理（仅存最优方案） + dest_dir = archive_root / category + dest_dir.mkdir(parents=True, exist_ok=True) + + best_speedup, best_item = find_best_version(dest_dir, category) + if best_item: + if speedup <= best_speedup and not args.force: + print(f'[ERROR] 当前加速比 {speedup:.4f}x 不优于已有最优版本 ' + f'v{best_item["ver"]} ({best_speedup:.4f}x)。' + f'如需强制覆盖，请加 --force', file=sys.stderr) + sys.exit(1) + # 新版本更优，清理所有旧版本物理文件 + remove_all_versions(dest_dir, category) + + # 3. 创建归档并复制文件 + base_name = f'{category}_v{version}_{date_str}' + src_files = { + work_dir / f'{op_name}_generated.py': dest_dir / f'{base_name}.py', + work_dir / 'report.md': dest_dir / f'{base_name}_report.md', + work_dir / 'summary.json': dest_dir / f'{base_name}_summary.json', + } + + for src, dst in src_files.items(): + shutil.copy2(src, dst) + print(f'[COPY] {src} -> {dst}') + + # 3. 更新 MEMORY.md + update_memory_md(memory_path, category, version, date_str, + speedup, passed_cases, total_cases) + print(f'[UPDATE] {memory_path}') + + print(f'\n[SUCCESS] {op_name} 已归档为 {base_name}.* (speedup={speedup:.4f}x)') + + # 4. 可选：自动创建经验文件模板 + if args.create_experience: + exp_path = root / '.claude' / 'memory' / f'kernel-opt-{category}.md' + if not exp_path.exists(): + import subprocess + result = subprocess.run( + [sys.executable, str(root / 'utils' / 'exp-init.py'), + category, '--op-name', op_name], + capture_output=True, text=True + ) + print(result.stdout) + if result.returncode != 0: + print(result.stderr, file=sys.stderr) + else: + print(f'[INFO] 经验文件已存在: {exp_path}，请手动更新 Layer 1-3 内容。') + + +if __name__ == '__main__': + main() diff --git a/utils/exp-check.py b/utils/exp-check.py new file mode 100644 index 00000000..1262b050 --- /dev/null +++ b/utils/exp-check.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +""" +归档与记忆规范检查工具：检查 archive 目录结构、MEMORY.md 索引、 +kernel-opt 经验文件是否符合四层隔离模型的规范要求。 + +用法: + python3 utils/exp-check.py [--root ] + +退出码: + 0 全部通过 + 1 存在失败项 +""" + +import argparse +import json +import os +import re +import sys +from pathlib import Path + + +class CheckReporter: + def __init__(self): + self.passed = 0 + self.failed = 0 + self.warnings = 0 + self.details = [] + + def ok(self, msg: str): + self.passed += 1 + self.details.append(f' [PASS] {msg}') + + def fail(self, msg: str): + self.failed += 1 + self.details.append(f' [FAIL] {msg}') + + def warn(self, msg: str): + self.warnings += 1 + self.details.append(f' [WARN] {msg}') + + def section(self, title: str): + self.details.append(f'\n{title}') + + def print_report(self): + for line in self.details: + print(line) + print('\n' + '='*60) + print(f'总计: {self.passed} 通过, {self.failed} 失败, {self.warnings} 警告') + if self.failed > 0: + print('结论: 不符合规范，请修复失败项') + elif self.warnings > 0: + print('结论: 基本符合规范，但存在警告') + else: + print('结论: 完全符合规范') + + +def check_archive_structure(root: Path, reporter: CheckReporter) -> dict: + """ + 检查 archive/ 目录结构，返回按类别组织的归档文件映射。 + 返回: {category: [{'py': Path, 'report': Path, 'summary': Path, 'ver': int, 'date': str}]} + """ + archive_root = root / '.claude' / 'memory' / 'archive' + reporter.section('## 检查 archive/ 目录结构') + + if not archive_root.exists(): + reporter.fail(f'archive 根目录不存在: {archive_root}') + return {} + + readme = archive_root / 'README.md' + if readme.exists(): + reporter.ok('archive/README.md 存在') + else: + reporter.fail('archive/README.md 缺失') + + # 遍历子目录 + categories = {} + for entry in sorted(archive_root.iterdir()): + if not entry.is_dir(): + continue + cat = entry.name + if cat.startswith('.') or cat == 'archive': + continue + + reporter.section(f'### 类别: {cat}') + if not re.match(r'^[a-z][a-z0-9_]*$', cat): + reporter.warn(f'目录名 "{cat}" 建议全小写且无特殊字符') + + py_files = list(entry.glob('*_v*_*.py')) + if not py_files: + reporter.fail(f'{cat}/ 下未找到符合命名规范的 .py 归档文件') + continue + + categories[cat] = [] + for py_file in sorted(py_files): + m = re.match(re.escape(cat) + r'_v(\d+)_(\d{8})\.py$', py_file.name) + if not m: + reporter.warn(f'文件名不符合规范: {py_file.name}') + continue + ver = int(m.group(1)) + date = m.group(2) + + report_file = py_file.with_suffix('').with_name(py_file.stem + '_report.md') + summary_file = py_file.with_suffix('').with_name(py_file.stem + '_summary.json') + + has_report = report_file.exists() + has_summary = summary_file.exists() + + if has_report and has_summary: + reporter.ok(f'{py_file.name} 配套文件完整 (v{ver}, {date})') + else: + if not has_report: + reporter.fail(f'{py_file.name} 缺少配套报告: {report_file.name}') + if not has_summary: + reporter.fail(f'{py_file.name} 缺少配套摘要: {summary_file.name}') + + categories[cat].append({ + 'py': py_file, + 'report': report_file if has_report else None, + 'summary': summary_file if has_summary else None, + 'ver': ver, + 'date': date, + }) + + return categories + + +def check_summary_json(categories: dict, reporter: CheckReporter): + """检查每个 summary.json 的字段与数值规范。""" + reporter.section('\n## 检查 summary.json 字段规范') + for cat, items in categories.items(): + for item in items: + sf = item['summary'] + if sf is None: + continue + try: + with open(sf, 'r', encoding='utf-8') as f: + data = json.load(f) + except json.JSONDecodeError as e: + reporter.fail(f'{cat}/{sf.name} JSON 解析失败: {e}') + continue + + # 关键字段 + for key in ('success', 'perf_data'): + if key not in data: + reporter.fail(f'{cat}/{sf.name} 缺少顶层字段: {key}') + break + else: + reporter.ok(f'{cat}/{sf.name} 包含必要顶层字段') + + perf = data.get('perf_data', {}) + total = perf.get('total_cases', 0) + passed = perf.get('passed_cases', 0) + speedup = perf.get('speedup_vs_torch') + + if passed == total and total > 0: + reporter.ok(f'{cat}/{sf.name} 精度全通过 ({passed}/{total})') + else: + reporter.fail(f'{cat}/{sf.name} 精度未全通过: {passed}/{total}') + + if speedup is not None and isinstance(speedup, (int, float)) and speedup > 0.8: + reporter.ok(f'{cat}/{sf.name} 加速比达标: {speedup:.4f}x') + else: + reporter.fail(f'{cat}/{sf.name} 加速比不达标: {speedup} (需 > 0.8x)') + + +def check_memory_md(root: Path, categories: dict, reporter: CheckReporter): + """检查 MEMORY.md 索引是否与 archive 目录一致。""" + reporter.section('\n## 检查 MEMORY.md 索引一致性') + memory_path = root / '.claude' / 'memory' / 'MEMORY.md' + if not memory_path.exists(): + reporter.fail(f'MEMORY.md 不存在: {memory_path}') + return + + content = memory_path.read_text(encoding='utf-8') + lines = content.splitlines() + + # 收集 archive 下所有预期条目 + indexed_cats = set() + for cat in categories: + # 查找是否包含指向该 category 的链接 + found = False + for line in lines: + if f'/{cat}/' in line: + found = True + # 检查行长度 + if len(line) > 150: + reporter.warn(f'MEMORY.md 行超长 ({len(line)} > 150): {line[:80]}...') + # 检查链接文件是否存在 + for link_match in re.finditer(r'\]\(([^)]+)\)', line): + link = link_match.group(1) + linked_path = (root / '.claude' / 'memory' / link).resolve() + if not linked_path.exists(): + reporter.fail(f'MEMORY.md 链接指向的文件不存在: {link}') + break + if found: + reporter.ok(f'MEMORY.md 包含 {cat} 的索引条目') + indexed_cats.add(cat) + else: + reporter.fail(f'MEMORY.md 缺少 {cat} 的索引条目') + + # 反向检查：MEMORY.md 中是否有指向不存在的 archive 条目的链接 + for line in lines: + if '/archive/' in line and line.strip().startswith('- '): + m = re.search(r'archive/([a-z][a-z0-9_]*)', line) + if m: + cat = m.group(1) + if cat not in categories: + reporter.warn(f'MEMORY.md 引用了 archive 中不存在的类别: {cat}') + + +def check_kernel_opt_refs(root: Path, categories: dict, reporter: CheckReporter): + """检查 kernel-opt-*.md 中 Layer 4 引用是否有效。""" + reporter.section('\n## 检查 kernel-opt-*.md 引用有效性') + memory_dir = root / '.claude' / 'memory' + for md_file in sorted(memory_dir.glob('kernel-opt-*.md')): + # 从文件名推断类别，如 kernel-opt-pad.md -> pad + cat = md_file.stem.replace('kernel-opt-', '') + if cat not in categories: + continue + + content = md_file.read_text(encoding='utf-8') + # 查找 Layer 4 或 archive 路径引用 + refs = re.findall(r'archive/[a-z][a-z0-9_/\-\.]+', content) + if not refs: + reporter.warn(f'{md_file.name} 的 Layer 4 部分未引用任何 archive 路径') + continue + + valid = True + for ref in refs: + full = memory_dir / ref + if not full.exists(): + reporter.fail(f'{md_file.name} 引用的路径不存在: {ref}') + valid = False + if valid: + reporter.ok(f'{md_file.name} Layer 4 引用全部有效 ({len(refs)} 处)') + + +def main(): + parser = argparse.ArgumentParser(description='检查归档与记忆规范') + parser.add_argument('--root', default='.', + help='项目根目录（默认当前目录）') + args = parser.parse_args() + + root = Path(args.root).resolve() + reporter = CheckReporter() + + categories = check_archive_structure(root, reporter) + if categories: + check_summary_json(categories, reporter) + check_memory_md(root, categories, reporter) + check_kernel_opt_refs(root, categories, reporter) + else: + reporter.fail('未检测到任何有效的归档类别，跳过后续检查') + + reporter.print_report() + sys.exit(1 if reporter.failed > 0 else 0) + + +if __name__ == '__main__': + main() diff --git a/utils/exp-edit.py b/utils/exp-edit.py new file mode 100644 index 00000000..ae34306d --- /dev/null +++ b/utils/exp-edit.py @@ -0,0 +1,509 @@ +#!/usr/bin/env python3 +""" +历史经验增删改查工具（exp-edit.py）。 + +对 .claude/memory/kernel-opt-{category}.md 中的 Layer 1-3 经验条目进行管理。 + +用法: + python3 utils/exp-edit.py list + python3 utils/exp-edit.py show [--layer N] [--id ID] + python3 utils/exp-edit.py add --layer N --title TITLE [--content TEXT] [--why TEXT] [--apply TEXT] + python3 utils/exp-edit.py update --id ID [--title TEXT] [--content TEXT] [--why TEXT] [--apply TEXT] + python3 utils/exp-edit.py remove --id ID + python3 utils/exp-edit.py undo + python3 utils/exp-edit.py history + +高危操作（update/remove）默认交互确认；Agent 可附加 --yes 跳过。 +""" + +import argparse +import difflib +import json +import os +import re +import shutil +import sys +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import List, Optional, Tuple + + +def get_project_root() -> Path: + return Path(__file__).resolve().parent.parent + + +@dataclass +class Entry: + id: str + layer: int + title: str + raw_title_line: str + body_lines: List[str] = field(default_factory=list) + start_line: int = 0 + end_line: int = 0 + + +class ExperienceStore: + def __init__(self, category: str): + self.category = category.lower() + self.root = get_project_root() + self.memory_dir = self.root / '.claude' / 'memory' + self.file_path = self.memory_dir / f'kernel-opt-{self.category}.md' + self.backup_dir = self.memory_dir / '.backup' + self.history_path = self.backup_dir / f'{self.category}_history.jsonl' + self.lines: List[str] = [] + self.frontmatter_end = -1 + self.entries: List[Entry] = [] + self._load() + + def exists(self) -> bool: + return self.file_path.exists() + + def _load(self): + if not self.exists(): + return + text = self.file_path.read_text(encoding='utf-8') + self.lines = text.splitlines() + self._parse() + + def _parse(self): + self.entries = [] + self.frontmatter_end = -1 + if len(self.lines) >= 2 and self.lines[0].strip() == '---': + for i in range(1, len(self.lines)): + if self.lines[i].strip() == '---': + self.frontmatter_end = i + break + + current_layer = 0 + i = self.frontmatter_end + 1 if self.frontmatter_end >= 0 else 0 + while i < len(self.lines): + line = self.lines[i] + # Detect Layer section + m = re.match(r'^## Layer\s+(\d+)\s*:', line) + if m: + current_layer = int(m.group(1)) + i += 1 + continue + + # Detect entry + if current_layer in (1, 2, 3) and line.startswith('### '): + entry = self._parse_entry(i, current_layer) + if entry: + self.entries.append(entry) + i = entry.end_line + continue + i += 1 + + def _parse_entry(self, start_idx: int, layer: int) -> Optional[Entry]: + title_line = self.lines[start_idx] + # Extract ID from + id_match = re.search(r'', title_line) + entry_id = id_match.group(1) if id_match else None + + # Extract title text + title_clean = re.sub(r'', '', title_line).strip() + title_clean = re.sub(r'^###\s+', '', title_clean) + title_clean = re.sub(r'^L\d+\.\d+\s+', '', title_clean) + + end_idx = start_idx + 1 + while end_idx < len(self.lines): + nxt = self.lines[end_idx] + if nxt.startswith('### ') or nxt.startswith('## '): + break + end_idx += 1 + + if entry_id is None: + # Auto-assign ID based on layer max sequence + seq = self._next_seq(layer) + entry_id = f'l{layer}-{seq:03d}' + # Inject ID into title line + new_title = title_line.rstrip() + f' ' + self.lines[start_idx] = new_title + title_line = new_title + + body = self.lines[start_idx + 1:end_idx] + return Entry( + id=entry_id, + layer=layer, + title=title_clean, + raw_title_line=title_line, + body_lines=body, + start_line=start_idx, + end_line=end_idx, + ) + + def _next_seq(self, layer: int) -> int: + seqs = [] + for e in self.entries: + if e.layer == layer: + m = re.match(rf'^l{layer}-(\d+)$', e.id) + if m: + seqs.append(int(m.group(1))) + return max(seqs, default=0) + 1 + + def _find_entry(self, entry_id: str) -> Optional[Entry]: + for e in self.entries: + if e.id == entry_id: + return e + return None + + def _ensure_backup_dir(self): + self.backup_dir.mkdir(parents=True, exist_ok=True) + + def _backup(self, reason: str): + self._ensure_backup_dir() + ts = time.strftime('%Y%m%d_%H%M%S') + bak = self.backup_dir / f'kernel-opt-{self.category}_{ts}.md' + shutil.copy2(self.file_path, bak) + self._log_history({'action': 'backup', 'reason': reason, 'file': str(bak), 'time': ts}) + return bak + + def _log_history(self, record: dict): + self._ensure_backup_dir() + with open(self.history_path, 'a', encoding='utf-8') as f: + f.write(json.dumps(record, ensure_ascii=False) + '\n') + + def get_last_backup(self) -> Optional[Path]: + if not self.backup_dir.exists(): + return None + files = sorted( + self.backup_dir.glob(f'kernel-opt-{self.category}_*.md'), + key=lambda p: p.stat().st_mtime + ) + return files[-1] if files else None + + def _save(self): + self.file_path.write_text('\n'.join(self.lines) + '\n', encoding='utf-8') + self._load() # re-parse + + # ------------------ Read operations ------------------ + + def list_categories(self) -> List[Tuple[str, int, str]]: + """Return [(category, entry_count, last_update)]""" + results = [] + for f in sorted(self.memory_dir.glob('kernel-opt-*.md')): + cat = f.stem.replace('kernel-opt-', '') + store = ExperienceStore(cat) + count = len(store.entries) + mtime = time.strftime('%Y-%m-%d', time.localtime(f.stat().st_mtime)) + results.append((cat, count, mtime)) + return results + + def show(self, layer: Optional[int] = None, entry_id: Optional[str] = None) -> str: + if not self.exists(): + return f'[ERROR] 经验文件不存在: {self.file_path}' + + if entry_id: + e = self._find_entry(entry_id) + if not e: + return f'[ERROR] 未找到 ID={entry_id}，可用条目: {", ".join(e.id for e in self.entries)}' + lines = [f'ID: {e.id} | Layer {e.layer} | {e.title}', '-' * 40] + lines.extend(e.body_lines) + return '\n'.join(lines) + + if layer: + filtered = [e for e in self.entries if e.layer == layer] + if not filtered: + return f'[INFO] Layer {layer} 下暂无条目' + lines = [f'Layer {layer} 条目列表 ({len(filtered)} 条):', ''] + for e in filtered: + preview = ' '.join(e.body_lines)[:60].replace('\n', ' ') + lines.append(f' {e.id}: {e.title}') + lines.append(f' {preview}...') + return '\n'.join(lines) + + # Show full file summary + lines = [f'# {self.category.capitalize()} 算子经验摘要', ''] + for lnum in (1, 2, 3): + filtered = [e for e in self.entries if e.layer == lnum] + lines.append(f'Layer {lnum}: {len(filtered)} 条') + for e in filtered: + lines.append(f' {e.id}: {e.title}') + lines.append('') + return '\n'.join(lines) + + # ------------------ Write operations ------------------ + + def add(self, layer: int, title: str, content: Optional[str] = None, + why: Optional[str] = None, apply: Optional[str] = None) -> str: + if not self.exists(): + return f'[ERROR] 经验文件不存在，请先运行: python3 utils/exp-init.py {self.category}' + + seq = self._next_seq(layer) + entry_id = f'l{layer}-{seq:03d}' + + # Build body + body_lines = [] + if content: + body_lines.append(f'- **必须** {content}' if layer == 1 else f'- {content}') + if why: + body_lines.append(f'- **Why:** {why}') + if apply: + body_lines.append(f'- **How to apply:** {apply}') + if not body_lines: + body_lines.append('- 待补充') + + # Build title line with L numbering + existing = [e for e in self.entries if e.layer == layer] + subnum = len(existing) + 1 + title_line = f'### L{layer}.{subnum} {title} ' + + # Find insert position: end of target layer + insert_pos = len(self.lines) + in_target_layer = False + for i, line in enumerate(self.lines): + m = re.match(r'^## Layer\s+(\d+)\s*:', line) + if m: + if in_target_layer and int(m.group(1)) != layer: + insert_pos = i + break + if int(m.group(1)) == layer: + in_target_layer = True + if in_target_layer and line.startswith('### '): + insert_pos = i + 1 + # skip to end of this entry + while insert_pos < len(self.lines) and not self.lines[insert_pos].startswith('### ') and not self.lines[insert_pos].startswith('## '): + insert_pos += 1 + + self._backup(f'add {entry_id}') + new_block = [title_line] + body_lines + self.lines = self.lines[:insert_pos] + new_block + self.lines[insert_pos:] + self._save() + self._log_history({'action': 'add', 'id': entry_id, 'layer': layer, 'title': title}) + return f'[ADD] {entry_id}: {title} (Layer {layer})' + + def update(self, entry_id: str, title: Optional[str] = None, + content: Optional[str] = None, why: Optional[str] = None, + apply: Optional[str] = None, dry_run: bool = False) -> Tuple[str, str, str]: + """Returns (diff_str, before_str, after_str). If dry_run=True, does not write.""" + e = self._find_entry(entry_id) + if not e: + available = ', '.join(e.id for e in self.entries) + raise ValueError(f'未找到 ID={entry_id}，可用条目: {available}') + + before_lines = [e.raw_title_line] + e.body_lines + after_lines = before_lines.copy() + + if title: + new_title = re.sub(r'^###\s+', '', after_lines[0]) + new_title = re.sub(r'', '', new_title).strip() + # preserve L numbering + prefix = re.match(r'^(L\d+\.\d+\s+)', new_title) + prefix_str = prefix.group(1) if prefix else '' + after_lines[0] = f'### {prefix_str}{title} ' + + # Rebuild body if any content fields provided + if content or why or apply: + new_body = [] + if content: + new_body.append(f'- **必须** {content}' if e.layer == 1 else f'- {content}') + if why: + new_body.append(f'- **Why:** {why}') + if apply: + new_body.append(f'- **How to apply:** {apply}') + after_lines = [after_lines[0]] + new_body + + diff = '\n'.join(difflib.unified_diff( + before_lines, after_lines, + fromfile=f'{entry_id} (before)', tofile=f'{entry_id} (after)', + lineterm='' + )) + + if not dry_run: + self._backup(f'update {entry_id}') + self.lines = ( + self.lines[:e.start_line] + + after_lines + + self.lines[e.end_line:] + ) + self._save() + self._log_history({'action': 'update', 'id': entry_id}) + return diff, '\n'.join(before_lines), '\n'.join(after_lines) + + def remove(self, entry_id: str, dry_run: bool = False) -> Tuple[str, str]: + """Returns (removed_title, removed_body). If dry_run=True, does not write.""" + e = self._find_entry(entry_id) + if not e: + available = ', '.join(e.id for e in self.entries) + raise ValueError(f'未找到 ID={entry_id}，可用条目: {available}') + + removed = '\n'.join(self.lines[e.start_line:e.end_line]) + if not dry_run: + self._backup(f'remove {entry_id}') + self.lines = self.lines[:e.start_line] + self.lines[e.end_line:] + self._save() + self._log_history({'action': 'remove', 'id': entry_id, 'removed': removed}) + return e.title, removed + + def undo(self) -> Optional[Path]: + last = self.get_last_backup() + if not last: + return None + shutil.copy2(last, self.file_path) + self._load() + self._log_history({'action': 'undo', 'restored_from': str(last)}) + return last + + def history(self) -> str: + if not self.history_path.exists(): + return '[INFO] 暂无修改历史' + lines = [] + with open(self.history_path, 'r', encoding='utf-8') as f: + for line in f: + rec = json.loads(line.strip()) + ts = rec.get('time', '?') + act = rec.get('action', '?') + detail = '' + if act == 'backup': + detail = f"reason={rec.get('reason')}" + elif act in ('add', 'update', 'remove'): + detail = f"id={rec.get('id')}" + elif act == 'undo': + detail = f"from={rec.get('restored_from')}" + lines.append(f' [{ts}] {act} {detail}') + return '\n'.join(lines) + + +# ------------------ CLI ------------------ + +def confirm(prompt: str, diff: str = '') -> bool: + if diff: + print('\n--- 变更 diff ---') + print(diff) + print('--- end diff ---\n') + resp = input(f'{prompt} (y/n/dry-run): ').strip().lower() + if resp == 'y': + return True + if resp == 'dry-run': + print('[DRY-RUN] 未执行写入') + return False + print('[CANCEL] 操作已取消') + return False + + +def main(): + parser = argparse.ArgumentParser(description='历史经验增删改查工具') + sub = parser.add_subparsers(dest='cmd', required=True) + + # list + sub.add_parser('list', help='列出所有经验类别') + + # show + p_show = sub.add_parser('show', help='查看经验内容') + p_show.add_argument('category', help='算子类别') + p_show.add_argument('--layer', type=int, choices=(1, 2, 3), help='仅查看指定 Layer') + p_show.add_argument('--id', help='查看指定条目') + + # add + p_add = sub.add_parser('add', help='新增经验条目') + p_add.add_argument('category', help='算子类别') + p_add.add_argument('--layer', type=int, required=True, choices=(1, 2, 3)) + p_add.add_argument('--title', required=True) + p_add.add_argument('--content') + p_add.add_argument('--why') + p_add.add_argument('--apply') + + # update + p_up = sub.add_parser('update', help='修改经验条目（高危）') + p_up.add_argument('category', help='算子类别') + p_up.add_argument('--id', required=True, help='条目 ID，如 l1-002') + p_up.add_argument('--title') + p_up.add_argument('--content') + p_up.add_argument('--why') + p_up.add_argument('--apply') + p_up.add_argument('--yes', action='store_true', help='跳过交互确认（Agent 模式）') + + # remove + p_rm = sub.add_parser('remove', help='删除经验条目（高危）') + p_rm.add_argument('category', help='算子类别') + p_rm.add_argument('--id', required=True, help='条目 ID') + p_rm.add_argument('--yes', action='store_true', help='跳过交互确认（Agent 模式）') + + # undo + p_undo = sub.add_parser('undo', help='回退上一次写操作') + p_undo.add_argument('category', help='算子类别') + + # history + p_hist = sub.add_parser('history', help='查看修改历史') + p_hist.add_argument('category', help='算子类别') + + args = parser.parse_args() + + if args.cmd == 'list': + store = ExperienceStore('dummy') + results = store.list_categories() + print(f'{"类别":<12} {"条目数":>6} {"最近更新":>12}') + print('-' * 32) + for cat, cnt, mtime in results: + print(f'{cat:<12} {cnt:>6} {mtime:>12}') + return + + store = ExperienceStore(args.category) + + if args.cmd == 'show': + print(store.show(layer=args.layer, entry_id=args.id)) + return + + if args.cmd == 'add': + result = store.add( + layer=args.layer, title=args.title, + content=args.content, why=args.why, apply=args.apply + ) + print(result) + return + + if args.cmd == 'update': + try: + diff, before, after = store.update( + args.id, title=args.title, + content=args.content, why=args.why, apply=args.apply, + dry_run=True + ) + except ValueError as e: + print(f'[ERROR] {e}', file=sys.stderr) + sys.exit(1) + if not args.yes and not confirm(f'确认更新 {args.id}?', diff=diff): + sys.exit(0) + # Execute for real + store.update( + args.id, title=args.title, + content=args.content, why=args.why, apply=args.apply, + dry_run=False + ) + print(f'[UPDATE] {args.id} 已更新') + print(f'[BACKUP] 原文件已备份') + return + + if args.cmd == 'remove': + try: + title, removed = store.remove(args.id, dry_run=True) + except ValueError as e: + print(f'[ERROR] {e}', file=sys.stderr) + sys.exit(1) + preview = removed[:200].replace('\n', ' ') + if not args.yes and not confirm( + f'确认删除 {args.id} ({title})?\n预览: {preview}...' + ): + sys.exit(0) + store.remove(args.id, dry_run=False) + print(f'[REMOVE] {args.id} ({title}) 已删除') + print(f'[BACKUP] 内容已备份') + return + + if args.cmd == 'undo': + bak = store.undo() + if bak: + print(f'[UNDO] 已恢复至备份: {bak.name}') + else: + print('[WARN] 未找到可回退的备份') + return + + if args.cmd == 'history': + print(store.history()) + return + + +if __name__ == '__main__': + main() diff --git a/utils/exp-init.py b/utils/exp-init.py new file mode 100644 index 00000000..014cfab1 --- /dev/null +++ b/utils/exp-init.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +""" +新算子类别经验文件初始化工具：基于模板创建 kernel-opt-{category}.md 并更新 MEMORY.md 索引。 + +用法: + python3 utils/init_experience.py [--op-name ] + +示例: + python3 utils/init_experience.py pad --op-name 15_Pad +""" + +import argparse +import re +import sys +from pathlib import Path + + +def get_project_root() -> Path: + return Path(__file__).resolve().parent.parent + + +def render_template(category: str, op_name: str) -> str: + """渲染 kernel-opt 模板。""" + category_cap = category.capitalize() + return f"""--- +name: kernel-opt-{category} +description: {category_cap} 算子的 Triton Ascend 四层隔离优化经验 +metadata: + type: reference +--- + +# {category_cap} 算子优化经验 + +**算子类别**: `待补充` +**典型特征**: 待补充 +**性能基准**: 待补充 + +--- + +## Layer 1: 设计约束（Agent 必须遵守） + +### L1.1 约束名称 +- **必须** / **禁止** 做某事 +- **Why:** 原因说明 +- **How to apply:** 适用场景 + +> 提示：每次算子探索后，将本次验证过的硬性约束追加到本层。 + +--- + +## Layer 2: 算法骨架（Agent 可参考架构） + +### L2.1 Host 侧分支决策树（伪代码） + +``` +待补充 +``` + +### L2.2 多核并行骨架模式 + +**模式 A - 按元素分配**: +``` +elements_per_core = cdiv(total_elements, num_cores) +core_start = pid * elements_per_core +... +``` + +**模式 B - 按行分配**: +``` +rows_per_core = cdiv(total_rows, num_cores) +row_start = pid * rows_per_core +... +``` + +--- + +## Layer 3: 关键技巧（Agent 可参考，但实现方式可不同） + +### L3.1 技巧名称 +```python +# 代码片段 +``` + +**可替代方向**: 说明其他可能的实现方式。 + +--- + +## Layer 4: 完整归档（Agent 默认不读取，仅人工复盘） + +> ⚠️ **Agent 注意**：以下仅为历史实现的路径记录。你**禁止**直接复制其代码结构、变量命名或 kernel 组织方式。 + +### 历史实现归档 + +| 版本 | 代码 | 报告 | 摘要 | 性能 | 特点 | +|------|------|------|------|------|------| +| 待补充 | 待补充 | 待补充 | 待补充 | 待补充 | 待补充 | + +### 完整归档路径（Layer 4） +``` +/home/zmm/OpAgent-Pad/.claude/memory/archive/{category}/ +``` + +### 原始工作目录 +``` +待补充 +``` + +### 性能基准（几何平均） + +| Shape 类型 | 典型加速比 | 说明 | +|-----------|-----------|------| +| 待补充 | 待补充 | 待补充 | + +**关键结论**：待补充 + +--- + +## 常见陷阱与避免方法 + +### 陷阱 1: 名称 +- **问题**: 描述 +- **解决**: 方案 +""" + + +def update_memory_md(memory_path: Path, category: str) -> None: + """在 MEMORY.md 中添加该算子类别的经验文件索引。""" + category_cap = category.capitalize() + exp_line = f'- [{category_cap} 算子优化经验](kernel-opt-{category}.md) — 待补充' + + if not memory_path.exists(): + memory_path.write_text('# Memory Index\n\n', encoding='utf-8') + + content = memory_path.read_text(encoding='utf-8') + lines = content.splitlines() + + # 查找是否已有该类别 + for line in lines: + if f'kernel-opt-{category}.md' in line: + return # 已存在 + + # 在 "算子优化经验" 节追加 + section_idx = None + insert_pos = len(lines) + for i, line in enumerate(lines): + if '算子优化经验' in line or '经验' in line: + section_idx = i + if section_idx is not None: + for j in range(section_idx + 1, len(lines)): + if lines[j].strip().startswith('- ') and 'kernel-opt-' in lines[j]: + insert_pos = j + 1 + elif lines[j].strip().startswith('## ') and j > section_idx: + insert_pos = j + break + lines.insert(insert_pos, exp_line) + else: + lines.append('') + lines.append('## 算子优化经验') + lines.append(exp_line) + + memory_path.write_text('\n'.join(lines) + '\n', encoding='utf-8') + + +def main(): + parser = argparse.ArgumentParser(description='初始化新算子类别经验文件') + parser.add_argument('category', help='算子类别名（小写，如 pad, softmax）') + parser.add_argument('--op-name', default=None, help='原始算子名称（如 15_Pad）') + args = parser.parse_args() + + category = args.category.lower() + op_name = args.op_name or category.capitalize() + root = get_project_root() + + exp_path = root / '.claude' / 'memory' / f'kernel-opt-{category}.md' + memory_path = root / '.claude' / 'memory' / 'MEMORY.md' + + if exp_path.exists(): + print(f'[SKIP] 经验文件已存在: {exp_path}') + print(' 如需更新，请直接编辑该文件。') + sys.exit(0) + + exp_path.write_text(render_template(category, op_name), encoding='utf-8') + print(f'[CREATE] {exp_path}') + + update_memory_md(memory_path, category) + print(f'[UPDATE] {memory_path}') + + print(f'\n[SUCCESS] {category} 类别经验文件已初始化。') + print(' 下一步：将本次探索提炼的 Layer 1-3 内容填入该文件。') + + +if __name__ == '__main__': + main() diff --git a/utils/exp_design.md b/utils/exp_design.md new file mode 100644 index 00000000..45426ab3 --- /dev/null +++ b/utils/exp_design.md @@ -0,0 +1,167 @@ +# 历史经验管理方案设计（增删改查） + +## 1. 设计目标 + +让 Agent 和人类用户都能对 `.claude/memory/kernel-opt-{category}.md` 中的历史经验进行增删改查，同时保证： +- **高危操作可回退**：更新/删除前必须确认，并保留备份 +- **多次操作一致性**：所有写操作后自动校验四层模型格式 +- **界面简洁**：命令语义清晰，输出聚焦核心信息 + +## 2. 命令接口设计 + +统一入口：`utils/exp-edit.py` + +```bash +# 查 +python3 utils/exp-edit.py list # 列出所有经验类别及条目数 +python3 utils/exp-edit.py show # 查看完整经验文件 +python3 utils/exp-edit.py show --layer 1 # 仅查看 Layer 1 +python3 utils/exp-edit.py show --id l1-002 # 查看指定条目 + +# 增 +python3 utils/exp-edit.py add --layer 1 \ + --title "坐标比较必须用 float32" \ + --content "禁止直接对整数坐标使用 tl.where(coord < 0, ...)" \ + --why "Ascend 后端整数比较可能行为异常" \ + --apply "所有边界映射 kernel" + +# 改（高危） +python3 utils/exp-edit.py update --id l1-002 \ + --content "必须先 .to(tl.float32) 再比较" + +# 删（高危） +python3 utils/exp-edit.py remove --id l1-002 + +# 工具 +python3 utils/exp-edit.py undo # 回退上一次写操作 +python3 utils/exp-edit.py history # 查看该类别的修改历史 +``` + +## 3. 数据模型与 ID 方案 + +### 3.1 Markdown 内嵌 ID + +在现有 `kernel-opt-*.md` 的 Markdown 结构中，为每个可独立操作的条目（L1.x、L3.x 等）末尾注入不可见的 HTML 注释 ID： + +```markdown +### L1.2 坐标比较必须用 float32 +- **必须**先 `.to(tl.float32)` 再比较 +- **Why:** Ascend 后端整数比较可能行为异常 +- **How to apply:** 所有边界映射 kernel +``` + +- `l1-002`：`l` + 层级 + `-` + 该层内序号 +- 序号按该层现有最大序号 + 1 自动生成 +- 注释对人类阅读完全透明，但对程序提供精确定位锚点 + +### 3.2 条目定位规则 + +| 场景 | 定位方式 | +|------|---------| +| 有 ID | 精确匹配 `` | +| 无 ID（历史遗留）| 按子标题 `### L{N}.{M} ...` 定位，操作后自动补 ID | +| 新增 | 在目标 Layer 最后一个条目后追加，自动生成 ID | + +## 4. 高危操作确认机制 + +### 4.1 交互确认（人类用户） + +`update` / `remove` 默认进入交互确认： + +``` +[UPDATE] pad / l1-002 +--- 当前内容 --- +- **必须**先 `.to(tl.float32)` 再比较 +--- 变更后内容 --- +- **必须**先 `.to(tl.float32)` 再比较，禁止直接对整数坐标使用 tl.where + +确认执行？ (y/n/dry-run): +``` + +- `y`：执行 +- `n`：取消 +- `dry-run`：仅打印 diff，不写入 + +### 4.2 Agent 非交互模式 + +Agent 调用时附加 `--yes` 参数跳过确认： + +```bash +python3 utils/exp-edit.py update pad --id l1-002 --content "..." --yes +``` + +**约束**：Agent 必须在调用前向用户展示变更摘要（diff），获得用户文字确认后方可加 `--yes` 执行。 + +### 4.3 删除特别保护 + +`remove` 操作无论是否 `--yes`，都会： +1. 将被删条目完整内容写入 `.claude/memory/.backup/{category}_removed_{timestamp}.md` +2. 在 `history` 中记录删除事件（含被删内容摘要） +3. 输出 `[REMOVED] ID={id} 内容已备份至 ...` + +## 5. 一致性保障机制 + +### 5.1 自动备份 + +每次写操作（add/update/remove）前： +- 自动将原文件复制到 `.claude/memory/.backup/kernel-opt-{category}_{timestamp}.md` +- 保留最近 20 个备份，超出的自动清理 + +### 5.2 自动校验 + +每次写操作后： +- 校验 Markdown 结构完整性（frontmatter、Layer 1-4 标题是否齐全） +- 校验所有条目是否都有 `exp-id` +- 若有缺失，自动补 ID 并告警 +- 调用 `exp-check.py` 做全局一致性检查 + +### 5.3 回退机制 + +```bash +python3 utils/exp-edit.py undo +``` + +- 将文件恢复为最近一次备份状态 +- 仅能回退一次（最近一次写操作） +- 回退后删除对应备份，防止连环 undo + +## 6. 简洁性设计 + +### 6.1 输出格式 + +- **查询**：仅显示条目标题 + ID + 前 80 字符摘要，不显示完整 Markdown +- **变更**：统一用 unified diff 格式展示，3 行上下文 +- **列表**：一行为一个类别，显示"类别名 / 总条目数 / 最近更新日期" + +### 6.2 错误提示 + +- 命令参数错误：直接输出用法示例（`Usage: ...`） +- ID 不存在：`[ERROR] pad 中不存在 ID=l1-099，可用条目: l1-001, l1-002, l3-001` +- 文件不存在：`[ERROR] kernel-opt-foo.md 不存在。如需创建，请先运行: python3 utils/exp-init.py foo` + +## 7. 实施计划 + +1. **Phase A**：实现 `show` / `list`（只读，无风险） +2. **Phase B**：实现 `add`（写操作，有备份和校验） +3. **Phase C**：实现 `update` / `remove` + 交互确认（高危操作） +4. **Phase D**：实现 `undo` / `history`（回退与审计） +5. **Phase E**：对现有 `kernel-opt-pad.md` 自动补全 `exp-id` + +## 8. 与现有工具的协作关系 + +``` +exp-init.py ----> 首次创建 kernel-opt-{category}.md（含模板结构） + | + v +exp-archive.py -> 归档代码到 Layer 4（更新 Layer 4 表格） + | + v + exp-edit.py ------------> 日常增删改查 Layer 1-3 经验条目 + | + v +exp-check.py -> 全局一致性最终校验 +``` + +--- + +**请确认本方案后，我按 Phase A→E 逐步实现。**