已合并
[fix]tiling_key_error #32340
cuiduo创建于 3月25日
[fix]tiling_key_error #32340
已合并
共 3 个文件变更+27-6
| @@ -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): |
| @@ -560,6 +560,21 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 560 | dtype = None | 560 | dtype = None |
| 561 | if axis is None: | 561 | if axis is None: |
| 562 | return None | 562 | return None |
| 563 | + | ||
| 564 | + def _lookup_dim(sym) -> Optional["IterationRangesEntryNPUIndex"]: | ||
| 565 | + dim = self.range_tree_nodes.get(sym) | ||
| 566 | + if dim is not None: | ||
| 567 | + return dim | ||
| 568 | + return self.range_tree_nodes_removed.get(sym) | ||
| 569 | + | ||
| 570 | + def _iter_candidate_syms(key): | ||
| 571 | + # indexing_map keys can be sympy.Symbol (common) or sympy.Expr (e.g. linearized indexing) | ||
| 572 | + # Try direct lookup first, then fall back to free_symbols for Expr. | ||
| 573 | + yield key | ||
| 574 | + if isinstance(key, sympy.Expr) and not isinstance(key, sympy.Symbol): | ||
| 575 | + for s in key.free_symbols: | ||
| 576 | + yield s | ||
| 577 | + | ||
| 563 | for node in self.node_schedule: | 578 | for node in self.node_schedule: |
| 564 | if node in (EnableReduction, DisableReduction): | 579 | if node in (EnableReduction, DisableReduction): |
| 565 | continue | 580 | continue |
| @@ -574,11 +589,14 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 574 | if node in (EnableReduction, DisableReduction): | 589 | if node in (EnableReduction, DisableReduction): |
| 575 | continue | 590 | continue |
| 576 | for key, _ in node._body.indexing_map.items(): | 591 | for key, _ in node._body.indexing_map.items(): |
| 577 | - if key in self.range_tree_nodes: | 592 | + dim = None |
| 578 | - dim = self.range_tree_nodes[key] | 593 | + for cand in _iter_candidate_syms(key): |
| 579 | - else: | 594 | + dim = _lookup_dim(cand) |
| 580 | - dim = self.range_tree_nodes_removed[key] | 595 | + if dim is not None: |
| 581 | - | 596 | + break |
| 597 | + if dim is None: | ||
| 598 | + continue | ||
| 599 | + | ||
| 582 | if dim.parent == axis.parent: | 600 | if dim.parent == axis.parent: |
| 583 | dtype = V.graph.get_dtype(node.node.name) | 601 | dtype = V.graph.get_dtype(node.node.name) |
| 584 | should_break_all = True | 602 | should_break_all = True |
| @@ -78,7 +78,9 @@ 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, | ||
| 82 | ] | 84 | ] |
| 83 | 85 | ||
| 84 | GENERATE_LIST2 = [ | 86 | GENERATE_LIST2 = [ |