Prune inductor cache during Triton extraction

[leon062112]

PR Category

Feature Enhancement

Description

当前 Triton kernel 提取流程会为每个计算图生成 TorchInductor 编译 cache，包括 output_code.py、Triton 编译中间产物 PTX、.best_config 等文件。完整 cache 体积较大，如果对完整 sample list 跑提取流程，占用空间太大。因此在每个 sample 编译完成后立即裁剪 cache，只保留后续提取和分析所需的最小文件集合。

同时修复了一个 PTX 提取问题：在当前环境中，部分 .best_config 的 triton_cache_hash 为 null，原逻辑无法从多个 PTX candidate 中选出 autotune 最终采用的 PTX。

改动内容

• 新增 cache_pruner.py，用于裁剪单个 sample 的 TorchInductor cache。
• 在 sample 编译完成后立即调用 cache pruner，降低批量处理时的磁盘峰值占用。
• 保留后续流程需要的文件：
  • test_compiler_log.log
  • original_graph/model.py
  • original_graph/graph_hash.txt
  • output_code.py
  • *.best_config
  • *.ptx
  • *.json
  • *.pruned_meta.json
• 删除不再需要的大体积中间产物：
  • *.source
  • *.cubin
  • *.ttir
  • *.ttgir
  • *.llir
  • 其他无关临时文件
• 删除 *.source 前提取其中的 block size 信息，并写入轻量的 *.pruned_meta.json。
• 增强 PTX 匹配逻辑：
  • 优先使用 .best_config 中的 triton_cache_hash
  • 当 triton_cache_hash 为 null 时，从 output_code.py 的 # kernel path: 找到对应
    .best_config
  • 使用 XBLOCK、num_warps、num_stages 等 metadata 匹配正确的 PTX candidate

验证

使用 2 个 typical subgraph sample测试：

Compiled: 2
Kept: 1
Discarded: 1
Extraction: 1 files, 3 kernels, 3 ptx, 1 samples

单个 sample dump 体积对比：
裁剪前: 14M
裁剪后: 1.9M
减少:   约 88.7%

[paddle-bot]

Thanks for your contribution!