已合并
dtype optimal pass upload #29585
dezheng889创建于 1月15日
dtype optimal pass upload #29585
已合并
共 4 个文件变更+228-144
| @@ -0,0 +1,171 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch.fx as fx | ||
| 3 | +from torch.testing._internal.common_utils import ( | ||
| 4 | + run_tests, parametrize, instantiate_parametrized_tests | ||
| 5 | +) | ||
| 6 | +from testutils import TestUtils | ||
| 7 | +from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import dtype_optimal_pass | ||
| 8 | + | ||
| 9 | + | ||
| 10 | +class DtypeOptimalPass(TestUtils): | ||
| 11 | + def create_gm(self, fn): | ||
| 12 | + """辅助函数: 创建 FX GraphModule (需启用 fake mode 以填充 meta['example_value'])""" | ||
| 13 | + gm = fx.symbolic_trace(fn) | ||
| 14 | + return gm | ||
| 15 | + | ||
| 16 | + def create_manual_gm_for_arange(self, start=0, end=2048, step=1, dtype=torch.int64): | ||
| 17 | + """手动创建 FX GM，避免常量折叠""" | ||
| 18 | + gm = fx.GraphModule({}, fx.Graph()) | ||
| 19 | + arange_node = gm.graph.call_function( | ||
| 20 | + torch.arange, | ||
| 21 | + args=(end,), | ||
| 22 | + kwargs={'start': start, 'step': step, 'dtype': dtype} | ||
| 23 | + ) | ||
| 24 | + gm.graph.output(arange_node) | ||
| 25 | + gm.recompile() | ||
| 26 | + return gm | ||
| 27 | + | ||
| 28 | + def apply_pass(self, gm): | ||
| 29 | + """应用 pass 并返回修改后的 gm""" | ||
| 30 | + dtype_optimal_pass(gm.graph) | ||
| 31 | + gm.recompile() | ||
| 32 | + return gm | ||
| 33 | + | ||
| 34 | + def test_static_safe_conversion_single_arg(self): | ||
| 35 | + """使用手动 GM 避免折叠""" | ||
| 36 | + gm = self.create_manual_gm_for_arange(end=2048) | ||
| 37 | + gm = self.apply_pass(gm) | ||
| 38 | + modified_readable = gm.print_readable() | ||
| 39 | + self.assertIn("dtype = torch.int32", modified_readable) | ||
| 40 | + self.assertIn("arange", modified_readable) # 确认无折叠 | ||
| 41 | + | ||
| 42 | + def test_static_safe_conversion_multi_arg(self): | ||
| 43 | + """测试静态多参数，范围安全: arange(1000, 3000, 2, dtype=int64) -> int32""" | ||
| 44 | + gm = self.create_manual_gm_for_arange(start=1000, end=3000, step=2) | ||
| 45 | + gm = self.apply_pass(gm) | ||
| 46 | + modified_readable = gm.print_readable() | ||
| 47 | + self.assertIn("dtype = torch.int32", modified_readable) | ||
| 48 | + | ||
| 49 | + def test_static_fallback_out_of_range(self): | ||
| 50 | + """测试静态范围超出: arange(2**31 + 1, dtype=int64) -> 保留 int64""" | ||
| 51 | + gm = self.create_manual_gm_for_arange(end=2**31 + 1) | ||
| 52 | + gm = self.apply_pass(gm) | ||
| 53 | + modified_readable = gm.print_readable() | ||
| 54 | + self.assertIn("dtype = torch.int64", modified_readable) # fallback | ||
| 55 | + | ||
| 56 | + def test_negative_step_safe(self): | ||
| 57 | + """测试负 step，范围安全: arange(2048, 0, -1, dtype=int64) -> int32""" | ||
| 58 | + gm = self.create_manual_gm_for_arange(start=2048, end=0, step=-1) | ||
| 59 | + gm = self.apply_pass(gm) | ||
| 60 | + modified_readable = gm.print_readable() | ||
| 61 | + self.assertIn("dtype = torch.int32", modified_readable) | ||
| 62 | + | ||
| 63 | + def test_zero_step_skip(self): | ||
| 64 | + """测试 zero step: 跳过转换""" | ||
| 65 | + gm = self.create_manual_gm_for_arange(start=0, end=10, step=0) | ||
| 66 | + gm = self.apply_pass(gm) | ||
| 67 | + modified_readable = gm.print_readable() | ||
| 68 | + self.assertIn("dtype = torch.int64", modified_readable) # 保留 | ||
| 69 | + | ||
| 70 | + def test_non_integer_step_fallback(self): | ||
| 71 | + """测试非整数 step: fallback""" | ||
| 72 | + gm = self.create_manual_gm_for_arange(start=0, end=10, step=0.5) | ||
| 73 | + gm = self.apply_pass(gm) | ||
| 74 | + modified_readable = gm.print_readable() | ||
| 75 | + self.assertIn("dtype = torch.int64", modified_readable) # fallback | ||
| 76 | + | ||
| 77 | + def test_non_int64_dtype_skip(self): | ||
| 78 | + """测试非 int64 dtype: 跳过""" | ||
| 79 | + gm = self.create_manual_gm_for_arange(end=2048, dtype=torch.float32) | ||
| 80 | + gm = self.apply_pass(gm) | ||
| 81 | + modified_readable = gm.print_readable() | ||
| 82 | + self.assertIn("dtype = torch.float32", modified_readable) # 未变 | ||
| 83 | + | ||
| 84 | + def test_conversion_float32_to_int64(self): | ||
| 85 | + """测试 float32 输入的 to(int64) -> to(int32)""" | ||
| 86 | + def fn(x): | ||
| 87 | + return x.to(torch.int64) # x 是 float32 | ||
| 88 | + | ||
| 89 | + gm = self.create_gm(fn) | ||
| 90 | + # 模拟 meta['example_value'] (在 fake mode 下自动填充) | ||
| 91 | + input_node = next(n for n in gm.graph.nodes if n.op == 'placeholder') | ||
| 92 | + input_node.meta['example_value'] = torch.tensor(1.0, dtype=torch.float32) # 模拟 float32 输入 | ||
| 93 | + | ||
| 94 | + original_readable = gm.print_readable() | ||
| 95 | + self.assertIn("to(torch.int64)", original_readable) | ||
| 96 | + | ||
| 97 | + gm = self.apply_pass(gm) | ||
| 98 | + modified_readable = gm.print_readable() | ||
| 99 | + self.assertIn("to(torch.int32)", modified_readable) | ||
| 100 | + | ||
| 101 | + def test_conversion_bool_to_int64(self): | ||
| 102 | + """测试 bool 输入的 to(int64) -> to(int32)""" | ||
| 103 | + def fn(x): | ||
| 104 | + return x.to(torch.int64) # x 是 bool | ||
| 105 | + | ||
| 106 | + gm = self.create_gm(fn) | ||
| 107 | + input_node = next(n for n in gm.graph.nodes if n.op == 'placeholder') | ||
| 108 | + input_node.meta['example_value'] = torch.tensor(True, dtype=torch.bool) # 模拟 bool 输入 | ||
| 109 | + | ||
| 110 | + gm = self.apply_pass(gm) | ||
| 111 | + modified_readable = gm.print_readable() | ||
| 112 | + self.assertIn("to(torch.int32)", modified_readable) | ||
| 113 | + | ||
| 114 | + def test_no_conversion_wrong_target_dtype(self): | ||
| 115 | + """测试 target_dtype 非 int64: 无转换""" | ||
| 116 | + def fn(x): | ||
| 117 | + return x.to(torch.float32) # 非 int64 | ||
| 118 | + | ||
| 119 | + gm = self.create_gm(fn) | ||
| 120 | + input_node = next(n for n in gm.graph.nodes if n.op == 'placeholder') | ||
| 121 | + input_node.meta['example_value'] = torch.tensor(1.0, dtype=torch.float32) | ||
| 122 | + | ||
| 123 | + gm = self.apply_pass(gm) | ||
| 124 | + modified_readable = gm.print_readable() | ||
| 125 | + self.assertIn("to(torch.float32)", modified_readable) # 未变 | ||
| 126 | + | ||
| 127 | + def test_no_conversion_wrong_target_method(self): | ||
| 128 | + """测试 target 非 'to': 无转换""" | ||
| 129 | + def fn(x): | ||
| 130 | + return x.view(1, -1) # 非 'to' | ||
| 131 | + | ||
| 132 | + gm = self.create_gm(fn) | ||
| 133 | + input_node = next(n for n in gm.graph.nodes if n.op == 'placeholder') | ||
| 134 | + input_node.meta['example_value'] = torch.tensor(1.0, dtype=torch.float32) | ||
| 135 | + | ||
| 136 | + gm = self.apply_pass(gm) | ||
| 137 | + modified_readable = gm.print_readable() | ||
| 138 | + self.assertIn("view", modified_readable) # 未变 | ||
| 139 | + | ||
| 140 | + def test_kwargs_dtype_handling(self): | ||
| 141 | + """测试 kwargs 中 dtype: to(int64) -> to(int32)""" | ||
| 142 | + def fn(x): | ||
| 143 | + return x.to(dtype=torch.int64) # 使用 kwargs | ||
| 144 | + | ||
| 145 | + gm = self.create_gm(fn) | ||
| 146 | + input_node = next(n for n in gm.graph.nodes if n.op == 'placeholder') | ||
| 147 | + input_node.meta['example_value'] = torch.tensor(1.0, dtype=torch.float32) | ||
| 148 | + gm = self.apply_pass(gm) | ||
| 149 | + modified_readable = gm.print_readable() | ||
| 150 | + self.assertIn("to(dtype = torch.int32)", modified_readable) | ||
| 151 | + | ||
| 152 | + def test_no_meta_example_value_skip(self): | ||
| 153 | + """测试无 meta['example_value']: 无转换 (安全 fallback)""" | ||
| 154 | + def fn(x): | ||
| 155 | + return x.to(torch.int64) | ||
| 156 | + | ||
| 157 | + gm = self.create_gm(fn) | ||
| 158 | + # 手动移除 meta 以模拟无 example_value | ||
| 159 | + input_node = next(n for n in gm.graph.nodes if n.op == 'placeholder') | ||
| 160 | + if 'example_value' in input_node.meta: | ||
| 161 | + del input_node.meta['example_value'] | ||
| 162 | + | ||
| 163 | + gm = self.apply_pass(gm) | ||
| 164 | + modified_readable = gm.print_readable() | ||
| 165 | + self.assertIn("to(torch.int64)", modified_readable) # 未变 | ||
| 166 | + | ||
| 167 | + | ||
| 168 | +instantiate_parametrized_tests(DtypeOptimalPass) | ||
| 169 | + | ||
| 170 | +if __name__ == "__main__": | ||
| 171 | + run_tests() | ||
| @@ -1,100 +0,0 @@ | |||
| 1 | -import torch | ||
| 2 | -from torch.fx.passes.shape_prop import ShapeProp | ||
| 3 | -from torch.testing._internal.common_utils import ( | ||
| 4 | - run_tests, parametrize, instantiate_parametrized_tests | ||
| 5 | -) | ||
| 6 | -from testutils import TestUtils | ||
| 7 | -from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import embedding_indice_i64_to_i32_pass | ||
| 8 | -from torch_npu._inductor.fx_passes.utils.check_op_util import check_embedding_op | ||
| 9 | - | ||
| 10 | - | ||
| 11 | - | ||
| 12 | -class TestEmbeddingIndiceI64ToI32Pass(TestUtils): | ||
| 13 | - # here, we use 'self.emb_table(input_ids)' | ||
| 14 | - class EmbeddingModel_X(torch.nn.Module): | ||
| 15 | - def __init__(self): | ||
| 16 | - super().__init__() | ||
| 17 | - self.emb_table = torch.nn.Embedding.from_pretrained( | ||
| 18 | - torch.normal( | ||
| 19 | - mean=0, std=0.1, size=(30522, 768) | ||
| 20 | - ) | ||
| 21 | - ).to('npu') | ||
| 22 | - | ||
| 23 | - def forward(self, input_ids): | ||
| 24 | - return self.emb_table(input_ids) | ||
| 25 | - | ||
| 26 | - | ||
| 27 | - | ||
| 28 | - def test_compile_case(self, shape, dtype): | ||
| 29 | - input_ids = self._generate_tensor(shape, dtype) | ||
| 30 | - model = self.EmbeddingModel_X() | ||
| 31 | - model.eval() | ||
| 32 | - | ||
| 33 | - with torch.no_grad(): | ||
| 34 | - compile_model = torch.compile(model, backend="inductor") | ||
| 35 | - compile_result = compile_model(input_ids) | ||
| 36 | - | ||
| 37 | - eager_result = model(input_ids) | ||
| 38 | - | ||
| 39 | - self.assertEqual(eager_result, compile_result, atol=1e-3, rtol=1e-3) | ||
| 40 | - | ||
| 41 | - # here, we use 'torch.nn.functional.embedding' | ||
| 42 | - class EmbeddingModel_Y(torch.nn.Module): | ||
| 43 | - def __init__(self): | ||
| 44 | - super().__init__() | ||
| 45 | - | ||
| 46 | - seed = 2026 | ||
| 47 | - torch.manual_seed(seed) | ||
| 48 | - torch.npu.manual_seed_all(seed) | ||
| 49 | - self.emb_table = torch.nn.Embedding(30522, 768, padding_idx=0).to('npu') | ||
| 50 | - torch.nn.init.uniform_(self.emb_table.weight, a=-1.0, b=1.0) | ||
| 51 | - | ||
| 52 | - def forward(self, input_ids): | ||
| 53 | - return torch.nn.functional.embedding(input_ids, self.emb_table.weight) | ||
| 54 | - | ||
| 55 | - | ||
| 56 | - | ||
| 57 | - def test_fx_model(self, shape, dtype): | ||
| 58 | - input_ids = self._generate_tensor(shape, dtype) | ||
| 59 | - model = self.EmbeddingModel_Y() | ||
| 60 | - model.eval() | ||
| 61 | - | ||
| 62 | - gm = torch.fx.symbolic_trace(model) | ||
| 63 | - ShapeProp(gm).propagate(input_ids) | ||
| 64 | - | ||
| 65 | - for node in gm.graph.nodes: | ||
| 66 | - if node.op == "placeholder": | ||
| 67 | - old_meta = node.meta['tensor_meta'] | ||
| 68 | - node.meta['tensor_meta'] = old_meta._replace(dtype=torch.int64) | ||
| 69 | - | ||
| 70 | - embedding_indice_i64_to_i32_pass(gm.graph) | ||
| 71 | - gm.recompile() | ||
| 72 | - | ||
| 73 | - # assert equal | ||
| 74 | - eager_result = model(input_ids) | ||
| 75 | - gm_result = gm(input_ids) | ||
| 76 | - self.assertEqual(eager_result, gm_result, atol=1e-3, rtol=1e-3) | ||
| 77 | - | ||
| 78 | - # test result | ||
| 79 | - embedding_node = None | ||
| 80 | - emb_input_is_cast = False | ||
| 81 | - found_cast = False | ||
| 82 | - | ||
| 83 | - for node in gm.graph.nodes: | ||
| 84 | - # check one of embedding-node's input-nodes is the inserted node | ||
| 85 | - if check_embedding_op(node): | ||
| 86 | - embedding_node = node | ||
| 87 | - if str(node.args[0].target) == "npu._npu_dtype_cast.default" or str(node.args[1].target) == "npu._npu_dtype_cast.default": | ||
| 88 | - emb_input_is_cast = True | ||
| 89 | - | ||
| 90 | - # check if cast node is inserted | ||
| 91 | - if "dtype_cast" in str(node.target) and node.kwargs.get('dtype') == torch.int32: | ||
| 92 | - found_cast = True | ||
| 93 | - | ||
| 94 | - self.assertTrue(found_cast, "cast-to-int32-node is not inserted by pass") | ||
| 95 | - self.assertIsNotNone(embedding_node) | ||
| 96 | - self.assertTrue(emb_input_is_cast, "embedding_node's input is not npu._npu_dtype_cast.default") | ||
| 97 | - | ||
| 98 | - | ||
| 99 | -if __name__ == "__main__": | ||
| 100 | - run_tests() | ||
| @@ -1,4 +1,5 @@ | |||
| 1 | import operator | 1 | import operator |
| 2 | +import math | ||
| 2 | import torch | 3 | import torch |
| 3 | import torch.fx | 4 | import torch.fx |
| 4 | from .register_custom_pass import register_custom_pass | 5 | from .register_custom_pass import register_custom_pass |
| @@ -21,7 +22,6 @@ from ..utils.check_op_util import ( | |||
| 21 | check_squeeze_op, | 22 | check_squeeze_op, |
| 22 | check_unsqueeze_op, | 23 | check_unsqueeze_op, |
| 23 | check_where_op, | 24 | check_where_op, |
| 24 | - check_embedding_op, | ||
| 25 | ) | 25 | ) |
| 26 | from ..utils.get_binary_fold_result import ( | 26 | from ..utils.get_binary_fold_result import ( |
| 27 | get_binary_fold_result, | 27 | get_binary_fold_result, |
| @@ -671,38 +671,63 @@ def fold_redundant_ops(graph: torch.fx.Graph): | |||
| 671 | if not any_removed: | 671 | if not any_removed: |
| 672 | break | 672 | break |
| 673 | eliminate_dead_code(graph, changed, fold_redundant_ops.__name__) | 673 | eliminate_dead_code(graph, changed, fold_redundant_ops.__name__) |
| 674 | - | 674 | + |
| 675 | - | 675 | + |
| 676 | -@register_custom_pass(PassType.POST) | 676 | +@register_custom_pass(PassType.PRE) |
W | |||
| 677 | -def embedding_indice_i64_to_i32_pass(graph: torch.fx.Graph) -> None: | 677 | +def dtype_optimal_pass(graph: torch.fx.Graph) -> None: |
| 678 | + int32_min, int32_max = -2**31, 2**31 - 1 | ||
| 679 | + cast_dtype_limit = [torch.float32, torch.int32, torch.bool, torch.int16, torch.int8] | ||
| 678 | changed = False | 680 | changed = False |
| 679 | - for node in graph.nodes: | 681 | + for node in list(graph.nodes): # 使用list避免修改时迭代问题 |
| 680 | - if not check_embedding_op(node): | 682 | + if node.op == 'call_function' and node.target == torch.arange \ |
| 681 | - continue | 683 | + and node.kwargs.get('dtype', None) == torch.int64: |
| 682 | - | 684 | + # 步骤1: 动态提取 start/end/step (处理不同 args 长度) |
| 683 | - indices_node = None | 685 | + args_len = len(node.args) |
| 684 | - args_id = -1 | 686 | + start = 0 |
| 685 | - if node.args[0].meta.get('tensor_meta') and node.args[0].meta.get('tensor_meta').dtype == torch.int64: | 687 | + end = None |
| 686 | - indices_node = node.args[0] | 688 | + step = 1 |
| 687 | - args_id = 0 | 689 | + if args_len == 1: |
| 688 | - elif node.args[1].meta.get('tensor_meta') and node.args[1].meta.get('tensor_meta').dtype == torch.int64: | 690 | + end = node.args[0] # arange(end) |
| 689 | - indices_node = node.args[1] | 691 | + elif args_len == 2: |
| 690 | - args_id = 1 | 692 | + start = node.args[0] |
| 691 | - | 693 | + end = node.args[1] # arange(start, end) |
| 692 | - if indices_node is not None: | 694 | + elif args_len >= 3: |
| 693 | - with graph.inserting_before(node): | 695 | + start = node.args[0] |
| 694 | - new_indices = graph.call_function( | 696 | + end = node.args[1] |
| 695 | - torch.ops.npu._npu_dtype_cast.default, | 697 | + step = node.args[2] # arange(start, end, step) |
| 696 | - args=(indices_node,), | 698 | + # 合并 kwargs 覆盖 (e.g., 用户指定 kwargs['start']) |
| 697 | - kwargs={"dtype": torch.int32} | 699 | + start = node.kwargs.get('start', start) |
| 698 | - ) | 700 | + end = node.kwargs.get('end', end) |
| 699 | - | 701 | + step = node.kwargs.get('step', step) |
| 700 | - new_args = list(node.args) | 702 | + # 如果 end 为 None，假设无限或跳过 (罕见，但安全) |
| 701 | - new_args[args_id] = new_indices | 703 | + if end is None: |
| 702 | - node.args = tuple(new_args) | 704 | + continue |
| 703 | - | 705 | + # 静态范围检查 (所有参数是常量) |
| 704 | - changed = True | 706 | + if all(isinstance(p, (int, float)) for p in [start, step, end]): |
| 705 | - eliminate_dead_code(graph, changed, embedding_indice_i64_to_i32_pass.__name__) | 707 | + if step == 0: |
| 708 | + continue | ||
| 709 | + # 如果 step 非整数且 dtype 是 int，警告 (arange 会自动转为 float) | ||
| 710 | + if not isinstance(step, int): | ||
| 711 | + continue | ||
| 712 | + # 计算序列长度和 min/max 值 | ||
| 713 | + num_elements = math.ceil((end - start) / step) if step > 0 else math.ceil((start - end) / -step) | ||
| 714 | + seq_min = min(start, start + (num_elements - 1) * step) | ||
| 715 | + seq_max = max(start, start + (num_elements - 1) * step) | ||
| 716 | + if seq_min > int32_min and seq_max < int32_max: | ||
| 717 | + node.kwargs = {**node.kwargs, 'dtype': torch.int32} | ||
| 718 | + changed = True | ||
| 719 | + if node.op == 'call_method': | ||
| 720 | + input_node = node.args[0] | ||
| 721 | + input_fake = input_node.meta.get('example_value', None) if hasattr(input_node, 'meta') else None | ||
| 722 | + input_dtype = input_fake.dtype if input_fake is not None else None | ||
| 723 | + target_dtype = node.args[1] if len(node.args) > 1 else node.kwargs.get('dtype', None) | ||
| 724 | + if input_dtype in cast_dtype_limit and node.target == 'to' and target_dtype == torch.int64: | ||
W dezheng889 1月26日 评论： 1月26日 评论： | |||
| 725 | + if len(node.args) > 1: | ||
| 726 | + node.args = (node.args[0], torch.int32) # 更新 positional dtype | ||
| 727 | + else: | ||
| 728 | + node.kwargs = {**node.kwargs, 'dtype': torch.int32} # 更新 kwargs dtype | ||
| 729 | + changed = True | ||
| 730 | + eliminate_dead_code(graph, changed, dtype_optimal_pass.__name__) | ||
| 706 | 731 | ||
| 707 | 732 | ||
| 708 | def eliminate_dead_code(graph, changed, fn_name): | 733 | def eliminate_dead_code(graph, changed, fn_name): |
| @@ -160,18 +160,6 @@ def check_div_op(node: fx.Node) -> bool: | |||
| 160 | return check_op(node, torch.ops.aten.div.Tensor) | 160 | return check_op(node, torch.ops.aten.div.Tensor) |
| 161 | 161 | ||
| 162 | 162 | ||
| 163 | -def check_embedding_op(node: fx.Node) -> bool: | ||
| 164 | - if node.op != "call_function": | ||
| 165 | - return False | ||
| 166 | - | ||
| 167 | - embedding_targets = { | ||
| 168 | - torch.nn.functional.embedding, | ||
| 169 | - torch.ops.aten.embedding.default, | ||
| 170 | - torch.embedding, | ||
| 171 | - } | ||
| 172 | - return node.target in embedding_targets | ||
| 173 | - | ||
| 174 | - | ||
| 175 | def check_op_by_targets(node: fx.Node, targets) -> bool: | 163 | def check_op_by_targets(node: fx.Node, targets) -> bool: |
| 176 | for target in targets: | 164 | for target in targets: |
| 177 | result = check_op(node, target) | 165 | result = check_op(node, target) |
这里为什么修改成PRE，POST有什么问题