[Autotuner] Fix fork-based autotuner to avoid re-initializing CUDA context in subprocess (#981)

Author: yf225

helion/autotuner/base_search.py

@@ -365,13 +365,18 @@ def start_precompile_and_check_for_hangs(
             )
             process.daemon = True
         else:
+            precompiler = _prepare_precompiler_for_fork(
+                fn, device_args, config, self.kernel, decorator
+            )
+            if precompiler is None:
+                return PrecompileFuture.skip(self, config, True)
             ctx = mp.get_context("fork")
             parent_conn, child_conn = ctx.Pipe()
             process = cast(
                 "mp.Process",
                 ctx.Process(
                     target=_run_kernel_in_subprocess_fork,
-                    args=(fn, device_args, config, self.kernel, child_conn, decorator),
+                    args=(precompiler, config, self.kernel, child_conn, decorator),
                 ),
             )
             process.daemon = True
@@ -1209,37 +1214,54 @@ def _run_kernel_in_subprocess_spawn(
         os._exit(status)
 
 
-def _run_kernel_in_subprocess_fork(
+def _prepare_precompiler_for_fork(
     fn: CompiledConfig,
     args: Sequence[object],
     config: Config,
     kernel: BoundKernel,
+    decorator: str,
+) -> Callable[[], None] | None:
+    def extract_launcher(
+        triton_kernel: object,
+        grid: tuple[int, ...],
+        *launch_args: object,
+        **launch_kwargs: object,
+    ) -> NoReturn:
+        raise _ExtractedLaunchArgs(triton_kernel, grid, launch_args, launch_kwargs)
+
+    try:
+        fn(*tuple(args), _launcher=extract_launcher)
+        raise RuntimeError("Expected _ExtractedLaunchArgs to be raised")
+    except _ExtractedLaunchArgs as extracted:
+        precompiler_factory = make_precompiler(
+            cast("Any", extracted.kernel),
+            config,
+            kernel,
+        )
+        precompiler = precompiler_factory(*extracted.args, **extracted.kwargs)
+        if precompiler is already_compiled:
+            return None
+        return precompiler
+    except Exception:
+        log.warning(
+            "Helion autotuner precompile error for %s\n\nGenerated Triton code:\n%s",
+            decorator,
+            kernel.to_triton_code(config),
+            exc_info=True,
+        )
+        raise
+
+
+def _run_kernel_in_subprocess_fork(
+    precompiler: Callable[[], None],
+    config: Config,
+    kernel: BoundKernel,
     conn: connection.Connection,
     decorator: str,
 ) -> None:
     status = 0
     try:
-
-        def extract_launcher(
-            triton_kernel: object,
-            grid: tuple[int, ...],
-            *launch_args: object,
-            **launch_kwargs: object,
-        ) -> NoReturn:
-            raise _ExtractedLaunchArgs(triton_kernel, grid, launch_args, launch_kwargs)
-
-        try:
-            fn(*tuple(args), _launcher=extract_launcher)
-            raise RuntimeError("Expected _ExtractedLaunchArgs to be raised")
-        except _ExtractedLaunchArgs as extracted:
-            precompiler_factory = make_precompiler(
-                cast("Any", extracted.kernel),
-                config,
-                kernel,
-            )
-            precompiler = precompiler_factory(*extracted.args, **extracted.kwargs)
-            if precompiler is not already_compiled:
-                precompiler()
+        precompiler()
         conn.send({"status": "ok"})
     except Exception as exc:
         status = 1

test/test_autotuner.py

@@ -5,12 +5,14 @@
 from contextlib import nullcontext
 import logging
 import math
+import multiprocessing as mp
 import os
 from pathlib import Path
 import pickle
 import random
 import tempfile
 from types import SimpleNamespace
+from typing import Callable
 import unittest
 from unittest import skip
 from unittest.mock import patch
@@ -70,14 +72,16 @@ def _autotune(self):
 
 
 class TestAutotuneIgnoreErrors(TestCase):
-    def _make_search(self, settings: Settings) -> BaseSearch:
+    def _make_search(
+        self, settings: Settings, *, args: tuple[object, ...] = ()
+    ) -> BaseSearch:
         search = BaseSearch.__new__(BaseSearch)
         search.settings = settings
         search.kernel = SimpleNamespace(
             format_kernel_decorator=lambda config, s: "decorator",
             to_triton_code=lambda config: "code",
         )
-        search.args = ()
+        search.args = args
         search.counters = collections.Counter()
         search.log = LambdaLogger(logging.CRITICAL)
         search._kernel_mutates_args = False
@@ -126,6 +130,53 @@ def bad_fn(*_args):
         self.assertEqual(result, float("inf"))
         warn.assert_not_called()
 
+    @pytest.mark.skipif(
+        "fork" not in mp.get_all_start_methods(),
+        reason="fork start method is unavailable on this platform",
+    )
+    def test_fork_precompile_avoids_cuda_reinit(self):
+        settings = Settings(
+            autotune_precompile="fork",
+            autotune_log_level=logging.CRITICAL,
+            autotune_compile_timeout=5,
+        )
+        search = self._make_search(settings, args=("arg0",))
+
+        parent_pid = os.getpid()
+        lazy_calls: list[int] = []
+
+        def fake_lazy_init() -> None:
+            lazy_calls.append(os.getpid())
+
+        def fake_make_precompiler(_kernel_obj, _config, _bound_kernel):
+            def binder(*_args: object, **_kwargs: object):
+                def run() -> None:
+                    return None
+
+                return run
+
+            return binder
+
+        def fake_compiled_fn(
+            *fn_args: object, _launcher: Callable[..., object]
+        ) -> None:
+            torch.cuda._lazy_init()
+            _launcher("fake_kernel", (1,), *fn_args)
+
+        with (
+            patch(
+                "helion.autotuner.base_search.make_precompiler",
+                side_effect=fake_make_precompiler,
+            ),
+            patch("torch.cuda._lazy_init", side_effect=fake_lazy_init),
+        ):
+            future = search.start_precompile_and_check_for_hangs(
+                "cfg", fake_compiled_fn
+            )
+            self.assertTrue(future())
+
+        self.assertEqual(set(lazy_calls), {parent_pid})
+
 
 class TestAutotuner(RefEagerTestDisabled, TestCase):
     def setUp(self):