已合并
[fix]tiling_key_error #32356
cuiduo创建于 3月25日
[fix]tiling_key_error #32356
已合并
共 4 个文件变更+29-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): |
| @@ -546,6 +546,21 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 546 | dtype = None | 546 | dtype = None |
| 547 | if axis is None: | 547 | if axis is None: |
| 548 | return None | 548 | return None |
| 549 | + | ||
| 550 | + def _lookup_dim(sym) -> Optional["IterationRangesEntryNPUIndex"]: | ||
| 551 | + dim = self.range_tree_nodes.get(sym) | ||
| 552 | + if dim is not None: | ||
| 553 | + return dim | ||
| 554 | + return self.range_tree_nodes_removed.get(sym) | ||
| 555 | + | ||
| 556 | + def _iter_candidate_syms(key): | ||
| 557 | + # indexing_map keys can be sympy.Symbol (common) or sympy.Expr (e.g. linearized indexing) | ||
| 558 | + # Try direct lookup first, then fall back to free_symbols for Expr. | ||
| 559 | + yield key | ||
| 560 | + if isinstance(key, sympy.Expr) and not isinstance(key, sympy.Symbol): | ||
| 561 | + for s in key.free_symbols: | ||
| 562 | + yield s | ||
| 563 | + | ||
| 549 | for node in self.node_schedule: | 564 | for node in self.node_schedule: |
| 550 | if node in (EnableReduction, DisableReduction): | 565 | if node in (EnableReduction, DisableReduction): |
| 551 | continue | 566 | continue |
| @@ -560,10 +575,13 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 560 | if node in (EnableReduction, DisableReduction): | 575 | if node in (EnableReduction, DisableReduction): |
| 561 | continue | 576 | continue |
| 562 | for key, _ in node._body.indexing_map.items(): | 577 | for key, _ in node._body.indexing_map.items(): |
| 563 | - if key in self.range_tree_nodes: | 578 | + dim = None |
| 564 | - dim = self.range_tree_nodes[key] | 579 | + for cand in _iter_candidate_syms(key): |
| 565 | - else: | 580 | + dim = _lookup_dim(cand) |
| 566 | - dim = self.range_tree_nodes_removed[key] | 581 | + if dim is not None: |
| 582 | + break | ||
| 583 | + if dim is None: | ||
| 584 | + continue | ||
| 567 | 585 | ||
| 568 | if dim.parent == axis.parent: | 586 | if dim.parent == axis.parent: |
| 569 | dtype = V.graph.get_dtype(node.node.name) | 587 | dtype = V.graph.get_dtype(node.node.name) |
| @@ -4,6 +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 | +from torch._inductor.pattern_matcher import stable_topological_sort | ||
| 7 | 8 | ||
| 8 | 9 | ||
| 9 | for _, module_name, _ in pkgutil.iter_modules(__path__): | 10 | for _, module_name, _ in pkgutil.iter_modules(__path__): |
| @@ -28,6 +29,7 @@ def run_register_pre_custom_passes(gm): | |||
| 28 | def run_register_post_custom_passes(gm): | 29 | def run_register_post_custom_passes(gm): |
| 29 | log.debug(f"before post_grad graph optimizer pass, graph is: {gm}") | 30 | log.debug(f"before post_grad graph optimizer pass, graph is: {gm}") |
| 30 | if is_inference_check(): | 31 | if is_inference_check(): |
| 32 | + stable_topological_sort(gm) | ||
| 31 | for level in sorted(FxPassLevel): | 33 | for level in sorted(FxPassLevel): |
| 32 | for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.POST][level]: | 34 | for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.POST][level]: |
| 33 | fn(gm) | 35 | fn(gm) |
| @@ -78,7 +78,10 @@ GENERATE_LIST = [ | |||
| 78 | aten.reciprocal, | 78 | aten.reciprocal, |
| 79 | aten._assert_scalar, | 79 | aten._assert_scalar, |
| 80 | triton_kernel_wrapper_mutation, | 80 | triton_kernel_wrapper_mutation, |
| 81 | - torch.ops.higher_order.invoke_subgraph | 81 | + torch.ops.higher_order.invoke_subgraph, |
| 82 | + torch.ops._inductor_test.realize, | ||
| 83 | + torch.ops._inductor_test.realize.default, | ||
| 84 | + | ||
| 82 | ] | 85 | ] |
| 83 | 86 | ||
| 84 | GENERATE_LIST2 = [ | 87 | GENERATE_LIST2 = [ |