diff --git a/Project.toml b/Project.toml index 5cb347a..df7199d 100644 --- a/Project.toml +++ b/Project.toml @@ -32,7 +32,7 @@ LinearAlgebra = "1.6" Metal = "1.9" PrecompileTools = "1" Random = "1.6" -StridedViews = "0.5.1" +StridedViews = "0.5.5" Test = "1.6" TupleTools = "1.6" julia = "1.10" diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 824e5c5..fb7fb93 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -297,7 +297,7 @@ function _mapreduce_threaded!( return nothing end -function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int) +function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int; unit::Bool = false) @assert N > 0 "Cannot generate 0-dim kernel" # many variables @@ -388,28 +388,19 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int) i = 1 if op == Nothing # Non-reducing (map!/permute/...): each destination element is written once, - # so skip the reduction-only `iszero` hoist below and emit just the - # unit-stride and generic `@simd` loops. + # so skip the reduction-only `iszero` hoist below. Whether the innermost loop + # has unit stride is loop invariant, and is selected by the caller via `unit`. ex = Expr(:(=), lhsex, fcallex) if N >= 1 + step1ex, step2ex = unit ? (unitstep1ex, unitstep2ex) : (stepstride1ex[i], stepstride2ex[i]) ex = quote - if $unitstridecond - @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i])) - $ex - $unitstep1ex - $unitstep2ex - end - $(returnstride1ex[i]) - $(returnstride2ex[i]) - else - @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i])) - $ex - $(stepstride1ex[i]) - $(stepstride2ex[i]) - end - $(returnstride1ex[i]) - $(returnstride2ex[i]) + @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i])) + $ex + $step1ex + $step2ex end + $(returnstride1ex[i]) + $(returnstride2ex[i]) end end else @@ -497,10 +488,12 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int) if initop === Nothing for outer i in 1:N ex = quote - for $(blockloopvars[i]) in 1:blocks[$i]:dims[$i] + $(blockloopvars[i]) = 1 + while $(blockloopvars[i]) <= dims[$i] $(blockdimvars[i]) = min(blocks[$i], dims[$i] - $(blockloopvars[i]) + 1) $ex $(outerstepstrideex[i]) + $(blockloopvars[i]) += blocks[$i] end $(outerreturnstrideex[i]) end @@ -516,11 +509,13 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int) for outer i in 1:N ex = quote $(initvars[i]) = $(initvars[i + 1]) - for $(blockloopvars[i]) in 1:blocks[$i]:dims[$i] + $(blockloopvars[i]) = 1 + while $(blockloopvars[i]) <= dims[$i] $(blockdimvars[i]) = min(blocks[$i], dims[$i] - $(blockloopvars[i]) + 1) $ex $(initvars[i]) &= $(stridevars[i, 1]) > 0 $(outerstepstrideex[i]) + $(blockloopvars[i]) += blocks[$i] end $(outerreturnstrideex[i]) end @@ -537,6 +532,42 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int) return ex end +# Trailing singleton dimensions, e.g. those left by fusing dimensions, still cost a loop each. +# Forward to the kernel of lower dimensionality instead, with the views rewrapped such that +# this kernel is shared with arrays that have that dimensionality to begin with. +function _mapreduce_trim_expr(N::Int, ex) + N > 1 || return ex + for k in 1:(N - 1) + ex = quote + if nactive == $k + dk = _firstk(dims, Val($k)) + sk = _firstk(strides, Val($k)) + return _mapreduce_kernel!( + f, op, initop, dk, _firstk(blocks, Val($k)), _rewrap(arrays, dk, sk), sk, offsets + ) + end + $ex + end + end + return quote + nactive = _lastnonsingleton(dims) + $ex + end +end +_firstk(t::Tuple, ::Val{k}) where {k} = ntuple(i -> t[i], Val(k)) +_firstk(ts::Tuple{Vararg{Tuple}}, v::Val) = map(t -> _firstk(t, v), ts) +function _lastnonsingleton(dims::Dims) + n = length(dims) + while n > 1 && dims[n] == 1 + n -= 1 + end + return n +end +# the kernel only indexes the parents, so the strides of the views need no normalization +_rewrap(arrays, dims, strides) = map(arrays, strides) do a, s + return StridedViews.unsafe_stridedview(eltype(a), parent(a), dims, s, offset(a), a.op) +end + @generated function _mapreduce_kernel!( (f), (op), (initop), dims::NTuple{N, Int}, @@ -545,7 +576,17 @@ end strides::NTuple{M, NTuple{N, Int}}, offsets::NTuple{M, Int} ) where {N, M} - return _mapreduce_kernel_expr(f, op, initop, N, M) + op === Nothing || return _mapreduce_trim_expr(N, _mapreduce_kernel_expr(f, op, initop, N, M)) + # select the unit-stride variant of the innermost loop once, instead of in every iteration + unitcond = Expr(:call, :all, :(==(1)), Expr(:tuple, (:(strides[$j][1]) for j in 1:M)...)) + ex = quote + if $unitcond + $(_mapreduce_kernel_expr(f, op, initop, N, M; unit = true)) + else + $(_mapreduce_kernel_expr(f, op, initop, N, M; unit = false)) + end + end + return _mapreduce_trim_expr(N, ex) end function indexorder(strides::NTuple{N, Int}) where {N}