From 3331e587657778e3203d33ff3c9f56f1cd3a102f Mon Sep 17 00:00:00 2001 From: d00845784 Date: Thu, 19 Jun 2025 11:58:37 +0800 Subject: [PATCH 1/2] support multiprocess compile --- torch_npu/_inductor/__init__.py | 6 ++++-- torch_npu/_inductor/config.py | 2 +- torch_npu/_inductor/npu_triton_heuristics.py | 6 +++--- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/torch_npu/_inductor/__init__.py b/torch_npu/_inductor/__init__.py index a13b5797995..d6c56d3e1e5 100644 --- a/torch_npu/_inductor/__init__.py +++ b/torch_npu/_inductor/__init__.py @@ -1,4 +1,5 @@ import os +import atexit import torch from torch._dynamo.device_interface import register_interface_for_device, get_interface_for_device @@ -6,6 +7,7 @@ from torch._inductor import lowering as inductor_lowering from torch._inductor.choices import InductorChoices from torch._inductor.codegen.common import register_backend_for_device, register_device_op_overrides from torch._inductor.runtime import autotune_cache +from torch._inductor.async_compile import shutdown_compile_workers from torch_npu.npu import device_count from torch_npu.utils._dynamo_device import NpuInterface, current_device, set_device from torch_npu.utils._inductor import NPUDeviceOpOverrides @@ -13,7 +15,7 @@ from torch_npu.utils._inductor import NPUDeviceOpOverrides from . import config as npu_config from . import codegen from .npu_fusion_attention_graph import register_fa_pass -from .config import aggresive_autotune, num_vector_core, set_compile_threads +from .config import aggresive_autotune, num_vector_core from .config import log as npulog from .decomposition import _register_npu_inductor_decompositons from .lowering import make_reduction, npu_make_fallback @@ -22,7 +24,7 @@ from .npu_device import NewNPUDeviceOpOverrides, NewNpuInterface from .runtime import _load_cached_autotuning from .utils import get_current_raw_stream -set_compile_threads() +atexit.register(shutdown_compile_workers) def _inductor_register_backend_for_device(): diff --git a/torch_npu/_inductor/config.py b/torch_npu/_inductor/config.py index f9bf23ee337..3a43ee8c255 100644 --- a/torch_npu/_inductor/config.py +++ b/torch_npu/_inductor/config.py @@ -95,7 +95,7 @@ inductor_static_mode = os.environ.get('INDUCTOR_STATIC_MODE', '0').lower() in (' profile_path = "./profile_result/" -def set_compile_threads(): +def set_compile_threads_to_1(): if "TORCHINDUCTOR_COMPILE_THREADS" in os.environ: torchinductor_compile_threads = int(os.environ["TORCHINDUCTOR_COMPILE_THREADS"]) if torchinductor_compile_threads == 1: diff --git a/torch_npu/_inductor/npu_triton_heuristics.py b/torch_npu/_inductor/npu_triton_heuristics.py index 19e2c230c4e..97f588a77c4 100644 --- a/torch_npu/_inductor/npu_triton_heuristics.py +++ b/torch_npu/_inductor/npu_triton_heuristics.py @@ -153,9 +153,9 @@ class NPUCachingAutotuner(CachingAutotuner): f"[thread {os.getpid()}][InductorNPU.precompile] Exception = {e}, kernel = {self.fn.__name__} config = {c}") # Skip the config if the compilation fails continue - if launcher is not None: - self.launchers.append(launcher) - compiled_binaries.append(compiled_binary) + + self.launchers.append(launcher) + compiled_binaries.append(compiled_binary) if len(self.launchers) == 0: raise RuntimeError( -- Gitee From 45ef87c7f909b09eea05d3760534ca01dea69a18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=9C=E6=89=BF=E6=98=86?= Date: Thu, 21 Aug 2025 07:31:31 +0000 Subject: [PATCH 2/2] update torch_npu/_inductor/npu_triton_heuristics.py. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 杜承昆 --- torch_npu/_inductor/npu_triton_heuristics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/torch_npu/_inductor/npu_triton_heuristics.py b/torch_npu/_inductor/npu_triton_heuristics.py index 97f588a77c4..19e2c230c4e 100644 --- a/torch_npu/_inductor/npu_triton_heuristics.py +++ b/torch_npu/_inductor/npu_triton_heuristics.py @@ -153,9 +153,9 @@ class NPUCachingAutotuner(CachingAutotuner): f"[thread {os.getpid()}][InductorNPU.precompile] Exception = {e}, kernel = {self.fn.__name__} config = {c}") # Skip the config if the compilation fails continue - - self.launchers.append(launcher) - compiled_binaries.append(compiled_binary) + if launcher is not None: + self.launchers.append(launcher) + compiled_binaries.append(compiled_binary) if len(self.launchers) == 0: raise RuntimeError( -- Gitee