run_once_fix-pytorch-AtomGit

已合并

run_once_fix #38396

run_once_fix #38396

已合并

cuiduo创建于 6月12日

共 5 个文件变更+21-18

	@@ -2,6 +2,7 @@ import unittest
2	from numbers import Number	2	from numbers import Number
3	import torch	3	import torch
4	from torch.testing import make_tensor	4	from torch.testing import make_tensor
		5	+import torch_npu
5	from torch_npu.testing.common_utils import SupportedDevices	6	from torch_npu.testing.common_utils import SupportedDevices
6	from torch_npu.testing.testcase import TestCase, run_tests	7	from torch_npu.testing.testcase import TestCase, run_tests
7		8

	@@ -65,7 +65,6 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
65	import math	65	import math
66	import random	66	import random
67	import os	67	import os
68	- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = 'mlir'
69	import tempfile	68	import tempfile
70	from math import inf, nan	69	from math import inf, nan
71	from torch._inductor.hooks import run_intermediate_hooks	70	from torch._inductor.hooks import run_intermediate_hooks

	@@ -11,7 +11,8 @@ from torch_npu._inductor.lowering_common import (
11	)	11	)
12		12
13	from .. import config	13	from .. import config
14	-from ..npu.utils import get_anir_mode, run_once	14	+from ..npu.utils import get_anir_mode
		15	+from torch_npu._inductor.lowering_common import run_once
15	from .utils import logger	16	from .utils import logger
16		17
17		18

	@@ -17,27 +17,12 @@ from torch._dynamo.symbolic_convert import break_graph_if_unsupported, Instructi
17	from torch._dynamo.exc import Unsupported	17	from torch._dynamo.exc import Unsupported
18	from torch._dynamo.variables.lists import TupleVariable	18	from torch._dynamo.variables.lists import TupleVariable
19	from torch._dynamo.variables.nn_module import NNModuleVariable	19	from torch._dynamo.variables.nn_module import NNModuleVariable
20	-	20	+from torch_npu._inductor.lowering_common import run_once
21		21
22	aten = torch.ops.aten	22	aten = torch.ops.aten
23	npu = torch.ops.npu	23	npu = torch.ops.npu
24		24
25		25
26	-def run_once(f):
27	- """Runs a function (successfully) only once.
28	- The running can be reset by setting the `has_run` attribute to False
29	- """
30	- @wraps(f)
31	- def wrapper(args, *kwargs):
32	- if not wrapper.has_run:
33	- result = f(args, *kwargs)
34	- wrapper.has_run = True
35	- return result
36	- return None
37	- wrapper.has_run = False
38	- return wrapper
39	-
40	-
41	npu_meta_table = {}	26	npu_meta_table = {}
42	break_fn_table = {}	27	break_fn_table = {}
43	break_mapping_table = {}	28	break_mapping_table = {}

	@@ -17,6 +17,7 @@ from __future__ import annotations
17	import inspect	17	import inspect
18	import sympy	18	import sympy
19	from functools import reduce	19	from functools import reduce
		20	+import functools
20	from dataclasses import dataclass	21	from dataclasses import dataclass
21	from typing import Any, Callable, Dict, List, Optional, Tuple, Union	22	from typing import Any, Callable, Dict, List, Optional, Tuple, Union
22		23
	@@ -29,6 +30,22 @@ from torch._inductor import ir
29	from torch._inductor.ir import ExpandView, IndexingConstant, TensorBox	30	from torch._inductor.ir import ExpandView, IndexingConstant, TensorBox
30	from torch._inductor.virtualized import V	31	from torch._inductor.virtualized import V
31		32
		33	+
		34	+def run_once(f):
		35	+ """Runs a function (successfully) only once.
		36	+ The running can be reset by setting the `has_run` attribute to False
		37	+ """
		38	+ @functools.wraps(f)
		39	+ def wrapper(args, *kwargs):
		40	+ if not wrapper.has_run:
		41	+ result = f(args, *kwargs)
		42	+ wrapper.has_run = True
		43	+ return result
		44	+ return None
		45	+ wrapper.has_run = False
		46	+ return wrapper
		47	+
		48	+
32	LOWERING_REGISTRY_ATTRS: tuple[str, ...] = (	49	LOWERING_REGISTRY_ATTRS: tuple[str, ...] = (
33	"lowerings",	50	"lowerings",
34	"_maybe_layout_constraints",	51	"_maybe_layout_constraints",