From 73e8da5cecf62d778553ace1303eeda9a6633700 Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Sun, 26 Jun 2022 21:18:54 +0800 Subject: [PATCH 1/2] Check the legality of zero-index based initialization. And add an optimized fallback for `keys::AbstractArray` (almost 2x faster with cheap `f`.) --- base/reducedim.jl | 84 +++++++++++++++++++++++++++++++++++++-------- test/offsetarray.jl | 8 +++++ test/reducedim.jl | 65 +++++++++++++++++++++-------------- 3 files changed, 117 insertions(+), 40 deletions(-) diff --git a/base/reducedim.jl b/base/reducedim.jl index dc34b4feb1f6a..749dfd4b54231 100644 --- a/base/reducedim.jl +++ b/base/reducedim.jl @@ -1025,21 +1025,25 @@ for (fname, op) in [(:sum, :add_sum), (:prod, :mul_prod), end ##### findmin & findmax ##### +# `iterate` based fallback for compatibility. # The initial values of Rval are not used if the corresponding indices in Rind are 0. -# -function findminmax!(f, op, Rval, Rind, A::AbstractArray{T,N}) where {T,N} +function _findminmax!(f, op, Rval, Rind, A::AbstractArray{T,N}, keys, init = true) where {T,N} (isempty(Rval) || isempty(A)) && return Rval, Rind lsiz = check_reducedims(Rval, A) for i = 1:N axes(Rval, i) == axes(Rind, i) || throw(DimensionMismatch("Find-reduction: outputs must have the same indices")) end + zi = zero(eltype(keys)) + zi in keys && throw(ArgumentError(LazyString("`keys` containing ", zi, " is not supported!"))) + if init + fill!(Rind, zi) + fill!(Rval, f(first(A))) + end # If we're reducing along dimension 1, for efficiency we can make use of a temporary. # Otherwise, keep the result in Rval/Rind so that we traverse A in storage order. indsAt, indsRt = safe_tail(axes(A)), safe_tail(axes(Rval)) keep, Idefault = Broadcast.shapeindexer(indsRt) - ks = keys(A) - y = iterate(ks) - zi = zero(eltype(ks)) + y = iterate(keys) if reducedim1(Rval, A) i1 = first(axes1(Rval)) @inbounds for IA in CartesianIndices(indsAt) @@ -1053,7 +1057,7 @@ function findminmax!(f, op, Rval, Rind, A::AbstractArray{T,N}) where {T,N} tmpRv = tmpAv tmpRi = k end - y = iterate(ks, kss) + y = iterate(keys, kss) end Rval[i1,IR] = tmpRv Rind[i1,IR] = tmpRi @@ -1070,7 +1074,57 @@ function findminmax!(f, op, Rval, Rind, A::AbstractArray{T,N}) where {T,N} Rval[i,IR] = tmpAv Rind[i,IR] = k end - y = iterate(ks, kss) + y = iterate(keys, kss) + end + end + end + Rval, Rind +end + +# Optimized fallback for `keys` with the same axes, e.g. `LinearIndices`, `CartesianIndices`. +# We initialize `Rval`/`Rind` via `map/copyfirst!` to support non-1 based `A`. +function _findminmax!(f, op, Rval, Rind, A::AbstractArray{T,N}, keys::AbstractArray{<:Any,N}, init = true) where {T,N} + axes(keys) == axes(A) || return @invoke _findminmax!(f, op, Rval, Rind, A::AbstractArray, keys::Any, init) + (isempty(Rval) || isempty(A)) && return Rval, Rind + lsiz = check_reducedims(Rval, A) + for i = 1:N + axes(Rval, i) == axes(Rind, i) || throw(DimensionMismatch("Find-reduction: outputs must have the same indices")) + end + if init + copyfirst!(Rind, keys) + mapfirst!(f, Rval, A) + end + # If we're reducing along dimension 1, for efficiency we can make use of a temporary. + # Otherwise, keep the result in Rval/Rind so that we traverse A in storage order. + indsAt, indsRt = safe_tail(axes(A)), safe_tail(axes(Rval)) + keep, Idefault = Broadcast.shapeindexer(indsRt) + if reducedim1(Rval, A) + i1 = first(axes1(Rval)) + @inbounds for IA in CartesianIndices(indsAt) + IR = Broadcast.newindex(IA, keep, Idefault) + tmpRv = Rval[i1,IR] + tmpRi = Rind[i1,IR] + for i in axes(A,1) + tmpAv = f(A[i,IA]) + if op(tmpRv, tmpAv) + tmpRv = tmpAv + tmpRi = keys[i,IA] + end + end + Rval[i1,IR] = tmpRv + Rind[i1,IR] = tmpRi + end + else + @inbounds for IA in CartesianIndices(indsAt) + IR = Broadcast.newindex(IA, keep, Idefault) + for i in axes(A, 1) + tmpAv = f(A[i,IA]) + tmpRv = Rval[i,IR] + tmpRi = Rind[i,IR] + if op(tmpRv, tmpAv) + Rval[i,IR] = tmpAv + Rind[i,IR] = keys[i,IA] + end end end end @@ -1086,7 +1140,7 @@ dimensions of `rval` and `rind`, and store the results in `rval` and `rind`. """ function findmin!(rval::AbstractArray, rind::AbstractArray, A::AbstractArray; init::Bool=true) - findminmax!(identity, isgreater, init && !isempty(A) ? fill!(rval, first(A)) : rval, fill!(rind,zero(eltype(keys(A)))), A) + _findminmax!(identity, isgreater, rval, rind, A, keys(A), init) end """ @@ -1142,9 +1196,9 @@ function _findmin(f, A, region) end similar(A, promote_op(f, eltype(A)), ri), zeros(eltype(keys(A)), ri) else - fA = f(first(A)) - findminmax!(f, isgreater, fill!(similar(A, _findminmax_inittype(f, A), ri), fA), - zeros(eltype(keys(A)), ri), A) + rval = similar(A, _findminmax_inittype(f, A), ri) + rind = similar(A, eltype(keys(A)), ri) + _findminmax!(f, isgreater, rval, rind, A, keys(A)) end end @@ -1157,7 +1211,7 @@ dimensions of `rval` and `rind`, and store the results in `rval` and `rind`. """ function findmax!(rval::AbstractArray, rind::AbstractArray, A::AbstractArray; init::Bool=true) - findminmax!(identity, isless, init && !isempty(A) ? fill!(rval, first(A)) : rval, fill!(rind,zero(eltype(keys(A)))), A) + _findminmax!(identity, isless, rval, rind, A, keys(A), init) end """ @@ -1213,9 +1267,9 @@ function _findmax(f, A, region) end similar(A, promote_op(f, eltype(A)), ri), zeros(eltype(keys(A)), ri) else - fA = f(first(A)) - findminmax!(f, isless, fill!(similar(A, _findminmax_inittype(f, A), ri), fA), - zeros(eltype(keys(A)), ri), A) + rval = similar(A, _findminmax_inittype(f, A), ri) + rind = similar(A, eltype(keys(A)), ri) + _findminmax!(f, isless, rval, rind, A, keys(A)) end end diff --git a/test/offsetarray.jl b/test/offsetarray.jl index bf5beab5c3437..25a5256258a2a 100644 --- a/test/offsetarray.jl +++ b/test/offsetarray.jl @@ -839,3 +839,11 @@ end # this is fixed in #40038, so the evaluation of its CartesianIndices should work @test CartesianIndices(A) == CartesianIndices(B) end + +# issue #38660 +@testset "`findmin/max` for OffsetArray" begin + ov = OffsetVector([-1, 1], 0:1) + @test @inferred(findmin(ov; dims = 1)) .|> first == (-1, 0) + ov = OffsetVector([-1, 1], -1:0) + @test @inferred(findmax(ov; dims = 1)) .|> first == (1, 0) +end diff --git a/test/reducedim.jl b/test/reducedim.jl index 5402376744e82..5cf368f657a02 100644 --- a/test/reducedim.jl +++ b/test/reducedim.jl @@ -273,33 +273,48 @@ end end end +struct WithIteratorsKeys{T,N} <: AbstractArray{T,N} + data::Array{T,N} +end +Base.size(a::WithIteratorsKeys) = size(a.data) +Base.getindex(a::WithIteratorsKeys, inds...) = a.data[inds...] +struct IteratorsKeys{T,A<:AbstractArray{T}} + iter::A + IteratorsKeys(iter) = new{eltype(iter),typeof(iter)}(iter) +end +Base.iterate(a::IteratorsKeys, state...) = Base.iterate(a.iter, state...) +Base.keys(a::WithIteratorsKeys) = IteratorsKeys(keys(a.data)) +Base.eltype(::IteratorsKeys{T}) where {T} = T + # findmin/findmax function arguments: output type inference @testset "findmin/findmax output type inference" begin - A = ["1" "22"; "333" "4444"] - for (tup, rval, rind) in [((1,), [1 2], [CartesianIndex(1, 1) CartesianIndex(1, 2)]), - ((2,), reshape([1, 3], 2, 1), reshape([CartesianIndex(1, 1), CartesianIndex(2, 1)], 2, 1)), - ((1,2), fill(1,1,1), fill(CartesianIndex(1,1),1,1))] - rval′, rind′ = findmin(length, A, dims=tup) - @test (rval, rind) == (rval′, rind′) - @test typeof(rval′) == Matrix{Int} - end - for (tup, rval, rind) in [((1,), [3 4], [CartesianIndex(2, 1) CartesianIndex(2, 2)]), - ((2,), reshape([2, 4], 2, 1), reshape([CartesianIndex(1, 2), CartesianIndex(2, 2)], 2, 1)), - ((1,2), fill(4,1,1), fill(CartesianIndex(2,2),1,1))] - rval′, rind′ = findmax(length, A, dims=tup) - @test (rval, rind) == (rval′, rind′) - @test typeof(rval) == Matrix{Int} - end - B = [1.5 1.0; 5.5 6.0] - for (tup, rval, rind) in [((1,), [3//2 1//1], [CartesianIndex(1, 1) CartesianIndex(1, 2)]), - ((2,), reshape([1//1, 11//2], 2, 1), reshape([CartesianIndex(1, 2), CartesianIndex(2, 1)], 2, 1)), - ((1,2), fill(1//1,1,1), fill(CartesianIndex(1,2),1,1))] - rval′, rind′ = findmin(Rational, B, dims=tup) - @test (rval, rind) == (rval′, rind′) - @test typeof(rval) == Matrix{Rational{Int}} - rval′, rind′ = findmin(Rational ∘ abs ∘ complex, B, dims=tup) - @test (rval, rind) == (rval′, rind′) - @test typeof(rval) == Matrix{Rational{Int}} + for wrapper in (identity, WithIteratorsKeys) + A = wrapper(["1" "22"; "333" "4444"]) + for (tup, rval, rind) in [((1,), [1 2], [CartesianIndex(1, 1) CartesianIndex(1, 2)]), + ((2,), reshape([1, 3], 2, 1), reshape([CartesianIndex(1, 1), CartesianIndex(2, 1)], 2, 1)), + ((1,2), fill(1,1,1), fill(CartesianIndex(1,1),1,1))] + rval′, rind′ = @inferred findmin(length, A, dims=tup) + @test (rval, rind) == (rval′, rind′) + @test typeof(rval′) == Matrix{Int} + end + for (tup, rval, rind) in [((1,), [3 4], [CartesianIndex(2, 1) CartesianIndex(2, 2)]), + ((2,), reshape([2, 4], 2, 1), reshape([CartesianIndex(1, 2), CartesianIndex(2, 2)], 2, 1)), + ((1,2), fill(4,1,1), fill(CartesianIndex(2,2),1,1))] + rval′, rind′ = @inferred findmax(length, A, dims=tup) + @test (rval, rind) == (rval′, rind′) + @test typeof(rval) == Matrix{Int} + end + B = wrapper([1.5 1.0; 5.5 6.0]) + for (tup, rval, rind) in [((1,), [3//2 1//1], [CartesianIndex(1, 1) CartesianIndex(1, 2)]), + ((2,), reshape([1//1, 11//2], 2, 1), reshape([CartesianIndex(1, 2), CartesianIndex(2, 1)], 2, 1)), + ((1,2), fill(1//1,1,1), fill(CartesianIndex(1,2),1,1))] + rval′, rind′ = @inferred findmin(Rational, B, dims=tup) + @test (rval, rind) == (rval′, rind′) + @test typeof(rval) == Matrix{Rational{Int}} + rval′, rind′ = @inferred findmin(Rational ∘ abs ∘ complex, B, dims=tup) + @test (rval, rind) == (rval′, rind′) + @test typeof(rval) == Matrix{Rational{Int}} + end end end From 64fcd547350ee6f6101488f83c69d98d2e10654d Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Tue, 12 Jul 2022 11:06:19 +0800 Subject: [PATCH 2/2] Adopt suggestions. 1. remove `check_reducedims` within `mapfirst!`. `init_array!` for other reductions also skip this check. it should be safe to remove it as we add no `@inbounds` 2. remove unneeded `getindex`. Co-Authored-By: Andrew Radcliffe <96091198+andrewjradcliffe@users.noreply.github.com> --- base/reducedim.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/base/reducedim.jl b/base/reducedim.jl index 749dfd4b54231..a7f19d75ffe9a 100644 --- a/base/reducedim.jl +++ b/base/reducedim.jl @@ -269,7 +269,6 @@ Extract first entry of slices of array A into existing array R. copyfirst!(R::AbstractArray, A::AbstractArray) = mapfirst!(identity, R, A) function mapfirst!(f::F, R::AbstractArray, A::AbstractArray{<:Any,N}) where {N, F} - lsiz = check_reducedims(R, A) t = _firstreducedslice(axes(R), axes(A)) map!(f, R, view(A, t...)) end @@ -1120,7 +1119,6 @@ function _findminmax!(f, op, Rval, Rind, A::AbstractArray{T,N}, keys::AbstractAr for i in axes(A, 1) tmpAv = f(A[i,IA]) tmpRv = Rval[i,IR] - tmpRi = Rind[i,IR] if op(tmpRv, tmpAv) Rval[i,IR] = tmpAv Rind[i,IR] = keys[i,IA]