fix: load env vars before

EnzymeAD · Jan 21, 2025 · 0c83b0a · 0c83b0a
1 parent c926ed1
commit 0c83b0a
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/src/XLA.jl b/src/XLA.jl
@@ -129,6 +129,18 @@ function __init__()
     backends["cpu"] = cpu
     default_backend[] = cpu
 
+    if haskey(ENV, "XLA_REACTANT_GPU_MEM_FRACTION")
+        XLA_REACTANT_GPU_MEM_FRACTION[] = parse(
+            Float64, ENV["XLA_REACTANT_GPU_MEM_FRACTION"]
+        )
+        @debug "XLA_REACTANT_GPU_MEM_FRACTION: " XLA_REACTANT_GPU_MEM_FRACTION[]
+    end
+
+    if haskey(ENV, "XLA_REACTANT_GPU_PREALLOCATE")
+        XLA_REACTANT_GPU_PREALLOCATE[] = parse(Bool, ENV["XLA_REACTANT_GPU_PREALLOCATE"])
+        @debug "XLA_REACTANT_GPU_PREALLOCATE: " XLA_REACTANT_GPU_PREALLOCATE[]
+    end
+
     @static if !Sys.isapple()
         if isfile("/usr/lib/libtpu.so")
             dataset_dir = @get_scratch!("libtpu")
@@ -160,18 +172,6 @@ function __init__()
         end
     end
 
-    if haskey(ENV, "XLA_REACTANT_GPU_MEM_FRACTION")
-        XLA_REACTANT_GPU_MEM_FRACTION[] = parse(
-            Float64, ENV["XLA_REACTANT_GPU_MEM_FRACTION"]
-        )
-        @debug "XLA_REACTANT_GPU_MEM_FRACTION: " XLA_REACTANT_GPU_MEM_FRACTION[]
-    end
-
-    if haskey(ENV, "XLA_REACTANT_GPU_PREALLOCATE")
-        XLA_REACTANT_GPU_PREALLOCATE[] = parse(Bool, ENV["XLA_REACTANT_GPU_PREALLOCATE"])
-        @debug "XLA_REACTANT_GPU_PREALLOCATE: " XLA_REACTANT_GPU_PREALLOCATE[]
-    end
-
     @ccall MLIR.API.mlir_c.RegisterEnzymeXLAGPUHandler()::Cvoid
 
     # This wasn't properly exported on macos, we'll remove the try once macOS JLL