PyTorch AutoHeuristic 实战指南pad_mm 启发式的数据收集与代码生成【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchpad_mm 是 PyTorch Inductor 中一项针对矩阵乘法torch.mm的自动填充padding优化它通过在矩阵维度不足对齐时补零来提升 Triton kernel 的执行效率但是否值得填充取决于 GPU 型号与具体形状。AutoHeuristic 是 PyTorch 仓库内置的框架它把自动调优autotuning的结果学习成可随编译器一起发布的决策树/回归树代码。本文以 torchgen/_autoheuristic/pad_mm/README.md 为核心完整讲解如何在本地重建已发布的 A100 启发式、如何为新 GPU 从零收集训练数据并生成 pad_mm 启发式同时结合 pad_mm 的源码实现 与 AutoHeuristic 工具链 说明每一步背后的工作原理。读完本文你将掌握一套可复制的基准测试 → 数据收集 → 训练 → 生成代码完整流程。pad_mm 优化与 AutoHeuristic 框架的关系什么是 pad_mmpad_mm 是 Inductor 图级别graph-level的一项融合优化在torch.mm的两个输入张量维度M、K、N未对齐时通过aten.constant_pad_nd补零到对齐长度使底层 GEMM kernel 能以更规整的 tile 布局执行从而换取吞吐提升。其入口位于 torch/_inductor/fx_passes/pad_mm.py核心安全判断逻辑设备、dtype、静态维度等集中在can_pad()中而是否值得填充则交由 AutoHeuristic 或传统 autotune 决策。其中对齐大小的定义在 pad_mm.py 的get_alignment_size_dtype()中float16/bfloat16对齐大小为 8float32对齐大小为 4其他 dtype返回 0不可填充。AutoHeuristic 如何为 pad_mm 决策在 pad_mm.py 的run_autoheuristic()中构造了一个两选一的AutoHeuristic实例choices [orig, pad]分别代表保持原始形状与填充后执行feedback是一个LocalFeedback对每个 choice 用do_bench实测耗时fallback返回autotune即当启发式不可用时回退到传统 autotuneaugment_contextpad_mm_operations()注入派生特征preconditionpad_mm_precondition决定哪些输入才值得交给 AutoHeuristic见下文特征工程一节。由此可以看出pad_mm 的 AutoHeuristic 场景是典型的本地自动调优feedback 函数能立刻返回实测耗时。学习得到的启发式会被编译进编译器在后续编译同一 GPU 上的同类形状时直接给出决策省去反复 benchmark 的开销。快速重建已发布的 A100 启发式如果只是想复现仓库已发布、基于已有数据训练的 pad_mm A100 启发式只需要依次执行 README.md 给出的两条脚本命令bash get_padmm_dataset.sh # Downloads A100 bash gen_pad_mm_a100.sh # Generates A100 heuristic第一步下载 A100 训练数据集get_padmm_dataset.sh 通过wget下载名为pad_mm_a100_data.zip的数据包再用unzip解压出pad_mm_a100_data.txt并清理压缩包。该 txt 文件即训练数据首行为 JSON 格式的设备元信息shared_memory、device_capa、name后续每一行是一次 autotune 采样记录。第二步训练并生成 A100 启发式gen_pad_mm_a100.sh 调用回归树训练脚本python train_regression_pad_mm.py pad_mm_a100_data.txt --heuristic-name PadMMA100训练完成后启发式代码会被写入torch/_inductor/autoheuristic/artifacts/_PadMMA100.py生成路径的逻辑见 train.py 的write_heuristic_to_file()。此后在编译 pad_mm 时Inductor 会自动加载并使用该启发式。为新 GPU 收集数据并生成启发式完整流程如果目标 GPU 不是 A100或希望用更新的数据重新训练就需要使用 generate_heuristic_pad_mm.sh 走完整的采集 → 训练流程。步骤 0修改脚本中的关键变量在运行脚本前必须先按注释修改脚本顶部的变量变量示例值含义GPU_DEVICE_IDS4,5参与数据采集的 CUDA 设备号逗号分隔多卡会被并行利用CONDA_ENVheuristic-pr运行前需要激活的 conda 环境名NUM_SAMPLES2000计划收集的样本总数会按 GPU 数量均分OUTPUT_DIRa100采集结果与中间文件的输出目录HEURISTIC_NAMEPadMMA100启发式名称若目标不是 A100必须修改此名同时建议修改BENCHMARK_SCRIPT/TRAIN_SCRIPT为对应 GPU 的采集与训练脚本如 H100 可参考仓库中的 gen_pad_mm_h100.shBENCHMARK_SCRIPTgen_data_pad_mm.py采集数据用的 benchmark 脚本TRAIN_SCRIPTtrain_regression_pad_mm.py训练启发式用的脚本脚本随后把上述参数转发给通用流程脚本 generate_heuristic.sh。步骤 1collect —— 跑 benchmark 并收集训练数据bash generate_heuristic_pad_mm.sh collect该命令会在随机输入上执行真实编译与基准测试。README 明确提示视 GPU 数量与算力而定这一步可能耗时约一天因为要对每个采样形状做torch.compiledo_bench实测。collect 模式内部由 collect_data.sh 实现并行化它依赖tmux为GPU_DEVICE_IDS中的每个设备创建一个 tmux pane执行形如conda activate CONDA_ENV python gen_data_pad_mm.py --device id \ -o OUTPUT_DIR/data_id.txt --num-samples NUM_SAMPLES_PER_GPU因此每个 GPU 会生成独立的data_id.txt文件互不干扰。步骤 2generate —— 用收集到的数据学习启发式bash generate_heuristic_pad_mm.sh generategenerate 模式在 generate_heuristic.sh 中完成两件事用 merge_data.py 把各 GPU 的data_id.txt合并为OUTPUT_DIR/HEURISTIC_NAME.txt调用训练脚本python train_regression_pad_mm.py OUTPUT_DIR/HEURISTIC_NAME.txt --heuristic-name HEURISTIC_NAME训练出的启发式同样输出到torch/_inductor/autoheuristic/artifacts/_HEURISTIC_NAME.py。数据收集的源码原理BenchmarkRunnerPadMM训练数据的质量直接决定启发式的泛化能力。pad_mm 的数据采集器 gen_data_pad_mm.py 中的BenchmarkRunnerPadMM类继承自 AutoHeuristic 的 BenchmarkRunner实现了create_input()与run_benchmark()两个核心方法并覆盖了run()来串联整条流水线。形状生成策略已知形状 随机形状形状生成器generate_mm_shapes()分两个阶段产生(m, k, n, dtype)第一阶段先穷举已知形状集合——包括通过--additional-shape-csv传入的外部 CSV格式为M,K,N,dtype支持float16/bfloat16/float32以及 collect_known_mm_shapes.py 收集的常见 mm 形状第二阶段无限生成随机形状其中一半概率使用均匀随机1 到 65536另一半使用 2 的幂与2 的幂之间的混合分布见get_random_dim()未对齐概率默认 0.25。关键过滤只采集未对齐形状因为 pad 只有在维度不齐时才有意义生成器对每个形状做对齐过滤源码 gen_data_pad_mm.pyalign_size get_alignment_size_dtype(dtype) # Skip if all dimensions are aligned (we need unaligned for padding to be relevant) if all(self.is_aligned(dim, align_size) for dim in [m, k, n]): continue # Check if it fits in memory if fits_in_memory(dtype, m, k, n): yield (m, k, n, dtype)同时用fits_in_memory()防止超大形状导致显存溢出。覆盖 padding 的真实成本run_benchmark()中值得一提的设计是prepadded变量概率默认 0.2它模拟输入已经预先填充好的场景通过torch.mm(a 1, b)等方式规避编译期形状传播从而把填充本身的开销排除在计时之外。这样训练数据能区分两种情形输入来自前一层填充成本已被吸收此时 pad 往往更划算输入形状固定填充会产生额外的拷贝开销。这正是启发式需要学习的关键权衡点之一。采集时的配置开关run()方法中通过torch._inductor.config控制采集行为源码 gen_data_pad_mm.pytorch._inductor.config.autoheuristic_use.pad_mm False torch._inductor.config.autoheuristic_collect.pad_mm True torch._inductor.config.autoheuristic_log_path args.o即采集模式下只记录 autotune 结果而不使用启发式。这与 AutoHeuristic 通用文档 README.md 中描述的环境变量用法一致——你也可以在任意触发 pad_mm 的用户脚本中用环境变量完成同样的控制# 采集模式把 autotune 结果写入 train.txt TORCHINDUCTOR_AUTOHEURISTIC_LOG_PATHtrain.txt \ TORCHINDUCTOR_AUTOHEURISTIC_COLLECTpad_mm python run.py # 使用模式加载已学习的启发式 TORCHINDUCTOR_AUTOHEURISTIC_USEpad_mm python run.py其中pad_mm即AutoHeuristic构造时传入的name。特征工程pad_mm 启发式依据什么做决策启发式的预测能力来自AHContext中收集的原始特征以及augment_context中的派生特征。原始特征在 pad_mm.py 的get_context()中注册包括m、k、n、mat1_stride_0/1、mat2_stride_0/1、m/k/n_padded_length、mat1/mat2_align_size、mat1/mat2_dtype类别特征、prepadded_mat1/mat2类别特征、using_tf32类别特征。派生特征定义在 torch/_inductor/autoheuristic/autoheuristic_utils.py 的pad_mm_operations()中训练与推理两侧通过 train_decision_pad_mm.py / train_regression_pad_mm.py 的add_new_features()保持一致地重建特征计算方式类型含义m*k、m*n、k*n维度两两相乘数值捕获计算量规模k/(m*n)k / (m*n)数值捕获瘦长 K 维的极端形态bfloat_perf_hitk m*1024 且 k n*1024 且 dtype 为 bfloat16类别bfloat16 下 K 极端大的性能敏感场景arith_intensitym*k*n / (m*k k*n m*n)数值算术强度衡量访存与计算比mat1_innermost_needs_padding/mat2_innermost_needs_padding由 stride 与 padded_length 推导类别最内层维度是否需要填充num_dims_needs_paddingm/k/n 需要填充的维度个数数值0~3填充工作量m_multiple_2/4/8/16/32等各维度对 2/4/8/16/32 的整除性类别对齐程度细粒度刻画mat1_iscontig/mat2_iscontigstride 与形状是否匹配类别张量连续性前置条件只在有价值的输入上启用pad_mm_precondition 通过元信息中的shared_memory与device_capa识别 GPU并为小形状直接返回 FalseA100shared_memory 166912device_capa (8, 0)要求 m、k、n 均 ≥ 512H100shared_memory 232448device_capa (9, 0)要求 m、k、n 均 ≥ 768其他 GPU默认返回 True。这保证 AutoHeuristic 不会把采样预算浪费在显然不需要 pad 的小矩阵上也解释了为什么 pad_mm.py 中当 precondition 不满足时不会收集数据。训练细节决策树与回归树两条路线pad_mm目录同时提供了两种训练入口train_decision_pad_mm.py继承 train_decision.py 的AHTrainDecisionTree学习决策树直接输出pad 还是 origtrain_regression_pad_mm.py继承 train_regression.py 的AHTrainRegressionTree学习回归树为每个 choice 预测一个分数分数高者胜出见 train_regression.py 的类注释。gen_pad_mm_a100.sh默认使用回归路线。无论哪条路线基类 train.py 都会完成读取数据集与元信息 → 用pad_mm_operations()补齐派生特征 → 对类别特征做 one-hot 编码pd.get_dummies→ 划分 train/val/test → 网格搜索超参 → 代码生成。以回归树为例train_regression.py 中的默认网格为max_depths [5, 6, 7]保持树与生成代码足够小min_samples_leafs [1, 2, 5, 10]。最终生成的 Python 启发式文件位于torch/_inductor/autoheuristic/artifacts/_HEURISTIC_NAME.py内含与训练元信息绑定的check_precondition()校验shared_memory与device_capa见 train.py以及树形predict逻辑随编译器一起发布使用。环境与注意事项运行前提数据采集依赖 CUDA GPU、PyTorch 源码环境含torch.compile与 Triton、conda 环境以及tmux并行采集需要见 collect_data.sh 的检查逻辑依赖清单可参考 requirements.txt。启发式与 GPU 绑定生成的启发式带shared_memory/device_capa前置条件不要跨 GPU 型号混用为新 GPU 生成时务必修改HEURISTIC_NAME以免覆盖旧文件。耗时预期collect 阶段在随机形状上做真实编译与计时README 提示可能需要约一天视 GPU 数量而定可用GPU_DEVICE_IDS增加并行度但样本会被均分。验证仓库还提供了 evaluate_pad_mm_heuristics.py 与 test_pad_mm.py 用于评估启发式效果与正确性可在生成后用于对比不同启发式在同一数据集上的命中率。总结从本仓库的 pad_mm 场景可以看到 AutoHeuristic 的完整闭环run_autoheuristic()在编译期构造AutoHeuristic实例 →BenchmarkRunnerPadMM在随机与已知形状上并行采集未对齐输入的实测数据 →train_regression_pad_mm.py结合pad_mm_operations()特征与网格搜索学习回归树 → 代码生成到artifacts/_PadMMA100.py随编译器发布。无论是两条命令重建 A100 启发式还是修改generate_heuristic_pad_mm.sh为新 GPU 从零采集训练这套流程都围绕同一原则用自动调优的离线结果换取编译期的即时决策。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyTorch AutoHeuristic 实战指南：pad_mm 启发式的数据收集与代码生成
PyTorch AutoHeuristic 实战指南pad_mm 启发式的数据收集与代码生成【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchpad_mm 是 PyTorch Inductor 中一项针对矩阵乘法torch.mm的自动填充padding优化它通过在矩阵维度不足对齐时补零来提升 Triton kernel 的执行效率但是否值得填充取决于 GPU 型号与具体形状。AutoHeuristic 是 PyTorch 仓库内置的框架它把自动调优autotuning的结果学习成可随编译器一起发布的决策树/回归树代码。本文以 torchgen/_autoheuristic/pad_mm/README.md 为核心完整讲解如何在本地重建已发布的 A100 启发式、如何为新 GPU 从零收集训练数据并生成 pad_mm 启发式同时结合 pad_mm 的源码实现 与 AutoHeuristic 工具链 说明每一步背后的工作原理。读完本文你将掌握一套可复制的基准测试 → 数据收集 → 训练 → 生成代码完整流程。pad_mm 优化与 AutoHeuristic 框架的关系什么是 pad_mmpad_mm 是 Inductor 图级别graph-level的一项融合优化在torch.mm的两个输入张量维度M、K、N未对齐时通过aten.constant_pad_nd补零到对齐长度使底层 GEMM kernel 能以更规整的 tile 布局执行从而换取吞吐提升。其入口位于 torch/_inductor/fx_passes/pad_mm.py核心安全判断逻辑设备、dtype、静态维度等集中在can_pad()中而是否值得填充则交由 AutoHeuristic 或传统 autotune 决策。其中对齐大小的定义在 pad_mm.py 的get_alignment_size_dtype()中float16/bfloat16对齐大小为 8float32对齐大小为 4其他 dtype返回 0不可填充。AutoHeuristic 如何为 pad_mm 决策在 pad_mm.py 的run_autoheuristic()中构造了一个两选一的AutoHeuristic实例choices [orig, pad]分别代表保持原始形状与填充后执行feedback是一个LocalFeedback对每个 choice 用do_bench实测耗时fallback返回autotune即当启发式不可用时回退到传统 autotuneaugment_contextpad_mm_operations()注入派生特征preconditionpad_mm_precondition决定哪些输入才值得交给 AutoHeuristic见下文特征工程一节。由此可以看出pad_mm 的 AutoHeuristic 场景是典型的本地自动调优feedback 函数能立刻返回实测耗时。学习得到的启发式会被编译进编译器在后续编译同一 GPU 上的同类形状时直接给出决策省去反复 benchmark 的开销。快速重建已发布的 A100 启发式如果只是想复现仓库已发布、基于已有数据训练的 pad_mm A100 启发式只需要依次执行 README.md 给出的两条脚本命令bash get_padmm_dataset.sh # Downloads A100 bash gen_pad_mm_a100.sh # Generates A100 heuristic第一步下载 A100 训练数据集get_padmm_dataset.sh 通过wget下载名为pad_mm_a100_data.zip的数据包再用unzip解压出pad_mm_a100_data.txt并清理压缩包。该 txt 文件即训练数据首行为 JSON 格式的设备元信息shared_memory、device_capa、name后续每一行是一次 autotune 采样记录。第二步训练并生成 A100 启发式gen_pad_mm_a100.sh 调用回归树训练脚本python train_regression_pad_mm.py pad_mm_a100_data.txt --heuristic-name PadMMA100训练完成后启发式代码会被写入torch/_inductor/autoheuristic/artifacts/_PadMMA100.py生成路径的逻辑见 train.py 的write_heuristic_to_file()。此后在编译 pad_mm 时Inductor 会自动加载并使用该启发式。为新 GPU 收集数据并生成启发式完整流程如果目标 GPU 不是 A100或希望用更新的数据重新训练就需要使用 generate_heuristic_pad_mm.sh 走完整的采集 → 训练流程。步骤 0修改脚本中的关键变量在运行脚本前必须先按注释修改脚本顶部的变量变量示例值含义GPU_DEVICE_IDS4,5参与数据采集的 CUDA 设备号逗号分隔多卡会被并行利用CONDA_ENVheuristic-pr运行前需要激活的 conda 环境名NUM_SAMPLES2000计划收集的样本总数会按 GPU 数量均分OUTPUT_DIRa100采集结果与中间文件的输出目录HEURISTIC_NAMEPadMMA100启发式名称若目标不是 A100必须修改此名同时建议修改BENCHMARK_SCRIPT/TRAIN_SCRIPT为对应 GPU 的采集与训练脚本如 H100 可参考仓库中的 gen_pad_mm_h100.shBENCHMARK_SCRIPTgen_data_pad_mm.py采集数据用的 benchmark 脚本TRAIN_SCRIPTtrain_regression_pad_mm.py训练启发式用的脚本脚本随后把上述参数转发给通用流程脚本 generate_heuristic.sh。步骤 1collect —— 跑 benchmark 并收集训练数据bash generate_heuristic_pad_mm.sh collect该命令会在随机输入上执行真实编译与基准测试。README 明确提示视 GPU 数量与算力而定这一步可能耗时约一天因为要对每个采样形状做torch.compiledo_bench实测。collect 模式内部由 collect_data.sh 实现并行化它依赖tmux为GPU_DEVICE_IDS中的每个设备创建一个 tmux pane执行形如conda activate CONDA_ENV python gen_data_pad_mm.py --device id \ -o OUTPUT_DIR/data_id.txt --num-samples NUM_SAMPLES_PER_GPU因此每个 GPU 会生成独立的data_id.txt文件互不干扰。步骤 2generate —— 用收集到的数据学习启发式bash generate_heuristic_pad_mm.sh generategenerate 模式在 generate_heuristic.sh 中完成两件事用 merge_data.py 把各 GPU 的data_id.txt合并为OUTPUT_DIR/HEURISTIC_NAME.txt调用训练脚本python train_regression_pad_mm.py OUTPUT_DIR/HEURISTIC_NAME.txt --heuristic-name HEURISTIC_NAME训练出的启发式同样输出到torch/_inductor/autoheuristic/artifacts/_HEURISTIC_NAME.py。数据收集的源码原理BenchmarkRunnerPadMM训练数据的质量直接决定启发式的泛化能力。pad_mm 的数据采集器 gen_data_pad_mm.py 中的BenchmarkRunnerPadMM类继承自 AutoHeuristic 的 BenchmarkRunner实现了create_input()与run_benchmark()两个核心方法并覆盖了run()来串联整条流水线。形状生成策略已知形状 随机形状形状生成器generate_mm_shapes()分两个阶段产生(m, k, n, dtype)第一阶段先穷举已知形状集合——包括通过--additional-shape-csv传入的外部 CSV格式为M,K,N,dtype支持float16/bfloat16/float32以及 collect_known_mm_shapes.py 收集的常见 mm 形状第二阶段无限生成随机形状其中一半概率使用均匀随机1 到 65536另一半使用 2 的幂与2 的幂之间的混合分布见get_random_dim()未对齐概率默认 0.25。关键过滤只采集未对齐形状因为 pad 只有在维度不齐时才有意义生成器对每个形状做对齐过滤源码 gen_data_pad_mm.pyalign_size get_alignment_size_dtype(dtype) # Skip if all dimensions are aligned (we need unaligned for padding to be relevant) if all(self.is_aligned(dim, align_size) for dim in [m, k, n]): continue # Check if it fits in memory if fits_in_memory(dtype, m, k, n): yield (m, k, n, dtype)同时用fits_in_memory()防止超大形状导致显存溢出。覆盖 padding 的真实成本run_benchmark()中值得一提的设计是prepadded变量概率默认 0.2它模拟输入已经预先填充好的场景通过torch.mm(a 1, b)等方式规避编译期形状传播从而把填充本身的开销排除在计时之外。这样训练数据能区分两种情形输入来自前一层填充成本已被吸收此时 pad 往往更划算输入形状固定填充会产生额外的拷贝开销。这正是启发式需要学习的关键权衡点之一。采集时的配置开关run()方法中通过torch._inductor.config控制采集行为源码 gen_data_pad_mm.pytorch._inductor.config.autoheuristic_use.pad_mm False torch._inductor.config.autoheuristic_collect.pad_mm True torch._inductor.config.autoheuristic_log_path args.o即采集模式下只记录 autotune 结果而不使用启发式。这与 AutoHeuristic 通用文档 README.md 中描述的环境变量用法一致——你也可以在任意触发 pad_mm 的用户脚本中用环境变量完成同样的控制# 采集模式把 autotune 结果写入 train.txt TORCHINDUCTOR_AUTOHEURISTIC_LOG_PATHtrain.txt \ TORCHINDUCTOR_AUTOHEURISTIC_COLLECTpad_mm python run.py # 使用模式加载已学习的启发式 TORCHINDUCTOR_AUTOHEURISTIC_USEpad_mm python run.py其中pad_mm即AutoHeuristic构造时传入的name。特征工程pad_mm 启发式依据什么做决策启发式的预测能力来自AHContext中收集的原始特征以及augment_context中的派生特征。原始特征在 pad_mm.py 的get_context()中注册包括m、k、n、mat1_stride_0/1、mat2_stride_0/1、m/k/n_padded_length、mat1/mat2_align_size、mat1/mat2_dtype类别特征、prepadded_mat1/mat2类别特征、using_tf32类别特征。派生特征定义在 torch/_inductor/autoheuristic/autoheuristic_utils.py 的pad_mm_operations()中训练与推理两侧通过 train_decision_pad_mm.py / train_regression_pad_mm.py 的add_new_features()保持一致地重建特征计算方式类型含义m*k、m*n、k*n维度两两相乘数值捕获计算量规模k/(m*n)k / (m*n)数值捕获瘦长 K 维的极端形态bfloat_perf_hitk m*1024 且 k n*1024 且 dtype 为 bfloat16类别bfloat16 下 K 极端大的性能敏感场景arith_intensitym*k*n / (m*k k*n m*n)数值算术强度衡量访存与计算比mat1_innermost_needs_padding/mat2_innermost_needs_padding由 stride 与 padded_length 推导类别最内层维度是否需要填充num_dims_needs_paddingm/k/n 需要填充的维度个数数值0~3填充工作量m_multiple_2/4/8/16/32等各维度对 2/4/8/16/32 的整除性类别对齐程度细粒度刻画mat1_iscontig/mat2_iscontigstride 与形状是否匹配类别张量连续性前置条件只在有价值的输入上启用pad_mm_precondition 通过元信息中的shared_memory与device_capa识别 GPU并为小形状直接返回 FalseA100shared_memory 166912device_capa (8, 0)要求 m、k、n 均 ≥ 512H100shared_memory 232448device_capa (9, 0)要求 m、k、n 均 ≥ 768其他 GPU默认返回 True。这保证 AutoHeuristic 不会把采样预算浪费在显然不需要 pad 的小矩阵上也解释了为什么 pad_mm.py 中当 precondition 不满足时不会收集数据。训练细节决策树与回归树两条路线pad_mm目录同时提供了两种训练入口train_decision_pad_mm.py继承 train_decision.py 的AHTrainDecisionTree学习决策树直接输出pad 还是 origtrain_regression_pad_mm.py继承 train_regression.py 的AHTrainRegressionTree学习回归树为每个 choice 预测一个分数分数高者胜出见 train_regression.py 的类注释。gen_pad_mm_a100.sh默认使用回归路线。无论哪条路线基类 train.py 都会完成读取数据集与元信息 → 用pad_mm_operations()补齐派生特征 → 对类别特征做 one-hot 编码pd.get_dummies→ 划分 train/val/test → 网格搜索超参 → 代码生成。以回归树为例train_regression.py 中的默认网格为max_depths [5, 6, 7]保持树与生成代码足够小min_samples_leafs [1, 2, 5, 10]。最终生成的 Python 启发式文件位于torch/_inductor/autoheuristic/artifacts/_HEURISTIC_NAME.py内含与训练元信息绑定的check_precondition()校验shared_memory与device_capa见 train.py以及树形predict逻辑随编译器一起发布使用。环境与注意事项运行前提数据采集依赖 CUDA GPU、PyTorch 源码环境含torch.compile与 Triton、conda 环境以及tmux并行采集需要见 collect_data.sh 的检查逻辑依赖清单可参考 requirements.txt。启发式与 GPU 绑定生成的启发式带shared_memory/device_capa前置条件不要跨 GPU 型号混用为新 GPU 生成时务必修改HEURISTIC_NAME以免覆盖旧文件。耗时预期collect 阶段在随机形状上做真实编译与计时README 提示可能需要约一天视 GPU 数量而定可用GPU_DEVICE_IDS增加并行度但样本会被均分。验证仓库还提供了 evaluate_pad_mm_heuristics.py 与 test_pad_mm.py 用于评估启发式效果与正确性可在生成后用于对比不同启发式在同一数据集上的命中率。总结从本仓库的 pad_mm 场景可以看到 AutoHeuristic 的完整闭环run_autoheuristic()在编译期构造AutoHeuristic实例 →BenchmarkRunnerPadMM在随机与已知形状上并行采集未对齐输入的实测数据 →train_regression_pad_mm.py结合pad_mm_operations()特征与网格搜索学习回归树 → 代码生成到artifacts/_PadMMA100.py随编译器发布。无论是两条命令重建 A100 启发式还是修改generate_heuristic_pad_mm.sh为新 GPU 从零采集训练这套流程都围绕同一原则用自动调优的离线结果换取编译期的即时决策。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考