Make Reactant loadable when the JLL isn't available (#1003)

* Make Reactant loadable when the JLL isn't available

* Make formatter happy

Author: giordano

deps/build_local.jl

@@ -194,6 +194,6 @@ using Preferences
 set_preferences!(
     joinpath(dirname(@__DIR__), "LocalPreferences.toml"),
     "Reactant_jll",
-    "libReactantExtra_path" => lib_path,
+    "libReactantExtra_path" => lib_path;
     force=true,
 )

src/Precompile.jl

@@ -57,29 +57,31 @@ function precompiling()
     return (@ccall jl_generating_output()::Cint) == 1
 end
 
-@setup_workload begin
-    initialize_dialect()
+if Reactant_jll.is_available()
+    @setup_workload begin
+        initialize_dialect()
 
-    if XLA.REACTANT_XLA_RUNTIME == "PJRT"
-        client = XLA.PJRT.CPUClient(; checkcount=false)
-    elseif XLA.REACTANT_XLA_RUNTIME == "IFRT"
-        client = XLA.IFRT.CPUClient(; checkcount=false)
-    else
-        error("Unsupported runtime: $(XLA.REACTANT_XLA_RUNTIME)")
-    end
+        if XLA.REACTANT_XLA_RUNTIME == "PJRT"
+            client = XLA.PJRT.CPUClient(; checkcount=false)
+        elseif XLA.REACTANT_XLA_RUNTIME == "IFRT"
+            client = XLA.IFRT.CPUClient(; checkcount=false)
+        else
+            error("Unsupported runtime: $(XLA.REACTANT_XLA_RUNTIME)")
+        end
 
-    @compile_workload begin
-        @static if precompilation_supported()
-            x = ConcreteRNumber(2.0; client)
-            Reactant.compile(sin, (x,); client, optimize=:all)
+        @compile_workload begin
+            @static if precompilation_supported()
+                x = ConcreteRNumber(2.0; client)
+                Reactant.compile(sin, (x,); client, optimize=:all)
 
-            y = ConcreteRArray([2.0]; client)
-            Reactant.compile(Base.sum, (y,); client, optimize=:all)
+                y = ConcreteRArray([2.0]; client)
+                Reactant.compile(Base.sum, (y,); client, optimize=:all)
+            end
         end
-    end
 
-    XLA.free_client(client)
-    client.client = C_NULL
-    deinitialize_dialect()
-    clear_oc_cache()
+        XLA.free_client(client)
+        client.client = C_NULL
+        deinitialize_dialect()
+        clear_oc_cache()
+    end
 end

src/Reactant.jl

@@ -252,8 +252,12 @@ function initialize_ptrs()
 end
 
 function __init__()
-    initialize_ptrs()
-    initialize_dialect()
+    if Reactant_jll.is_available()
+        initialize_ptrs()
+        initialize_dialect()
+    else
+        @warn "Reactant_jll isn't availble for your platform $(Reactant_jll.host_platform)"
+    end
     return nothing
 end
 

src/mlir/MLIR.jl

@@ -7,7 +7,11 @@ module API
     using Preferences
     using Reactant_jll
 
-    const mlir_c = Reactant_jll.libReactantExtra
+    const mlir_c = if Reactant_jll.is_available()
+        Reactant_jll.libReactantExtra
+    else
+        ""
+    end
 
     # MLIR C API
     let

src/xla/XLA.jl

@@ -104,53 +104,58 @@ function update_global_state!(args...; kwargs...)
 end
 
 function __init__()
-    # This must be the very first thing initialized (otherwise we can't throw errors)
-    errptr = cglobal((:ReactantThrowError, MLIR.API.mlir_c), Ptr{Ptr{Cvoid}})
-    unsafe_store!(errptr, @cfunction(reactant_err, Cvoid, (Cstring,)))
-
-    initLogs = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, "InitializeLogs")
-    ccall(initLogs, Cvoid, ())
-    # Add most log level
-    # SetLogLevel(0)
-
-    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[]
-        if XLA_REACTANT_GPU_MEM_FRACTION[] > 1 || XLA_REACTANT_GPU_MEM_FRACTION[] < 0
-            error("XLA_REACTANT_GPU_MEM_FRACTION must be between 0 and 1")
+    if Reactant_jll.is_available()
+        # This must be the very first thing initialized (otherwise we can't throw errors)
+        errptr = cglobal((:ReactantThrowError, MLIR.API.mlir_c), Ptr{Ptr{Cvoid}})
+        unsafe_store!(errptr, @cfunction(reactant_err, Cvoid, (Cstring,)))
+
+        initLogs = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, "InitializeLogs")
+        ccall(initLogs, Cvoid, ())
+        # Add most log level
+        # SetLogLevel(0)
+
+        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[]
+            if XLA_REACTANT_GPU_MEM_FRACTION[] > 1 || XLA_REACTANT_GPU_MEM_FRACTION[] < 0
+                error("XLA_REACTANT_GPU_MEM_FRACTION must be between 0 and 1")
+            end
         end
-    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
+        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
 
-    if haskey(ENV, "REACTANT_VISIBLE_GPU_DEVICES")
-        global_state.local_gpu_device_ids =
-            parse.(Int, split(ENV["REACTANT_VISIBLE_GPU_DEVICES"], ","))
-        @debug "REACTANT_VISIBLE_GPU_DEVICES: " global_state.local_gpu_device_ids
-    end
+        if haskey(ENV, "REACTANT_VISIBLE_GPU_DEVICES")
+            global_state.local_gpu_device_ids =
+                parse.(Int, split(ENV["REACTANT_VISIBLE_GPU_DEVICES"], ","))
+            @debug "REACTANT_VISIBLE_GPU_DEVICES: " global_state.local_gpu_device_ids
+        end
 
-    @debug "REACTANT_XLA_RUNTIME: " REACTANT_XLA_RUNTIME
+        @debug "REACTANT_XLA_RUNTIME: " REACTANT_XLA_RUNTIME
 
-    @ccall MLIR.API.mlir_c.RegisterEnzymeXLACPUHandler()::Cvoid
-    @ccall MLIR.API.mlir_c.RegisterEnzymeXLAGPUHandler()::Cvoid
+        @ccall MLIR.API.mlir_c.RegisterEnzymeXLACPUHandler()::Cvoid
+        @ccall MLIR.API.mlir_c.RegisterEnzymeXLAGPUHandler()::Cvoid
 
-    @static if !Sys.isapple()
-        lljit = Enzyme.LLVM.JuliaOJIT()
-        jd_main = Enzyme.LLVM.JITDylib(lljit)
+        @static if !Sys.isapple()
+            lljit = Enzyme.LLVM.JuliaOJIT()
+            jd_main = Enzyme.LLVM.JITDylib(lljit)
 
-        for name in ("XLAExecute", "XLAExecuteSharded", "ifrt_loaded_executable_execute")
-            ptr = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, name)
-            Enzyme.LLVM.define(
-                jd_main,
-                Enzyme.Compiler.JIT.absolute_symbol_materialization(
-                    Enzyme.LLVM.mangle(lljit, name), ptr
-                ),
-            )
+            for name in
+                ("XLAExecute", "XLAExecuteSharded", "ifrt_loaded_executable_execute")
+                ptr = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, name)
+                Enzyme.LLVM.define(
+                    jd_main,
+                    Enzyme.Compiler.JIT.absolute_symbol_materialization(
+                        Enzyme.LLVM.mangle(lljit, name), ptr
+                    ),
+                )
+            end
         end
     end