-
-
Notifications
You must be signed in to change notification settings - Fork 5.6k
WIP: The great pairwise reduction refactor #58418
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.
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
base: master
Are you sure you want to change the base?
The head ref may contain hidden characters: "mb/ReduceReuse\u267B"
Conversation
…]Int`. (cherry picked from commit c923b1d)
(cherry picked from commit 4bee191)
for accumulate inference (cherry picked from commit 4fadd8f)
Previously, Julia tried to guess at initial values for empty dimensional reductions in a slightly different way than whole-array reductions. This change unifies those behaviors, such that `mapreduce_empty` is called for empty dimensional reductions, just like it is called for empty whole-array reductions. Beyond the unification (which is valuable in its own right), this change enables some downstream simplifications to dimensional reductions and is likely to be the "most breaking" public behavior in a refactoring targeting more correctness improvements. (cherry picked from commit 94f24b8)
(cherry picked from commit 0e924cf)
they already supported axes through the default fallback that uses size to construct OneTos
and this was causing ambiguities
(cherry picked from commit 32fc16b)
Fantastic stuff, this is really exciting to see! 🥳 |
length(t::ImmutableDict) = count(Returns(true), t) | ||
# length is defined in terms of iteration; using a higher-order function to do the iteration | ||
# is likely to call `length` again, so this is a manual for loop: | ||
function length(t::ImmutableDict) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
won't this be horribly slow?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I mean, wasn't it already doing this? This shouldn't be any worse. It's definitely faster than infinite recursion :)
The trouble is that ImmutableDict
has an IteratorSize
of HasLength
. And so the new reduction machinery happily calls length
to determine the pairwise splits. So if length
itself uses reductions (like count) to calculate length, we're in trouble.
On nightly:
julia> using BenchmarkTools
julia> d = Base.ImmutableDict((i => i+1 for i in 1:200)...);
julia> @benchmark length($d)
BenchmarkTools.Trial: 10000 samples with 641 evaluations per sample.
Range (min … max): 193.512 ns … 392.225 ns ┊ GC (min … max): 0.00% … 0.00%
Time (median): 193.772 ns ┊ GC (median): 0.00%
Time (mean ± σ): 196.869 ns ± 11.655 ns ┊ GC (mean ± σ): 0.00% ± 0.00%
█▁ ▁▃▂▄▂ ▁
███████████▇██▇▇▆▇▆▆▆▅▅▇▆▆▅▆▅▅▆▅▅▅▅▅▅▅▄▄▅▅▅▅▂▄▃▃▄▄▄▄▃▂▂▂▃▄▂▅▄ █
194 ns Histogram: log(frequency) by time 241 ns <
Memory estimate: 0 bytes, allocs estimate: 0.
julia> @eval Base function length(t::ImmutableDict)
r = 0
for _ in t
r+=1
end
return r
end
length (generic function with 95 methods)
julia> @benchmark length($d)
BenchmarkTools.Trial: 10000 samples with 631 evaluations per sample.
Range (min … max): 193.542 ns … 366.482 ns ┊ GC (min … max): 0.00% … 0.00%
Time (median): 193.740 ns ┊ GC (median): 0.00%
Time (mean ± σ): 196.829 ns ± 11.533 ns ┊ GC (mean ± σ): 0.00% ± 0.00%
█ ▁ ▄ ▂ ▁
███████▇███▇██▇▇▇▇▆▇▆▆▅▆▆▆▇▆▆▆▅▅▅▅▅▅▄▅▆▅▄▅▅▅▅▄▄▅▄▄▄▅▅▄▄▃▄▄▃▃▄ █
194 ns Histogram: log(frequency) by time 238 ns <
Memory estimate: 0 bytes, allocs estimate: 0.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ah that's unfortunate...
Worth seeing if it also closes #43501 |
StaticArrays explicitly requests the fastmath version of fma here: https://github.com/JuliaArrays/StaticArrays.jl/blob/563adebc264a23d38da97bf5d9e5d3c6333bdf26/src/matrix_multiply_add.jl#L428 |
I have benchmarks! They're... ok. A bit of a mixed bag. The codeusing Chairmarks: Chairmarks, @be
using CSV: CSV
using ProgressMeter: Progress, next!
struct LengthUnknown{T}
itr::T
end
@inline Base.iterate(x::LengthUnknown) = iterate(x.itr)
@inline Base.iterate(x::LengthUnknown, s) = iterate(x.itr, s)
Base.IteratorSize(::LengthUnknown) = Base.SizeUnknown()
Base.IteratorEltype(::LengthUnknown) = Base.HasEltype()
Base.eltype(u::LengthUnknown) = Base.eltype(u.itr)
using Random: Random
const rng = Random.MersenneTwister(1155)
rand(sz...) = Random.rand(rng, sz...)
rand(::Type{T}, sz...) where {T} = Random.rand(rng, T, sz...)
haslength(x) = Iterators.take(x,length(x))
function cartesianvec(x)
n = length(x)
n1 = nextpow(2, sqrt(n))
return @view reshape(x, n1, :)[begin:end, begin:end]
end
function bench_wholearray(Ns, Ts)
results = (; style=String[], eltype=String[], init=Bool[], dims=Union{Missing,String}[], n=Int[], time=Float64[])
P = Progress(4*4*2*length(Ts))
for (xf,type) in [(identity, "IndexLinear"), (cartesianvec, "IndexCartesian"), (haslength, "HasLength"), (LengthUnknown, "SizeUnknown")],
(xf,type) in [(xf, type),
(x->skipmissing(xf(Array{Union{Missing, eltype(x)}}(x))), "SkipMissing $type 0%"),
(x->skipmissing(xf(Array{Union{Missing, eltype(x)}}([rand() < .5 ? missing : x for x in x]))), "SkipMissing $type 50%"),
(x->skipmissing(xf(Array{Union{Missing, eltype(x)}}(fill(missing, size(x))))), "SkipMissing $type 100%"),],
T in Ts,
init in (nothing, zero(T))
ts = [begin
A = xf(rand(T, N÷sizeof(T)))
if isempty(A) && isnothing(init) && try mapreduce(identity, +, A); false; catch; true; end
# Empty SkipMissings don't define reduce_empty; skip it.
[(;time = NaN,)]
elseif isnothing(init)
@be(mapreduce(identity, +, $A))
else
@be(mapreduce(identity, +, $A; init=$init))
end
end for N in Ns]
append!(results.time, getproperty.(minimum.(ts), :time))
append!(results.n, Ns)
append!(results.dims, fill(missing, length(ts)))
append!(results.eltype, fill(string(T), length(ts)))
append!(results.style, fill(type, length(ts)))
append!(results.init, fill(!isnothing(init), length(ts)))
next!(P)
end
results
end
function bench_dims(Ns, Ts)
results = (; style=String[], eltype=String[], init=Bool[], dims=Union{Missing,String}[], n=Int[], time=Float64[])
P = Progress(2*8*2*length(Ts))
for (xf,type) in [(identity, "IndexLinear"), (x->@view(x[1:end,1:end,1:end,1:end]), "IndexCartesian")],
(szf, dims) in [((n,_,_)->(n÷4,4), 1), ((n,_,_)->(4,n÷4), 2),
((n,n2,_)->(n2,n÷n2÷4,4), (1,2)),
((n,n2,_)->(n2,4,n÷n2÷4), (1,3)),
((n,n2,_)->(4,n2,n÷n2÷4), (2,3)),
((n,_,n3)->(n3,n3,n÷n3÷n3÷4,4), (1,2,3)),
((n,_,n3)->(n3,n3,4,n÷n3÷n3÷4), (1,2,4)),
((n,_,n3)->(n3,4,n3,n÷n3÷n3÷4), (1,3,4)),
((n,_,n3)->(4,n3,n3,n÷n3÷n3÷4), (2,3,4))],
T in Ts,
init in (nothing, zero(T))
ts = [begin
n = N÷sizeof(T)
A = xf(rand(T, szf(n, nextpow(2, sqrt(n)), nextpow(2, cbrt(n)))))
if isnothing(init)
@be(mapreduce(identity, +, $A; dims=$dims))
else
@be(mapreduce(identity, +, $A; dims=$dims, init=$init))
end
end for N in Ns]
append!(results.time, getproperty.(minimum.(ts), :time))
append!(results.n, Ns)
append!(results.eltype, fill(string(T), length(ts)))
append!(results.style, fill(type, length(ts)))
append!(results.init, fill(!isnothing(init), length(ts)))
append!(results.dims, fill(string(dims), length(ts)))
next!(P)
end
return results
end
whole = bench_wholearray(2 .^ (3:22), (Int32,Int64,Float32,Float64))
dims = bench_dims(2 .^ (5:22), (Int32,Int64,Float32,Float64))
foreach(Base.splat(append!), zip(whole, dims))
CSV.write(Base.GIT_VERSION_INFO.commit_short * ".csv", whole) This gives us 744e3fd976.csv vs c5df01807e3.csv (nightly from a few days ago). So we can break this down over a wide variety of cases. ![]() ![]() Most of the slow cases here are the "small n": ![]() ![]() The variance is pretty high on SkipMissing — Julia's inference really doesn't like iteration's tuple of |
OK, I've implemented specialized kernels for Here are the latest benchmarks: 144c253358.csv. The remaining big regressions are in mid-size dimensional reductions, and in particular |
and slightly improve perf for abstract iterables
they ended up losing type stability and it isn't hard to just define our own pairwise impl
they are just different enough to break some assumptions about this function
OK, I must confess I've never really processed how SIMD doesn't really re-associate things. It fundamentally changes the entire ordering. It goes far above and beyond what we promise It's not uniformly a win, however. Hard-coding such a re-ordering means that we no longer let
A bunch of the 2x regressions are the 32-bit ones, because we're no longer saturating my M1's simd width. And those that are much worse than that are mostly small n and/or non-concrete. I think I can add some threshold tuning to improve some of that. |
some_exception(op) = try return (Some(op()), nothing); catch ex; return (nothing, ex); end | ||
reduced_shape(sz, dims) = ntuple(d -> d in dims ? 1 : sz[d], length(sz)) | ||
|
||
@testset "Ensure that calling, e.g., sum(empty; dims) has the same behavior as sum(empty)" begin |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
fyi this testset is repeated from L579
@nexprs 8 N->begin | ||
it = iterate(itr, s) | ||
if it === nothing | ||
N > 7 && (v_7 = op(v_7, f(a_7))) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
N > 7 && (v_7 = op(v_7, f(a_7))) | |
@nexprs 7 M->((N > M) && (v_{8-M} = op(v_{8-M}, f(a_{8-M})))) |
|
||
# This special internal method must have at least 8 indices and allows passing | ||
# optional scalar leading and trailing dimensions | ||
function _mapreduce_kernel_commutative(f, op, A, init, inds, leading=(), trailing=()) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
potentially profitable to unroll it out to @nexprs 16
when on a 64bit platform but sizeof(T) <= 4
|
||
@deprecate (reducedim_initarray(A::Union{Base.AbstractBroadcasted, AbstractArray}, region, init, ::Type{T}) where {T}) fill!(mapreduce_similar(A,T,reduced_indices(A,region)), init) false | ||
@deprecate reducedim_initarray(A::Union{Base.AbstractBroadcasted, AbstractArray}, region, init) fill!(mapreduce_similar(A,typeof(init),reduced_indices(A,region)), init) false | ||
const _dep_message_reducedim_init = ", these internals have been removed. To customize the array returned by dimensional reductions, implement mapreduce_similar instead" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
const _dep_message_reducedim_init = ", these internals have been removed. To customize the array returned by dimensional reductions, implement mapreduce_similar instead" |
dup. on L567
mapreduce(f::F, op::G, A::AbstractArrayOrBroadcasted; init=_InitialValue(), dims=(:)) where {F,G} = mapreducedim(f, op, A, init, dims) | ||
mapreduce(f::F, op::G, A::AbstractArrayOrBroadcasted, As::AbstractArrayOrBroadcasted...; init=_InitialValue(), dims=(:)) where {F,G} = | ||
reduce(op, map(f, A, As...); init, dims) | ||
mapreducedim(f::F, op::G, A, init, ::Colon) where {F,G} = mapreduce_pairwise(f, op, A, init) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I wonder (suspect?) if smaller pairwise_blocksize
is better when the cost of f
exceeds the cost of op
? e.g. to start with:
pairwise_blocksize(::typeof(sqrt), op) = 64
pairwise_blocksize(::typeof(sin), op) = 64
pairwise_blocksize(::typeof(cos), op) = 64
pairwise_blocksize(::typeof(exp), op) = 64
pairwise_blocksize(::typeof(log), op) = 64
|
||
# Reduction kernels that explicitly encode simplistic SIMD-ish reorderings | ||
const CommutativeOps = Union{typeof(+),typeof(Base.add_sum),typeof(min),typeof(max),typeof(Base._extrema_rf),typeof(|),typeof(&)} | ||
is_commutative_op(::CommutativeOps) = true |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
maybe two arg is better with a default?
is_commutative_op(op, ::Type) = is_commutative_op(op)
so we can say like *
on Float64
commutes (for the purposes of reduce
) but not on Matrix
or quaternions or whatever
This is the grand culmination of (and core motivation for) all the reductions work I've slowly been chipping away at over the past year. It implements the following major changes:
reduce
-based reductions use a pairwise re-association by default. This includes seven different implementations:foldl
s without SIMD)foldl
s without SIMD)foldl
s without SIMD)mapreduce_kernel
function (with specific signatures for arrays and iterators) after performing the pairwise splits and/or dimensional selection. This could eventually become a public API that should help alleviate some of the disruption to those who had been extending the internals before. See, for example, how LinearAlgebra replaced its internal and tricky_sum
overloads with this method: JuliaLang/LinearAlgebra.jl@e70953e.init
to start every reduction chain as we decided the documentation should direct implementations to do in Document that reduce-like functions useinit
one or more times #53945. They also consistently usemapreduce_first
.This currently stands atop the yet-to-be-merged #58374 and #55628.
As the successor to #55318, this carries all its tests and addresses all those issues in the dimensional reductions. Closes: #45566, closes #47231, closes #55213, closes #31427, closes #41054, closes #54920, closes #42259, closes #38660, closes #21097, closes #32366, closes #54875, closes #43731, closes #45562.
In addition, this further brings pairwise reassociations (and SIMD!) to the other class of reduction. Closes #29883, closes #52365, closes #30421, closes #52457, closes #39385.
This requires some changes to the stdlibs: It requires JuliaStats/Statistics.jl#172 (which fixes another four issues: JuliaStats/Statistics.jl#7, JuliaStats/Statistics.jl#126, JuliaStats/Statistics.jl#160, JuliaStats/Statistics.jl#183). This also requires changes to SparseArrays and LinearAlgebra.
There are other pull requests that this supersedes that were inspirational for my work here. In particular, this supersedes and would close #45651 (pairwise whole-array IndexCartesian), closes #52397 (pairwise whole-array Iterators.SizeUnknown). It also supersedes a few other PRs; would close #45822 (findmin/max for offset arrays; I've added those tests here), close #58241 (my PR for whole-array reductions
init
refactor that snowballed into this one).Yet to do:
sum
function for custom arrays withIndexCartesian
#33928, Reduction on array view is slow #28848,sum(f, itr; init)
slower thansum(f, itr)
#47216, Reducing over bool array fails to optimize somehow #43501. Unfortunately it looks like this makes performance regression in 1.12.0-beta2 when summing over iterator #58370 wayyyy worse.--depwarn=error
— all it needed was amapreduce_similar
definition to satisfy some type checks!~~ But because the internals changed, the sparse optimizations are no longer taking effect — they're just adding methods to vestigial deprecated methods that aren't called.mapreduce_kernel
on the entire dimension/array. This would be helpful for integers and other fully associative operations