已合并
aot_autograd_pass_fo_FA #30156
Lu_G创建于 1月28日
aot_autograd_pass_fo_FA #30156
已合并
共 3 个文件变更+102-1
| @@ -0,0 +1,77 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch.fx as fx | ||
| 3 | +from torch.fx.passes.shape_prop import ShapeProp | ||
| 4 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 5 | +from testutils import TestUtils | ||
| 6 | +import torch_npu | ||
| 7 | +import torch_npu._inductor | ||
| 8 | + | ||
| 9 | + | ||
| 10 | +class FusionAttentionUnchangeModel(torch.nn.Module): | ||
| 11 | + def forward(self, primals_1, primals_2, primals_3): | ||
| 12 | + npu_fusion_attention = torch.ops.npu.npu_fusion_attention.default(primals_1, primals_2, primals_3, 8, 'BNSD', None, None, None, 0.125, 0.9) | ||
| 13 | + getitem = npu_fusion_attention[0] | ||
| 14 | + getitem_1 = npu_fusion_attention[1] | ||
| 15 | + getitem_2 = npu_fusion_attention[2] | ||
| 16 | + getitem_3 = npu_fusion_attention[3] | ||
| 17 | + getitem_4 = npu_fusion_attention[4] | ||
| 18 | + getitem_5 = npu_fusion_attention[5] | ||
| 19 | + return {"getitem": getitem, "getitem_1": getitem_1, "getitem_2": getitem_2, "getitem_3": getitem_3, "getitem_4": getitem_4, "getitem_5": getitem_5} | ||
| 20 | + | ||
| 21 | + | ||
| 22 | +class TestFusionAttentionUnchangePass(TestUtils): | ||
| 23 | + def op_calc(self, primals_1, primals_2, primals_3): | ||
| 24 | + npu_fusion_attention_v3 = torch.ops.npu.npu_fusion_attention_v3.default(primals_1, primals_2, primals_3, 8, 'BNSD', None, None, None, 0.125, 0.9) | ||
| 25 | + getitem = npu_fusion_attention_v3[0] | ||
| 26 | + getitem_1 = npu_fusion_attention_v3[1] | ||
| 27 | + getitem_2 = npu_fusion_attention_v3[2] | ||
| 28 | + getitem_3 = npu_fusion_attention_v3[3] | ||
| 29 | + getitem_4 = npu_fusion_attention_v3[4] | ||
| 30 | + getitem_5 = npu_fusion_attention_v3[5] | ||
| 31 | + return {"getitem": getitem, "getitem_1": getitem_1, "getitem_2": getitem_2, "getitem_3": getitem_3, "getitem_4": getitem_4, "getitem_5": getitem_5} | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + def test_compile_cases(self): | ||
| 36 | + primals_1 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu") | ||
| 37 | + primals_2 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu") | ||
| 38 | + primals_3 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu") | ||
| 39 | + | ||
| 40 | + state = torch.npu.get_rng_state() | ||
| 41 | + std_result = self.op_calc(primals_1, primals_2, primals_3) | ||
| 42 | + torch.npu.set_rng_state(state) | ||
| 43 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor") | ||
| 44 | + inductor_result = compiled_op_calc(primals_1, primals_2, primals_3) | ||
| 45 | + self.assertIsInstance(inductor_result["getitem_4"], torch.Tensor, "The output parameter 'seed' of the npu_fusion_attention_v3 should be Tensor") | ||
| 46 | + self.assertIsInstance(inductor_result["getitem_5"], torch.Tensor, "The output parameter 'offset' of the npu_fusion_attention_v3 should be Tensor") | ||
| 47 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 48 | + | ||
| 49 | + | ||
| 50 | + def test_ut_cases(self): | ||
| 51 | + primals_1 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu") | ||
| 52 | + primals_2 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu") | ||
| 53 | + primals_3 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu") | ||
| 54 | + model = FusionAttentionUnchangeModel() | ||
| 55 | + graph_module = fx.symbolic_trace(model) | ||
| 56 | + ShapeProp(graph_module).propagate(primals_1, primals_2, primals_3) | ||
| 57 | + | ||
| 58 | + # 应用优化 Pass | ||
| 59 | + from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fusion_attention_v3_pass | ||
| 60 | + fusion_attention_v3_pass(graph_module.graph) | ||
| 61 | + graph_module.recompile() | ||
| 62 | + | ||
| 63 | + # 验证输出是否一致 | ||
| 64 | + state = torch.npu.get_rng_state() | ||
| 65 | + std_result = model(primals_1, primals_2, primals_3) | ||
| 66 | + torch.npu.set_rng_state(state) | ||
| 67 | + inductor_result = graph_module(primals_1, primals_2, primals_3) | ||
| 68 | + self.assertIsInstance(inductor_result["getitem_4"], torch.Tensor, "After the pass, the output parameter 'seed' of the npu_fusion_attention should be Tensor") | ||
| 69 | + self.assertIsInstance(inductor_result["getitem_5"], torch.Tensor, "After the pass, the output parameter 'offset' of the npu_fusion_attention should be Tensor") | ||
| 70 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 71 | + | ||
| 72 | + | ||
| 73 | +instantiate_parametrized_tests(TestFusionAttentionUnchangePass) | ||
| 74 | + | ||
| 75 | + | ||
| 76 | +if __name__ == "__main__": | ||
| 77 | + run_tests() | ||
| @@ -17,7 +17,12 @@ def run_register_pre_custom_passes(gm): | |||
| 17 | for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.PRE][level]: | 17 | for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.PRE][level]: |
| 18 | fn(gm) | 18 | fn(gm) |
| 19 | 19 | ||
| 20 | - log.debug(f"after pre_grad graph optimizer pass, graph is: {gm}") | 20 | + for level in sorted(FxPassLevel): |
| 21 | + for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.PRE][level]: | ||
| 22 | + if fn.__name__ == "fusion_attention_v3_pass": | ||
| 23 | + fn(gm) | ||
| 24 | + | ||
| 25 | + log.debug(f"after pre_grad graph optimizer pass, graph is: {gm}") | ||
D Ddezheng8892月3日   | |||
| 21 | 26 | ||
| 22 | 27 | ||
| 23 | def run_register_post_custom_passes(gm): | 28 | def run_register_post_custom_passes(gm): |
| @@ -805,6 +805,25 @@ def unfold_dual_reduction_pass(graph: torch.fx.Graph) -> None: | |||
| 805 | eliminate_dead_code(graph, changed, unfold_dual_reduction_pass.__name__) | 805 | eliminate_dead_code(graph, changed, unfold_dual_reduction_pass.__name__) |
| 806 | 806 | ||
| 807 | 807 | ||
| 808 | + | ||
| 809 | +def fusion_attention_v3_pass(graph: torch.fx.Graph) -> None: | ||
| 810 | + changed = False | ||
| 811 | + for node in list(graph.nodes): # 使用list避免迭代时修改图结构 | ||
| 812 | + if node.op == 'call_function' and node.target == torch.ops.npu.npu_fusion_attention.default: | ||
| 813 | + # 创建新节点调用v3版本 | ||
| 814 | + with graph.inserting_before(node): | ||
| 815 | + new_node = graph.call_function( | ||
| 816 | + torch.ops.npu.npu_fusion_attention_v3.default, | ||
| 817 | + args=node.args, | ||
| 818 | + kwargs=node.kwargs | ||
| 819 | + ) | ||
| 820 | + new_node.meta.update(node.meta) | ||
| 821 | + node.replace_all_uses_with(new_node) | ||
D 最好是有个ut 作为测试用例 | |||
| 822 | + graph.erase_node(node) | ||
| 823 | + changed = True | ||
| 824 | + eliminate_dead_code(graph, changed, fusion_attention_v3_pass.__name__, False) | ||
| 825 | + | ||
| 826 | + | ||
| 808 | def eliminate_dead_code(graph, changed, fn_name, POST=True): | 827 | def eliminate_dead_code(graph, changed, fn_name, POST=True): |
| 809 | if changed: | 828 | if changed: |
| 810 | if POST: | 829 | if POST: |