Faster mapreduce for Broadcasted #31020
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -12,6 +12,9 @@ else | |
| const SmallUnsigned = Union{UInt8,UInt16,UInt32} | ||
| end | ||
|
|
||
| abstract type AbstractBroadcasted end | ||
| const AbstractArrayOrBroadcasted = Union{AbstractArray, AbstractBroadcasted} | ||
|
There was a problem hiding this comment. I needed to define this abstract type AbstractIndexable end
const Indexable = Union{AbstractArray, AbstractIndexable}or even abstract type Indexable{N} end
abstract type AbstractArray{T,N} <: Indexable{N} endso that it's much easier for downstream projects to use indexing-based There was a problem hiding this comment. Having an abstract type or a trait like Though a trait would probably be better than an abstract type, since it would be more flexible. In particular, an object returned by There was a problem hiding this comment. If we go with trait direction, one option is to use struct NonIndexable <: IndexStyle end
IndexStyle(::Type) = NonIndexable()There was a problem hiding this comment. That's a clever idea. I think the union here is sensible for now — that trait could possibly come later. There was a problem hiding this comment. Differing the whole designing work for the trait makes sense. There was a problem hiding this comment. @vtjnash mentioned there may be another way around this bootstrapping issue. There was a problem hiding this comment. Given that the 1.4 feature freeze is approaching, I support merging the PR with the FWIW, this PR is giving us incredibly good performance to be able to reuse the pairwise summation algorithm with weights at JuliaStats/StatsBase.jl#518. Great work! There was a problem hiding this comment. Thanks a lot for testing this PR in a real-world scenario! |
||
|
|
||
| """ | ||
| Base.add_sum(x, y) | ||
|
|
||
|
|
@@ -227,7 +230,8 @@ foldr(op, itr; kw...) = mapfoldr(identity, op, itr; kw...) | |
|
|
||
| # This is a generic implementation of `mapreduce_impl()`, | ||
| # certain `op` (e.g. `min` and `max`) may have their own specialized versions. | ||
| @noinline function mapreduce_impl(f, op, A::AbstractArrayOrBroadcasted, | ||
| ifirst::Integer, ilast::Integer, blksize::Int) | ||
| if ifirst == ilast | ||
| @inbounds a1 = A[ifirst] | ||
| return mapreduce_first(f, op, a1) | ||
|
|
@@ -250,7 +254,7 @@ foldr(op, itr; kw...) = mapfoldr(identity, op, itr; kw...) | |
| end | ||
| end | ||
|
|
||
| mapreduce_impl(f, op, A::AbstractArrayOrBroadcasted, ifirst::Integer, ilast::Integer) = | ||
| mapreduce_impl(f, op, A, ifirst, ilast, pairwise_blocksize(f, op)) | ||
|
|
||
| """ | ||
|
|
@@ -383,13 +387,13 @@ The default is `reduce_first(op, f(x))`. | |
| """ | ||
| mapreduce_first(f, op, x) = reduce_first(op, f(x)) | ||
|
|
||
| _mapreduce(f, op, A::AbstractArrayOrBroadcasted) = _mapreduce(f, op, IndexStyle(A), A) | ||
|
|
||
| function _mapreduce(f, op, ::IndexLinear, A::AbstractArrayOrBroadcasted) | ||
| inds = LinearIndices(A) | ||
| n = length(inds) | ||
| if n == 0 | ||
| return mapreduce_empty_iter(f, op, A, IteratorEltype(A)) | ||
| elseif n == 1 | ||
| @inbounds a1 = A[first(inds)] | ||
| return mapreduce_first(f, op, a1) | ||
|
|
@@ -410,7 +414,7 @@ end | |
|
|
||
| mapreduce(f, op, a::Number) = mapreduce_first(f, op, a) | ||
|
|
||
| _mapreduce(f, op, ::IndexCartesian, A::AbstractArrayOrBroadcasted) = mapfoldl(f, op, A) | ||
|
|
||
| """ | ||
| reduce(op, itr; [init]) | ||
|
|
@@ -560,7 +564,7 @@ isgoodzero(::typeof(max), x) = isbadzero(min, x) | |
| isgoodzero(::typeof(min), x) = isbadzero(max, x) | ||
|
|
||
| function mapreduce_impl(f, op::Union{typeof(max), typeof(min)}, | ||
| A::AbstractArrayOrBroadcasted, first::Int, last::Int) | ||
| a1 = @inbounds A[first] | ||
| v1 = mapreduce_first(f, op, a1) | ||
| v2 = v3 = v4 = v1 | ||
|
|
@@ -856,7 +860,7 @@ function count(pred, itr) | |
| end | ||
| return n | ||
| end | ||
| function count(pred, a::AbstractArrayOrBroadcasted) | ||
| n = 0 | ||
| for i in eachindex(a) | ||
| @inbounds n += pred(a[i])::Bool | ||
|
|
||
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.
It turned out my previous approach (combine
IndexStyleof each argument) was wrong since it didn't work with, e.g.,broadcasted(+, zeros(5), zeros(1, 4)). Each argument can beIndexLineareven though the broadcasted indexing is not (that's the main job of broadcasted).So now the new approach is to return an equivalent of
IndexStyle(bc.axes[1])iflength(bc.axes) == 1and then returnIndexCartesian()otherwise. As I needAxestype parameter now, this implementation requires theBroadcastedto beinstantiated first. I guess this is OK since that's the usual broadcasting pipeline.Is this implementation fine? Is
IndexStyle(bc.axes[1]) == IndexLinear() && length(bc.axes) == 1the only possibleIndexLinear()case that is detectable in a type-stable manner? That is to say, I'm ignoring the case likebroadcasted(+, zeros(1, 4), zeros(1, 4))as I need to look at the value to check that it supports linear indexing.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.
This doesn't really require the
Broadcastedto beinstantiated to get anIndexStyle— it just means that if it's not instantiated that it falls back to a cartesian implementation. AllBroadcasteds should support cartesian indexing; only one-dimensional ones will allow straight integers.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.
Yes, that's much more accurate way to put it. I wanted to say that "for the broadcasted to be dispatched to the optimized method", it has to be
instantiated first.