已合并
[fix]tiling_key_error #32336
cuiduo创建于 3月25日
[fix]tiling_key_error #32336
已合并
共 4 个文件变更+27-5
| @@ -6,6 +6,7 @@ from torch.testing._internal.common_utils import ( | |||
| 6 | ) | 6 | ) |
| 7 | from testutils import TestUtils | 7 | from testutils import TestUtils |
| 8 | import torch_npu | 8 | import torch_npu |
| 9 | +import torch_npu._inductor | ||
| 9 | 10 | ||
| 10 | 11 | ||
| 11 | class TestResizeAs(TestUtils): | 12 | class TestResizeAs(TestUtils): |
| @@ -786,6 +786,21 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 786 | dtype = None | 786 | dtype = None |
| 787 | if axis is None: | 787 | if axis is None: |
| 788 | return None | 788 | return None |
| 789 | + | ||
| 790 | + def _lookup_dim(sym) -> Optional["IterationRangesEntryNPUIndex"]: | ||
| 791 | + dim = self.range_tree_nodes.get(sym) | ||
| 792 | + if dim is not None: | ||
| 793 | + return dim | ||
| 794 | + return self.range_tree_nodes_removed.get(sym) | ||
| 795 | + | ||
| 796 | + def _iter_candidate_syms(key): | ||
| 797 | + # indexing_map keys can be sympy.Symbol (common) or sympy.Expr (e.g. linearized indexing) | ||
| 798 | + # Try direct lookup first, then fall back to free_symbols for Expr. | ||
| 799 | + yield key | ||
| 800 | + if isinstance(key, sympy.Expr) and not isinstance(key, sympy.Symbol): | ||
| 801 | + for s in key.free_symbols: | ||
| 802 | + yield s | ||
| 803 | + | ||
| 789 | for node in self.node_schedule: | 804 | for node in self.node_schedule: |
| 790 | if node in (EnableReduction, DisableReduction): | 805 | if node in (EnableReduction, DisableReduction): |
| 791 | continue | 806 | continue |
| @@ -800,10 +815,13 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 800 | if node in (EnableReduction, DisableReduction): | 815 | if node in (EnableReduction, DisableReduction): |
| 801 | continue | 816 | continue |
| 802 | for key, _ in node._body.indexing_map.items(): | 817 | for key, _ in node._body.indexing_map.items(): |
| 803 | - if key in self.range_tree_nodes: | 818 | + dim = None |
 | |||
| 804 | - dim = self.range_tree_nodes[key] | 819 | + for cand in _iter_candidate_syms(key): |
| 805 | - else: | 820 | + dim = _lookup_dim(cand) |
| 806 | - dim = self.range_tree_nodes_removed[key] | 821 | + if dim is not None: |
| 822 | + break | ||
| 823 | + if dim is None: | ||
| 824 | + continue | ||
| 807 | 825 | ||
| 808 | if dim.parent == axis.parent: | 826 | if dim.parent == axis.parent: |
| 809 | dtype = V.graph.get_dtype(node.node.name) | 827 | dtype = V.graph.get_dtype(node.node.name) |
| @@ -4,7 +4,7 @@ from ..utils.check_mode import is_inference_check | |||
| 4 | from .register_custom_pass import ASCEND_CUSTOME_PASS_REGISTER | 4 | from .register_custom_pass import ASCEND_CUSTOME_PASS_REGISTER |
| 5 | from ..utils.fx_pass_level import FxPassLevel, PassType | 5 | from ..utils.fx_pass_level import FxPassLevel, PassType |
| 6 | from ...config import log | 6 | from ...config import log |
| 7 | - | 7 | +from torch._inductor.pattern_matcher import stable_topological_sort |
| 8 | 8 | ||
| 9 | for _, module_name, _ in pkgutil.iter_modules(__path__): | 9 | for _, module_name, _ in pkgutil.iter_modules(__path__): |
| 10 | importlib.import_module(f"{__name__}.{module_name}") | 10 | importlib.import_module(f"{__name__}.{module_name}") |
| @@ -28,6 +28,7 @@ def run_register_pre_custom_passes(gm): | |||
| 28 | def run_register_post_custom_passes(gm): | 28 | def run_register_post_custom_passes(gm): |
| 29 | log.debug(f"before post_grad graph optimizer pass, graph is: {gm}") | 29 | log.debug(f"before post_grad graph optimizer pass, graph is: {gm}") |
| 30 | if is_inference_check(): | 30 | if is_inference_check(): |
| 31 | + stable_topological_sort(gm) | ||
| 31 | for level in sorted(FxPassLevel): | 32 | for level in sorted(FxPassLevel): |
| 32 | for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.POST][level]: | 33 | for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.POST][level]: |
| 33 | fn(gm) | 34 | fn(gm) |
可将if dim is None修改为else，不用预定义dim=None