已合并
多项式问题修复 #44410
zzll创建于 20 天前
多项式问题修复 #44410
已合并
共 2 个文件变更+171-14
| @@ -3,9 +3,83 @@ import torch | |||
| 3 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | 3 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests |
| 4 | from testutils import TestUtils | 4 | from testutils import TestUtils |
| 5 | import torch_npu | 5 | import torch_npu |
| 6 | +import sympy | ||
| 7 | +from sympy import Symbol | ||
| 8 | +from torch.utils._sympy.functions import FloorDiv, ModularIndexing | ||
| 9 | +from torch_npu._inductor.codegen.ir import analyze_floordiv_expression, analyze_modular_expression | ||
| 10 | +from torch._inductor.virtualized import V | ||
| 11 | +from torch.utils._sympy.value_ranges import ValueRanges, int_oo | ||
| 12 | +from types import SimpleNamespace | ||
| 13 | +from unittest.mock import patch | ||
| 14 | +from unittest.mock import MagicMock | ||
| 6 | 15 | ||
| 7 | 16 | ||
| 8 | -class TestUnifiedAxis(TestUtils): | 17 | +class MockRangeNode: |
| 18 | + def __init__(self, length): | ||
| 19 | + self.length = length | ||
| 20 | + | ||
| 21 | + | ||
| 22 | +def make_range_tree_nodes(mapping): | ||
| 23 | + """mapping: {symbol: length}""" | ||
| 24 | + return {sym: MockRangeNode(length) for sym, length in mapping.items()} | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +class TestLinearDynamic(TestUtils): | ||
| 28 | + | ||
| 29 | + def test_analyze_floordiv_expression(self): | ||
| 30 | + x0 = Symbol("x0", integer=True, nonnegative=True) | ||
| 31 | + x1 = Symbol("x1", integer=True, nonnegative=True) | ||
| 32 | + s0 = Symbol("s0", integer=True, positive=True) | ||
| 33 | + s1 = Symbol("s1", integer=True, positive=True) | ||
| 34 | + s2 = Symbol("s2", integer=True, positive=True) | ||
| 35 | + | ||
| 36 | + nodes = make_range_tree_nodes({x0: s0, x1: s1}) | ||
| 37 | + expr = FloorDiv(s1 * x0 + x1, s0*s1) | ||
| 38 | + res = analyze_floordiv_expression(expr, nodes) | ||
| 39 | + print("\n[Case3] symbolic max contains divisor symbols → remainder 0 path") | ||
| 40 | + print(res) | ||
| 41 | + assert "can_split" in res and res["can_split"] == True | ||
| 42 | + | ||
| 43 | + | ||
| 44 | + def test_analyze_modular_expression(self): | ||
| 45 | + x0 = Symbol("x0", integer=True, nonnegative=True) | ||
| 46 | + x1 = Symbol("x1", integer=True, nonnegative=True) | ||
| 47 | + s0 = Symbol("s0", integer=True, positive=True) | ||
| 48 | + s1 = Symbol("s1", integer=True, positive=True) | ||
| 49 | + s2 = Symbol("s2", integer=True, positive=True) | ||
| 50 | + | ||
| 51 | + nodes = make_range_tree_nodes({x0: s0, x1: s1}) | ||
| 52 | + expr = ModularIndexing(s1 * x0 + x1, 1, s0*s1) | ||
| 53 | + res = analyze_modular_expression(expr, nodes) | ||
| 54 | + print("\n[Case3] symbolic max contains divisor symbols → remainder 0 path") | ||
| 55 | + print(res) | ||
| 56 | + assert "can_split" in res and res["can_split"] == True | ||
| 57 | + | ||
| 58 | + | ||
| 59 | + def test_analyze_modular_expression_mod_is_symbol(self): | ||
| 60 | + x0 = Symbol("x0", integer=True, nonnegative=True) | ||
| 61 | + x1 = Symbol("x1", integer=True, nonnegative=True) | ||
| 62 | + s0 = Symbol("s0", integer=True, positive=True) | ||
| 63 | + s1 = Symbol("s1", integer=True, positive=True) | ||
| 64 | + s2 = Symbol("s2", integer=True, positive=True) | ||
| 65 | + s3 = Symbol("s3", integer=True, positive=True) | ||
| 66 | + mock_kernel = MagicMock() | ||
| 67 | + | ||
| 68 | + mock_kernel.symbol_range_map = { | ||
| 69 | + "s0": MagicMock(lower=1), | ||
| 70 | + "s1": MagicMock(lower=1), | ||
| 71 | + "s2": MagicMock(lower=1), | ||
| 72 | + "s3": MagicMock(lower=1), | ||
| 73 | + } | ||
| 74 | + | ||
| 75 | + with V.set_kernel_handler(mock_kernel): | ||
| 76 | + nodes = make_range_tree_nodes({x0: s3*s2, x1: s2}) | ||
| 77 | + expr = ModularIndexing((x0 + x1*s3), s2, s3) # 这里 divisor=s0，第一项 coeff*length = s1*s0 | ||
| 78 | + res = analyze_modular_expression(expr, nodes) | ||
| 79 | + print("\n[Case3] symbolic max contains divisor symbols → remainder 0 path") | ||
| 80 | + print(res) | ||
| 81 | + assert "can_split" in res and res["can_split"] == True | ||
| 82 | + | ||
| 9 | 83 | ||
| 10 | def op_calc_dynamic(self, x, y, batch_size, seq_len, hidden1, hidden2, dim1, dim2): | 84 | def op_calc_dynamic(self, x, y, batch_size, seq_len, hidden1, hidden2, dim1, dim2): |
| 11 | view_1 = x.view(batch_size, seq_len, hidden1, dim1).permute(0, 1, 3, 2).reshape(batch_size, seq_len, hidden1*dim1) | 85 | view_1 = x.view(batch_size, seq_len, hidden1, dim1).permute(0, 1, 3, 2).reshape(batch_size, seq_len, hidden1*dim1) |
| @@ -75,7 +149,7 @@ class TestUnifiedAxis(TestUtils): | |||
| 75 | self.assertEqual(std_result, inductor_result, atol=1e-2, rtol=1e-2) | 149 | self.assertEqual(std_result, inductor_result, atol=1e-2, rtol=1e-2) |
| 76 | 150 | ||
| 77 | 151 | ||
| 78 | -instantiate_parametrized_tests(TestUnifiedAxis) | 152 | +instantiate_parametrized_tests(TestLinearDynamic) |
| 79 | 153 | ||
| 80 | 154 | ||
| 81 | if __name__ == "__main__": | 155 | if __name__ == "__main__": |
| @@ -556,7 +556,6 @@ def calculate_max_remainder(coeff, length, divisor_or_mod): | |||
| 556 | # 证明divisor_symbols真包含max_length_value，即divisor_or_mod是coeff*length的整数倍，所以(coeff*length) % divisor_or_mod的结果是(coeff*length) | 556 | # 证明divisor_symbols真包含max_length_value，即divisor_or_mod是coeff*length的整数倍，所以(coeff*length) % divisor_or_mod的结果是(coeff*length) |
| 557 | return coeff*length | 557 | return coeff*length |
| 558 | else: | 558 | else: |
| 559 | - breakpoint() | ||
| 560 | if ( | 559 | if ( |
| 561 | isinstance(max_length_value, sympy.Symbol) | 560 | isinstance(max_length_value, sympy.Symbol) |
| 562 | and isinstance(divisor_or_mod, sympy.Symbol) | 561 | and isinstance(divisor_or_mod, sympy.Symbol) |
| @@ -745,11 +744,30 @@ def analyze_floordiv_expression(expr, range_tree_nodes: dict) -> dict: | |||
| 745 | result["reason"] = ( | 744 | result["reason"] = ( |
| 746 | f"expr can not split, max_remainder_sum {max_remainder_sum} >= divisor {divisor}" | 745 | f"expr can not split, max_remainder_sum {max_remainder_sum} >= divisor {divisor}" |
| 747 | ) | 746 | ) |
| 747 | + elif ( | ||
| 748 | + max_remainder_sum == 0 | ||
| 749 | + and isinstance(divisor, sympy.Symbol) | ||
| 750 | + and str(divisor) in V.kernel.symbol_range_map.keys() | ||
| 751 | + ): | ||
| 752 | + lower_value = V.kernel.symbol_range_map.get(str(divisor)).lower | ||
| 753 | + if lower_value >= max_remainder_sum: | ||
| 754 | + result["can_split"] = True | ||
| 755 | + result["reason"] = ( | ||
| 756 | + f"expr can split, max_remainder_sum {max_remainder_sum} < divisor {divisor}, divisor lower_value={lower_value}" | ||
| 757 | + ) | ||
| 758 | + split_terms = [] | ||
| 759 | + for term in add_terms: | ||
| 760 | + split_terms.append(f"({term} // {divisor})") | ||
| 761 | + result["split_form"] = " + ".join(split_terms) | ||
| 762 | + else: | ||
| 763 | + # 不能拆分 | ||
| 764 | + result["reason"] = ( | ||
| 765 | + f"expr can not split, max_remainder_sum {max_remainder_sum} >= divisor {divisor}" | ||
| 766 | + ) | ||
| 748 | elif ( | 767 | elif ( |
| 749 | hasattr(max_remainder_sum, "free_symbols") | 768 | hasattr(max_remainder_sum, "free_symbols") |
| 750 | and hasattr(divisor, "free_symbols") | 769 | and hasattr(divisor, "free_symbols") |
| 751 | ): | 770 | ): |
| 752 | - # breakpoint() | ||
| 753 | max_remainder_sum_str = str(max_remainder_sum) | 771 | max_remainder_sum_str = str(max_remainder_sum) |
| 754 | divisor_str = str(divisor) | 772 | divisor_str = str(divisor) |
| 755 | if max_remainder_sum_str in divisor_str: | 773 | if max_remainder_sum_str in divisor_str: |
| @@ -855,10 +873,34 @@ def analyze_modular_expression(expr, range_tree_nodes: dict) -> dict: | |||
| 855 | add_terms = expr_to_mod.args | 873 | add_terms = expr_to_mod.args |
| 856 | # 优先判断expr_to_mod//lower是否可拆分 | 874 | # 优先判断expr_to_mod//lower是否可拆分 |
| 857 | term_details = [] | 875 | term_details = [] |
| 858 | - # breakpoint() | ||
| 859 | sub_expr = FloorDiv(expr_to_mod, lower) | 876 | sub_expr = FloorDiv(expr_to_mod, lower) |
| 860 | result = analyze_floordiv_expression(sub_expr, range_tree_nodes) | 877 | result = analyze_floordiv_expression(sub_expr, range_tree_nodes) |
| 861 | - return result | 878 | + if ("can_split" in result |
| 879 | + and "split_form" in result | ||
| 880 | + and result["can_split"] | ||
| 881 | + ): | ||
| 882 | + sub_exprs = [] | ||
| 883 | + for term in add_terms: | ||
| 884 | + sub_exprs.append(FloorDiv(term, lower)) | ||
| 885 | + dividen_expression = sympy.Add(*sub_exprs) | ||
| 886 | + lower = sympy.Integer(1) | ||
| 887 | + expression = ModularIndexing(dividen_expression, lower, upper) | ||
| 888 | + result = { | ||
| 889 | + "expression": str(expression), | ||
| 890 | + "type": "ModularIndexing", | ||
| 891 | + "can_split": False, | ||
| 892 | + "reason": "", | ||
| 893 | + "details": {}, | ||
| 894 | + "split_form": "", | ||
| 895 | + } | ||
| 896 | + result["details"]["expr_to_mod"] = dividen_expression | ||
| 897 | + result["details"]["lower"] = lower | ||
| 898 | + result["details"]["upper"] = upper | ||
| 899 | + free_symbols = dividen_expression.free_symbols | ||
| 900 | + num_symbols = len(free_symbols) | ||
| 901 | + result["details"]["num_symbols"] = num_symbols | ||
| 902 | + result["details"]["symbols"] = list(free_symbols) | ||
| 903 | + expr_to_mod = dividen_expression | ||
 | |||
| 862 | # For ModularIndexing, the split condition is: | 904 | # For ModularIndexing, the split condition is: |
| 863 | # (expr1 % mod) + (expr2 % mod) < mod | 905 | # (expr1 % mod) + (expr2 % mod) < mod |
| 864 | # where mod = upper - lower + 1 | 906 | # where mod = upper - lower + 1 |
| @@ -894,7 +936,12 @@ def analyze_modular_expression(expr, range_tree_nodes: dict) -> dict: | |||
| 894 | # Extract coefficient and symbol | 936 | # Extract coefficient and symbol |
| 895 | coeff = 1 | 937 | coeff = 1 |
| 896 | symbol = None | 938 | symbol = None |
| 897 | - | 939 | + if isinstance(term, FloorDiv): |
| 940 | + numerator, denominator = term.args | ||
| 941 | + term = sympy.Mul( | ||
| 942 | + numerator, | ||
| 943 | + sympy.Pow(denominator, -1), | ||
| 944 | + ) | ||
| 898 | if isinstance(term, sympy.Symbol): | 945 | if isinstance(term, sympy.Symbol): |
| 899 | symbol = term | 946 | symbol = term |
| 900 | elif isinstance(term, sympy.Mul): | 947 | elif isinstance(term, sympy.Mul): |
| @@ -965,11 +1012,30 @@ def analyze_modular_expression(expr, range_tree_nodes: dict) -> dict: | |||
| 965 | result["reason"] = ( | 1012 | result["reason"] = ( |
| 966 | f"expr can not split, max_remainder_sum {max_remainder_sum} >= mod {mod}" | 1013 | f"expr can not split, max_remainder_sum {max_remainder_sum} >= mod {mod}" |
| 967 | ) | 1014 | ) |
| 1015 | + elif ( | ||
| 1016 | + max_remainder_sum == 0 | ||
| 1017 | + and isinstance(mod, sympy.Symbol) | ||
| 1018 | + and str(mod) in V.kernel.symbol_range_map.keys() | ||
| 1019 | + ): | ||
| 1020 | + lower_value = V.kernel.symbol_range_map.get(str(mod)).lower | ||
| 1021 | + if lower_value >= max_remainder_sum: | ||
| 1022 | + result["can_split"] = True | ||
| 1023 | + result["reason"] = ( | ||
| 1024 | + f"expr can split, max_remainder_sum {max_remainder_sum} < mod {mod}, mod lower_value={lower_value}" | ||
| 1025 | + ) | ||
| 1026 | + split_terms = [] | ||
| 1027 | + for term in add_terms: | ||
| 1028 | + split_terms.append(f"ModularIndexing({term}, {lower}, {upper})") | ||
| 1029 | + result["split_form"] = " + ".join(split_terms) | ||
| 1030 | + else: | ||
| 1031 | + # 不能拆分 | ||
| 1032 | + result["reason"] = ( | ||
| 1033 | + f"expr can not split, max_remainder_sum {max_remainder_sum} >= mod {mod}" | ||
| 1034 | + ) | ||
| 968 | elif ( | 1035 | elif ( |
| 969 | hasattr(max_remainder_sum, "free_symbols") | 1036 | hasattr(max_remainder_sum, "free_symbols") |
| 970 | and hasattr(mod, "free_symbols") | 1037 | and hasattr(mod, "free_symbols") |
| 971 | ): | 1038 | ): |
| 972 | - # breakpoint() | ||
| 973 | max_remainder_sum_str = str(max_remainder_sum) | 1039 | max_remainder_sum_str = str(max_remainder_sum) |
| 974 | mod_str = str(mod) | 1040 | mod_str = str(mod) |
| 975 | if max_remainder_sum_str in mod_str: | 1041 | if max_remainder_sum_str in mod_str: |
| @@ -1205,12 +1271,29 @@ def split_expression(expr): | |||
| 1205 | # Split: ModularIndexing(a+b, lower, upper) -> | 1271 | # Split: ModularIndexing(a+b, lower, upper) -> |
| 1206 | # ModularIndexing(a, lower, upper) + ModularIndexing(b, lower, upper) | 1272 | # ModularIndexing(a, lower, upper) + ModularIndexing(b, lower, upper) |
| 1207 | split_terms = [] | 1273 | split_terms = [] |
| 1208 | - for term in expr_to_mod.args: | 1274 | + if ( |
| 1209 | - new_mod = ModularIndexing(term, lower, upper) | 1275 | + isinstance(lower, sympy.Symbol) |
| 1210 | - # ModularIndexing(16*z0, 1, 128) -> 16*ModularIndexing(z0, 1, 8) | 1276 | + or ( |
| 1211 | - new_mod = eliminate_modular(new_mod) | 1277 | + hasattr(lower, "free_symbols") |
| 1212 | - new_mod = extract_modular_indexing_coefficient(new_mod) | 1278 | + and len(lower["free_symbols"]) > 0 |
| 1213 | - split_terms.append(new_mod) | 1279 | + ) |
| 1280 | + ): | ||
| 1281 | + temp_splits = [] | ||
| 1282 | + for term in expr_to_mod.args: | ||
| 1283 | + temp_new_mod = FloorDiv(term, lower) | ||
| 1284 | + temp_splits.append(temp_new_mod) | ||
| 1285 | + for term in temp_splits: | ||
| 1286 | + new_mod = ModularIndexing(term, 1, upper) | ||
| 1287 | + new_mod = eliminate_modular(new_mod) | ||
| 1288 | + new_mod = extract_modular_indexing_coefficient(new_mod) | ||
| 1289 | + split_terms.append(new_mod) | ||
| 1290 | + else: | ||
| 1291 | + for term in expr_to_mod.args: | ||
| 1292 | + new_mod = ModularIndexing(term, lower, upper) | ||
| 1293 | + # ModularIndexing(16*z0, 1, 128) -> 16*ModularIndexing(z0, 1, 8) | ||
| 1294 | + new_mod = eliminate_modular(new_mod) | ||
| 1295 | + new_mod = extract_modular_indexing_coefficient(new_mod) | ||
| 1296 | + split_terms.append(new_mod) | ||
| 1214 | return sympy.Add(*split_terms) | 1297 | return sympy.Add(*split_terms) |
| 1215 | else: | 1298 | else: |
| 1216 | new_mod = ModularIndexing(expr_to_mod, lower, upper) | 1299 | new_mod = ModularIndexing(expr_to_mod, lower, upper) |
🟡 Medium Priority
在
analyze_modular_expression中，当lower为符号类型时（第 869-906 行），第 880 行将result覆写为analyze_floordiv_expression的返回值（其type为"FloorDiv"）。旧代码在此处直接return result，本次 diff 删除了该return并新增了条件块（第 881-906 行）：当can_split=True时正确重建了type: "ModularIndexing"的 result；但当can_split=False时跳过该块，使 result 仍保留 FloorDiv 的 type/expression 字段，随后落入第 907 行之后的模块化分析逻辑（lines 1008-1087），result 被混合修改后返回。调用方（第 508 行）将该结果追加到modular_expressions列表中，但 dict 内type仍为"FloorDiv"，形成类型错乱。changed line: 880 行
result = analyze_floordiv_expression(...)覆写了analyze_modular_expression的 result，diff 新增的条件块仅在 can_split=True 时修复，can_split=False 路径残留 FloorDiv type → 调用方collect_expressions（第 508 行）将 type="FloorDiv" 的结果放入modular_expressions→ 下游消费者若依赖 type 字段判断将得到错误分类 → 修复：在进入第 880 行前用局部变量保存 floordiv 分析结果，或在该分支末尾无条件重建 ModularIndexing 类型的 result。建议：将第 880 行的 floordiv 结果保存到局部变量（如
floordiv_result）而非直接覆写result，然后在条件块内根据 floordiv_result 判定；若 can_split=False 且需要继续模块化分析，应先重建 ModularIndexing 类型的 result 再继续后续逻辑。