Skip to content

perf(jit): cache AST structure analysis for eager calls - #2953

Merged
YunjiQin merged 1 commit into
hw-native-sys:feat/kernel-mode-integration-testfrom
mihudaner:perf/jit-ast-structure-cache
Oct 8, 2026
Merged

perf(jit): cache AST structure analysis for eager calls#2953
YunjiQin merged 1 commit into
hw-native-sys:feat/kernel-mode-integration-testfrom
mihudaner:perf/jit-ast-structure-cache

Conversation

@mihudaner

@mihudaner mihudaner commented Oct 8, 2026

Copy link
Copy Markdown

修改内容

在 Eager 调用中，JIT 即使命中已编译的 kernel，仍会重复遍历函数 AST，以查找引用的常量名称和依赖函数调用位置。本次将这两类结构信息缓存，减少每次调用的 Host 开销。

只缓存源码结构，不缓存常量的当前值。 每次请求仍重新读取全局变量和闭包中的常量，保留原有的缓存失效判断。新增测试覆盖缓存命中后不再扫描、AST 替换后重新分析、调用顺序以及返回列表隔离。

部分测试用例验证

基于 feat/kernel-mode-integration-test@e3ed3a0c，验证结果如下：

  • test_decorator.py 和 test_cache.py：251 项通过。
  • Ruff 0.14.8、格式、语法及 diff 检查通过。
  • Debug 构建退出时有 nanobind 引用泄漏诊断，测试退出码为 0。

性能优化结果

优化前耗时：
image

优化后耗时：
image

本次修改的两个文件与此前在 3e87a84 环境验证的版本完全一致。此前单卡 TMR PTO Eager、B1／128K／S6／TP1 测试中，无 profiler 的 Host 调用中位数从 107.85 ms 降至 11.18 ms，减少约 89.6%。同输入下，优化前后的输出、Top-K／分数及五组
cache 更新逐元素完全一致，重复执行和越界写检查均通过。

@mihudaner
perf(jit): cache AST structure analysis for eager calls
a4bf67f 
Cache referenced constant names and direct call sites by read-only parsed
function definition instead of walking the same AST on every eager call.
Bound each cache to 512 entries and return a fresh call-site list to avoid
mutating the cached index.

Continue resolving current global and closure values on every request;
only source structure is cached, preserving value-sensitive invalidation.
Add regression tests for warm lookup without rescanning, replacement ASTs,
source ordering, and returned-list isolation.

Base: feat/kernel-mode-integration-test at e3ed3a0.

Validation on this base:
- Built the Python extension in the PR worktree (Debug, two build jobs).
- test_decorator.py and test_cache.py: 251 passed with two pytest workers.
  Debug nanobind reports reference leaks at shutdown; test exit code is 0.
- Ruff 0.14.8, formatting, syntax, and diff checks pass.

The modified files are byte-identical to the previously validated 3e87a84
integration checkout. There, single-card TMR PTO Eager B1/128K/S6/TP1 host
call median dropped from 107.85 ms to 11.18 ms without profiler. Same-input
outputs, Top-K/scores, and all five cache groups compared exactly; repeated
calls and write guards passed. These measurements are from that checkout,
not a new whole-model performance run on this branch.
@coderabbitai

coderabbitai Bot commented Oct 8, 2026

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration
  • Configuration used: Organization UI
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: 7e5236f6-fb6a-4d68-8f3f-a848f9dc8b6f

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
  • Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@YunjiQin
