已合并
lazy_init #39539
cuiduo创建于 6月29日
lazy_init #39539
已合并
共 7 个文件变更+201-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 |
 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8,ruff，请Committer检视其合理性。   | |||
| 6 | 6 | ||
| 7 | class TestNpuStream(TestCase): | 7 | class TestNpuStream(TestCase): |
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 | |||
| 8 | 8 | ||
| @@ -9,7 +9,7 @@ AsyncCompile.warm_pool() | |||
| 9 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD | 9 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD |
| 10 | 10 | ||
| 11 | import os | 11 | import os |
| 12 | - | 12 | +from torch_npu.utils._dynamo import _dynamo_register_interface_for_device, patch_SkipFunctionVariable, patch_TensorVariable_call_method |
| 13 | # all backends need register npu/cpu/mps device_op_overrides | 13 | # all backends need register npu/cpu/mps device_op_overrides |
| 14 | from .graph import patch_codegen_with_cpp_wrapper | 14 | from .graph import patch_codegen_with_cpp_wrapper |
| 15 | from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, get_current_raw_stream | 15 | from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, get_current_raw_stream |
| @@ -19,6 +19,7 @@ from .shape_handling import NPUShapeHandling, patch_shape_handling | |||
| 19 | from ._npu_meta_registration import npu_patch_meta | 19 | from ._npu_meta_registration import npu_patch_meta |
| 20 | 20 | ||
| 21 | npu_patch_meta() | 21 | npu_patch_meta() |
| 22 | +_dynamo_register_interface_for_device() | ||
| 22 | register_device_op_overrides_npu() | 23 | register_device_op_overrides_npu() |
| 23 | patch_has_triton() | 24 | patch_has_triton() |
| 24 | patch_is_gpu() | 25 | 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 |
| 69 | # Dynamo's global rules maps. | 67 | # Dynamo's global rules maps. |
| @@ -4,6 +4,7 @@ import logging | |||
| 4 | import sys | 4 | import sys |
| 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 | if TYPE_CHECKING: | 9 | if TYPE_CHECKING: |
| 9 | from torch._dynamo.symbolic_convert import InstructionTranslator | 10 | from torch._dynamo.symbolic_convert import InstructionTranslator |
| @@ -11,107 +12,76 @@ if TYPE_CHECKING: | |||
| 11 | import torch | 12 | import torch |
| 12 | import torch_npu | 13 | import torch_npu |
| 13 | from torch import _TorchCompileWrapper | 14 | from torch import _TorchCompileWrapper |
| 14 | -from torch._dynamo import optimize | ||
| 15 | -from torch._dynamo.utils import tensortype_to_dtype | ||
| 16 | -from torch._dynamo.variables.base import VariableTracker | ||
| 17 | -from torch._dynamo.variables.constant import ConstantVariable | ||
| 18 | -from torch._dynamo.variables.ctx_manager import AutocastModeVariable | ||
| 19 | -from torch._dynamo.variables.functions import SkipFunctionVariable | ||
| 20 | -from torch._dynamo.variables.lists import TupleVariable | ||
| 21 | -from torch._dynamo.variables.streams import StreamContextVariable, StreamVariable | ||
| 22 | -from torch._dynamo.variables.tensor import TensorVariable | ||
| 23 | -from torch._dynamo.variables.torch import ( | ||
| 24 | - TorchCtxManagerClassVariable, | ||
| 25 | - TorchInGraphFunctionVariable, | ||
| 26 | -) | ||
| 27 | -from torch._dynamo.variables.user_defined import UserDefinedClassVariable | ||
| 28 | -from torch_npu.dynamo import _get_global_npu_backend | ||
| 29 | 15 | ||
| 30 | 16 | ||
| 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: |
| @@ -217,6 +187,7 @@ def patch_inductor_wrapper(): | |||
| 217 | return ori_dict | 187 | return ori_dict |
| 218 | 188 | ||
| 219 | def new_init(self, mode, options, dynamic): | 189 | def new_init(self, mode, options, dynamic): |
| 190 | + add_dynamo_methods_init() | ||
| 220 | src_init(self, mode, options, dynamic) | 191 | src_init(self, mode, options, dynamic) |
| 221 | backend = _resolve_npu_backend_from_wrapper(self) | 192 | backend = _resolve_npu_backend_from_wrapper(self) |
| 222 | if backend=="mlir": | 193 | if backend=="mlir": |
| @@ -240,6 +211,8 @@ def patch_inductor_wrapper(): | |||
| 240 | 211 | ||
| 241 | 212 | ||
| 242 | def patch_dynamo_optimize(): | 213 | def patch_dynamo_optimize(): |
| 214 | + from torch._dynamo import optimize | ||
| 215 | + from torch_npu.dynamo import _get_global_npu_backend | ||
| 243 | src_optimize = optimize | 216 | src_optimize = optimize |
| 244 | 217 | ||
| 245 | def npu_optimize(*args, **kwargs): | 218 | def npu_optimize(*args, **kwargs): |
| @@ -288,75 +261,77 @@ def patch_event_variable_python_type(): | |||
| 288 | torch._dynamo.variables.streams.EventVariable.python_type = python_type | 261 | torch._dynamo.variables.streams.EventVariable.python_type = python_type |
| 289 | 262 | ||
| 290 | 263 | ||
| 291 | -class NpuStreamContextVariable(StreamContextVariable): | ||
| 292 | - """This represents NPU stream context with FX graph set_stream node creation.""" | ||
| 293 | - | ||
| 294 | - | ||
| 295 | - def create( | ||
| 296 | - tx: "InstructionTranslator", | ||
| 297 | - stream_to_enter: "StreamVariable", | ||
| 298 | - **kwargs: dict[str, Any], | ||
| 299 | - ) -> "NpuStreamContextVariable": | ||
| 300 | - from torch._dynamo.device_interface import get_interface_for_device | ||
| 301 | - from torch._dynamo.variables.builder import wrap_fx_proxy_cls | ||
| 302 | - | ||
| 303 | - device_interface = get_interface_for_device(stream_to_enter.device) | ||
| 304 | - current_stream_var = wrap_fx_proxy_cls( | ||
| 305 | - StreamVariable, | ||
| 306 | - tx, | ||
| 307 | - tx.output.create_proxy( | ||
| 308 | - "call_function", | ||
| 309 | - device_interface.current_stream, | ||
| 310 | - (None,), | ||
| 311 | - {}, | ||
| 312 | - ), | ||
| 313 | - ) | ||
| 314 | - | ||
| 315 | - return NpuStreamContextVariable( | ||
| 316 | - stream_to_enter, | ||
| 317 | - current_stream=current_stream_var, | ||
| 318 | - device_interface=device_interface, | ||
| 319 | - **kwargs, | ||
| 320 | - ) | ||
| 321 | - | ||
| 322 | - def __init__( | ||
| 323 | - self, | ||
| 324 | - stream: Optional["StreamVariable"], | ||
| 325 | - current_stream: Optional["StreamVariable"] = None, | ||
| 326 | - device_interface: Any | None = None, | ||
| 327 | - **kwargs: Any, | ||
| 328 | - ) -> None: | ||
| 329 | - self.current_stream = current_stream | ||
| 330 | - self.device_interface = device_interface | ||
| 331 | - super().__init__(stream, **kwargs) | ||
| 332 | - | ||
| 333 | - def enter( | ||
| 334 | - self, tx: "InstructionTranslator", *args: VariableTracker | ||
| 335 | - ) -> VariableTracker: | ||
| 336 | - if self.get_stream(): | ||
| 337 | - tx.output.create_proxy( | ||
| 338 | - "call_function", | ||
| 339 | - self.device_interface.set_stream, | ||
| 340 | - (self.get_stream().as_proxy(),), | ||
| 341 | - {}, | ||
| 342 | - ) | ||
| 343 | - return super().enter(tx) | ||
| 344 | - | ||
| 345 | - def exit( | ||
| 346 | - self, tx: "InstructionTranslator", *args: VariableTracker | ||
| 347 | - ) -> VariableTracker: | ||
| 348 | - if self.get_stream(): | ||
| 349 | - tx.output.create_proxy( | ||
| 350 | - "call_function", | ||
| 351 | - self.device_interface.set_stream, | ||
| 352 | - (self.current_stream.as_proxy(),), | ||
| 353 | - {}, | ||
| 354 | - ) | ||
| 355 | - return super().exit(tx, *args) | ||
| 356 | - | ||
| 357 | - | ||
| 358 | def patch_npu_stream_context(): | 264 | def patch_npu_stream_context(): |
| 359 | from torch._dynamo.device_interface import get_interface_for_device | 265 | from torch._dynamo.device_interface import get_interface_for_device |
| 266 | + from torch._dynamo.variables.base import VariableTracker | ||
| 267 | + from torch._dynamo.variables.streams import StreamContextVariable, StreamVariable | ||
| 268 | + from torch._dynamo.variables.torch import TorchInGraphFunctionVariable | ||
| 269 | + | ||
| 270 | + class NpuStreamContextVariable(StreamContextVariable): | ||
| 271 | + """This represents NPU stream context with FX graph set_stream node creation.""" | ||
| 272 | + | ||
| 273 | + | ||
| 274 | + def create( | ||
| 275 | + tx: "InstructionTranslator", | ||
| 276 | + stream_to_enter: "StreamVariable", | ||
| 277 | + **kwargs: dict[str, Any], | ||
| 278 | + ) -> "NpuStreamContextVariable": | ||
| 279 | + from torch._dynamo.device_interface import get_interface_for_device | ||
| 280 | + from torch._dynamo.variables.builder import wrap_fx_proxy_cls | ||
| 281 | + | ||
| 282 | + device_interface = get_interface_for_device(stream_to_enter.device) | ||
| 283 | + current_stream_var = wrap_fx_proxy_cls( | ||
| 284 | + StreamVariable, | ||
| 285 | + tx, | ||
| 286 | + tx.output.create_proxy( | ||
| 287 | + "call_function", | ||
| 288 | + device_interface.current_stream, | ||
| 289 | + (None,), | ||
| 290 | + {}, | ||
| 291 | + ), | ||
| 292 | + ) | ||
| 293 | + | ||
| 294 | + return NpuStreamContextVariable( | ||
| 295 | + stream_to_enter, | ||
| 296 | + current_stream=current_stream_var, | ||
| 297 | + device_interface=device_interface, | ||
| 298 | + **kwargs, | ||
| 299 | + ) | ||
| 300 | + | ||
| 301 | + def __init__( | ||
| 302 | + self, | ||
| 303 | + stream: Optional["StreamVariable"], | ||
| 304 | + current_stream: Optional["StreamVariable"] = None, | ||
| 305 | + device_interface: Any | None = None, | ||
| 306 | + **kwargs: Any, | ||
| 307 | + ) -> None: | ||
| 308 | + self.current_stream = current_stream | ||
| 309 | + self.device_interface = device_interface | ||
| 310 | + super().__init__(stream, **kwargs) | ||
| 311 | + | ||
| 312 | + def enter( | ||
| 313 | + self, tx: "InstructionTranslator", *args: VariableTracker | ||
| 314 | + ) -> VariableTracker: | ||
| 315 | + if self.get_stream(): | ||
| 316 | + tx.output.create_proxy( | ||
| 317 | + "call_function", | ||
| 318 | + self.device_interface.set_stream, | ||
| 319 | + (self.get_stream().as_proxy(),), | ||
| 320 | + {}, | ||
| 321 | + ) | ||
| 322 | + return super().enter(tx) | ||
| 323 | + | ||
| 324 | + def exit( | ||
| 325 | + self, tx: "InstructionTranslator", *args: VariableTracker | ||
| 326 | + ) -> VariableTracker: | ||
| 327 | + if self.get_stream(): | ||
| 328 | + tx.output.create_proxy( | ||
| 329 | + "call_function", | ||
| 330 | + self.device_interface.set_stream, | ||
| 331 | + (self.current_stream.as_proxy(),), | ||
| 332 | + {}, | ||
| 333 | + ) | ||
| 334 | + return super().exit(tx, *args) | ||
| 360 | 335 | ||
| 361 | def _handle_npu_device_interface_stream(self, tx, stream): | 336 | def _handle_npu_device_interface_stream(self, tx, stream): |
| 362 | return NpuStreamContextVariable.create(tx, stream) | 337 | return NpuStreamContextVariable.create(tx, stream) |
| @@ -396,8 +371,13 @@ def patch_record_stream(): | |||
| 396 | 371 | ||
| 397 | def patch_user_defined_class_variable(): | 372 | def patch_user_defined_class_variable(): |
| 398 | import functools | 373 | import functools |
| 399 | - | 374 | + from torch._dynamo.variables.user_defined import UserDefinedClassVariable |
| 375 | + from torch._dynamo.variables.torch import TorchCtxManagerClassVariable | ||
| 376 | + from torch._dynamo.variables.torch import TorchInGraphFunctionVariable | ||
| 400 | original_method = UserDefinedClassVariable._in_graph_classes | 377 | original_method = UserDefinedClassVariable._in_graph_classes |
| 378 | + class NPUTorchCtxManagerClassVariable(TorchCtxManagerClassVariable): | ||
| 379 | + def call_function(self, tx, args, kwargs): | ||
| 380 | + return _create_npu_autocast_mode_variable(self.value, args, kwargs) | ||
| 401 | 381 | ||
| 402 | 382 | ||
| 403 | 383 | ||
| @@ -407,20 +387,71 @@ def patch_user_defined_class_variable(): | |||
| 407 | result.add(torch.npu.Stream) | 387 | result.add(torch.npu.Stream) |
| 408 | return result | 388 | return result |
| 409 | 389 | ||
| 390 | + def UserDefinedClassVariable__new__(cls, value, **kwargs): | ||
| 391 | + if value in [ | ||
| 392 | + torch.npu.amp.autocast, | ||
| 393 | + torch_npu.npu.amp.autocast, | ||
| 394 | + torch.npu.amp.autocast_mode.autocast, | ||
| 395 | + torch_npu.npu.amp.autocast_mode.autocast, | ||
| 396 | + ]: | ||
| 397 | + return NPUTorchCtxManagerClassVariable(value, **kwargs) | ||
| 398 | + elif value in [ | ||
| 399 | + torch_npu.npu.BoolTensor, | ||
| 400 | + torch_npu.npu.ByteTensor, | ||
| 401 | + torch_npu.npu.CharTensor, | ||
| 402 | + torch_npu.npu.DoubleTensor, | ||
| 403 | + torch_npu.npu.FloatTensor, | ||
| 404 | + torch_npu.npu.HalfTensor, | ||
| 405 | + torch_npu.npu.IntTensor, | ||
| 406 | + torch_npu.npu.LongTensor, | ||
| 407 | + torch_npu.npu.ShortTensor, | ||
| 408 | + torch_npu.npu.BFloat16Tensor, | ||
| 409 | + ]: | ||
| 410 | + return TorchInGraphFunctionVariable(value, **kwargs) | ||
| 411 | + return cls.__new__raw(cls) | ||
| 412 | + | ||
| 410 | UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes | 413 | UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes |
| 411 | - | ||
| 412 | - | ||
| 413 | -def add_dynamo_methods(): | ||
| 414 | UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__ | 414 | UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__ |
| 415 | UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__ | 415 | UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__ |
| 416 | - SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__ | 416 | + |
| 417 | - SkipFunctionVariable.__new__ = SkipFunctionVariable__new__ | 417 | + |
| 418 | - TensorVariable.call_method_raw = TensorVariable.call_method | 418 | +def run_once(f): |
| 419 | - TensorVariable.call_method = TensorVariable_call_method | 419 | + """Runs a function (successfully) only once. |
| 420 | - patch_dynamo_optimize() | 420 | + The running can be reset by setting the `has_run` attribute to False |
| 421 | - patch_inductor_wrapper() | 421 | + """ |
| 422 | + | ||
| 423 | + def wrapper(*args, **kwargs): | ||
| 424 | + if not wrapper.has_run: | ||
| 425 | + result = f(*args, **kwargs) | ||
| 426 | + wrapper.has_run = True | ||
| 427 | + return result | ||
| 428 | + return None | ||
| 429 | + wrapper.has_run = False | ||
| 430 | + return wrapper | ||
| 431 | + | ||
| 432 | + | ||
| 433 | + | ||
| 434 | +def _dynamo_register_interface_for_device(): | ||
| 435 | + from torch._dynamo.device_interface import register_interface_for_device | ||
| 436 | + from torch_npu.utils._dynamo_device import NpuInterface | ||
| 437 | + | ||
| 438 | + register_interface_for_device("npu", NpuInterface) | ||
| 439 | + for i in range(32): | ||
| 440 | + | ||
| 441 | + register_interface_for_device(f"npu:{i}", NpuInterface) | ||
| 442 | + | ||
| 443 | + | ||
| 444 | +def add_dynamo_methods_init(): | ||
| 445 | + _dynamo_register_interface_for_device() | ||
| 446 | + patch_SkipFunctionVariable() | ||
| 447 | + patch_TensorVariable_call_method() | ||
| 448 | + patch_user_defined_class_variable() | ||
| 449 | + patch_record_stream() | ||
| 422 | patch_event_variable_python_type() | 450 | patch_event_variable_python_type() |
| 423 | patch_builtin_variable() | 451 | patch_builtin_variable() |
| 424 | patch_npu_stream_context() | 452 | patch_npu_stream_context() |
| 425 | - patch_record_stream() | 453 | + |
| 426 | - patch_user_defined_class_variable() | 454 | + |
| 455 | +def add_dynamo_methods(): | ||
| 456 | + patch_dynamo_optimize() | ||
| 457 | + patch_inductor_wrapper() | ||
| @@ -1,7 +1,6 @@ | |||
| 1 | from typing import Optional | 1 | from typing import Optional |
| 2 | 2 | ||
| 3 | import torch | 3 | import torch |
| 4 | -from torch._inductor.codegen.common import DeviceOpOverrides, register_device_op_overrides | ||
| 5 | from torch._prims_common import TensorLike | 4 | from torch._prims_common import TensorLike |
| 6 | from torch._prims.rng_prims import register_rng_prim | 5 | from torch._prims.rng_prims import register_rng_prim |
| 7 | 6 | ||
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8,ruff，请Committer检视其合理性。