hotfix: HIP supports parallel BEAM search

2025-03-29 11:49:53 +08:00
parent 162f286a0e
commit 77f0d09ecf
1 changed files with 1 additions and 1 deletions
@@ -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