已合并
lazy_init #39541
cuiduo创建于 6月29日
lazy_init #39541
已合并
共 6 个文件变更+202-192
| @@ -9,8 +9,6 @@ from torch.testing._internal.common_utils import ( | |||
| 9 | instantiate_parametrized_tests, | 9 | instantiate_parametrized_tests, |
| 10 | ) | 10 | ) |
| 11 | from testutils import TestUtils | 11 | from testutils import TestUtils |
| 12 | -import torch_npu | ||
| 13 | -import torch_npu._inductor | ||
| 14 | 12 | ||
| 15 | 13 | ||
| 16 | class TestWrapTriton(TestUtils): | 14 | class TestWrapTriton(TestUtils): |
| @@ -2,7 +2,7 @@ import torch | |||
| 2 | 2 | ||
| 3 | import torch_npu | 3 | import torch_npu |
| 4 | from torch_npu.testing.testcase import TestCase, run_tests | 4 | from torch_npu.testing.testcase import TestCase, run_tests |
| 5 | - | 5 | +import torch_npu._inductor # noqa: F401 |
 | |||
| 6 | 6 | ||
| 7 | class TestNpuStream(TestCase): | 7 | class TestNpuStream(TestCase): |
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。   | |||
| 8 | 8 | ||
| @@ -1,5 +1,5 @@ | |||
| 1 | import os | 1 | import os |
| 2 | - | 2 | +from torch_npu.utils._dynamo import _dynamo_register_interface_for_device, patch_SkipFunctionVariable, patch_TensorVariable_call_method |
| 3 | # All backends need npu/cpu/mps device_op_overrides. | 3 | # All backends need npu/cpu/mps device_op_overrides. |
| 4 | from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system | 4 | from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system |
| 5 | from .graph import patch_codegen_with_cpp_wrapper | 5 | from .graph import patch_codegen_with_cpp_wrapper |
| @@ -8,6 +8,7 @@ from .cpp_builder import patch_get_optimization_cflags | |||
| 8 | from ._npu_meta_registration import npu_patch_meta | 8 | from ._npu_meta_registration import npu_patch_meta |
| 9 | 9 | ||
| 10 | npu_patch_meta() | 10 | npu_patch_meta() |
| 11 | +_dynamo_register_interface_for_device() | ||
| 11 | register_device_op_overrides_npu() | 12 | register_device_op_overrides_npu() |
| 12 | patch_has_triton() | 13 | patch_has_triton() |
| 13 | patch_is_gpu() | 14 | patch_is_gpu() |
| @@ -1,27 +1,9 @@ | |||
| 1 | -from torch._dynamo.device_interface import register_interface_for_device | ||
| 2 | - | ||
| 3 | -from torch_npu.utils._dynamo_device import NpuInterface | ||
| 4 | - | ||
| 5 | - | ||
| 6 | -def _dynamo_register_interface_for_device(): | ||
| 7 | - register_interface_for_device("npu", NpuInterface) | ||
| 8 | - for i in range(32): | ||
| 9 | - register_interface_for_device(f"npu:{i}", NpuInterface) | ||
| 10 | - | ||
| 11 | - | ||
| 12 | def register_dynamo_backends(): | 1 | def register_dynamo_backends(): |
| 13 | from torch_npu.dynamo import _register_backends | 2 | from torch_npu.dynamo import _register_backends |
| 14 | 3 | ||
| 15 | _register_backends() | 4 | _register_backends() |
| 16 | 5 | ||
| 17 | 6 | ||
| 18 | -def register_dynamo_device_interface(): | ||
| 19 | - """ | ||
| 20 | - Register NPU device interface for Dynamo | ||
| 21 | - """ | ||
| 22 | - _dynamo_register_interface_for_device() | ||
| 23 | - | ||
| 24 | - | ||
| 25 | def register_dynamo_trace_rules(): | 7 | def register_dynamo_trace_rules(): |
| 26 | """ | 8 | """ |
| 27 | # Support stream into Dynamo charts. Enable Dynamo to recognize NPU | 9 | # Support stream into Dynamo charts. Enable Dynamo to recognize NPU |
| @@ -58,12 +58,10 @@ def _register_dynamo(): | |||
| 58 | """ | 58 | """ |
| 59 | from torch_npu._init.registry.dynamo import ( | 59 | from torch_npu._init.registry.dynamo import ( |
| 60 | register_dynamo_backends, | 60 | register_dynamo_backends, |
| 61 | - register_dynamo_device_interface, | ||
| 62 | register_dynamo_trace_rules, | 61 | register_dynamo_trace_rules, |
| 63 | ) | 62 | ) |
| 64 | 63 | ||
| 65 | register_dynamo_backends() | 64 | register_dynamo_backends() |
| 66 | - register_dynamo_device_interface() | ||
| 67 | 65 | ||
| 68 | # Do not repeat this call for register_dynamo_trace_rules appends rules into | 66 | # Do not repeat this call for register_dynamo_trace_rules appends rules into |
 🟠 High Priority 变更从 现在 device interface 仅在以下时机注册：
