fix: apply init values after reduce (#896)

Author: avik-pal

src/Ops.jl

@@ -391,7 +391,7 @@ end
 end
 
 # shape ops
-function reshape(x::TracedRArray, dims...; kwargs...)
+function reshape(x::TracedRArray, dims::Integer...; kwargs...)
     return reshape(x, collect(dims); kwargs...)
 end
 
@@ -2392,9 +2392,9 @@ end
 """
     reduce(
         x::TracedRArray{T},
-        init_values::TracedRNumber{T},
+        init_values::Union{Nothing,TracedRNumber{T}},
         dimensions::Vector{Int},
-        fn::Function,
+        fn::Function;
         location=mlir_stacktrace("rand", @__FILE__, @__LINE__),
     )
 
@@ -2433,18 +2433,36 @@ Applies a reduction function `fn` along the specified `dimensions` of input `x`,
 """
 @noinline function reduce(
     x::TracedRArray{T},
-    init_values::TracedRNumber{T},
+    init_values::Union{TracedRNumber{T},Nothing},
     dimensions::Vector{Int},
-    fn::Function,
+    fn::Function;
     location=mlir_stacktrace("reduce", @__FILE__, @__LINE__),
 ) where {T}
+    elT = T
+    if init_values === nothing
+        if fn === min || fn === Base.FastMath.min_fast
+            init = typemax(elT)
+        elseif fn === max || fn === Base.FastMath.max_fast
+            init = typemin(elT)
+        else
+            init = Base.reduce_empty(Base.BottomRF(fn), elT)
+        end
+
+        initT = unwrapped_eltype(typeof(init))
+        if initT != elT # Bool, etc. reductions
+            elT = promote_type(initT, elT)
+            x = elT.(x)
+        end
+        init_values = Reactant.TracedUtils.promote_to(TracedRNumber{elT}, init)
+    end
+
     reduced_shape = Tuple(deleteat!(collect(size(x)), dimensions))
 
-    result_type = mlir_type(TracedRArray{T,length(reduced_shape)}, reduced_shape)
+    result_type = mlir_type(TracedRArray{elT,length(reduced_shape)}, reduced_shape)
 
     sample_inputs = [
-        Reactant.TracedUtils.promote_to(TracedRNumber{T}, 0),
-        Reactant.TracedUtils.promote_to(TracedRNumber{T}, 0),
+        Reactant.TracedUtils.promote_to(TracedRNumber{elT}, 0),
+        Reactant.TracedUtils.promote_to(TracedRNumber{elT}, 0),
     ]
 
     func =
@@ -2458,14 +2476,8 @@ Applies a reduction function `fn` along the specified `dimensions` of input `x`,
             return_dialect=:stablehlo,
         ).f
     @assert MLIR.IR.nregions(func) == 1
-    fn_name = String(
-        MLIR.IR.attr(func, String(MLIR.API.mlirSymbolTableGetSymbolAttributeName()))
-    )
-    ftype_attr = MLIR.IR.attr(func, "function_type")
-    ftype = MLIR.IR.Type(ftype_attr)
-    @assert MLIR.IR.result(ftype) == MLIR.IR.TensorType((), MLIR.IR.Type(T)) error (
-        "$fn return type is not tensor<i1>"
-    )
+    ftype = MLIR.IR.Type(MLIR.IR.attr(func, "function_type"))
+    @assert MLIR.IR.result(ftype) == MLIR.IR.TensorType((), MLIR.IR.Type(elT)) "$fn return type is not tensor<i1>"
     fn = MLIR.IR.Region()
     MLIR.API.mlirRegionTakeBody(fn, MLIR.IR.region(func, 1))
     MLIR.IR.rmfromparent!(func)
@@ -2483,7 +2495,7 @@ Applies a reduction function `fn` along the specified `dimensions` of input `x`,
         ),
     )
 
-    return TracedRArray{T,length(reduced_shape)}((), res, reduced_shape)
+    return TracedRArray{elT,length(reduced_shape)}((), res, reduced_shape)
 end
 
 end # module Ops

src/TracedRArray.jl

@@ -468,100 +468,29 @@ function Base.mapreduce(
     dims=:,
     init=nothing,
 ) where {T,N}
-    A = materialize_traced_array(A)
+    inp = broadcast(f, materialize_traced_array(A))
 
-    if dims isa Int
-        dims = [dims]
-    end
-
-    op_in_T = Core.Compiler.return_type(f, Tuple{T})
+    dims isa Number && (dims = (dims,))
 