YunjiQin merged commit b63d203 into hw-native-sys:feat/kernel-mode-integration-test Oct 8, 2026
12 of 14 checks passed
@mihudaner mihudaner mentioned this pull request Oct 10, 2026
Merged
YunjiQin pushed a commit that referenced this pull request Oct 10, 2026
@mihudaner
perf(jit): reduce host overhead for eager single-operator calls (#2954)
3cc5c61 
## 问题与目的

延续 #2953 的 非aclgraph场景下host开销性能优化。已从 **无优化 103.49 ms -----------> #2953
优化后 11.25 ms**

当前即使 Eager 调用命中已编译 kernel，JIT 仍会在不同 helper 中重复解析同一个函数签名，并把已经取得的普通常量值先构造成
AST、再转换为源码文本。本 PR 减少这两部分 Host 准备开销。**验证后耗时优化到 7.784 ms。**

包含两个独立 commit，分别对应以下改动， **都是逻辑等价优化；不改变 kernel 计算、调度器或提交协议。**

## 改动一：请求内复用函数签名

`python/pypto/jit/decorator.py` 新增 `_function_signature`，使用已有的
`@cache_in_snapshot`；六处签名查询统一经由该 helper。

```text
原来：同一请求的参数名 / constexpr / layout / 参数绑定等 helper
        ----> 各自 inspect.signature(func)
现在：本请求首次查询 ----> inspect.signature(func)
      本请求再次查询 ----> 复用 Signature
      请求结束       ----> 丢弃快照，下次请求重新读取
```

参数绑定、默认值应用和注解求值仍按原路径执行，不缓存实参绑定结果或整个 specialization。跨请求修改默认值、注解和自定义
`__signature__` 后会重新读取；嵌套调用共享当前快照，独立线程隔离，异常退出也清理快照。与现有请求快照语义一致。

## 改动二：普通常量文本直接生成
python/pypto/jit/specializer.py 的 constant_source 对精确的 None/bool/int/str
直接使用 repr；free_name_source 读取当前值后复用同一转换逻辑。

  原来：当前值 128 ----> ast.Constant(128) ----> ast.unparse ----> "128"
  现在：当前值 128 ----> repr(128)                          ----> "128"
默认 AST unparser 对这些类型原本就使用 repr，因此仅省去 AST 构造和遍历。精确类型用 is
检查，避免子类或元类自定义相等比较干扰。浮点数（含 NaN/Inf、负零）、子类、容器、枚举和 DSL
符号保留原转换流程。仍读取当前绑定，不新增跨请求缓存，不缓存表达式的 eval 结果或 Python 编译对象。
所以这个优化直白简单描述一下是干了什么


• 就是把常量变成源码字符串时，省掉中间绕路。

## 性能验证

单卡 **TMR PTO Eager，B1 / 128K 历史 / S6 / TP1**，BF16 hidden、正式 C4 层 W8A8
权重，NZ2，`atomic_add=0`、`deterministic_level=0`，合成独立物理 Cache 历史。以下为无
profiler 的 Host 调用中位数，各版每段 8 次：

| 独立 A/B | 优化前 | 优化后 | Host 降幅 |
|---|---:|---:|---:|
| 改动1：`b63d2031` → 请求内签名复用 | 11.323 ms | 8.805 ms | 22.2% |
| 改动2：`80336c7d` → 常量文本快速生成 | 8.888 ms | 7.784 ms | 12.4% |

两组为不同轮次，不将百分比直接相加。Phase 3 插桩后的另一段无 profiler 测量为 8.765 → 7.772 ms，方向一致。

轻量插桩分段如下，单位 ms、各版 8 次均值；父子边界为包含关系，不能相加：

| 边界 | 改动一 前 → 后 | 改动二 前 → 后 |
|---|---:|---:|
| `_resolve_constexpr_bindings` | 4.109 → 2.342 | 2.378 → 2.124 |
| `_get_source_hash` | 2.095 → 2.070 | 2.081 → 1.292 |
| `_resolve_specialization` | 2.682 → 1.914 | 1.935 → 1.917 |

独立 cProfile 的 3 次调用中，Phase 2 的 `inspect.signature` 次数从 645 降至
168；cProfile 绝对耗时不作为性能结论。Phase 3 的主要收益来自常量文本处理，未优化参数特化。

两阶段均保存优化前后 CANN Level1 CPU/NPU profiling。Phase 3 的 CANN Host 包络为 9.498 →
7.966 ms，设备 kernel 中位数为 483.71 → 470.84 μs，每版 8 个完整设备任务。设备仍约 0.5
ms，独立轮次差异不归因为本次核内提速。本 PR 改动位于 HBG/TMR 共用准备逻辑，但本轮只测 TMR 单算子，不将该比例外推到
HBG、ACLGraph replay 或整模型吞吐。

## 正确性与回归

- 两阶段分别验证同输入、同初态下的优化前后等价性：输出、Top-K/分数、五组 Cache 更新逐元素一致，最大绝对差
0；重复执行一致，保护区检查通过。此项是 Host 优化前后比较，不是 Native/PTO 精度验收。
- 两版生成的 208 份 PTO/C++ 文件一致，仅归一化自动 inline 编号及构建目录。
- 完整 `tests/ut/jit/`：**833 passed**；最后补强新字面量用例后定向回归：**26 passed**。
- 新增测试覆盖签名共享、默认值/注解/自定义签名刷新、异常清理、线程隔离，以及常量文本一致性、快速路径、全局重绑定、容器修改、子类和元类行为。
- 四个变更文件的 Ruff 检查、格式检查及 `git diff --check` 通过；提交前确认源码与已完成 NPU 验证的快照一致。
- Debug 扩展退出时有已知 nanobind 引用泄漏提示；验证任务退出码为 0。
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@YunjiQin YunjiQin YunjiQin approved these changes

Assignees

No one assigned

Labels

None yet

Projects

pto project
Status: No status

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@mihudaner @YunjiQin