feat: introduce a special (::dev)(...) syntax to compile reactant models

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.2.0"
+version = "1.3.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

ext/LuxReactantExt/LuxReactantExt.jl

@@ -9,6 +9,8 @@ using Static: False
 using Lux: Lux, LuxOps, Training
 using Lux.Training: TrainingBackendCache, ReactantBackend
 
+Lux.is_extension_loaded(::Val{:Reactant}) = true
+
 include("training.jl")
 
 end

src/Lux.jl

@@ -17,6 +17,7 @@ using Optimisers: Optimisers
 using Random: Random, AbstractRNG
 using Static: StaticBool, StaticInt, StaticSymbol, True, False, static, known, dynamic
 using Reexport: Reexport, @reexport
+using Setfield: @set!
 using Statistics: mean
 
 import LuxCore: AbstractLuxLayer, AbstractLuxContainerLayer, AbstractLuxWrapperLayer,
@@ -69,6 +70,7 @@ include("helpers/losses.jl")
 include("helpers/recursive_ops.jl")
 include("helpers/match_eltype.jl")
 include("helpers/size_propagator.jl")
+include("helpers/compile.jl")
 
 # AutoDiff
 include("autodiff/api.jl")

src/helpers/compile.jl

@@ -0,0 +1,32 @@
+function compile_reactant_model end
+
+function (dev::AbstractDevice)(model::AbstractExplicitLayer, x, ps, st)
+    return model, dev(x), dev(ps), dev(st)
+end
+
+function (dev::XLADevice)(model::AbstractExplicitLayer, x, ps, st)
+    if !is_extension_loaded(Val(:Reactant))
+        error("Reactant.jl needs to be loaded to use XLA compilation")
+    end
+    return compile_reactant_model(model, x, ps, st)
+end
+
+rewrite_conv_as_cross_correlation(model, ps, st, ::Functors.KeyPath) = model, ps, st
+function rewrite_conv_as_cross_correlation(model::Conv, ps, st, ::Functors.KeyPath)
+    if model.cross_correlation isa False
+        @set! model.cross_correlation = True()
+        @allowscalar ps_reversed = reverse(
+            ps.weight; dims=ntuple(i -> i, ndims(ps.weight) - 2)
+        )
+        @set! ps.weight = ps_reversed
+    end
+    return model, ps, st
+end
+
+function (dev::AMDGPUDevice)(model::AbstractExplicitLayer, x, ps, st)
+    # Conv is slow in MIOpen, so we rewrite it as a cross-correlation
+    model, ps, st = Experimental.layer_map(
+        rewrite_conv_as_cross_correlation, model, ps, st
+    )
+    return model, dev(x), dev(ps), dev(st)
+end