-    if init === nothing
-        if op === min
-            init = typemax(op_in_T)
-        elseif op === max
-            init = typemin(op_in_T)
-        else
-            init = Base.reduce_empty(Base.BottomRF(op), op_in_T)
-        end
-
-        if typeof(init) != op_in_T
-            op_in_T = typeof(init)
-            A = typeof(init).(A)
-        end
+    if init !== nothing && typeof(init) != unwrapped_eltype(inp)
+        inp = typeof(init).(inp)
     end
 
-    init = [TracedUtils.broadcast_to_size(init, ()).mlir_data]
-
-    inp = [broadcast(f, A).mlir_data]
+    rdims = dims == (:) ? collect(Int64, 1:N) : collect(Int64, dims)
 
-    rdims = Int64[]
+    reduction_result = Ops.reduce(inp, nothing, rdims, op)
 
-    if dims == (:)
-        for i in 0:(N - 1)
-            push!(rdims, i)
-        end
+    reduction_result = if dims != (:)
+        Ops.reshape(reduction_result, Int64[i ∈ rdims ? 1 : size(A, i) for i in 1:N])
     else
-        for i in dims
-            push!(rdims, i - 1)
-        end
-    end
-
-    in_tys = [
-        MLIR.IR.TensorType(Int64[], eltype(MLIR.IR.type(inp[1]))),
-        MLIR.IR.TensorType(Int64[], eltype(MLIR.IR.type(init[1]))),
-    ]
-
-    fnbody = MLIR.IR.Block(in_tys, [MLIR.IR.Location(), MLIR.IR.Location()])
-
-    args = (
-        TracedRNumber{Reactant.unwrapped_eltype(op_in_T)}((), MLIR.IR.argument(fnbody, 1)),
-        TracedRNumber{Reactant.unwrapped_eltype(op_in_T)}((), MLIR.IR.argument(fnbody, 2)),
-    )
-
-    resty = MLIR.IR.block!(fnbody) do
-        tmp = TracedUtils.broadcast_to_size(op(args...), ())
-        Ops.return_(tmp)
-        return eltype(MLIR.IR.type(tmp.mlir_data))
+        TracedRNumber{unwrapped_eltype(reduction_result)}((), reduction_result.mlir_data)
     end
 
-    toonedims = Int[]
-    outdims = Int[]
-    for i in 1:N
-        tmp = if in(i - 1, rdims)
-            1
-        else
-            sz = size(A, i)
-            push!(outdims, sz)
-            sz
-        end
-        push!(toonedims, tmp)
-    end
-
-    TT = MLIR.IR.Type[MLIR.IR.TensorType(outdims, resty)]
-
-    body = MLIR.IR.Region()
-    push!(body, fnbody)
-    red = MLIR.Dialects.stablehlo.reduce(
-        inp, init; result_0=TT, dimensions=MLIR.IR.DenseArrayAttribute(rdims), body
-    )
-
-    red = MLIR.IR.result(red, 1)
-    redT = eltype(MLIR.IR.julia_type(MLIR.IR.type(red)))
-
-    if dims != (:)
-        red = Ops.reshape(TracedRArray(red), toonedims...)
-    else
-        if length(outdims) == 0
-            red = TracedRNumber{redT}((), red)
-        else
-            red = TracedRArray{redT,length(outdims)}((), red, (outdims...,))
-        end
-    end
-    return red
+    init === nothing && return reduction_result
+    return broadcast(op, reduction_result, init)
 end
 
-function Base.mapreducedim!(
+function Base._mapreducedim!(
     @nospecialize(f),
     @nospecialize(op),
     @nospecialize(R::AnyTracedRArray),
@@ -573,9 +502,9 @@ function Base.mapreducedim!(
         @assert sR == 1
         return i
     end
+    isempty(A) && return R
     tmp = mapreduce(f, op, A; dims=filter(!isnothing, dims))
-    # set_mlir_data!(R, get_mlir_data(tmp))
-    R .= op.(R, tmp) # match native Julia's behavior
+    R .= op.(R, tmp)
     return R
 end
 

test/basic.jl

@@ -939,6 +939,18 @@ end
     )
 end
 
+@testset "mapreduce with init" begin
+    x = reshape(collect(Float32, 1:12), 3, 4)
+    x_ra = Reactant.to_rarray(x)
+
+    init = 3.0
+    init_ra = Reactant.to_rarray(init; track_numbers=Number)
+
+    fn(x, init; kwargs...) = sum(x; init, kwargs...)
+
+    @test @jit(fn(x_ra, init_ra; dims=2)) ≈ fn(x, init; dims=2)
+end
+
 @testset "map!" begin
     x = randn(Float32, 2, 3)
     y = zeros(Float32, 2, 3)