Use the REACTANT_DEFAULT_DEVICE environment variable

Author: giordano

src/xla/Client.jl

@@ -16,7 +16,8 @@ function platform_name end
 """
     DEFAULT_DEVICE :: Ref{Int}
 
-0-based index of default device to use, by default 0 (first available device).
+0-based index of default device to use.
+By default, the value of the environment variable `REACTANT_DEFAULT_DEVICE` is used when set to a non-negative integer, otherwise it is set to 0 (first available device).
 """
 const DEFAULT_DEVICE = Ref{Int}(0)
 

src/xla/XLA.jl

@@ -158,6 +158,13 @@ function __init__()
                 1
         end
 
+        if haskey(ENV, "REACTANT_DEFAULT_DEVICE")
+            DEFAULT_DEVICE[] = max(
+                0, something(tryparse(Int, ENV["REACTANT_DEFAULT_DEVICE"]), 0)
+            )
+            @debug "REACTANT_DEFAULT_DEVICE: " DEFAULT_DEVICE[] maxlog = 1
+        end
+
         @debug "REACTANT_XLA_RUNTIME: " REACTANT_XLA_RUNTIME maxlog = 1
 
         @ccall MLIR.API.mlir_c.RegisterEnzymeXLACPUHandler()::Cvoid