Star887
1399
代码
代码
Issues1807
Pull Requests1791
Actions
项目讨论
Wiki
分析
Star887
1399
已合并
Register custom meta function for _to_copy to handle NPU-specific tensor conversions #34572
du-jin-hang创建于 4月28日
Register custom meta function for _to_copy to handle NPU-specific tensor conversions #34572
已合并
du-jin-hang创建于 4月28日
从du-jin-hang/pytorchv2.11.0合入到Ascend/pytorchv2.11.0
共 6 个文件变更+364-148
Atest/_inductor/test_to_copy_stride.py+57-0
@@ -0,0 +1,57 @@
1+# Owner(s): ["module: inductor"]
2+import unittest
3+ 
4+import torch
5+from torch._subclasses.fake_tensor import FakeTensorMode
6+from torch.testing._internal.common_utils import run_tests, TestCase
7+ 
8+ 
9+RUN_NPU = torch.npu.is_available()
10+SOURCE_SIZE = (4, 2048, 16, 192)
11+SOURCE_STRIDE = (6291456, 192, 393216, 1)
12+CONTIGUOUS_STRIDE = (6291456, 3072, 192, 1)
13+ 
14+ 
15+def _to_copy_strided_float32(x):
16+ return torch.ops.aten._to_copy.default(
17+ x,
18+ dtype=torch.float32,
19+ layout=torch.strided,
20+ )
21+ 
22+ 
23+def _make_source_tensor():
24+ device = torch.device("npu:0")
25+ base_tensor = torch.empty(SOURCE_SIZE, dtype=torch.float32, device=device)
26+ return base_tensor.as_strided(
27+ size=SOURCE_SIZE,
28+ stride=SOURCE_STRIDE,
29+ )
30+ 
31+ 
32+@unittest.skipIf(not RUN_NPU, "requires npu")
33+class TestToCopyStride(TestCase):
34+ def test_to_copy_fake_tensor_mode_stride(self):
35+ source = _make_source_tensor()
36+ 
37+ fake_mode = FakeTensorMode()
38+ with fake_mode:
39+ fake_source = fake_mode.from_tensor(source)
40+ fake_copied = _to_copy_strided_float32(fake_source)
41+ self.assertEqual(fake_copied.shape, SOURCE_SIZE)
42+ self.assertEqual(fake_copied.stride(), CONTIGUOUS_STRIDE)
43+ 
44+ def test_to_copy_compile_stride(self):
45+ source = _make_source_tensor()
46+ real_copied = _to_copy_strided_float32(source)
47+ self.assertEqual(real_copied.shape, SOURCE_SIZE)
48+ self.assertEqual(real_copied.stride(), CONTIGUOUS_STRIDE)
49+ 
50+ compiled_to_copy = torch.compile(_to_copy_strided_float32, backend="inductor")
51+ compiled_copied = compiled_to_copy(source)
52+ self.assertEqual(compiled_copied.shape, SOURCE_SIZE)
53+ self.assertEqual(compiled_copied.stride(), CONTIGUOUS_STRIDE)
54+ 
55+ 
56+if __name__ == "__main__":
57+ run_tests()
Mtorch_npu/_inductor/__init__.py+55-30
@@ -1,54 +1,74 @@
1import os1import os
2 2 
3from .codegen.common import register_device_op_overrides_npu3from .codegen.common import register_device_op_overrides_npu
4+ 
5+ 
4register_device_op_overrides_npu()6register_device_op_overrides_npu()
5 7 
6-if os.getenv('TORCHINDUCTOR_NPU_BACKEND', 'default') == 'mlir':8+if os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default") == "mlir":
7 try:9 try:
8 import torch_mlir10 import torch_mlir
9 from torch_mlir import ir11 from torch_mlir import ir
10- except:12+ except ImportError as err:
11- raise ImportError("torch_mlir is not installed, install it first.")13+ raise ImportError("torch_mlir is not installed, install it first.") from err
12 from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin14 from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin
13- from .utils import patch_is_gpu,patch_has_triton15+ from .utils import patch_has_triton, patch_is_gpu
16+ 
14 patch_is_gpu()17 patch_is_gpu()
15 patch_has_triton()18 patch_has_triton()
16else:19else:
17 import torch20 import torch
18- from torch._dynamo.device_interface import register_interface_for_device, get_interface_for_device21+ from torch._dynamo.device_interface import (
22+ get_interface_for_device,
23+ register_interface_for_device,
24+ )
19 from torch._inductor import lowering as inductor_lowering25 from torch._inductor import lowering as inductor_lowering
20- from torch._inductor.lowering import make_fallback as ori_make_fallback
21 from torch._inductor.choices import InductorChoices26 from torch._inductor.choices import InductorChoices
22- from torch._inductor.codegen.common import register_backend_for_device, register_device_op_overrides27+ from torch._inductor.codegen.common import (
28+ register_backend_for_device,
29+ register_device_op_overrides,
30+ )
31+ from torch._inductor.lowering import make_fallback as ori_make_fallback
23 from torch._inductor.runtime import autotune_cache32 from torch._inductor.runtime import autotune_cache
24 from torch_npu.npu import device_count33 from torch_npu.npu import device_count
25- from torch_npu.utils._dynamo_device import NpuInterface, current_device, set_device34+ from torch_npu.utils._dynamo_device import current_device, NpuInterface, set_device
26 from torch_npu.utils._inductor import NPUDeviceOpOverrides35 from torch_npu.utils._inductor import NPUDeviceOpOverrides
27 36 
28- from . import config as npu_config37+ from . import codegen, config as npu_config
29- from . import codegen38+ from .codecache import patch_aot_code_compiler_compile, patch_cache_base_get_system
30- from .npu_fusion_attention_graph import register_fa_pass39+ from .config import (
31- from .config import aggresive_autotune, num_vector_core, set_compile_threads, disable_comprehensive_padding40+ aggresive_autotune,
32- from .config import log as npulog41+ disable_comprehensive_padding,
42+ log as npulog,
43+ num_vector_core,
44+ set_compile_threads,
45+ )
46+ from .cpp_builder import patch_get_optimization_cflags
33 from .decomposition import _register_npu_inductor_decompositons47 from .decomposition import _register_npu_inductor_decompositons
34 from .lowering import make_reduction, npu_make_fallback48 from .lowering import make_reduction, npu_make_fallback
35 from .npu_choices import should_use_persistent_reduction49 from .npu_choices import should_use_persistent_reduction
36 from .npu_device import NewNPUDeviceOpOverrides50 from .npu_device import NewNPUDeviceOpOverrides
51+ from .npu_fusion_attention_graph import register_fa_pass
37 from .runtime import _load_cached_autotuning52 from .runtime import _load_cached_autotuning
38- from .utils import get_current_raw_stream, patch_is_gpu, patch_has_triton, disable_foreach, patch_fx_node_is_input_dependent_cudagraph_unsafe53+ from .utils import (
39- from .codecache import patch_aot_code_compiler_compile, patch_cache_base_get_system54+ disable_foreach,
40- from .cpp_builder import patch_get_optimization_cflags55+ get_current_raw_stream,
56+ patch_fx_node_is_input_dependent_cudagraph_unsafe,
57+ patch_has_triton,
58+ patch_is_gpu,
59+ )
41 60 
42 set_compile_threads()61 set_compile_threads()
43 disable_comprehensive_padding()62 disable_comprehensive_padding()
44 63 
45- 
46 def _inductor_register_backend_for_device():64 def _inductor_register_backend_for_device():
65+ from .codegen.cpp_wrapper import CppWrapperNpu
47 from .codegen.scheduling import NPUTritonScheduling66 from .codegen.scheduling import NPUTritonScheduling
48 from .codegen.wrapper import NPUWrapperCodeGen67 from .codegen.wrapper import NPUWrapperCodeGen
49- from .codegen.cpp_wrapper import CppWrapperNpu
50- register_backend_for_device('npu', NPUTritonScheduling, NPUWrapperCodeGen, CppWrapperNpu)
51 68 
69+ register_backend_for_device(
70+ "npu", NPUTritonScheduling, NPUWrapperCodeGen, CppWrapperNpu
71+ )
52 72 
53 _inductor_register_backend_for_device()73 _inductor_register_backend_for_device()
54 74 
@@ -57,14 +77,16 @@ else:
57 inductor_lowering.make_reduction = make_reduction77 inductor_lowering.make_reduction = make_reduction
58 inductor_lowering.make_fallback = npu_make_fallback78 inductor_lowering.make_fallback = npu_make_fallback
59 79 
60- 
61 def patch_torch_for_aoti():80 def patch_torch_for_aoti():
62- from .graph import patch_codegen_with_cpp_wrapper
63- from .cpp_builder import patch_get_cpp_torch_device_options
64 from .codegen.cpp_utils import patch_device_to_aten81 from .codegen.cpp_utils import patch_device_to_aten
65- from .utils import patch_is_same_tensor82+ from .cpp_builder import patch_get_cpp_torch_device_options
66 from .fx_passes.joint_graph import patch_constant_fold_uniform_value83 from .fx_passes.joint_graph import patch_constant_fold_uniform_value
67- from .ir import patch_fallback_kernel_codegen84+ from .graph import patch_codegen_with_cpp_wrapper
85+ from .ir import (
86+ patch_extern_kernel_codegen_size_asserts,
87+ patch_fallback_kernel_codegen,
88+ )
89+ from .utils import patch_is_same_tensor
68 90 
69 patch_codegen_with_cpp_wrapper()91 patch_codegen_with_cpp_wrapper()
70 patch_get_cpp_torch_device_options()92 patch_get_cpp_torch_device_options()
@@ -72,24 +94,26 @@ else:
72 patch_is_same_tensor()94 patch_is_same_tensor()
73 patch_constant_fold_uniform_value()95 patch_constant_fold_uniform_value()
74 patch_fallback_kernel_codegen()96 patch_fallback_kernel_codegen()
97+ patch_extern_kernel_codegen_size_asserts()
75 98 
76 patch_aot_code_compiler_compile()99 patch_aot_code_compiler_compile()
77 100 
78- 
79 if os.environ.get("DISABLE_AOTI_PATCH", "0") != "1":101 if os.environ.get("DISABLE_AOTI_PATCH", "0") != "1":
80 patch_torch_for_aoti()102 patch_torch_for_aoti()
81 103 
82- 
83 if npu_config.dump_fx_graph:104 if npu_config.dump_fx_graph:
84 from .codegen.ir_fx import _patch_npu_inductor_ir105 from .codegen.ir_fx import _patch_npu_inductor_ir
85 106 
86 _patch_npu_inductor_ir()107 _patch_npu_inductor_ir()
87 108 
88- from .lowering import _register_npu_inductor_fallbacks, _enable_full_lowering_fallback109+ from .lowering import (
110+ _enable_full_lowering_fallback,
111+ _register_npu_inductor_fallbacks,
112+ )
89 113 
90 _register_npu_inductor_decompositons()114 _register_npu_inductor_decompositons()
91 115 
92- if npu_config.enable_full_lowering_fallback.strip()=='allfallback':116+ if npu_config.enable_full_lowering_fallback.strip() == "allfallback":
93 _enable_full_lowering_fallback()117 _enable_full_lowering_fallback()
94 else:118 else:
95 _register_npu_inductor_fallbacks()119 _register_npu_inductor_fallbacks()
@@ -97,11 +121,12 @@ else:
97 # register fx_pass should be put behind of _register_npu_inductor_decompositons121 # register fx_pass should be put behind of _register_npu_inductor_decompositons
98 def _replace_benchmark_all_configs():122 def _replace_benchmark_all_configs():
99 from torch._inductor.triton_heuristics import CachingAutotuner123 from torch._inductor.triton_heuristics import CachingAutotuner
124+ 
100 from .npu_triton_heuristics import benchmark_all_configs125 from .npu_triton_heuristics import benchmark_all_configs
126+ 
101 CachingAutotuner.benchmark_all_configs = benchmark_all_configs127 CachingAutotuner.benchmark_all_configs = benchmark_all_configs
102 128 
103- 129+ if aggresive_autotune:
104- if (aggresive_autotune):
105 _replace_benchmark_all_configs()130 _replace_benchmark_all_configs()
106 import os131 import os
107 132 
Mtorch_npu/_inductor/codegen/wrapper.py+55-59
@@ -1,25 +1,25 @@
1-import os
2import copy1import copy
3-from typing import Any, Callable, Optional, TYPE_CHECKING, Union
4import hashlib2import hashlib
3+import os
4+ 
5import sympy5import sympy
6 6 
7import torch7import torch
8-from torch._inductor import config
9-from torch._inductor.codegen.wrapper import PythonWrapperCodegen, SymbolicCallArg, SubgraphPythonWrapperCodegen
10-from torch._inductor.runtime import triton_heuristics
11-from torch._inductor.utils import (
12- cache_on_self,
13-)
14-from torch._inductor.virtualized import V
15-from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
16-from torch.utils._sympy.singleton_int import SingletonInt
17-from torch._inductor.ir import GraphPartitionSignature
18- 
19-from torch_npu._inductor import config as npu_config
20import torch_npu.npu.aclnn8import torch_npu.npu.aclnn
9+from torch._inductor import config
10+from torch._inductor.codegen.wrapper import (
11+ PythonWrapperCodegen,
12+ SubgraphPythonWrapperCodegen,
13+ SymbolicCallArg,
14+)
15+from torch._inductor.ir import GraphPartitionSignature
16+from torch._inductor.runtime import triton_heuristics
17+from torch._inductor.utils import cache_on_self
18+from torch._inductor.virtualized import V
19+from torch._subclasses.fake_tensor import FakeTensor
20+from torch.utils._sympy.singleton_int import SingletonInt
21+from torch_npu._inductor import config as npu_config
21from torch_npu._inductor.codegen.triton import NPUIndexTritonKernel22from torch_npu._inductor.codegen.triton import NPUIndexTritonKernel
22-from torch_npu._inductor.npu_triton_heuristics import PrecomputedGridNpu, user_autotune_npu
23 23 
24 24 
25class NPUWrapperCodeGen(PythonWrapperCodegen):25class NPUWrapperCodeGen(PythonWrapperCodegen):
@@ -31,16 +31,18 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
31 is_subgraph: bool,31 is_subgraph: bool,
32 subgraph_name: str,32 subgraph_name: str,
33 parent_wrapper: PythonWrapperCodegen,33 parent_wrapper: PythonWrapperCodegen,
34- partition_signatures: Optional[GraphPartitionSignature] = None,34+ partition_signatures: GraphPartitionSignature | None = None,
35 ):35 ):
36 if is_subgraph:36 if is_subgraph:
37- return SubgraphPythonWrapperCodegen(subgraph_name, parent_wrapper, partition_signatures)37+ return SubgraphPythonWrapperCodegen(
38+ subgraph_name, parent_wrapper, partition_signatures
39+ )
38 return NPUWrapperCodeGen()40 return NPUWrapperCodeGen()
39 41 
40 def write_header(self) -> None:42 def write_header(self) -> None:
41 super().write_header()43 super().write_header()
42 self.imports.splice(44 self.imports.splice(
43- f"""45+ """
44 import torch_npu46 import torch_npu
45 """,47 """,
46 strip=True,48 strip=True,
@@ -86,10 +88,6 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
86 # it suffices as a type hint for the purposes of producing the correct code for this type.88 # it suffices as a type hint for the purposes of producing the correct code for this type.
87 return SymbolicCallArg(expr, numel_expr)89 return SymbolicCallArg(expr, numel_expr)
88 90 
89- # don't free anything
90- def make_buffer_free(self, buffer):
91- return ""
92- 
93 # don't assert91 # don't assert
94 def codegen_input_size_asserts(self) -> None:92 def codegen_input_size_asserts(self) -> None:
95 pass93 pass
@@ -104,7 +102,7 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
104 """102 """
105 if not config.benchmark_harness:103 if not config.benchmark_harness:
106 return None104 return None
107- 105+ 
108 if npu_config.aot_inductor.debug_kernel:106 if npu_config.aot_inductor.debug_kernel:
109 return self.add_npu_repro(output)107 return self.add_npu_repro(output)
110 108 
@@ -127,7 +125,7 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
127 "print(result)",125 "print(result)",
128 ]126 ]
129 )127 )
130- 128+ 
131 def add_repro_func(self, output):129 def add_repro_func(self, output):
132 seen_constants = set()130 seen_constants = set()
133 131 
@@ -140,11 +138,11 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
140 )138 )
141 139 
142 def get_hash(name):140 def get_hash(name):
143- byte = name.encode('utf-8')141+ byte = name.encode("utf-8")
144- sha1 = hashlib.sha1()142+ sha1 = hashlib.sha1(usedforsecurity=False)
145 sha1.update(byte)143 sha1.update(byte)
146 return sha1.hexdigest()144 return sha1.hexdigest()
147- 145+ 
148 def save_tensor(tensor, path):146 def save_tensor(tensor, path):
149 dirname = os.path.dirname(path)147 dirname = os.path.dirname(path)
150 if not os.path.exists(dirname):148 if not os.path.exists(dirname):
@@ -154,21 +152,22 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
154 def add_real_tensor(name, tensor):152 def add_real_tensor(name, tensor):
155 tensor_dir = npu_config.aot_inductor.repro_tensor_path153 tensor_dir = npu_config.aot_inductor.repro_tensor_path
156 if isinstance(tensor, FakeTensor):154 if isinstance(tensor, FakeTensor):
157- raise RuntimeError(f"Could not generate repro func because detected {name} is FakeTensor "155+ raise RuntimeError(
158- f"when trying to dump it. Set repro and debug_kernel false to avoid it.")156+ f"Could not generate repro func because detected {name} is FakeTensor "
157+ f"when trying to dump it. Set repro and debug_kernel false to avoid it."
158+ )
159 hash_name = get_hash(name)159 hash_name = get_hash(name)
160 tensor_path = os.path.join(os.getcwd(), tensor_dir, f"{hash_name}.pt")160 tensor_path = os.path.join(os.getcwd(), tensor_dir, f"{hash_name}.pt")
161 if name not in seen_constants:161 if name not in seen_constants:
162 save_tensor(tensor, tensor_path)162 save_tensor(tensor, tensor_path)
163 seen_constants.add(name)163 seen_constants.add(name)
164- output.writeline(164+ output.writeline(f"{name} = torch.load('{tensor_path}')")
165- f"{name} = torch.load('{tensor_path}')"
166- )
167 165 
168 def add_torchbind_input(name, value):166 def add_torchbind_input(name, value):
169 import pickle167 import pickle
170 168 
171 output.writeline(f"{name} = pickle.loads({pickle.dumps(value)!r})")169 output.writeline(f"{name} = pickle.loads({pickle.dumps(value)!r})")
170+ 
172 output.writelines(171 output.writelines(
173 ["", "", f"def repro_run({', '.join(V.graph.graph_inputs.keys())}):"]172 ["", "", f"def repro_run({', '.join(V.graph.graph_inputs.keys())}):"]
174 )173 )
@@ -193,11 +192,11 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
193 # these 'global var_name' lines192 # these 'global var_name' lines
194 output.writeline(f"global {name}")193 output.writeline(f"global {name}")
195 add_torchbind_input(name, torchbind_obj)194 add_torchbind_input(name, torchbind_obj)
196- 195+ 
197 call_str = f"call([{', '.join(V.graph.graph_inputs.keys())}])"196 call_str = f"call([{', '.join(V.graph.graph_inputs.keys())}])"
198 output.writeline(f"fn = lambda: {call_str}")197 output.writeline(f"fn = lambda: {call_str}")
199 output.writeline("return fn()")198 output.writeline("return fn()")
200- 199+ 
201 def add_benchmark_func(self, output):200 def add_benchmark_func(self, output):
202 def add_fake_input(name, shape, stride, device, dtype):201 def add_fake_input(name, shape, stride, device, dtype):
203 output.writeline(202 output.writeline(
@@ -244,7 +243,7 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
244 value.get_device(),243 value.get_device(),
245 value.get_dtype(),244 value.get_dtype(),
246 )245 )
247- 246+ 
248 call_str = f"repro_run({', '.join(V.graph.graph_inputs.keys())})"247 call_str = f"repro_run({', '.join(V.graph.graph_inputs.keys())})"
249 output.writeline(f"fn = lambda: {call_str}")248 output.writeline(f"fn = lambda: {call_str}")
250 output.writeline("return fn()")249 output.writeline("return fn()")
@@ -254,53 +253,50 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
254 if torch_npu.npu.aclnn._use_static_aclnn_kernel:253 if torch_npu.npu.aclnn._use_static_aclnn_kernel:
255 self.prefix.do_indent()254 self.prefix.do_indent()
256 with self.prefix.indent():255 with self.prefix.indent():
257- self.prefix.writeline('global has_initialized')256+ self.prefix.writeline("global has_initialized")
258- self.prefix.writeline('if not has_initialized:')257+ self.prefix.writeline("if not has_initialized:")
259 self.prefix.do_indent()258 self.prefix.do_indent()
260 with self.prefix.indent():259 with self.prefix.indent():
261- self.prefix.writeline('from torch_npu._inductor.npu_static_kernel import StaticKernelCompiler')260+ self.prefix.writeline(
262- self.prefix.writeline('static_kernel_compiler = StaticKernelCompiler()')261+ "from torch_npu._inductor.npu_static_kernel import StaticKernelCompiler"
263- self.prefix.writeline('static_kernel_compiler.__enter__()')262+ )
264- self.prefix.writeline('has_initialized = True')263+ self.prefix.writeline("static_kernel_compiler = StaticKernelCompiler()")
264+ self.prefix.writeline("static_kernel_compiler.__enter__()")
265+ self.prefix.writeline("has_initialized = True")
265 self.prefix.do_indent()266 self.prefix.do_indent()
266 267 
267 def generate_return(self, output_refs: list[str]) -> None:268 def generate_return(self, output_refs: list[str]) -> None:
268 if torch_npu.npu.aclnn._use_static_aclnn_kernel:269 if torch_npu.npu.aclnn._use_static_aclnn_kernel:
269 self.wrapper_call.do_unindent()270 self.wrapper_call.do_unindent()
270 with self.wrapper_call.indent():271 with self.wrapper_call.indent():
271- self.wrapper_call.writeline('if not has_initialized:')272+ self.wrapper_call.writeline("if not has_initialized:")
272 self.wrapper_call.do_indent()273 self.wrapper_call.do_indent()
273 with self.wrapper_call.indent():274 with self.wrapper_call.indent():
274- self.wrapper_call.writeline('exc_info=(None, None, None)')275+ self.wrapper_call.writeline("exc_info=(None, None, None)")
275- self.wrapper_call.writeline('static_kernel_compiler.__exit__(*exc_info)')276+ self.wrapper_call.writeline(
277+ "static_kernel_compiler.__exit__(*exc_info)"
278+ )
276 super().generate_return(output_refs)279 super().generate_return(output_refs)
277 280 
278 def define_kernel(281 def define_kernel(
279 self,282 self,
280 kernel_name: str,283 kernel_name: str,
281 kernel_body: str,284 kernel_body: str,
282- metadata: Optional[str] = None,285+ metadata: str | None = None,
283 gpu: bool = True,286 gpu: bool = True,
284- cpp_definition: Optional[str] = None,287+ cpp_definition: str | None = None,
285 ):288 ):
286 # 重写父类逻辑，将triton_heuristics.user_autotune替换为npu_triton_heuristics.user_autotune_npu，289 # 重写父类逻辑，将triton_heuristics.user_autotune替换为npu_triton_heuristics.user_autotune_npu，
287 # 将PrecomputedGrid替换为PrecomputedGridNpu，以适配NPU设备，避免core dump错误。290 # 将PrecomputedGrid替换为PrecomputedGridNpu，以适配NPU设备，避免core dump错误。
288 if "user_autotune" in kernel_body and "user_autotune_npu" not in kernel_body:291 if "user_autotune" in kernel_body and "user_autotune_npu" not in kernel_body:
289 kernel_body = kernel_body.replace(292 kernel_body = kernel_body.replace(
290 "triton_heuristics.user_autotune(",293 "triton_heuristics.user_autotune(",
291- "npu_triton_heuristics.user_autotune_npu("294+ "npu_triton_heuristics.user_autotune_npu(",
292 )295 )
296+ kernel_body = kernel_body.replace("PrecomputedGrid", "PrecomputedGridNpu")
297+ kernel_body = kernel_body.replace("FixedGrid", "FixedGridNpu")
298+ # import npu_triton_heuristicsd相关头文件
293 kernel_body = kernel_body.replace(299 kernel_body = kernel_body.replace(
294- "PrecomputedGrid",300+ "'''\n", "'''\n" + NPUIndexTritonKernel.gen_triton_ext_imports() + "\n"
295- "PrecomputedGridNpu"
296 )301 )
297- kernel_body = kernel_body.replace(302+ super().define_kernel(kernel_name, kernel_body, metadata, gpu, cpp_definition)
298- "FixedGrid",
299- "FixedGridNpu"
300- )
301- #import npu_triton_heuristicsd相关头文件
302- kernel_body = kernel_body.replace(
303- "'''\n",
304- "'''\n" + NPUIndexTritonKernel.gen_triton_ext_imports() + "\n"
305- )
306- super().define_kernel(kernel_name, kernel_body, metadata, gpu, cpp_definition)
Mtorch_npu/_inductor/config.py+50-28
@@ -1,9 +1,11 @@
1import logging1import logging
2import os # noqa: C1012import os # noqa: C101
3-from typing import Any, Callable, Dict, Optional, TYPE_CHECKING3+ 
4+from triton.runtime.driver import driver
5+ 
4import torch6import torch
5from torch._inductor import config7from torch._inductor import config
6-from triton.runtime.driver import driver8+ 
7 9 
8enable_npu_indexing = True10enable_npu_indexing = True
9 11 
@@ -23,14 +25,14 @@ prop = driver.active.utils.get_device_properties(device)
23num_cube_core = prop["num_aicore"]25num_cube_core = prop["num_aicore"]
24num_vector_core = prop["num_aicore"]26num_vector_core = prop["num_aicore"]
25 27 
26-# unit byte 28+# unit byte
27npu_block = 3229npu_block = 32
28 30 
29 31 
30# For debug32# For debug
31class aot_inductor:33class aot_inductor:
32 # If debug_kernel is set, codegen in python wrapper (output_code.py) and cpp wrapper (model.pt2)34 # If debug_kernel is set, codegen in python wrapper (output_code.py) and cpp wrapper (model.pt2)
33- # will be modified to dump fx graph and weights. Meanwhile, generate repro func in output_code.py. 35+ # will be modified to dump fx graph and weights. Meanwhile, generate repro func in output_code.py.
34 # Then, run aoti and output_code.py will dump tensor args before and after each triton kernel,36 # Then, run aoti and output_code.py will dump tensor args before and after each triton kernel,
35 # which can be used to detect which kernel is incorrect.37 # which can be used to detect which kernel is incorrect.
36 debug_kernel = os.environ.get("AOTI_ASCEND_DEBUG_KERNEL", False)38 debug_kernel = os.environ.get("AOTI_ASCEND_DEBUG_KERNEL", False)
@@ -40,7 +42,9 @@ class aot_inductor:
40 debug_kernel_in_run = False42 debug_kernel_in_run = False
41 43 
42 # Path that to be used for dump weights in aoti to reproduce when debug_kernel is set.44 # Path that to be used for dump weights in aoti to reproduce when debug_kernel is set.
43- repro_tensor_path = os.environ.get("AOTI_ASCEND_REPRO_TENSOR_PATH", "aoti_repro_tensors")45+ repro_tensor_path = os.environ.get(
46+ "AOTI_ASCEND_REPRO_TENSOR_PATH", "aoti_repro_tensors"
47+ )
44 48 
45 # Path that to be used for dump tensor args before and after triton kernel in aoti execute49 # Path that to be used for dump tensor args before and after triton kernel in aoti execute
46 # when debug_kernel is set.50 # when debug_kernel is set.
@@ -63,7 +67,7 @@ class _npugraph_trees:
63 @disable_cpu_input_check.setter67 @disable_cpu_input_check.setter
64 def disable_cpu_input_check(self, value):68 def disable_cpu_input_check(self, value):
65 self._disable_cpu_input_check = bool(value)69 self._disable_cpu_input_check = bool(value)
66- # When disable_cpu_input_check is True, set slow_path_cudagraph_asserts to True to skip the CPU check. 70+ # When disable_cpu_input_check is True, set slow_path_cudagraph_asserts to True to skip the CPU check.
67 if value:71 if value:
68 torch._inductor.config.triton.slow_path_cudagraph_asserts = False72 torch._inductor.config.triton.slow_path_cudagraph_asserts = False
69 73 
@@ -81,9 +85,11 @@ auto_fallback = os.environ.get("INDUCTOR_ASCEND_AUTO_FALLBACK", True)
81fallback_warning = os.environ.get("INDUCTOR_ASCEND_FALLBACK_WARNING", False)85fallback_warning = os.environ.get("INDUCTOR_ASCEND_FALLBACK_WARNING", False)
82 86 
83# Trace fx graph when lowering and dump.87# Trace fx graph when lowering and dump.
84-dump_fx_graph = os.environ.get("INDUCTOR_ASCEND_DUMP_FX_GRAPH", False) \88+dump_fx_graph = (
85- or check_accuracy \89+ os.environ.get("INDUCTOR_ASCEND_DUMP_FX_GRAPH", False)
86- or aot_inductor.debug_kernel90+ or check_accuracy
91+ or aot_inductor.debug_kernel
92+)
87# Specify kernel ids that to be force fallback to fx graph call.93# Specify kernel ids that to be force fallback to fx graph call.
88# Usage: `torch_npu._inductor.config.force_fallback_kernel_id = 'all' `94# Usage: `torch_npu._inductor.config.force_fallback_kernel_id = 'all' `
89# or `torch_npu._inductor.config.force_fallback_kernel_id = [1, 2, 10] `95# or `torch_npu._inductor.config.force_fallback_kernel_id = [1, 2, 10] `
@@ -91,33 +97,45 @@ dump_fx_graph = os.environ.get("INDUCTOR_ASCEND_DUMP_FX_GRAPH", False) \
91# (2) [1, 2, 10] means try to fallback kernel like triton_xxx_1, triton_xxx_2 and triton_xxx_1097# (2) [1, 2, 10] means try to fallback kernel like triton_xxx_1, triton_xxx_2 and triton_xxx_10
92force_fallback_kernel_id = []98force_fallback_kernel_id = []
93 99 
100+# Control whether to skip stride assertions for ops that may change stride
101+# at runtime (like _to_copy on NPU forcing Contiguous memory format).
102+#
103+# Usage:
104+# - Skip specific ops: skip_specific_stride_asserts = [torch.ops.aten._to_copy.default, ...]
105+# - Disable skip: skip_specific_stride_asserts = [] (default)
106+skip_specific_stride_asserts = []
107+ 
94acc_comp_tol = {108acc_comp_tol = {
95- torch.float32: {'rtol': 1.3e-6, 'atol': 1e-5},109+ torch.float32: {"rtol": 1.3e-6, "atol": 1e-5},
96- torch.float16: {'rtol': 1e-3, 'atol': 1e-5},110+ torch.float16: {"rtol": 1e-3, "atol": 1e-5},
97- torch.bfloat16: {'rtol': 1.6e-2, 'atol': 1e-5},111+ torch.bfloat16: {"rtol": 1.6e-2, "atol": 1e-5},
98- "default": {'rtol': 1.3e-6, 'atol': 1e-5},112+ "default": {"rtol": 1.3e-6, "atol": 1e-5},
99}113}
100 114 
101-if ("Ascend910B" in target.arch):115+if "Ascend910B" in target.arch:
102 num_vector_core = num_cube_core * 2116 num_vector_core = num_cube_core * 2
103 117 
104-log_level_env = os.getenv('INDUCTOR_ASCEND_LOG_LEVEL', 'WARNING').upper()118+log_level_env = os.getenv("INDUCTOR_ASCEND_LOG_LEVEL", "WARNING").upper()
105log_level_mapping = {119log_level_mapping = {
106- 'DEBUG': logging.DEBUG,120+ "DEBUG": logging.DEBUG,
107- 'INFO': logging.INFO,121+ "INFO": logging.INFO,
108- 'WARNING': logging.WARNING,122+ "WARNING": logging.WARNING,
109- 'ERROR': logging.ERROR,123+ "ERROR": logging.ERROR,
110- 'CRITICAL': logging.CRITICAL124+ "CRITICAL": logging.CRITICAL,
111}125}
112log_level = log_level_mapping.get(log_level_env.upper(), logging.INFO)126log_level = log_level_mapping.get(log_level_env.upper(), logging.INFO)
113-logging.basicConfig(127+logging.basicConfig(level=log_level, format="%(asctime)s - %(levelname)s - %(message)s")
114- level=log_level,
115- format='%(asctime)s - %(levelname)s - %(message)s'
116-)
117log = logging.getLogger(__name__)128log = logging.getLogger(__name__)
118 129 
119-aggresive_autotune = os.getenv("INDUCTOR_ASCEND_AGGRESSIVE_AUTOTUNE", '0').lower() in ('1', 'true')130+aggresive_autotune = os.getenv("INDUCTOR_ASCEND_AGGRESSIVE_AUTOTUNE", "0").lower() in (
120-inductor_static_mode = os.environ.get('INDUCTOR_STATIC_MODE', '0').lower() in ('1', 'yes', 'true')131+ "1",
132+ "true",
133+)
134+inductor_static_mode = os.environ.get("INDUCTOR_STATIC_MODE", "0").lower() in (
135+ "1",
136+ "yes",
137+ "true",
138+)
121profile_path = "./profile_result/"139profile_path = "./profile_result/"
122 140 
123 141 
@@ -126,14 +144,18 @@ def set_compile_threads():
126 torchinductor_compile_threads = int(os.environ["TORCHINDUCTOR_COMPILE_THREADS"])144 torchinductor_compile_threads = int(os.environ["TORCHINDUCTOR_COMPILE_THREADS"])
127 if torchinductor_compile_threads == 1:145 if torchinductor_compile_threads == 1:
128 return146 return
129- log.warning(f"TORCHINDUCTOR_COMPILE_THREADS is set to {torchinductor_compile_threads}, "147+ log.warning(
130- "but currently only support 1. It will be modified to 1.")148+ "TORCHINDUCTOR_COMPILE_THREADS is set to %s, "
149+ "but currently only support 1. It will be modified to 1.",
150+ torchinductor_compile_threads,
151+ )
131 152 
132 os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"153 os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
133 torch._inductor.config.compile_threads = 1154 torch._inductor.config.compile_threads = 1
134 155 
135 def get_env_num_workers():156 def get_env_num_workers():
136 return 1157 return 1
158+ 
137 torch._inductor.select_algorithm.get_env_num_workers = get_env_num_workers159 torch._inductor.select_algorithm.get_env_num_workers = get_env_num_workers
138 160 
139 161 
Mtorch_npu/_inductor/ir.py+40-4
@@ -1,8 +1,11 @@
1import itertools1import itertools
2+ 
2import torch3import torch
3-from torch._inductor.virtualized import ops, OpsValue, V
4-from torch._inductor.ir import log, Layout
5from torch._inductor import config4from torch._inductor import config
5+from torch._inductor.ir import ExternKernel, Layout, log
6+from torch._inductor.virtualized import V
7+ 
8+from . import config as npu_config
6 9 
7 10 
8def patch_fallback_kernel_codegen():11def patch_fallback_kernel_codegen():
@@ -10,10 +13,13 @@ def patch_fallback_kernel_codegen():
10 kernel = self.op_overload13 kernel = self.op_overload
11 if kernel.namespace == "aten": # type: ignore[union-attr]14 if kernel.namespace == "aten": # type: ignore[union-attr]
12 if not isinstance(kernel, torch._ops.OpOverload):15 if not isinstance(kernel, torch._ops.OpOverload):
13- raise AssertionError(f"kernel should be OpOverload, but got {type(kernel)}")16+ raise AssertionError(
17+ f"kernel should be OpOverload, but got {type(kernel)}"
18+ )
14 if V.graph.cpp_wrapper:19 if V.graph.cpp_wrapper:
15 # Fallback all npu op to proxy executor and warn when gpu do not.20 # Fallback all npu op to proxy executor and warn when gpu do not.
16 from torchgen.aoti.fallback_ops import inductor_fallback_ops21 from torchgen.aoti.fallback_ops import inductor_fallback_ops
22+ 
17 self.use_runtime_dispatch = True23 self.use_runtime_dispatch = True
18 if str(kernel) in inductor_fallback_ops:24 if str(kernel) in inductor_fallback_ops:
19 log.warning(25 log.warning(
@@ -85,4 +91,34 @@ def patch_fallback_kernel_codegen():
85 self.codegen_unbacked_symbol_defs(wrapper)91 self.codegen_unbacked_symbol_defs(wrapper)
86 92 
87 from torch._inductor.ir import FallbackKernel93 from torch._inductor.ir import FallbackKernel
88- FallbackKernel.codegen = codegen_npu94+ 
95+ FallbackKernel.codegen = codegen_npu
96+ 
97+ 
98+def patch_extern_kernel_codegen_size_asserts():
99+ original_codegen_size_asserts = ExternKernel.codegen_size_asserts
100+ 
101+ def npu_codegen_size_asserts(self, wrapper):
102+ fx_node = getattr(self, "fx_node", None)
103+ 
104+ should_skip = False
105+ if fx_node and fx_node.target:
106+ skip_config = npu_config.skip_specific_stride_asserts
107+ 
108+ # Only skip ops that are in the configured list
109+ if isinstance(skip_config, (list, tuple)):
110+ should_skip = fx_node.target in skip_config
111+ 
112+ if should_skip:
113+ if config.size_asserts and not V.graph.cpp_wrapper:
114+ from torch._inductor.utils import sympy_product
115+ 
116+ if sympy_product(self.get_size()) == 0:
117+ return
118+ wrapper.writeline(
119+ f"# NPU: Skipping stride assertion for {fx_node.target} (stride may change at runtime)"
120+ )
121+ else:
122+ original_codegen_size_asserts(self, wrapper)
123+ 
124+ ExternKernel.codegen_size_asserts = npu_codegen_size_asserts
Mtorch_npu/utils/_npu_meta_registration.py+107-27
@@ -1,32 +1,37 @@
1-import os
2-import sys
3import operator1import operator
4-from functools import wraps, reduce, lru_cache2+from collections.abc import Callable
5-from typing import Callable, Optional3+from functools import lru_cache, reduce, wraps
4+ 
6import torch5import torch
6+import torch_npu # noqa: F401
7from torch import Tensor7from torch import Tensor
8-from torch._ops import OpOverload, OpOverloadPacket
9-from torch._subclasses import fake_tensor as _subclasses_fake_tensor
10from torch._C import DispatchKey8from torch._C import DispatchKey
11-from torch._refs import div as refs_div, _broadcast_shapes9+from torch._decomp import decomposition_table, meta_table
12from torch._inductor import decomposition as inductor_decompo10from torch._inductor import decomposition as inductor_decompo
13-from torch._prims_common import corresponding_real_dtype, corresponding_complex_dtype11+from torch._ops import OpOverload, OpOverloadPacket
14from torch._prims_common.wrappers import out_wrapper12from torch._prims_common.wrappers import out_wrapper
15-from torch._decomp import decomposition_table, decompositions_for_rng, get_decompositions13+from torch._subclasses import fake_tensor as _subclasses_fake_tensor
16-from torch._dynamo.symbolic_convert import break_graph_if_unsupported, InstructionTranslatorBase, stack_op14+ 
17-from torch._dynamo.exc import Unsupported
18-from torch._dynamo.variables.lists import TupleVariable
19-from torch._dynamo.variables.nn_module import NNModuleVariable
20-import torch_npu
21 15 
22aten = torch.ops.aten16aten = torch.ops.aten
23npu = torch.ops.npu17npu = torch.ops.npu
24 18 
19+META_BLACKLIST = {
20+ "aten::empty_strided", # causing infinite recursion, test_meta.py
21+ "aten::clone", # causing infinite recursion
22+ "aten::_to_copy", # causing infinite recursion, test_serialization.py -k test_tensor_subclass_getstate_overwrite # noqa: B950
23+ "aten::copy_", # Exception not raised, test_torch.py -k test_storage_meta_errors_cpu_int64 # noqa: B950
24+ "aten::constant_pad_nd", # requires_grad mismatch, test_ops.py -k test_fake_crossref_backward_amp_istft_cuda_float32 # noqa: B950
25+ "aten::rot90", # requires_grad mismatch! test_ops.py -k test_fake_crossref_backward_amp_rot90_cuda_float32 # noqa: B950
26+ "aten::as_strided_scatter", # requires_grad mismatch, test_ops.py -k test_fake_crossref_backward_no_amp_as_strided_scatter_cuda_float32 # noqa: B950
27+}
28+ 
25 29 
26def run_once(f):30def run_once(f):
27 """Runs a function (successfully) only once.31 """Runs a function (successfully) only once.
28 The running can be reset by setting the `has_run` attribute to False32 The running can be reset by setting the `has_run` attribute to False
29 """33 """
34+ 
30 @wraps(f)35 @wraps(f)
31 def wrapper(*args, **kwargs):36 def wrapper(*args, **kwargs):
32 if not wrapper.has_run:37 if not wrapper.has_run:
@@ -34,6 +39,7 @@ def run_once(f):
34 wrapper.has_run = True39 wrapper.has_run = True
35 return result40 return result
36 return None41 return None
42+ 
37 wrapper.has_run = False43 wrapper.has_run = False
38 return wrapper44 return wrapper
39 45 
@@ -56,7 +62,9 @@ def _add_op_to_meta_table(op, fn, avoid_fallback_flag=False, inductor_decomp=Fal
56 62 
57 for op_overload in overloads:63 for op_overload in overloads:
58 if op_overload in npu_meta_table:64 if op_overload in npu_meta_table:
59- raise RuntimeError(f"duplicate registrations for npu_meta_table {op_overload}")65+ raise RuntimeError(
66+ f"duplicate registrations for npu_meta_table {op_overload}"
67+ )
60 npu_meta_table[op_overload] = fn68 npu_meta_table[op_overload] = fn
61 if avoid_fallback_flag:69 if avoid_fallback_flag:
62 avoid_make_fallback_table.append(op_overload)70 avoid_make_fallback_table.append(op_overload)
@@ -65,31 +73,36 @@ def _add_op_to_meta_table(op, fn, avoid_fallback_flag=False, inductor_decomp=Fal
65 73 
66 74 
67def patch_torch_inductor_decompositions():75def patch_torch_inductor_decompositions():
68- '''76+ """
69 TorchInductor traces compiled backward with its own decomposition table.77 TorchInductor traces compiled backward with its own decomposition table.
70 Only patch ops that explicitly opted in via inductor_decomp=True so we78 Only patch ops that explicitly opted in via inductor_decomp=True so we
71 don't accidentally overwrite unrelated inductor decompositions.79 don't accidentally overwrite unrelated inductor decompositions.
72- '''80+ """
73 import torch._inductor.decomposition as inductor_decomposition81 import torch._inductor.decomposition as inductor_decomposition
74 82 
75 for op_overload in inductor_decomp_table:83 for op_overload in inductor_decomp_table:
76 if op_overload in npu_meta_table:84 if op_overload in npu_meta_table:
77- inductor_decomposition.decompositions[op_overload] = npu_meta_table[op_overload]85+ inductor_decomposition.decompositions[op_overload] = npu_meta_table[
86+ op_overload
87+ ]
78 88 
79 89 
80def patch_torch_decomp_decompositions():90def patch_torch_decomp_decompositions():
81- '''91+ """
82 Because source torch_decomp_decompositions only enable the decompositions in92 Because source torch_decomp_decompositions only enable the decompositions in
83 torch/_decomp/decompositions.py. Patch it to make decompositions in this file work.93 torch/_decomp/decompositions.py. Patch it to make decompositions in this file work.
84- '''94+ """
85 src_func = _subclasses_fake_tensor.torch_decomp_decompositions95 src_func = _subclasses_fake_tensor.torch_decomp_decompositions
86 96 
87 @lru_cache(None)97 @lru_cache(None)
88 def torch_decomp_decompositions_new(func):98 def torch_decomp_decompositions_new(func):
89- if func in npu_meta_table.keys():99+ if func in npu_meta_table:
90 return True100 return True
91 return src_func(func)101 return src_func(func)
92- _subclasses_fake_tensor.torch_decomp_decompositions = torch_decomp_decompositions_new102+ 
103+ _subclasses_fake_tensor.torch_decomp_decompositions = (
104+ torch_decomp_decompositions_new
105+ )
93 106 
94 107 
95def register_meta_npu(op, avoid_fallback_flag=False, inductor_decomp=False):108def register_meta_npu(op, avoid_fallback_flag=False, inductor_decomp=False):
@@ -102,26 +115,40 @@ def register_meta_npu(op, avoid_fallback_flag=False, inductor_decomp=False):
102 115 
103@run_once116@run_once
104def npu_patch_meta():117def npu_patch_meta():
105- '''118+ """
106 Torch official register decompostions and meta func for some aten ops,119 Torch official register decompostions and meta func for some aten ops,
107 which will raise conflict when npu outputs' dtype and shape are different120 which will raise conflict when npu outputs' dtype and shape are different
108 from native impl. Delete decompositions and meta func of these ops and add121 from native impl. Delete decompositions and meta func of these ops and add
109 npu decompositions and meta func.122 npu decompositions and meta func.
110- '''123+ """
124+ _meta_library = torch.library.Library("aten", "IMPL", "Meta")
111 for op_overload, fn in npu_meta_table.items():125 for op_overload, fn in npu_meta_table.items():
112 if not isinstance(op_overload, OpOverload):126 if not isinstance(op_overload, OpOverload):
113 raise AssertionError("op_overload must be instance of OpOverload.")127 raise AssertionError("op_overload must be instance of OpOverload.")
128+ op_name = op_overload.name()
114 if op_overload not in avoid_make_fallback_table:129 if op_overload not in avoid_make_fallback_table:
115 decomposition_table[op_overload] = fn130 decomposition_table[op_overload] = fn
116 op_overload.py_kernels.pop(DispatchKey.Meta, None)131 op_overload.py_kernels.pop(DispatchKey.Meta, None)
117 op_overload.py_impl(DispatchKey.Meta)(fn)132 op_overload.py_impl(DispatchKey.Meta)(fn)
118 133 
134+ if op_name not in META_BLACKLIST:
135+ meta_table[op_overload] = fn
136+ key = (
137+ _meta_library.ns
138+ + "/"
139+ + op_name.split("::")[-1]
140+ + "/"
141+ + _meta_library.dispatch_key
142+ )
143+ if key in torch.library._impls:
144+ torch.library._impls.remove(key)
145+ _meta_library.impl(op_overload, fn)
146+ 
119 inductor_decompo.fast_random_decomps.cache_clear()147 inductor_decompo.fast_random_decomps.cache_clear()
120 patch_torch_decomp_decompositions()148 patch_torch_decomp_decompositions()
121 patch_torch_inductor_decompositions()149 patch_torch_inductor_decompositions()
122 150 
123 151 
124- 
125@register_meta_npu(aten.index_put.default)152@register_meta_npu(aten.index_put.default)
126def meta_index_put_patch(self, indices, values, accumulate=False):153def meta_index_put_patch(self, indices, values, accumulate=False):
127 return self.new_empty(self.shape)154 return self.new_empty(self.shape)
@@ -129,7 +156,7 @@ def meta_index_put_patch(self, indices, values, accumulate=False):
129 156 
130@register_meta_npu(aten.native_dropout, inductor_decomp=True)157@register_meta_npu(aten.native_dropout, inductor_decomp=True)
131@out_wrapper("out0", "out1")158@out_wrapper("out0", "out1")
132-def meta_native_dropout_patch(tensor_input: Tensor, p: float, train: Optional[bool]):159+def meta_native_dropout_patch(tensor_input: Tensor, p: float, train: bool | None):
133 if torch._inductor.config.fallback_random:160 if torch._inductor.config.fallback_random:
134 if train and p != 0:161 if train and p != 0:
135 if tensor_input.is_meta:162 if tensor_input.is_meta:
@@ -144,6 +171,7 @@ def meta_native_dropout_patch(tensor_input: Tensor, p: float, train: Optional[bo
144 return (tensor_input, torch.ones_like(tensor_input, dtype=torch.bool))171 return (tensor_input, torch.ones_like(tensor_input, dtype=torch.bool))
145 else:172 else:
146 from torch._decomp.decompositions import native_dropout173 from torch._decomp.decompositions import native_dropout
174+ 
147 return native_dropout(tensor_input, p, train)175 return native_dropout(tensor_input, p, train)
148 176 
149 177 
@@ -157,4 +185,56 @@ def meta_native_dropout_backward_patch(grad_output: Tensor, mask: Tensor, scale:
157 return torch.ops.npu.npu_dropout_backward(grad_output, mask, p)185 return torch.ops.npu.npu_dropout_backward(grad_output, mask, p)
158 else:186 else:
159 from torch._decomp.decompositions import native_dropout_backward187 from torch._decomp.decompositions import native_dropout_backward
160- return native_dropout_backward(grad_output, mask, scale)188+ 
189+ return native_dropout_backward(grad_output, mask, scale)
190+ 
191+ 
192+@register_meta_npu(aten._to_copy.default, inductor_decomp=True)
193+def meta_to_copy_default(
194+ x,
195+ *,
196+ dtype: torch.dtype | None = None,
197+ layout=None,
198+ device: torch.device | None = None,
199+ pin_memory: bool = False,
200+ non_blocking: bool = False,
201+ memory_format: torch.memory_format | None = None,
202+):
203+ if layout and layout != torch.strided:
204+ raise AssertionError(f"Only strided layout is supported, got {layout}")
205+ if pin_memory:
206+ raise AssertionError("pin_memory is not supported")
207+ if not isinstance(x, (torch.Tensor, int, float, bool, complex)):
208+ raise AssertionError(f"x must be Tensor or scalar type, got {type(x)}")
209+ 
210+ out_memory_format = (
211+ memory_format if memory_format is not None else torch.contiguous_format
212+ )
213+ 
214+ if device is None and dtype is None and memory_format is None:
215+ if isinstance(x, torch.Tensor):
216+ return x.clone(memory_format=out_memory_format)
217+ else:
218+ return x
219+ dtype_converted = False
220+ 
221+ if isinstance(x, torch.Tensor):
222+ x_tensor = x
223+ else:
224+ x_tensor = torch.scalar_tensor(x)
225+ 
226+ if device is not None and device != x_tensor.device:
227+ # avoid conversions on cpu
228+ if dtype is not None and device.type == "cpu":
229+ x_tensor = torch._prims.convert_element_type(x_tensor, dtype)
230+ dtype_converted = True
231+ x_tensor = torch._prims.device_put(x_tensor, device, non_blocking)
232+ 
233+ if dtype is not None and not dtype_converted:
234+ x_tensor = torch._prims.convert_element_type(x_tensor, dtype)
235+ dtype_converted = True
236+ 
237+ if memory_format is not None: # no ref/prim for memory format
238+ return torch.clone(x_tensor, memory_format=memory_format)
239+ else:
240+ return torch.clone(x_tensor, memory_format=out_memory_format)