Revert "fix: apply init values after reduce (#896)"

This reverts commit 44910bf.

Author: giordano

src/Ops.jl

@@ -391,7 +391,7 @@ end
 end
 
 # shape ops
-function reshape(x::TracedRArray, dims::Integer...; kwargs...)
+function reshape(x::TracedRArray, dims...; kwargs...)
     return reshape(x, collect(dims); kwargs...)
 end
 
@@ -2392,9 +2392,9 @@ end
 """
     reduce(
         x::TracedRArray{T},
-        init_values::Union{Nothing,TracedRNumber{T}},
+        init_values::TracedRNumber{T},
         dimensions::Vector{Int},
-        fn::Function;
+        fn::Function,
         location=mlir_stacktrace("rand", @__FILE__, @__LINE__),
     )
 
@@ -2433,36 +2433,18 @@ Applies a reduction function `fn` along the specified `dimensions` of input `x`,
 """
 @noinline function reduce(
     x::TracedRArray{T},
-    init_values::Union{TracedRNumber{T},Nothing},
+    init_values::TracedRNumber{T},
     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{elT,length(reduced_shape)}, reduced_shape)
+    result_type = mlir_type(TracedRArray{T,length(reduced_shape)}, reduced_shape)
 
     sample_inputs = [
-        Reactant.TracedUtils.promote_to(TracedRNumber{elT}, 0),
-        Reactant.TracedUtils.promote_to(TracedRNumber{elT}, 0),
+        Reactant.TracedUtils.promote_to(TracedRNumber{T}, 0),
+        Reactant.TracedUtils.promote_to(TracedRNumber{T}, 0),
     ]
 
     func =
@@ -2476,8 +2458,14 @@ Applies a reduction function `fn` along the specified `dimensions` of input `x`,
             return_dialect=:stablehlo,
         ).f
     @assert MLIR.IR.nregions(func) == 1
-    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_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>"
+    )
     fn = MLIR.IR.Region()
     MLIR.API.mlirRegionTakeBody(fn, MLIR.IR.region(func, 1))
     MLIR.IR.rmfromparent!(func)
@@ -2495,7 +2483,7 @@ Applies a reduction function `fn` along the specified `dimensions` of input `x`,
         ),
     )
 
-    return TracedRArray{elT,length(reduced_shape)}((), res, reduced_shape)
+    return TracedRArray{T,length(reduced_shape)}((), res, reduced_shape)
 end
 
 end # module Ops

src/TracedRArray.jl

@@ -468,29 +468,100 @@ function Base.mapreduce(
     dims=:,
     init=nothing,
 ) where {T,N}
-    inp = broadcast(f, materialize_traced_array(A))
+    A = materialize_traced_array(A)
 
-    dims isa Number && (dims = (dims,))
+    if dims isa Int
+        dims = [dims]
+    end
+
+    op_in_T = Core.Compiler.return_type(f, Tuple{T})
 
-    if init !== nothing && typeof(init) != unwrapped_eltype(inp)
-        inp = typeof(init).(inp)
+    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
     end
 
-    rdims = dims == (:) ? collect(Int64, 1:N) : collect(Int64, dims)
+    init = [TracedUtils.broadcast_to_size(init, ()).mlir_data]
+
+    inp = [broadcast(f, A).mlir_data]
 
-    reduction_result = Ops.reduce(inp, nothing, rdims, op)
+    rdims = Int64[]
 
-    reduction_result = if dims != (:)
-        Ops.reshape(reduction_result, Int64[i ∈ rdims ? 1 : size(A, i) for i in 1:N])
+    if dims == (:)
+        for i in 0:(N - 1)
+            push!(rdims, i)
+        end
     else
-        TracedRNumber{unwrapped_eltype(reduction_result)}((), reduction_result.mlir_data)
+        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))
     end
 
-    init === nothing && return reduction_result
-    return broadcast(op, reduction_result, init)
+    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
 end
 
-function Base._mapreducedim!(
+function Base.mapreducedim!(
     @nospecialize(f),
     @nospecialize(op),
     @nospecialize(R::AnyTracedRArray),
@@ -502,9 +573,9 @@ function Base._mapreducedim!(
         @assert sR == 1
         return i
     end
-    isempty(A) && return R
     tmp = mapreduce(f, op, A; dims=filter(!isnothing, dims))
-    R .= op.(R, tmp)
+    # set_mlir_data!(R, get_mlir_data(tmp))
+    R .= op.(R, tmp) # match native Julia's behavior
     return R
 end
 

test/basic.jl

@@ -939,18 +939,6 @@ 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)