diff --git a/src/broadcast.jl b/src/broadcast.jl index bdc6bd6..083ba9f 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -52,7 +52,7 @@ Base.dotview(a::StridedView{<:Any, N}, I::Vararg{SliceIndex, N}) where {N} = get # promote StridedView to have same size, by giving artificial zero strides stridedargs = promoteshape(dims, capturestridedargs(bc)...) c = make_capture(bc) - _mapreduce_fuse!(c, nothing, nothing, dims, (dest, stridedargs...)) + _mapreduce_order!(c, nothing, nothing, dims, (dest, stridedargs...)) return dest end diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 7a48290..150be8b 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -54,7 +54,7 @@ function Base.map!( any(isequal(0), dims) && return b # don't do anything - _mapreduce_fuse!(f, nothing, nothing, dims, (b, a1, A...)) + _mapreduce_order!(f, nothing, nothing, dims, (b, a1, A...)) return b end @@ -101,18 +101,35 @@ function _mapreducedim!( map!(initop, arrays[1], arrays[1]) end else - _mapreduce_fuse!(f, op, initop, dims, promoteshape(dims, arrays...)) + _mapreduce_order!(f, op, initop, dims, promoteshape(dims, arrays...)) end return arrays[1] end -function _mapreduce_fuse!( - @nospecialize(f), @nospecialize(op), @nospecialize(initop), - dims::Dims, arrays::Tuple{Vararg{StridedView}} - ) - # Fuse dimensions if possible: assume that at least one array, e.g. the output array in - # arrays[1], has its strides sorted - allstrides = map(strides, arrays) +# Cache-friendly reordering of the loop dimensions: put the most important +# (smallest-stride, most contiguous) dimension first and size-1 dimensions last. +# The importance of each dimension is modelled from the `indexorder` of every +# array's strides, with the output array (`strides[1]`) weighted by a factor 2. +# `strides` is the per-array tuple-of-tuples; returns the reordered `(dims, strides)`. +function _sortdims(dims, strides) + M = length(strides) + N = length(dims) + # ceil(Int, log2(M+2)) # to account for the fact that there are M arrays, where the first one is counted with a factor 2 + g = 8 * sizeof(Int) - leading_zeros(M + 1) + importance = 2 .* (1 .<< (g .* (N .- indexorder(strides[1])))) + for k in 2:M + importance = importance .+ (1 .<< (g .* (N .- indexorder(strides[k])))) + end + + importance = importance .* (dims .> 1) # put dims 1 at the back + p = TupleTools.sortperm(importance; rev = true) + return TupleTools.getindices(dims, p), broadcast(TupleTools.getindices, strides, (p,)) +end + +# Fuse dimensions if possible: merge dimension `i` into `i-1` when the two are +# contiguous in *every* array (`s[i] == dims[i-1] * s[i-1]`). +# Merged dimensions fold their extent into `i-1` and become size 1. +function _fusedims(dims, allstrides) @inbounds for i in length(dims):-1:2 merge = true for s in allstrides @@ -126,31 +143,29 @@ function _mapreduce_fuse!( dims = TupleTools.setindex(dims, 1, i) end end - return _mapreduce_order!(f, op, initop, dims, allstrides, arrays) + return dims end +function order_and_fuse_dims(dims, allstrides) + dims, allstrides = _sortdims(dims, allstrides) + dims = _fusedims(dims, allstrides) + return _sortdims(dims, allstrides) +end + +# Per-dimension cost used by the blocking and thread-splitting heuristics, +# derived from the minimum (doubled) stride across arrays. +_computecosts(strides) = map(a -> ifelse(iszero(a), 1, a << 1), map(min, strides...)) + +# Pipeline entry point: order → fuse → order → block → kernel. +# Fusing needs ordered entries, and final order pass brings remaining dim 1 to end function _mapreduce_order!( @nospecialize(f), @nospecialize(op), @nospecialize(initop), - dims, strides, arrays + dims::Dims, arrays::Tuple{Vararg{StridedView}} ) - M = length(arrays) - N = length(dims) - # sort order of loops/dimensions by modelling the importance of each dimension - g = 8 * sizeof(Int) - leading_zeros(M + 1) # ceil(Int, log2(M+2)) # to account for the fact - # that there are M arrays, where the first one is counted with a factor 2 - importance = 2 .* (1 .<< (g .* (N .- indexorder(strides[1])))) # first array is output - # and is more important by a factor 2 - for k in 2:M - importance = importance .+ (1 .<< (g .* (N .- indexorder(strides[k])))) - end - - importance = importance .* (dims .> 1) # put dims 1 at the back - p = TupleTools.sortperm(importance; rev = true) - dims = TupleTools.getindices(dims, p) - strides = broadcast(TupleTools.getindices, strides, (p,)) + dims, allstrides = order_and_fuse_dims(dims, map(strides, arrays)) offsets = map(offset, arrays) - costs = map(a -> ifelse(iszero(a), 1, a << 1), map(min, strides...)) - return _mapreduce_block!(f, op, initop, dims, strides, offsets, costs, arrays) + costs = _computecosts(allstrides) + return _mapreduce_block!(f, op, initop, dims, allstrides, offsets, costs, arrays) end const MINTHREADLENGTH = 1 << 15 # minimal length before any kind of threading is applied