Skip to content

WIP: a more principled take on dimensional reduction inits #55318

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 22 commits into from
Closed

WIP: a more principled take on dimensional reduction inits #55318

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
b9bdd63
a more principled take on dimensional reduction inits
mbauman Jul 31, 2024
2b000b1
add support for broadcasteds
mbauman Jul 31, 2024
07c1da1
add tests for complex reductions
mbauman Jul 31, 2024
2179df3
add tests for bitwise operators on integers
mbauman Jul 31, 2024
344ed43
fixup! add tests for bitwise operators on integers
mbauman Jul 31, 2024
8eb7d91
restore performance for linear fast dim=1+
mbauman Jul 31, 2024
11e98ed
fixup! add support for broadcasteds
mbauman Jul 31, 2024
c32c6ed
add more tests
mcabbott Jul 31, 2024
58d2b27
perf: add loop switching
mbauman Jul 31, 2024
2d595bb
perf: slice up the slow part
mbauman Jul 31, 2024
b7adc89
use builtin incremental widening slowly when needed
mbauman Aug 1, 2024
e59a4d8
fixup! use builtin incremental widening slowly when needed
mbauman Aug 2, 2024
6af4850
add tests for 44906
mbauman Aug 2, 2024
d31cba6
also revamp findmin/findmax
mbauman Aug 4, 2024
6025380
minor cleanups, slightly better naming and hoist some merges
mbauman Aug 6, 2024
3e9e9d0
fixup findmin/max/etc and add support for custom allocators
mbauman Aug 27, 2024
6bfd6c1
more tests
mbauman Aug 29, 2024
1635026
better comment and slight cleanups
mbauman Aug 29, 2024
64a27ea
Merge branch 'master' into mapreducedim-init-adienes-testing
adienes Dec 16, 2024
ce4c0cd
Merge remote-tracking branch 'origin' into mapreducedim-init-adienes-…
adienes Dec 21, 2024
48b4ad0
couple tuneups to dimensional mapreduce refactor
adienes Dec 21, 2024
aacdcee
Merge branch 'master' into mb/mapreducedim-init
mbauman May 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions base/fastmath.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -408,9 +408,10 @@ maximum!_fast(r::AbstractArray, A::AbstractArray; kw...) =
minimum!_fast(r::AbstractArray, A::AbstractArray; kw...) =
minimum!_fast(identity, r, A; kw...)

maximum!_fast(f::Function, r::AbstractArray, A::AbstractArray; init::Bool=true) =
Base.mapreducedim!(f, max_fast, Base.initarray!(r, f, max, init, A), A)
minimum!_fast(f::Function, r::AbstractArray, A::AbstractArray; init::Bool=true) =
Base.mapreducedim!(f, min_fast, Base.initarray!(r, f, min, init, A), A)
maximum!_fast(f::Function, r::AbstractArray, A::AbstractArray; init::Bool=true) = init ?
Base.mapreduce!(f, max_fast, r, A) : Base.mapreducedim!(f, max_fast, r, A)

minimum!_fast(f::Function, r::AbstractArray, A::AbstractArray; init::Bool=true) = init ?
Base.mapreduce!(f, min_fast, r, A) : Base.mapreducedim!(f, min_fast, r, A)

end
61 changes: 22 additions & 39 deletions base/reducedim.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,16 +129,13 @@ function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArrayOrBroadcasted,
return R
end

if isdefined(Core, :Compiler)
_mapreduce_might_widen(_, _, _, _) = false
function _mapreduce_might_widen(f::F, op::G, A::T, ::Nothing) where {F,G,T}
return !isconcretetype(Core.Compiler.return_type(x->op(f(first(x)), f(first(x))), Tuple{T}))
end
else
_mapreduce_might_widen(_, _, _, ::Nothing) = true
_mapreduce_might_widen(_, _, _, _) = false

_mapreduce_might_widen(_, _, _, _) = false
function _mapreduce_might_widen(f::F, op::G, A::T, ::Nothing) where {F,G,T}
return !isconcretetype(Base._return_type(x->op(f(first(x)), f(first(x))), Tuple{T}))
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@adienes I'm coming back to this and am curious — do you remember why you changed this?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe (if I remember correctly), that it was causing some test failures in Statistics.cov because it skipped some pairwise summation path which changed the accuracy of the reduction. unfortunately I don't remember which test exactly --- this change might not be necessary in the end

end


