From 77f0d09ecf95d4bb77f48555cb519a7b777c2ece Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sat, 29 Mar 2025 11:49:53 +0800 Subject: [PATCH] hotfix: HIP supports parallel BEAM search --- tinygrad/engine/search.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/engine/search.py b/tinygrad/engine/search.py index 368f9963d9..c2906a29bf 100644 --- a/tinygrad/engine/search.py +++ b/tinygrad/engine/search.py @@ -142,7 +142,7 @@ def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True, beam: list[tuple[Kernel, float]] = [(lin, float("inf"))] seen_libs = set() - default_parallel = multiprocessing.cpu_count() if lin.opts.device in {"CUDA", "AMD", "NV", "METAL"} else 0 + default_parallel = multiprocessing.cpu_count() if lin.opts.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0 if beam_pool is None and (workers := getenv("PARALLEL", default_parallel)): beam_pool = multiprocessing.get_context("spawn").Pool(workers, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16)) @atexit.register