这会导致任何在 受影响的已知代码： | |||
| 69 | # Dynamo's global rules maps. | 67 | # Dynamo's global rules maps. |
| @@ -4,25 +4,11 @@ import sys | |||
| 4 | import logging | 4 | import logging |
| 5 | from typing import Any, Optional, TYPE_CHECKING | 5 | from typing import Any, Optional, TYPE_CHECKING |
| 6 | import importlib | 6 | import importlib |
| 7 | +import functools | ||
| 7 | 8 | ||
| 8 | import torch | 9 | import torch |
| 9 | import torch_npu | 10 | import torch_npu |
| 10 | from torch import _TorchCompileWrapper | 11 | from torch import _TorchCompileWrapper |
| 11 | -from torch._dynamo import optimize | ||
| 12 | -from torch._dynamo.utils import tensortype_to_dtype | ||
| 13 | -from torch._dynamo.variables.base import VariableTracker | ||
| 14 | -from torch._dynamo.variables.constant import ConstantVariable | ||
| 15 | -from torch._dynamo.variables.ctx_manager import AutocastModeVariable | ||
| 16 | -from torch._dynamo.variables.functions import SkipFunctionVariable | ||
| 17 | -from torch._dynamo.variables.lists import TupleVariable | ||
| 18 | -from torch._dynamo.variables.streams import StreamContextVariable, StreamVariable | ||
| 19 | -from torch._dynamo.variables.tensor import TensorVariable | ||
| 20 | -from torch._dynamo.variables.torch import ( | ||
| 21 | - TorchCtxManagerClassVariable, | ||
| 22 | - TorchInGraphFunctionVariable, | ||
| 23 | -) | ||
| 24 | -from torch._dynamo.variables.user_defined import UserDefinedClassVariable | ||
| 25 | -from torch_npu.dynamo import _get_global_npu_backend | ||
| 26 | 12 | ||
| 27 | 13 | ||
| 28 | if TYPE_CHECKING: | 14 | if TYPE_CHECKING: |
| @@ -31,87 +17,71 @@ if TYPE_CHECKING: | |||
| 31 | use_jit_script = False | 17 | use_jit_script = False |
| 32 | log = logging.getLogger(__name__) | 18 | log = logging.getLogger(__name__) |
| 33 | 19 | ||
| 34 | -class NPUTorchCtxManagerClassVariable(TorchCtxManagerClassVariable): | 20 | +def _create_npu_autocast_mode_variable(func, args, kwargs): |
| 35 | - def call_function(self, tx, args, kwargs): | 21 | + from torch._dynamo.variables.ctx_manager import AutocastModeVariable |
| 36 | - return NPUAutocastModeVariable.create(self.value, args, kwargs) | 22 | + from torch._dynamo.variables.base import VariableTracker |
| 23 | + bound_args = inspect.signature(func).bind(*args, **kwargs) | ||
| 24 | + bound_args.apply_defaults() | ||
| 25 | + target_values = [] | ||
| 26 | + kwargs.clear() | ||
| 37 | 27 | ||
| 28 | + for key in ["device_type", "dtype", "enabled", "cache_enabled"]: | ||
| 29 | + if key == "device_type" and func in [ | ||
| 30 | + torch_npu.npu.amp.autocast, | ||
| 31 | + ]: | ||
| 32 | + arg = "npu" if func is torch_npu.npu.amp.autocast else "cpu" | ||
| 33 | + else: | ||
| 34 | + arg = bound_args.arguments[key] | ||
| 35 | + if isinstance(arg, VariableTracker): | ||
| 36 | + target_values.append(arg.as_python_constant()) | ||
| 37 | + else: | ||
| 38 | + target_values.append(arg) | ||
| 38 | 39 | ||
| 39 | -class NPUAutocastModeVariable(AutocastModeVariable): | 40 | + var = AutocastModeVariable(target_values, initial_values=None, **kwargs) |
| 40 | - @staticmethod | 41 | + return var |
| 41 | - def create(func, args, kwargs): | ||
| 42 | - bound_args = inspect.signature(func).bind(*args, **kwargs) | ||
| 43 | - bound_args.apply_defaults() | ||
| 44 | - target_values = [] | ||
| 45 | - kwargs.clear() | ||
| 46 | 42 | ||
| 47 | - for key in ["device_type", "dtype", "enabled", "cache_enabled"]: | 43 | +def patch_SkipFunctionVariable(): |
| 48 | - if key == "device_type" and func in [ | 44 | + from torch._dynamo.variables.functions import SkipFunctionVariable |
| 49 | - torch_npu.npu.amp.autocast, | 45 | + from torch._dynamo.variables.torch import TorchInGraphFunctionVariable |
| 50 | - ]: | 46 | + def SkipFunctionVariable__new__(cls, value, reason=None, **kwargs): |
| 51 | - arg = "npu" if func is torch_npu.npu.amp.autocast else "cpu" | 47 | + if value in [ |
| 52 | - else: | 48 | + torch.npu.stream, |
| 53 | - arg = bound_args.arguments[key] | 49 | + torch_npu.npu.stream, |
| 54 | - if isinstance(arg, VariableTracker): | 50 | + torch_npu.npu.utils.stream, |
| 55 | - target_values.append(arg.as_python_constant()) | 51 | + ]: |
| 56 | - else: | 52 | + return TorchInGraphFunctionVariable(value, **kwargs) |
| 57 | - target_values.append(arg) | 53 | + return cls.__new__raw(cls) |
| 58 | 54 | ||
| 59 | - var = AutocastModeVariable(target_values, initial_values=None, **kwargs) | 55 | + SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__ |
| 60 | - return var | 56 | + SkipFunctionVariable.__new__ = SkipFunctionVariable__new__ |
| 61 | 57 | ||
| 58 | +def patch_TensorVariable_call_method(): | ||
| 59 | + from torch._dynamo.variables.tensor import TensorVariable | ||
| 60 | + from torch._dynamo.utils import tensortype_to_dtype | ||
| 61 | + from torch._dynamo.variables.constant import ConstantVariable | ||
| 62 | + from torch._dynamo.variables.lists import TupleVariable | ||
| 62 | 63 | ||
| 63 | -def UserDefinedClassVariable__new__(cls, value, **kwargs): | 64 | + def TensorVariable_call_method(self, tx, name, args, kwargs): |
| 64 | - if value in [ | 65 | + if ( |
| 65 | - torch.npu.amp.autocast, | 66 | + name == "type" |
| 66 | - torch_npu.npu.amp.autocast, | 67 | + and self.dtype is not None |
| 67 | - torch.npu.amp.autocast_mode.autocast, | 68 | + and len(args) == 0 |
| 68 | - torch_npu.npu.amp.autocast_mode.autocast, | 69 | + and isinstance(self.device, torch.device) |
| 69 | - ]: | 70 | + and self.device.type == "npu" |
| 70 | - return NPUTorchCtxManagerClassVariable(value, **kwargs) | 71 | + ): |
| 71 | - elif value in [ | 72 | + tensortype = next(k for k, v in tensortype_to_dtype.items() if self.dtype in v) |
| 72 | - torch_npu.npu.BoolTensor, | 73 | + constant_result = ConstantVariable.create(f"torch.npu.{tensortype.__name__}") |
| 73 | - torch_npu.npu.ByteTensor, | ||
| 74 | - torch_npu.npu.CharTensor, | ||
| 75 | - torch_npu.npu.DoubleTensor, | ||
| 76 | - torch_npu.npu.FloatTensor, | ||
| 77 | - torch_npu.npu.HalfTensor, | ||
| 78 | - torch_npu.npu.IntTensor, | ||
| 79 | - torch_npu.npu.LongTensor, | ||
| 80 | - torch_npu.npu.ShortTensor, | ||
| 81 | - torch_npu.npu.BFloat16Tensor, | ||
| 82 | - ]: | ||
| 83 | - return TorchInGraphFunctionVariable(value, **kwargs) | ||
| 84 | - return cls.__new__raw(cls) | ||
| 85 | 74 | ||
| 75 | + if len(args) == 1: | ||
| 76 | + return constant_result.getitem_const(args[0]) | ||
| 77 | + elif args: | ||
| 78 | + return TupleVariable([constant_result.getitem_const(a) for a in args]) | ||
| 79 | + return constant_result | ||
| 80 | + else: | ||
| 81 | + return TensorVariable.call_method_raw(self, tx, name, args, kwargs) | ||
| 86 | 82 | ||
| 87 | -def SkipFunctionVariable__new__(cls, value, reason=None, **kwargs): | 83 | + TensorVariable.call_method_raw = TensorVariable.call_method |
| 88 | - if value in [ | 84 | + TensorVariable.call_method = TensorVariable_call_method |
| 89 | - torch.npu.stream, | ||
| 90 | - torch_npu.npu.stream, | ||
| 91 | - torch_npu.npu.utils.stream, | ||
| 92 | - ]: | ||
| 93 | - return TorchInGraphFunctionVariable(value, **kwargs) | ||
| 94 | - return cls.__new__raw(cls) | ||
| 95 | - | ||
| 96 | - | ||
| 97 | -def TensorVariable_call_method(self, tx, name, args, kwargs): | ||
| 98 | - if ( | ||
| 99 | - name == "type" | ||
| 100 | - and self.dtype is not None | ||
| 101 | - and len(args) == 0 | ||
| 102 | - and isinstance(self.device, torch.device) | ||
| 103 | - and self.device.type == "npu" | ||
| 104 | - ): | ||
| 105 | - tensortype = next(k for k, v in tensortype_to_dtype.items() if self.dtype in v) | ||
| 106 | - constant_result = ConstantVariable.create(f"torch.npu.{tensortype.__name__}") | ||
| 107 | - | ||
| 108 | - if len(args) == 1: | ||
| 109 | - return constant_result.getitem_const(args[0]) | ||
| 110 | - elif args: | ||
| 111 | - return TupleVariable([constant_result.getitem_const(a) for a in args]) | ||
| 112 | - return constant_result | ||
| 113 | - else: | ||
| 114 | - return TensorVariable.call_method_raw(self, tx, name, args, kwargs) | ||
| 115 | 85 | ||
| 116 | 86 | ||
| 117 | class _InductorNpuRegistry: | 87 | class _InductorNpuRegistry: |
| @@ -223,6 +193,7 @@ def patch_inductor_wrapper(): | |||
| 223 | return ori_dict | 193 | return ori_dict |
| 224 | 194 | ||
| 225 | def new_init(self, mode, options, dynamic, name=None): | 195 | def new_init(self, mode, options, dynamic, name=None): |
| 196 | + add_dynamo_methods_init() | ||
| 226 | if name is not None: | 197 | if name is not None: |
| 227 | src_init(self, mode, options, dynamic, name) | 198 | src_init(self, mode, options, dynamic, name) |
| 228 | else: | 199 | else: |
| @@ -244,6 +215,8 @@ def patch_inductor_wrapper(): | |||
| 244 | 215 | ||
| 245 | 216 | ||
| 246 | def patch_dynamo_optimize(): | 217 | def patch_dynamo_optimize(): |
| 218 | + from torch._dynamo import optimize | ||
| 219 | + from torch_npu.dynamo import _get_global_npu_backend | ||
| 247 | src_optimize = optimize | 220 | src_optimize = optimize |
| 248 | 221 | ||
| 249 | def npu_optimize(*args, **kwargs): | 222 | def npu_optimize(*args, **kwargs): |
| @@ -298,78 +271,80 @@ def patch_stream_event_variable_python_type(): | |||
| 298 | streams.EventVariable.python_type = python_type | 271 | streams.EventVariable.python_type = python_type |
| 299 | 272 | ||
| 300 | 273 | ||
| 301 | -class NpuStreamContextVariable(StreamContextVariable): | ||
| 302 | - """This represents NPU stream context with FX graph set_stream node creation.""" | ||
| 303 | - | ||
| 304 | - | ||
| 305 | - def create( | ||
| 306 | - tx: "InstructionTranslator", | ||
| 307 | - stream_to_enter: "StreamVariable", | ||
| 308 | - **kwargs: dict[str, Any], | ||
| 309 | - ) -> "NpuStreamContextVariable": | ||
| 310 | - from torch._dynamo.device_interface import get_interface_for_device | ||
| 311 | - from torch._dynamo.variables.builder import wrap_fx_proxy_cls | ||
| 312 | - | ||
| 313 | - device_interface = get_interface_for_device(stream_to_enter.device) | ||
| 314 | - current_stream_var = wrap_fx_proxy_cls( | ||
| 315 | - StreamVariable, | ||
| 316 | - tx, | ||
| 317 | - tx.output.create_proxy( | ||
| 318 | - "call_function", | ||
| 319 | - device_interface.current_stream, | ||
| 320 | - (None,), | ||
| 321 | - {}, | ||
| 322 | - ), | ||
| 323 | - ) | ||
| 324 | - | ||
| 325 | - return NpuStreamContextVariable( | ||
| 326 | - stream_to_enter, | ||
| 327 | - current_stream=current_stream_var, | ||
| 328 | - device_interface=device_interface, | ||
| 329 | - **kwargs, | ||
| 330 | - ) | ||
| 331 | - | ||
| 332 | - def __init__( | ||
| 333 | - self, | ||
| 334 | - stream: Optional["StreamVariable"], | ||
| 335 | - current_stream: Optional["StreamVariable"] = None, | ||
| 336 | - device_interface: Any | None = None, | ||
| 337 | - **kwargs: Any, | ||
| 338 | - ) -> None: | ||
| 339 | - self.current_stream = current_stream | ||
| 340 | - self.device_interface = device_interface | ||
| 341 | - super().__init__(stream, **kwargs) | ||
| 342 | - | ||
| 343 | - def enter( | ||
| 344 | - self, tx: "InstructionTranslator", *args: VariableTracker | ||
| 345 | - ) -> VariableTracker: | ||
| 346 | - # Create set_stream node to switch to self.stream | ||
| 347 | - if self.get_stream(): | ||
| 348 | - tx.output.create_proxy( | ||
| 349 | - "call_function", | ||
| 350 | - self.device_interface.set_stream, | ||
| 351 | - (self.get_stream().as_proxy(),), | ||
| 352 | - {}, | ||
| 353 | - ) | ||
| 354 | - return super().enter(tx) | ||
| 355 | - | ||
| 356 | - def exit( | ||
| 357 | - self, tx: "InstructionTranslator", *args: VariableTracker | ||
| 358 | - ) -> VariableTracker: | ||
| 359 | - # First exit the symbolic stream state | ||
| 360 | - # Create set_stream node to restore current_stream | ||
| 361 | - if self.get_stream(): | ||
| 362 | - tx.output.create_proxy( | ||
| 363 | - "call_function", | ||
| 364 | - self.device_interface.set_stream, | ||
| 365 | - (self.current_stream.as_proxy(),), | ||
| 366 | - {}, | ||
| 367 | - ) | ||
| 368 | - return super().exit(tx, *args) | ||
| 369 | - | ||
| 370 | - | ||
| 371 | def patch_npu_stream_context(): | 274 | def patch_npu_stream_context(): |
| 372 | from torch._dynamo.device_interface import get_interface_for_device | 275 | from torch._dynamo.device_interface import get_interface_for_device |
| 276 | + from torch._dynamo.variables.base import VariableTracker | ||
| 277 | + from torch._dynamo.variables.streams import StreamContextVariable, StreamVariable | ||
| 278 | + from torch._dynamo.variables.torch import TorchInGraphFunctionVariable | ||
| 279 | + | ||
| 280 | + class NpuStreamContextVariable(StreamContextVariable): | ||
| 281 | + """This represents NPU stream context with FX graph set_stream node creation.""" | ||
| 282 | + | ||
| 283 | + | ||
| 284 | + def create( | ||
| 285 | + tx: "InstructionTranslator", | ||
| 286 | + stream_to_enter: "StreamVariable", | ||
| 287 | + **kwargs: dict[str, Any], | ||
| 288 | + ) -> "NpuStreamContextVariable": | ||
| 289 | + from torch._dynamo.device_interface import get_interface_for_device | ||
| 290 | + from torch._dynamo.variables.builder import wrap_fx_proxy_cls | ||
| 291 | + | ||
| 292 | + device_interface = get_interface_for_device(stream_to_enter.device) | ||
| 293 | + current_stream_var = wrap_fx_proxy_cls( | ||
| 294 | + StreamVariable, | ||
| 295 | + tx, | ||
| 296 | + tx.output.create_proxy( | ||
| 297 | + "call_function", | ||
| 298 | + device_interface.current_stream, | ||
| 299 | + (None,), | ||
| 300 | + {}, | ||
| 301 | + ), | ||
| 302 | + ) | ||
| 303 | + | ||
| 304 | + return NpuStreamContextVariable( | ||
| 305 | + stream_to_enter, | ||
| 306 | + current_stream=current_stream_var, | ||
| 307 | + device_interface=device_interface, | ||
| 308 | + **kwargs, | ||
| 309 | + ) | ||
| 310 | + | ||
| 311 | + def __init__( | ||
| 312 | + self, | ||
| 313 | + stream: Optional["StreamVariable"], | ||
| 314 | + current_stream: Optional["StreamVariable"] = None, | ||
| 315 | + device_interface: Any | None = None, | ||
| 316 | + **kwargs: Any, | ||
| 317 | + ) -> None: | ||
| 318 | + self.current_stream = current_stream | ||
| 319 | + self.device_interface = device_interface | ||
| 320 | + super().__init__(stream, **kwargs) | ||
| 321 | + | ||
| 322 | + def enter( | ||
| 323 | + self, tx: "InstructionTranslator", *args: VariableTracker | ||
| 324 | + ) -> VariableTracker: | ||
| 325 | + # Create set_stream node to switch to self.stream | ||
| 326 | + if self.get_stream(): | ||
| 327 | + tx.output.create_proxy( | ||
| 328 | + "call_function", | ||
| 329 | + self.device_interface.set_stream, | ||
| 330 | + (self.get_stream().as_proxy(),), | ||
| 331 | + {}, | ||
| 332 | + ) | ||
| 333 | + return super().enter(tx) | ||
| 334 | + | ||
| 335 | + def exit( | ||
| 336 | + self, tx: "InstructionTranslator", *args: VariableTracker | ||
| 337 | + ) -> VariableTracker: | ||
| 338 | + # First exit the symbolic stream state | ||
| 339 | + # Create set_stream node to restore current_stream | ||
| 340 | + if self.get_stream(): | ||
| 341 | + tx.output.create_proxy( | ||
| 342 | + "call_function", | ||
| 343 | + self.device_interface.set_stream, | ||
| 344 | + (self.current_stream.as_proxy(),), | ||
| 345 | + {}, | ||
| 346 | + ) | ||
| 347 | + return super().exit(tx, *args) | ||
| 373 | 348 | ||
| 374 | def _handle_npu_device_interface_stream(self, tx, stream): | 349 | def _handle_npu_device_interface_stream(self, tx, stream): |
| 375 | return NpuStreamContextVariable.create(tx, stream) | 350 | return NpuStreamContextVariable.create(tx, stream) |
| @@ -381,14 +356,20 @@ def patch_npu_stream_context(): | |||
| 381 | 356 | ||
| 382 | def patch_npu_current_stream(): | 357 | def patch_npu_current_stream(): |
| 383 | """Reuse PT handle_current_stream so current_stream gets user_object_index.""" | 358 | """Reuse PT handle_current_stream so current_stream gets user_object_index.""" |
| 359 | + from torch._dynamo.variables.torch import TorchInGraphFunctionVariable | ||
| 384 | handlers = TorchInGraphFunctionVariable._get_handlers() | 360 | handlers = TorchInGraphFunctionVariable._get_handlers() |
| 385 | handlers[torch.npu.current_stream] = handlers[torch.accelerator.current_stream] | 361 | handlers[torch.npu.current_stream] = handlers[torch.accelerator.current_stream] |
| 386 | 362 | ||
| 387 | 363 | ||
| 388 | def patch_user_defined_class_variable(): | 364 | def patch_user_defined_class_variable(): |
| 389 | import functools | 365 | import functools |
| 390 | - | 366 | + from torch._dynamo.variables.user_defined import UserDefinedClassVariable |
| 367 | + from torch._dynamo.variables.torch import TorchCtxManagerClassVariable | ||
| 368 | + from torch._dynamo.variables.torch import TorchInGraphFunctionVariable | ||
| 391 | original_method = UserDefinedClassVariable._in_graph_classes | 369 | original_method = UserDefinedClassVariable._in_graph_classes |
| 370 | + class NPUTorchCtxManagerClassVariable(TorchCtxManagerClassVariable): | ||
| 371 | + def call_function(self, tx, args, kwargs): | ||
| 372 | + return _create_npu_autocast_mode_variable(self.value, args, kwargs) | ||
| 392 | 373 | ||
| 393 | 374 | ||
| 394 | 375 | ||
| @@ -398,20 +379,70 @@ def patch_user_defined_class_variable(): | |||
| 398 | result.add(torch.npu.Stream) | 379 | result.add(torch.npu.Stream) |
| 399 | return result | 380 | return result |
| 400 | 381 | ||
| 382 | + def UserDefinedClassVariable__new__(cls, value, **kwargs): | ||
| 383 | + if value in [ | ||
| 384 | + torch.npu.amp.autocast, | ||
| 385 | + torch_npu.npu.amp.autocast, | ||
| 386 | + torch.npu.amp.autocast_mode.autocast, | ||
| 387 | + torch_npu.npu.amp.autocast_mode.autocast, | ||
| 388 | + ]: | ||
| 389 | + return NPUTorchCtxManagerClassVariable(value, **kwargs) | ||
| 390 | + elif value in [ | ||
| 391 | + torch_npu.npu.BoolTensor, | ||
| 392 | + torch_npu.npu.ByteTensor, | ||
| 393 | + torch_npu.npu.CharTensor, | ||
| 394 | + torch_npu.npu.DoubleTensor, | ||
| 395 | + torch_npu.npu.FloatTensor, | ||
| 396 | + torch_npu.npu.HalfTensor, | ||
| 397 | + torch_npu.npu.IntTensor, | ||
| 398 | + torch_npu.npu.LongTensor, | ||
| 399 | + torch_npu.npu.ShortTensor, | ||
| 400 | + torch_npu.npu.BFloat16Tensor, | ||
| 401 | + ]: | ||
| 402 | + return TorchInGraphFunctionVariable(value, **kwargs) | ||
| 403 | + return cls.__new__raw(cls) | ||
| 404 | + | ||
| 401 | UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes | 405 | UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes |
| 402 | - | ||
| 403 | - | ||
| 404 | -def add_dynamo_methods(): | ||
| 405 | UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__ | 406 | UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__ |
| 406 | UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__ | 407 | UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__ |
| 407 | - SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__ | 408 | + |
| 408 | - SkipFunctionVariable.__new__ = SkipFunctionVariable__new__ | 409 | +def run_once(f): |
| 409 | - TensorVariable.call_method_raw = TensorVariable.call_method | 410 | + """Runs a function (successfully) only once. |
| 410 | - TensorVariable.call_method = TensorVariable_call_method | 411 | + The running can be reset by setting the `has_run` attribute to False |
| 411 | - patch_dynamo_optimize() | 412 | + """ |
| 412 | - patch_inductor_wrapper() | 413 | + @functools.wraps(f) |
| 414 | + def wrapper(*args, **kwargs): | ||
| 415 | + if not wrapper.has_run: | ||
| 416 | + result = f(*args, **kwargs) | ||
| 417 | + wrapper.has_run = True | ||
| 418 | + return result | ||
| 419 | + return None | ||
| 420 | + wrapper.has_run = False | ||
| 421 | + return wrapper | ||
| 422 | + | ||
| 423 | + | ||
| 424 | + | ||
| 425 | +def _dynamo_register_interface_for_device(): | ||
| 426 | + from torch._dynamo.device_interface import register_interface_for_device | ||
| 427 | + from torch_npu.utils._dynamo_device import NpuInterface | ||
| 428 | + | ||
| 429 | + register_interface_for_device("npu", NpuInterface) | ||
| 430 | + for i in range(32): | ||
| 431 | + | ||
| 432 | + register_interface_for_device(f"npu:{i}", NpuInterface) | ||
| 433 | + | ||
| 434 | + | ||
| 435 | +def add_dynamo_methods_init(): | ||
| 436 | + _dynamo_register_interface_for_device() | ||
| 437 | + patch_SkipFunctionVariable() | ||
| 438 | + patch_TensorVariable_call_method() | ||
| 413 | patch_stream_event_variable_python_type() | 439 | patch_stream_event_variable_python_type() |
| 414 | patch_builtin_variable() | 440 | patch_builtin_variable() |
| 415 | patch_npu_stream_context() | 441 | patch_npu_stream_context() |
| 416 | patch_npu_current_stream() | 442 | patch_npu_current_stream() |
| 417 | patch_user_defined_class_variable() | 443 | patch_user_defined_class_variable() |
| 444 | + | ||
| 445 | + | ||
| 446 | +def add_dynamo_methods(): | ||
| 447 | + patch_dynamo_optimize() | ||
| 448 | + patch_inductor_wrapper() | ||
此条代码评论区间4至+5
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8,ruff，请Committer检视其合理性。