fix: more dispatches for mapreduce (#1521)

Author: avik-pal

src/Ops.jl

@@ -1064,6 +1064,7 @@ function broadcast_in_dim(
     location=mlir_stacktrace("broadcast_in_dim", @__FILE__, @__LINE__),
 ) where {T,N}
     @assert length(dims) == N
+    @assert length(result_size) ≥ N
 
     res = MLIR.IR.result(
         stablehlo.broadcast_in_dim(

src/TracedRArray.jl

@@ -543,7 +543,14 @@ end
 function overloaded_mapreduce(
     @nospecialize(f), @nospecialize(op), @nospecialize(A); dims=:, init=nothing
 )
-    return overloaded_mapreduce(identity, op, unwrapped_broadcast(f, A); dims, init)
+    res = unwrapped_broadcast(f, A)
+    # This means we are unable to use the optimized dispatches. For now we will
+    # unroll the mapreduce.
+    if typeof(res) == typeof(A)
+        @assert dims == Colon() "dims not supported for mapreduce currently."
+        return foldl(op, res; init)
+    end
+    return overloaded_mapreduce(identity, op, res; dims=:, init)
 end
 
 function overloaded_mapreduce(
@@ -1436,16 +1443,27 @@ struct BroadcastIterator{F}
     f::F
 end
 
-(fn::BroadcastIterator)(args...) = fn.f((args...,))
+(fn::BroadcastIterator)(args...) = Reactant.call_with_reactant(fn.f, (args...,))
 
 function unwrapped_broadcast(f::F, x::Base.Iterators.Zip) where {F}
     min_length = Base.inferencebarrier(minimum)(length, x.is)
     itrs = [length(itr) > min_length ? itr[1:min_length] : itr for itr in x.is]
-    return (BroadcastIterator(f)).(itrs...)
+    if any(Base.Fix2(isa, AnyTracedRArray), itrs)
+        return (BroadcastIterator(f)).(itrs...)
+    else
+        fn = BroadcastIterator(f)
+        return [fn(Base.Fix2(getindex, i).(itrs)...) for i in 1:min_length]
+    end
 end
 
 function unwrapped_broadcast(f::F, x::Base.Iterators.Enumerate) where {F}
-    return (BroadcastIterator(f)).(1:length(x.itr), x.itr)
+    if x.itr isa AnyTracedRArray
+        return (BroadcastIterator(f)).(1:length(x.itr), x.itr)
+    else
+        return [f((i, x.itr[i])) for i in 1:length(x.itr)]
+    end
 end
 
+unwrapped_broadcast(f::F, xs::Vector) where {F} = [f(x) for x in xs]
+
 end

test/basic.jl

@@ -1444,6 +1444,20 @@ end
 zip_iterator(a, b) = mapreduce(splat(*), +, zip(a, b))
 enumerate_iterator(a) = mapreduce(splat(*), +, enumerate(a))
 
+function nested_mapreduce_zip(x, y)
+    return mapreduce(+, zip(eachcol(x), eachcol(y)); init=0.0f0) do (x, y)
+        return sum(abs2, x) + sum(abs2, y)
+    end
+end
+
+function nested_mapreduce_hcat(x, y)
+    return mapreduce(
+        hcat, zip(eachcol(x), eachcol(y)); init=similar(x, size(x, 1), 0)
+    ) do (x, y)
+        return x .+ y
+    end
+end
+
 @testset "Base.Iterators" begin
     @testset "zip" begin
         N = 10
@@ -1460,4 +1474,24 @@ enumerate_iterator(a) = mapreduce(splat(*), +, enumerate(a))
 
         @test @jit(enumerate_iterator(x_ra)) ≈ enumerate_iterator(x)
     end
+
+    @testset "nested mapreduce" begin
+        x = rand(Float32, 4, 3)
+        y = rand(Float32, 4, 3)
+
+        x_ra = Reactant.to_rarray(x)
+        y_ra = Reactant.to_rarray(y)
+
+        @test @jit(nested_mapreduce_zip(x_ra, y_ra)) ≈ nested_mapreduce_zip(x, y)
+    end
+
+    @testset "nested mapreduce hcat" begin
+        x = rand(Float32, 4, 3)
+        y = rand(Float32, 4, 3)
+
+        x_ra = Reactant.to_rarray(x)
+        y_ra = Reactant.to_rarray(y)
+
+        @test @jit(nested_mapreduce_hcat(x_ra, y_ra)) ≈ nested_mapreduce_hcat(x, y)
+    end
 end