Use temporary files to work around mktempdir slowness.

JuliaGPU · May 10, 2021 · 80b4084 · 80b4084
1 parent 9c6572e
commit 80b4084
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl
@@ -340,8 +340,11 @@ function cufunction_compile(@nospecialize(job::CompilerJob))
     end
 
     # compile to machine code
-    image = mktempdir() do dir
-        ptx_input = joinpath(dir, "input.ptx")
+    # NOTE: we use tempname since mktemp doesn't support suffixes, and mktempdir is slow
+    ptx_input = tempname(cleanup=false) * ".ptx"
+    ptxas_output = tempname(cleanup=false) * ".cubin"
+    nvlink_output = tempname(cleanup=false) * ".cubin"
+    image = try
         write(ptx_input, code)
 
         arch = "sm_$(job.target.cap.major)$(job.target.cap.minor)"
@@ -353,7 +356,6 @@ function cufunction_compile(@nospecialize(job::CompilerJob))
         #    version, which is hard to correlate to PTX JIT improvements;
         # 3. if we want to be able to use newer (minor upgrades) of the CUDA toolkit on an
         #    older driver, we should use the newer compiler to ensure compatibility.
-        ptxas_output = joinpath(dir, "ptxas.cubin")
         proc, log = run_and_collect(
             ```$ptxas_cmd --gpu-name $arch
                           --verbose
@@ -371,7 +373,6 @@ function cufunction_compile(@nospecialize(job::CompilerJob))
         # TODO: try LTO, `--link-time-opt --nvvmpath /opt/cuda/nvvm`.
         #       fails with `Ignoring -lto option because no LTO objects found`
         if needs_cudadevrt
-            nvlink_output = joinpath(dir, "nvlink.cubin")
             proc, log = run_and_collect(
                 ```$nvlink_cmd --arch $arch
                                --library-path $(dirname(libcudadevrt()))
@@ -388,6 +389,10 @@ function cufunction_compile(@nospecialize(job::CompilerJob))
         else
             read(ptxas_output)
         end
+    finally
+        rm(ptx_input)
+        ispath(ptxas_output) && rm(ptxas_output)
+        ispath(nvlink_output) && rm(nvlink_output)
     end
 
     return (image, entry=LLVM.name(kernel), external_gvars)