chore: updates for new JLL (#1190)

* fix: minor fixes

* fix: run passes 3rd time

* chore: bump jll

* test: simple RNN test for whileOp autodiff

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.74"
+version = "0.2.75"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -87,7 +87,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.9"
-Reactant_jll = "0.0.146"
+Reactant_jll = "0.0.147"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/Compiler.jl

@@ -779,6 +779,7 @@ function optimization_passes(;
                 "transpose_wrap",
                 "transpose_extend",
                 "transpose_rotate",
+                "transpose_dynamic_slice",
             ],
         )
         if AGGRESSIVE_PROPAGATION[]
@@ -1248,6 +1249,7 @@ function compile_mlir!(
                     "enzyme-batch",
                     opt_passes2,
                     enzyme_pass,
+                    opt_passes2,
                     "canonicalize",
                     "remove-unnecessary-enzyme-ops",
                     "enzyme-simplify-math",
@@ -1269,6 +1271,7 @@ function compile_mlir!(
                     "enzyme-batch",
                     opt_passes2,
                     enzyme_pass,
+                    opt_passes2,
                     "canonicalize",
                     "remove-unnecessary-enzyme-ops",
                     "enzyme-simplify-math",
@@ -1287,6 +1290,7 @@ function compile_mlir!(
                     "enzyme-batch",
                     opt_passes2,
                     enzyme_pass,
+                    opt_passes2,
                     "canonicalize",
                     "remove-unnecessary-enzyme-ops",
                     "enzyme-simplify-math",
@@ -1307,6 +1311,7 @@ function compile_mlir!(
                     "enzyme-batch",
                     opt_passes2,
                     enzyme_pass,
+                    opt_passes2,
                     "canonicalize",
                     "remove-unnecessary-enzyme-ops",
                     "enzyme-simplify-math",
@@ -1326,6 +1331,7 @@ function compile_mlir!(
                     "enzyme-batch",
                     opt_passes2,
                     enzyme_pass,
+                    opt_passes2,
                     "canonicalize",
                     "remove-unnecessary-enzyme-ops",
                     "enzyme-simplify-math",
@@ -1397,16 +1403,7 @@ function compile_mlir!(
         error("Invalid optimize option: $(Meta.quot(optimize))")
     end
 
-    # HACK: remove with next JLL
     if !(optimize isa String)
-        if transpose_propagate === :up
-            run_pass_pipeline!(
-                mod,
-                "enzyme-hlo-generate-td{patterns=transpose_while},transform-interpreter,enzyme-hlo-remove-transform",
-                "transpose_while",
-            )
-        end
-
         if optimize ∉ (:none, :just_batch, :canonicalize) &&
             (transpose_propagate === :up || reshape_propagate === :up)
             # We tried propagating reshapes and transposes up. If at this point we are left with
@@ -1425,7 +1422,6 @@ function compile_mlir!(
     end
 
     # Now we resolve paddings if `optimize_then_pad`
-    prepad_fnname = fnname
     if optimize_then_pad
         padded_inputs = IdDict()
         has_padded_inputs = false

src/stdlibs/Random.jl

@@ -14,8 +14,8 @@ using ..Reactant:
     Reactant,
     TracedUtils,
     Ops,
-    ConcretePJRTArray,
-    ConcretePJRTNumber,
+    AbstractConcreteArray,
+    AbstractConcreteNumber,
     unwrapped_eltype
 using Random: Random, AbstractRNG
 
@@ -48,29 +48,29 @@ end
 end
 
 @noinline function Random.seed!(rng::ConcreteRNG, seed::Number)
-    seed isa ConcretePJRTNumber && (seed = unwrapped_eltype(seed)(seed))
+    seed isa AbstractConcreteNumber && (seed = unwrapped_eltype(seed)(seed))
     seed = reinterpret(UInt64, Random.hash_seed(seed))
-    return Random.seed!(rng, ConcretePJRTArray(seed))
+    return Random.seed!(rng, Reactant.to_rarray(seed))
 end
 
 @noinline function Random.seed!(rng::ConcreteRNG, seed::AbstractVector{<:Integer})
     return Random.seed!(rng, seed)
 end
 
 @noinline function Random.seed!(rng::ConcreteRNG, seed::AbstractVector{UInt64})
-    return Random.seed!(rng, ConcretePJRTArray(seed))
+    return Random.seed!(rng, Reactant.to_rarray(seed))
 end
 
-@noinline function Random.seed!(rng::ConcreteRNG, seed::ConcretePJRTArray{UInt64,1})
+@noinline function Random.seed!(rng::ConcreteRNG, seed::AbstractConcreteArray{UInt64,1})
     rng.seed = seed
     return rng
 end
 
 Base.copy(rng::ConcreteRNG) = ConcreteRNG(copy(rng.seed), rng.algorithm)
 Base.copy(rng::TracedRNG) = TracedRNG(copy(rng.seed), rng.algorithm)
 
-@noinline ConcreteRNG() = ConcreteRNG(ConcretePJRTArray(make_seed()))
-@noinline ConcreteRNG(seed::ConcretePJRTArray{UInt64,1}) = ConcreteRNG(seed, "DEFAULT")
+@noinline ConcreteRNG() = ConcreteRNG(Reactant.to_rarray(make_seed()))
+@noinline ConcreteRNG(seed::AbstractConcreteArray{UInt64,1}) = ConcreteRNG(seed, "DEFAULT")
 
 @noinline default_rng() = ConcreteRNG()
 

test/nn/lux.jl

@@ -5,6 +5,11 @@ function loss_function(model, x, y, ps, st)
     return CrossEntropyLoss()(y_hat, y)
 end
 
+function loss_function(model, x, ps, st)
+    y_hat, _ = model(x, ps, st)
+    return sum(abs2, y_hat)
+end
+
 function gradient_loss_function(model, x, y, ps, st)
     dps = Enzyme.make_zero(ps)
     _, res = Enzyme.autodiff(
@@ -20,6 +25,20 @@ function gradient_loss_function(model, x, y, ps, st)
     return res, dps
 end
 
+function gradient_loss_function(model, x, ps, st)
+    dps = Enzyme.make_zero(ps)
+    _, res = Enzyme.autodiff(
+        set_runtime_activity(ReverseWithPrimal),
+        loss_function,
+        Active,
+        Const(model),
+        Const(x),
+        Duplicated(ps, dps),
+        Const(st),
+    )
+    return res, dps
+end
+
 @testset "Lux.jl Integration" begin
     # Generate some data for the XOR problem: vectors of length 2, as columns of a matrix:
     noisy = rand(Float32, 2, 1000)                                        # 2×1000 Matrix{Float32}
@@ -67,3 +86,16 @@ end
         @test_skip dps1 ≈ dps2 atol = 1e-3 rtol = 1e-2
     end
 end
+
+@testset "RNN Integration" begin
+    using Reactant, Lux, Enzyme, Random
+
+    model = Recurrence(RNNCell(4 => 4); ordering=BatchLastIndex())
+    ps, st = Reactant.to_rarray(Lux.setup(Random.default_rng(), model))
+
+    x = Reactant.to_rarray(rand(Float32, 4, 16, 12))
+
+    # This test requires running optimizations between the enzyme autodiff passes
+    res, ∂ps = @jit gradient_loss_function(model, x, ps, st)
+    @test res isa Reactant.ConcreteRNumber
+end