mapreduce_similar(A, ::Type{T}, dims) where {T} = similar(A, T, dims)
function _mapreduce_similar_curried(A)
return function(::Type{T}, dims) where {T}
Expand Down Expand Up @@ -190,7 +187,6 @@ end
_mapreducedim_impl(f::Type{F}, op, init, A::AbstractArrayOrBroadcasted, raxs, out) where {F} =
_mapreducedim_impl(x->F(x), op, init, A::AbstractArrayOrBroadcasted, raxs, out)
function _mapreducedim_impl(f, op, init, A::AbstractArrayOrBroadcasted, raxs, out)
# @show f, op, init, A, raxs, out
outer_inds = CartesianIndices(raxs)
a_inds = CartesianIndices(axes(A))
Aaxs = axes(A)
Expand All @@ -200,8 +196,6 @@ function _mapreducedim_impl(f, op, init, A::AbstractArrayOrBroadcasted, raxs, ou
n = prod(map((b, ax)->ifelse(b, 1, length(ax)), is_outer_dim, Aaxs))

# We have special support for returning empty arrays
# isempty(outer_inds) && return collect_allocator(_mapreduce_similar_curried(A),
# (_mapreduce_naive_inner_loop(f, op, init, A, is_outer_dim, outer_ind, a_inds) for outer_ind in outer_inds))
(n == 0 || isempty(A)) && return _mapreduce_empty_array(f, op, init, A, raxs, out)
n == 1 && return _mapreduce_one_array(f, op, init, A, raxs, out)

Expand All @@ -210,13 +204,13 @@ function _mapreducedim_impl(f, op, init, A::AbstractArrayOrBroadcasted, raxs, ou

# Create the result vector with the first step of the reduction, using the passed output if given
# note that this is using a (potentially) bad cache ordering, so we don't want to continue like this
i1 = first(outer_inds)
inner1 = mergeindices(is_outer_dim, i1, a_inds)
o1 = first(outer_inds)
inner1 = mergeindices(is_outer_dim, o1, a_inds)
i1, i2 = inner1
@inbounds a1, a2 = A[i1], A[i2]
v = _mapreduce_start(f, op, init, a1, a2)
R = isnothing(out) ? mapreduce_similar(A, typeof(v), raxs) : out
@inbounds R[i1] = v
@inbounds R[o1] = v
for i in Iterators.drop(outer_inds, 1)
innerj = mergeindices(is_outer_dim, i, a_inds)
j1, j2 = innerj
Expand Down Expand Up @@ -246,35 +240,24 @@ function _mapreducedim_impl(f, op, init, A::AbstractArrayOrBroadcasted, raxs, ou
end
end
else
# Take care of the smallest cartesian "rectangle" that includes our two peeled elements
small_inner, large_inner = _split2(inner1)
# and we may as well try to do this small part in a cache-advantageous manner
if something(findfirst((>)(1), size(inner1))) < something(findfirst((>)(1), size(outer_inds)), typemax(Int))
for iR in outer_inds
@inbounds r = R[iR]
inner_inds = mergeindices(is_outer_dim, CartesianIndices(map(sliceat, Aaxs, iR.I)), small_inner)
for iA in Iterators.drop(inner_inds, 2)
r = op(r, f(@inbounds(A[iA])))
end
@inbounds R[iR] = r
end
else
for iA in Iterators.drop(small_inner, 2)
for iR in outer_inds
iA = mergeindices(is_outer_dim, iR, iA)
v = op(@inbounds(R[iR]), f(@inbounds(A[iA])))
@inbounds R[iR] = v
end
# Fall back to a straightforward loop for the remaining elements
# We already handled the first two items in each slice above.
# So we skip them and keep accumulating in R[iR].
# Critically, we do *not* call _split2 after dropping 2 items,
# because that leftover set isn’t rectangular for ND.
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm still struggling to see what I got wrong here in my initial take, but I'd really like to find it. I'd like to add it to the tests!

If this was wrong, then there must be some A for which I missed or double-counted some elements — and thus sum(A; dims) !≈ mapslices(sum, A; dims), yeah? But I can't find it!

julia> versioninfo()
Julia Version 1.12.0-DEV.970
Commit 163502626b (2024-08-29 22:35 UTC)

julia> using Test, Combinatorics

julia> function test()
           @testset "size $sz" for sz in Iterators.product(Iterators.repeated(1:5, 5)...)
               @testset "dims $dims" for dims in map(Base.splat(tuple), Iterators.flatten(combinations.((1:5,), 1:5)))
                   for f in (identity, x->view(x, range.(1, sz)...)) # both IndexLinear and IndexCartesian
                       A = f(rand(Float64, sz))
                       @test sum(A; dims=dims) ≈ mapslices(sum, A; dims=dims)
                   end
               end
           end
        end
test (generic function with 3 methods)

julia> @testset test();
Test Summary: |   Pass   Total  Time
test          | 193750  193750  3.7s

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let me see... it's been a bit. I'll try to generate some relevant test cases (it's also possible I was wrong)

for o in outer_inds
# The partial result for the first 2 items is in R[o]
@inbounds r = R[o]
# Now process the rest of that slice, skipping the first 2
fullslice = mergeindices(is_outer_dim, o, a_inds)
for iA in Iterators.drop(fullslice, 2)
r = op(r, f(@inbounds(A[iA])))
end
end
# And then if there's anything left, defer to _mapreducedim!
if !isempty(large_inner)
large_inds = mergeindices(is_outer_dim, outer_inds, large_inner)
_mapreducedim!(f, op, R, A, large_inds.indices)
@inbounds R[o] = r
end
end
return R
end
end

function _mapreduce_naive_inner_loop(f, op, init, A, is_outer_dim, outer_ind, a_inds)
inner_inds = mergeindices(is_outer_dim, outer_ind, a_inds)
Expand Down
26 changes: 0 additions & 26 deletions contrib/juliac-buildscript.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,32 +97,6 @@ end
end
end
show_type_name(io::IO, tn::Core.TypeName) = print(io, tn.name)

mapreduce(f::F, op::F2, A::AbstractArrayOrBroadcasted; dims=:, init=_InitialValue()) where {F, F2} =
_mapreduce_dim(f, op, init, A, dims)
mapreduce(f::F, op::F2, A::AbstractArrayOrBroadcasted...; kw...) where {F, F2} =
reduce(op, map(f, A...); kw...)

_mapreduce_dim(f::F, op::F2, nt, A::AbstractArrayOrBroadcasted, ::Colon) where {F, F2} =
mapfoldl_impl(f, op, nt, A)

_mapreduce_dim(f::F, op::F2, ::_InitialValue, A::AbstractArrayOrBroadcasted, ::Colon) where {F, F2} =
_mapreduce(f, op, IndexStyle(A), A)

_mapreduce_dim(f::F, op::F2, nt, A::AbstractArrayOrBroadcasted, dims) where {F, F2} =
mapreducedim!(f, op, reducedim_initarray(A, dims, nt), A)

_mapreduce_dim(f::F, op::F2, ::_InitialValue, A::AbstractArrayOrBroadcasted, dims) where {F,F2} =
mapreducedim!(f, op, reducedim_init(f, op, A, dims), A)

mapreduce_empty_iter(f::F, op::F2, itr, ItrEltype) where {F, F2} =
reduce_empty_iter(MappingRF(f, op), itr, ItrEltype)
mapreduce_first(f::F, op::F2, x) where {F,F2} = reduce_first(op, f(x))

_mapreduce(f::F, op::F2, A::AbstractArrayOrBroadcasted) where {F,F2} = _mapreduce(f, op, IndexStyle(A), A)
mapreduce_empty(::typeof(identity), op::F, T) where {F} = reduce_empty(op, T)
mapreduce_empty(::typeof(abs), op::F, T) where {F} = abs(reduce_empty(op, T))
mapreduce_empty(::typeof(abs2), op::F, T) where {F} = abs2(reduce_empty(op, T))
end
@eval Base.Unicode begin
function utf8proc_map(str::Union{String,SubString{String}}, options::Integer, chartransform::F = identity) where F
Expand Down
6 changes: 3 additions & 3 deletions test/reducedim.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,7 @@ end
end
end

@testset "reducedim_init min/max unorderable handling" begin
@testset "mapreducedim min/max unorderable handling" begin
x = Any[1.0, NaN]
y = [1, missing]
for (v, rval1, rval2) in [(x, [NaN], x),
Expand Down Expand Up @@ -609,7 +609,7 @@ end
try
#=@test_throws BoundsError=# mapreduce(identity, op, arr; dims)
catch ex
@test_broken ex isa BoundsError
@test ex isa BoundsError
end
@test_throws BoundsError mapreduce(identity, op, arr; dims, init=0)

Expand Down Expand Up @@ -938,7 +938,7 @@ end
a = collect(1.0:4)
f(x) = x < 3 ? missing : x
@test ismissing(minimum(f, a))
@test ismissing(minimum(f, a); dims = 1)
@test ismissing(only(minimum(f, a; dims = 1)))
end

@testset "zero indices are not special, issue #38660" begin
Expand Down