已合并
[inductor] add fx_pass embedding_indice_i64_to_i32_pass #29405
jimmycao9929创建于 1月11日
[inductor] add fx_pass embedding_indice_i64_to_i32_pass #29405
已合并
共 4 个文件变更+180-3
| @@ -202,3 +202,5 @@ torch_npu/csrc/aten/RegisterCPU.cpp | |||
| 202 | torch_npu/csrc/aten/RegisterNPU.cpp | 202 | torch_npu/csrc/aten/RegisterNPU.cpp |
| 203 | torch_npu/csrc/aten/python_custom_functions.cpp | 203 | torch_npu/csrc/aten/python_custom_functions.cpp |
| 204 | torch_npu/_op_plugin_docs.py | 204 | torch_npu/_op_plugin_docs.py |
| 205 | + | ||
| 206 | +.codemate | ||
S | |||
| @@ -0,0 +1,100 @@ | |||
| 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() | ||
| @@ -2,8 +2,39 @@ import operator | |||
| 2 | import torch | 2 | import torch |
| 3 | import torch.fx | 3 | import torch.fx |
| 4 | from .register_custom_pass import register_custom_pass | 4 | from .register_custom_pass import register_custom_pass |
| 5 | -from ..utils.check_op_util import is_zero_like, check_op_by_targets, is_one_like, try_match, is_cast_node, normalize_dtype, get_node_dtype, get_cast_dtype, check_cat_op, get_input_node, get_input_kw_node, check_support_op, _get_tensor_meta, check_view, check_act_op, check_squeeze_op, check_unsqueeze_op, check_where_op | 5 | +from ..utils.check_op_util import ( |
| 6 | -from ..utils.get_binary_fold_result import get_binary_fold_result, get_node_meta, get_node_shape, get_node_unique_id, has_storage_or_layout, _get_fold_result, _fold_slice, _fold_slice_scatter, get_slice_dim, get_pad_dim_and_size | 6 | + is_zero_like, |
| 7 | + check_op_by_targets, | ||
| 8 | + is_one_like, | ||
| 9 | + try_match, | ||
| 10 | + is_cast_node, | ||
| 11 | + normalize_dtype, | ||
| 12 | + get_node_dtype, | ||
| 13 | + get_cast_dtype, | ||
| 14 | + check_cat_op, | ||
| 15 | + get_input_node, | ||
| 16 | + get_input_kw_node, | ||
| 17 | + check_support_op, | ||
| 18 | + _get_tensor_meta, | ||
| 19 | + check_view, | ||
| 20 | + check_act_op, | ||
| 21 | + check_squeeze_op, | ||
| 22 | + check_unsqueeze_op, | ||
| 23 | + check_where_op, | ||
| 24 | + check_embedding_op, | ||
| 25 | +) | ||
| 26 | +from ..utils.get_binary_fold_result import ( | ||
| 27 | + get_binary_fold_result, | ||
| 28 | + get_node_meta, | ||
| 29 | + get_node_shape, | ||
| 30 | + get_node_unique_id, | ||
| 31 | + has_storage_or_layout, | ||
| 32 | + _get_fold_result, | ||
| 33 | + _fold_slice, | ||
| 34 | + _fold_slice_scatter, | ||
| 35 | + get_slice_dim, | ||
| 36 | + get_pad_dim_and_size, | ||
| 37 | +) | ||
| 7 | from ..utils.fx_pass_level import PassType | 38 | from ..utils.fx_pass_level import PassType |
| 8 | from ...config import log | 39 | from ...config import log |
| 9 | 40 | ||
| @@ -640,10 +671,42 @@ def fold_redundant_ops(graph: torch.fx.Graph): | |||
| 640 | if not any_removed: | 671 | if not any_removed: |
| 641 | break | 672 | break |
| 642 | eliminate_dead_code(graph, changed, fold_redundant_ops.__name__) | 673 | eliminate_dead_code(graph, changed, fold_redundant_ops.__name__) |
| 674 | + | ||
| 675 | + | ||
| 676 | + | ||
| 677 | +def embedding_indice_i64_to_i32_pass(graph: torch.fx.Graph) -> None: | ||
| 678 | + changed = False | ||
| 679 | + for node in graph.nodes: | ||
| 680 | + if not check_embedding_op(node): | ||
| 681 | + continue | ||
| 682 | + | ||
| 683 | + indices_node = None | ||
| 684 | + args_id = -1 | ||
| 685 | + if node.args[0].meta.get('tensor_meta') and node.args[0].meta.get('tensor_meta').dtype == torch.int64: | ||
| 686 | + indices_node = node.args[0] | ||
| 687 | + args_id = 0 | ||
| 688 | + elif node.args[1].meta.get('tensor_meta') and node.args[1].meta.get('tensor_meta').dtype == torch.int64: | ||
| 689 | + indices_node = node.args[1] | ||
| 690 | + args_id = 1 | ||
| 691 | + | ||
| 692 | + if indices_node is not None: | ||
| 693 | + with graph.inserting_before(node): | ||
| 694 | + new_indices = graph.call_function( | ||
| 695 | + torch.ops.npu._npu_dtype_cast.default, | ||
| 696 | + args=(indices_node,), | ||
| 697 | + kwargs={"dtype": torch.int32} | ||
W jimmycao9929 1月14日 评论： 1月14日 评论： | |||
| 698 | + ) | ||
| 699 | + | ||
| 700 | + new_args = list(node.args) | ||
| 701 | + new_args[args_id] = new_indices | ||
| 702 | + node.args = tuple(new_args) | ||
| 703 | + | ||
| 704 | + changed = True | ||
| 705 | + eliminate_dead_code(graph, changed, embedding_indice_i64_to_i32_pass.__name__) | ||
| 643 | 706 | ||
| 644 | 707 | ||
| 645 | def eliminate_dead_code(graph, changed, fn_name): | 708 | def eliminate_dead_code(graph, changed, fn_name): |
| 646 | if changed: | 709 | if changed: |
| 647 | graph.lint() | 710 | graph.lint() |
| 648 | graph.eliminate_dead_code() | 711 | graph.eliminate_dead_code() |
| 649 | - log.info(f"{fn_name} pass works") | 712 | + log.info(f"[inductor_fx_pas] {fn_name} works") |
| @@ -160,6 +160,18 @@ 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 | + | ||
| 163 | def check_op_by_targets(node: fx.Node, targets) -> bool: | 175 | def check_op_by_targets(node: fx.Node, targets) -> bool: |
| 164 | for target in targets: | 176 | for target in targets: |
| 165 | result = check_op(node, target) | 177 | result = check_op(node, target) |
codemate是不是应该不加呀