From 1d20c2a7726a4286e8e05ef8d6e88a7cf9afd506 Mon Sep 17 00:00:00 2001 From: Kino Date: Fri, 14 Aug 2026 17:05:46 -1000 Subject: [PATCH] fix #161: XQA cubin generation ignores resolved output dir under spawn/forkserver gen_cubins.py's worker pool (multiprocessing.Pool.map) reads the module-scope `cubin_dir`/`nvcc_bin`/`clean_cubin` globals. These are only rebound to their CLI-resolved values inside `if __name__ == "__main__":`, which worker processes never execute when the multiprocessing start method re-imports the module (spawn/forkserver) instead of forking. Under fork (Linux's current default) this is masked because children inherit the parent's post-__main__ state, but it silently falls back to the relative module-level default elsewhere, matching the ptxas "could not be opened" failure reported in #161. Pass the resolved values into each worker explicitly via Pool(initializer=...) so behavior no longer depends on the start method. Signed-off-by: Kino --- kernelSrcs/xqa/gen_cubins.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/kernelSrcs/xqa/gen_cubins.py b/kernelSrcs/xqa/gen_cubins.py index cf88b8ce8..243a0e2d7 100755 --- a/kernelSrcs/xqa/gen_cubins.py +++ b/kernelSrcs/xqa/gen_cubins.py @@ -454,6 +454,18 @@ def convert_cubin_cpp_np(cubin_file_name: str): return cpp_array, cubin_size +def init_cubin_gen_worker(worker_cubin_dir: str, worker_nvcc_bin: str, + worker_clean_cubin: bool): + # Workers under a start method that re-imports this module (spawn/forkserver) + # never execute the `if __name__ == "__main__":` block, so the module-scope + # defaults for these globals would otherwise be used instead of the values + # resolved from CLI args. Set them explicitly in each worker process. + global cubin_dir, nvcc_bin, clean_cubin + cubin_dir = worker_cubin_dir + nvcc_bin = worker_nvcc_bin + clean_cubin = worker_clean_cubin + + def run_cubin_gen(arch_micro_file_list: CompileArchMacrosAndFile): nvcc_command, xxd_command, cubin_file_name = build_commands( build_func_name_prefix, arch_micro_file_list.arch, @@ -833,7 +845,10 @@ def get_generated_output_paths( cpu_count = os.cpu_count() thread_count = cpu_count // 2 if cpu_count >= 2 else cpu_count - with multiprocessing.Pool(processes=thread_count) as pool: + with multiprocessing.Pool( + processes=thread_count, + initializer=init_cubin_gen_worker, + initargs=(cubin_dir, nvcc_bin, clean_cubin)) as pool: name_size_list = pool.map(run_cubin_gen, arch_macro_lists) header_file_contents = generate_header_file_contents( arch_macro_lists, name_size_list)