当前编译栈是什么?
仓库的编译器由 Python 图捕获、Graph/Pass IR、Stax 原生执行器和 Triton/TVM 后端组成。它不试图一次复制完整工业编译器,而是先把每层的输入、输出和失败边界写清楚。
现有计划按 L1 到 L7 展开:捕获前端、IR 与 pass、符号形状、AOT 自动微分、代码生成、guard 管理和 export IR。当前已落地的部分包括 PassManager、常量折叠、死代码消除、ShapeProp、元数据特化、shape guards 和部分 AOT 原生路径;更远的层仍然是路线图,不应被包装成已完成能力。
图捕获:哪些值可以特化?
捕获期可以读取 shape、dtype、device、ndim 和 len 等元数据,并将安全的 Python 控制流特化进图。张量数据本身仍保持符号;如果分支依赖真实数据,fullgraph 模式应报告 GraphCaptureError,而不是生成一个只对当前输入碰巧正确的图。
dynamic 模式还需要 guard:如果图读取了 x.shape[0] 并据此选择分支,缓存键必须记录这个触碰过的维度。否则不同尺寸复用同一份特化图,就会把性能优化变成静默的正确性错误。
if x.shape[0] > 2:
y = x.relu()
else:
y = x.sigmoid()
shape read -> guard -> specialized graph
data-dependent branch -> explicit capture errorStax 的性能工作:先把成本暴露出来
路线图把 codegen cache、autotune、cudagraphs、分解表和更宽的融合面列为后续里程碑。它们分别影响编译重复成本、kernel 配置选择、启动开销、算子语义和跨算子优化;在这些能力完成并有基准结果之前,不应笼统宣称“编译后更快”。
性能验收应该回到仓库的 benchmark/benchmark_resnet_classification.py:相同权重、相同数据顺序、相同输入和相同计时口径,同时记录首次编译成本、稳态吞吐、p50/p95 延迟、logits parity 和 codegen。这样每个优化都能说明它改善了哪一项成本。
- 编译成本:wrapper 与 first pass 的时间。
- 运行成本:稳态 p50/p95、吞吐和内存行为。
- 正确性成本:eager 与编译路径的 logits 和预测一致性。
- 执行证据:native、Triton、TVM 或 fallback 的实际 codegen。
现在能准确说什么?
- 可以说:执行路径被拆成可检查的捕获、pass、lowering 和 kernel 边界。
- 可以说:不支持的图结构会显式失败或按契约 fallback,并保留数值验证。
- 可以说:性能正在通过 GEMM、ResNet 和编译对照逐配置测量。
- 不能说:所有模型、GPU、dtype 和训练任务都比成熟框架更快。
