feat: precompile common workloads (#1485)

* feat: precompile common workloads

* fix: try reactivating precompile

* fix: add behind a check

* perf: run compilation time experiments on each commit

* fix: add all deps

* fix: only pre-compile forward pass for now

* test: disable some testing (drop me)

* fix: try moving to module

* fix: bad rebase

* fix: clear oc cache

Author: avik-pal

Project.toml

@@ -56,10 +56,10 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [sources]
-LuxCore = {path = "lib/LuxCore"}
-LuxLib = {path = "lib/LuxLib"}
-MLDataDevices = {path = "lib/MLDataDevices"}
-WeightInitializers = {path = "lib/WeightInitializers"}
+LuxCore = { path = "lib/LuxCore" }
+LuxLib = { path = "lib/LuxLib" }
+MLDataDevices = { path = "lib/MLDataDevices" }
+WeightInitializers = { path = "lib/WeightInitializers" }
 
 [extensions]
 LuxComponentArraysExt = "ComponentArrays"
@@ -71,7 +71,16 @@ LuxMLUtilsExt = "MLUtils"
 LuxMPIExt = "MPI"
 LuxMPINCCLExt = ["CUDA", "MPI", "NCCL"]
 LuxMooncakeExt = "Mooncake"
-LuxReactantExt = ["Enzyme", "Reactant"]
+LuxReactantExt = [
+    "Enzyme",
+    "Reactant",
+    "ReactantCore",
+    "LuxLib",
+    "MLDataDevices",
+    "Optimisers",
+    "NNlib",
+    "Statistics",
+]
 LuxReverseDiffExt = ["FunctionWrappers", "ReverseDiff"]
 LuxSimpleChainsExt = "SimpleChains"
 LuxTrackerExt = "Tracker"

ext/LuxReactantExt/LuxReactantExt.jl

@@ -1,8 +1,10 @@
 module LuxReactantExt
 
+using ADTypes: AutoEnzyme
 using Enzyme: Enzyme, Active, Const, Duplicated
 using Functors: Functors
 using Preferences: load_preference
+using Random: Random
 using Optimisers: Optimisers
 using Reactant:
     Reactant,
@@ -23,7 +25,8 @@ using Lux: Lux, LuxOps, Training, Utils, StatefulLuxLayer
 using Lux.Training: TrainingBackendCache, ReactantBackend
 using Lux: get_time_dimension, time_dimension_size, init_recurrent_state
 using LuxCore: LuxCore, AbstractLuxLayer
-using MLDataDevices: MLDataDevices, ReactantDevice, get_device
+using LuxLib: LuxLib
+using MLDataDevices: MLDataDevices, ReactantDevice, reactant_device, get_device
 
 Lux.is_extension_loaded(::Val{:Reactant}) = true
 
@@ -61,4 +64,6 @@ include("layers.jl")
 include("tracing.jl")
 include("saved_model.jl")
 
+include("precompile_workloads.jl")
+
 end

ext/LuxReactantExt/precompile_workloads.jl

@@ -0,0 +1,69 @@
+using PrecompileTools: @setup_workload, @compile_workload
+
+module PrecompileWorkloads
+
+function sumabs2attnloss(model, ps, st, data)
+    (y, _), stₙ = model(data, ps, st)
+    return sum(abs2, y), stₙ, NamedTuple()
+end
+
+function sumabs2loss(model, ps, st, data)
+    y, stₙ = model(data, ps, st)
+    return sum(abs2, y), stₙ, NamedTuple()
+end
+
+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
+
+        dev = reactant_device(; force=true)
+
+        # 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
+
+        # 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
+
+        @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)),
+            )
+
+            @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
+
+        Reactant.clear_oc_cache()
+        Reactant.set_default_backend(orig_backend)
+    end
+end

src/Lux.jl

@@ -123,7 +123,7 @@ include("serialization/serialization.jl")
 # Deprecations for v2
 include("deprecations.jl")
 
-# Precompile common workloads
+# Precompile Common Workloads
 include("precompile_workloads.jl")
 
 # Layers