变更摘要

此 PR 将动态 shape 场景下「移除未使用的升维轴（dim-up axes）」的逻辑从 ir.py 的 rebuild_flattened_dims 函数中提取并重构，迁移至 triton.py 的 NPUIndexTritonKernel 类中，作为独立的 _remove_unused_dim_up_axes 方法。重构后的逻辑在代码生成流程中更精确的位置（_transform_schedule_indexing 之后、_record_store_unified_indexing 之前）执行，并增加了对 schedule 节点的遍历、range_tree_nodes 的正确清理（通过 node.parent.remove_entry）、以及残留替换候选的清理（_drop_dead_substitution_candidates），从而避免因过早判定而误删仍被 Store 索引引用的轴。

主要改动

• 将升维轴清理逻辑从 ir.py 迁移至 triton.py：删除了 rebuild_flattened_dims 中基于 kernel.expr_substituted 和 kernel.range_tree_nodes 直接删除未使用升维轴的代码块（共 11 行），改为在 NPUIndexTritonKernel 中新增 _remove_unused_dim_up_axes 方法统一处理。
• 新增 _index_symbol_users 方法：遍历 node_schedule 中所有节点的 _body.indexing，建立符号到索引键的映射；若某节点缺少变换后索引则返回 None，避免基于不完整信息做出错误删除判断。
• 新增 _drop_dead_substitution_candidates 方法：在移除未使用升维轴后，清理 range_tree_nodes_substituted 中那些重建已删除轴的父扩展候选，防止代码生成引用不存在的变量。
• 改进 range_tree_nodes 的删除方式：原实现直接 del kernel.range_tree_nodes[var]，新实现改为通过 node.parent.remove_entry(var) 正确地从父节点中移除条目，确保数据结构一致性。
• 调整执行时机：_remove_unused_dim_up_axes 被插入到 _transform_schedule_indexing 与 _record_store_unified_indexing 之间，确保父轴已展开后再判定使用关系，避免广播维度被误删。

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=42762&platform=gitcode&codeHostingPlatformFlag=gitcode

代码审查

审查完毕

已审查文件

共 2 个文件，均已审查：

发现摘要

总体风险评估：中等。P2 问题在特定场景下（一个 var 有多个替换候选，且仅部分候选引用了被移除的升维轴）会触发，可能导致生成的 Triton kernel 引用未定义变量。修复简单（加一行 else 赋值），建议在合入前修复。

💬 仅评论

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.