fix: reactant precompilation + throw error

Author: avik-pal

Project.toml

@@ -122,7 +122,7 @@ Optimisers = "0.4.6"
 PrecompileTools = "1.2.1"
 Preferences = "1.4.3"
 Random = "1.10"
-Reactant = "0.2.200"
+Reactant = "0.2.203"
 ReactantCore = "0.1.16"
 Reexport = "1.2.2"
 ReverseDiff = "1.15"

docs/Project.toml

@@ -68,7 +68,7 @@ NPZ = "0.4.3"
 Optimisers = "0.4.6"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.200"
+Reactant = "0.2.203"
 StableRNGs = "1"
 StaticArrays = "1"
 WeightInitializers = "1.3"

ext/ReactantExt/ReactantExt.jl

@@ -59,6 +59,6 @@ include("tracing.jl")
 include("saved_model.jl")
 include("batched_jacobian.jl")
 
-# include("precompile_workloads.jl")
+include("precompile_workloads.jl")
 
 end

ext/ReactantExt/precompile_workloads.jl

@@ -16,54 +16,60 @@ end
 
 if Reactant.Reactant_jll.is_available()
     @setup_workload begin
-        orig_backend = Reactant.XLA.default_backend()
-        Reactant.set_default_backend("cpu") # always precompile on CPU
+        try
+            orig_backend = Reactant.XLA.default_backend()
+            Reactant.set_default_backend("cpu") # always precompile on CPU
 
-        dev = reactant_device(; force=true)
+            dev = reactant_device(; force=true)
 
-        # attention model
-        mha = Lux.MultiHeadAttention(4; nheads=2)
-        ps_mha, st_mha = Lux.setup(Random.default_rng(), mha) |> dev
+            # attention model
+            mha = Lux.MultiHeadAttention(4; nheads=2)
+            ps_mha, st_mha = Lux.setup(Random.default_rng(), mha) |> dev
 
-        q = rand(Float32, (4, 3, 2)) |> dev
-        k = rand(Float32, (4, 3, 2)) |> dev
-        v = rand(Float32, (4, 3, 2)) |> dev
+            q = rand(Float32, (4, 3, 2)) |> dev
+            k = rand(Float32, (4, 3, 2)) |> dev
+            v = rand(Float32, (4, 3, 2)) |> dev
 
-        # convolution + dense model
-        conv_model = Lux.Chain(
-            Lux.Conv((3, 3), 3 => 32),
-            Lux.Conv((3, 3), 32 => 64),
-            Lux.GlobalMaxPool(),
-            Lux.FlattenLayer(),
-            Lux.Dense(64 => 10),
-        )
-        ps_conv_model, st_conv_model = Lux.setup(Random.default_rng(), conv_model) |> dev
+            # convolution + dense model
+            conv_model = Lux.Chain(
+                Lux.Conv((3, 3), 3 => 32),
+                Lux.Conv((3, 3), 32 => 64),
+                Lux.GlobalMaxPool(),
+                Lux.FlattenLayer(),
+                Lux.Dense(64 => 10),
+            )
+            ps_conv_model, st_conv_model = Lux.setup(Random.default_rng(), conv_model) |> dev
 
-        x = rand(Float32, (28, 28, 3, 2)) |> dev
+            x = rand(Float32, (28, 28, 3, 2)) |> dev
 
-        @compile_workload begin
-            @compile mha((q, k, v), ps_mha, LuxCore.testmode(st_mha))
+            @compile_workload begin
+                @compile mha((q, k, v), ps_mha, LuxCore.testmode(st_mha))
 
-            Lux.Training.single_train_step(
-                AutoEnzyme(),
-                PrecompileWorkloads.sumabs2attnloss,
-                (q, k, v),
-                Lux.Training.TrainState(mha, ps_mha, st_mha, Optimisers.Adam(0.001f0)),
-            )
+                Lux.Training.single_train_step(
+                    AutoEnzyme(),
+                    PrecompileWorkloads.sumabs2attnloss,
+                    (q, k, v),
+                    Lux.Training.TrainState(mha, ps_mha, st_mha, Optimisers.Adam(0.001f0)),
+                )
 
-            @compile conv_model(x, ps_conv_model, LuxCore.testmode(st_conv_model))
+                @compile conv_model(x, ps_conv_model, LuxCore.testmode(st_conv_model))
 
-            Lux.Training.single_train_step(
-                AutoEnzyme(),
-                PrecompileWorkloads.sumabs2loss,
-                x,
-                Lux.Training.TrainState(
-                    conv_model, ps_conv_model, st_conv_model, Optimisers.Adam(0.001f0)
-                ),
-            )
-        end
+                Lux.Training.single_train_step(
+                    AutoEnzyme(),
+                    PrecompileWorkloads.sumabs2loss,
+                    x,
+                    Lux.Training.TrainState(
+                        conv_model, ps_conv_model, st_conv_model, Optimisers.Adam(0.001f0)
+                    ),
+                )
+            end
 
-        Reactant.clear_oc_cache()
-        Reactant.set_default_backend(orig_backend)
+            Reactant.clear_oc_cache()
+            Reactant.set_default_backend(orig_backend)
+        catch err
+            if !(err isa Reactant.ReactantPrecompilationException)
+                rethrow(err)
+            end
+        end
     end
 end

lib/LuxCore/Project.toml

@@ -42,7 +42,7 @@ Flux = "0.16.3"
 Functors = "0.5"
 MLDataDevices = "1.17"
 Random = "1.10"
-Reactant = "0.2.200"
+Reactant = "0.2.203"
 ReverseDiff = "1.15"
 SciMLPublic = "1.0.0"
 Setfield = "1"

lib/MLDataDevices/Project.toml

@@ -67,7 +67,7 @@ Metal = "1"
 OneHotArrays = "0.2.10"
 Preferences = "1.4.3"
 Random = "1.10"
-Reactant = "0.2.200"
+Reactant = "0.2.203"
 RecursiveArrayTools = "3.8"
 ReverseDiff = "1.15"
 SciMLPublic = "1.0.0"

test/Project.toml

@@ -80,7 +80,7 @@ Preferences = "1.4.3"
 PythonCall = "0.9"
 Random = "1.10"
 ReTestItems = "1.24.0"
-Reactant = "0.2.170"
+Reactant = "0.2.203"
 Reexport = "1.2.2"
 ReverseDiff = "1.15.3"
 Setfield = "1.1.1"