feat: new scatter / gather optimization patterns (#1397)

* feat: new scatter / gather optimization patterns

* fix: use iota

* Update test/integration/onehotarrays.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

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.133"
+version = "0.2.134"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -90,7 +90,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.12"
-Reactant_jll = "0.0.200"
+Reactant_jll = "0.0.201"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

ext/ReactantOneHotArraysExt.jl

@@ -2,6 +2,7 @@ module ReactantOneHotArraysExt
 
 using OneHotArrays
 using Reactant
+using Reactant: TracedRArray, TracedRNumber, TracedUtils, Ops
 
 function Reactant.traced_type_inner(
     @nospecialize(_::Type{OneHotArrays.OneHotArray{T,N,Np1,I}}),
@@ -21,16 +22,24 @@ function Reactant.traced_type_inner(
 end
 
 # OneHotArray is a <: AbstractArray{Bool, M} so our usual dispatches don't work
-function Reactant.TracedUtils.broadcast_to_size(
+function TracedUtils.broadcast_to_size(
     r::OneHotArrays.OneHotArray{T,N,Np1,<:Reactant.TracedRArray}, rsize
 ) where {T,N,Np1}
-    return Reactant.TracedUtils.broadcast_to_size(
-        Reactant.TracedUtils.materialize_traced_array(r), rsize
-    )
+    return TracedUtils.broadcast_to_size(TracedUtils.materialize_traced_array(r), rsize)
 end
 
-function Reactant.TracedUtils.materialize_traced_array(r::OneHotArrays.OneHotArray)
-    return reshape(r.indices, 1, size(r.indices)...) .== 1:(r.nlabels)
+function TracedUtils.materialize_traced_array(r::OneHotArrays.OneHotArray)
+    indices = vec(r.indices)
+    N = r.nlabels
+    B = length(indices)
+
+    linear_indices =
+        TracedUtils.promote_to(TracedRArray{Int64,ndims(r.indices)}, indices) .+
+        Ops.iota(Int64, [B]; iota_dimension=1) .* N
+
+    z = Ops.fill(false, (N, B))
+    z[linear_indices] = fill(true, length(linear_indices))
+    return reshape(z, size(r))
 end
 
 function Base.Array(

src/Compiler.jl

@@ -928,6 +928,8 @@ function optimization_passes(;
         "split_convolution_into_reverse_convolution",
         # TODO we want to enable but may cause an infinite compile time
         # "concat_to_onedim_dusslice",
+        "scatter_multiply_simplify",
+        "unary_elementwise_scatter_simplify",
     ]
 
     # constant prop patterns
@@ -980,6 +982,7 @@ function optimization_passes(;
             "concat_const_prop<1>($max_constant_threshold)",
             "dynamic_update_slice_const_prop($max_constant_threshold)",
             "scatter_update_computation_const_prop",
+            "gather_const_prop",
         ],
     )
 

test/integration/onehotarrays.jl

@@ -16,9 +16,18 @@ end
     x = rand(Float32, 4, 5)
     x_ra = Reactant.to_rarray(x)
 
-    res_ra = @jit r_m .+ x_ra
-    res = m .+ x
-    @test res_ra ≈ res
+    @testset "addition" begin
+        res_ra = @jit r_m .+ x_ra
+        res = m .+ x
+        @test res_ra ≈ res
 
-    @test Array(r_m) isa Matrix{Bool}
+        @test Array(r_m) isa Matrix{Bool}
+    end
+
+    @testset "multiplication" begin
+        # Broadcasting a multiplication has special passes
+        res_ra = @jit r_m .* x_ra
+        res = m .* x
+        @test res_ra ≈ res
+    end
 end