Skip to content

Commit

Permalink
mapreduce: don't inbounds unknown functions (#55329)
Browse files Browse the repository at this point in the history
More finely scope the `@inbounds` annotations to ensure neither `f` nor
`op` are erroneously `@inbounds`ed.

(cherry picked from commit 1dffd77)
  • Loading branch information
mbauman authored and KristofferC committed Sep 12, 2024
1 parent 6f3fdf7 commit 21ccfc0
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 23 deletions.
10 changes: 5 additions & 5 deletions base/reduce.jl
Original file line number Diff line number Diff line change
Expand Up @@ -649,11 +649,11 @@ function mapreduce_impl(f, op::Union{typeof(max), typeof(min)},
start = first + 1
simdstop = start + chunk_len - 4
while simdstop <= last - 3
@inbounds for i in start:4:simdstop
v1 = _fast(op, v1, f(A[i+0]))
v2 = _fast(op, v2, f(A[i+1]))
v3 = _fast(op, v3, f(A[i+2]))
v4 = _fast(op, v4, f(A[i+3]))
for i in start:4:simdstop
v1 = _fast(op, v1, f(@inbounds(A[i+0])))
v2 = _fast(op, v2, f(@inbounds(A[i+1])))
v3 = _fast(op, v3, f(@inbounds(A[i+2])))
v4 = _fast(op, v4, f(@inbounds(A[i+3])))
end
checkbounds(A, simdstop+3)
start += chunk_len
Expand Down
37 changes: 19 additions & 18 deletions base/reducedim.jl
Original file line number Diff line number Diff line change
Expand Up @@ -302,19 +302,20 @@ function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArrayOrBroadcasted)
if reducedim1(R, A)
# keep the accumulator as a local variable when reducing along the first dimension
i1 = first(axes1(R))
@inbounds for IA in CartesianIndices(indsAt)
for IA in CartesianIndices(indsAt)
IR = Broadcast.newindex(IA, keep, Idefault)
r = R[i1,IR]
@inbounds r = R[i1,IR]
@simd for i in axes(A, 1)
r = op(r, f(A[i, IA]))
r = op(r, f(@inbounds(A[i, IA])))
end
R[i1,IR] = r
@inbounds R[i1,IR] = r
end
else
@inbounds for IA in CartesianIndices(indsAt)
for IA in CartesianIndices(indsAt)
IR = Broadcast.newindex(IA, keep, Idefault)
@simd for i in axes(A, 1)
R[i,IR] = op(R[i,IR], f(A[i,IA]))
v = op(@inbounds(R[i,IR]), f(@inbounds(A[i,IA])))
@inbounds R[i,IR] = v
end
end
end
Expand Down Expand Up @@ -1058,33 +1059,33 @@ function findminmax!(f, op, Rval, Rind, A::AbstractArray{T,N}) where {T,N}
zi = zero(eltype(ks))
if reducedim1(Rval, A)
i1 = first(axes1(Rval))
@inbounds for IA in CartesianIndices(indsAt)
for IA in CartesianIndices(indsAt)
IR = Broadcast.newindex(IA, keep, Idefault)
tmpRv = Rval[i1,IR]
tmpRi = Rind[i1,IR]
@inbounds tmpRv = Rval[i1,IR]
@inbounds tmpRi = Rind[i1,IR]
for i in axes(A,1)
k, kss = y::Tuple
tmpAv = f(A[i,IA])
tmpAv = f(@inbounds(A[i,IA]))
if tmpRi == zi || op(tmpRv, tmpAv)
tmpRv = tmpAv
tmpRi = k
end
y = iterate(ks, kss)
end
Rval[i1,IR] = tmpRv
Rind[i1,IR] = tmpRi
@inbounds Rval[i1,IR] = tmpRv
@inbounds Rind[i1,IR] = tmpRi
end
else
@inbounds for IA in CartesianIndices(indsAt)
for IA in CartesianIndices(indsAt)
IR = Broadcast.newindex(IA, keep, Idefault)
for i in axes(A, 1)
k, kss = y::Tuple
tmpAv = f(A[i,IA])
tmpRv = Rval[i,IR]
tmpRi = Rind[i,IR]
tmpAv = f(@inbounds(A[i,IA]))
@inbounds tmpRv = Rval[i,IR]
@inbounds tmpRi = Rind[i,IR]
if tmpRi == zi || op(tmpRv, tmpAv)
Rval[i,IR] = tmpAv
Rind[i,IR] = k
@inbounds Rval[i,IR] = tmpAv
@inbounds Rind[i,IR] = k
end
y = iterate(ks, kss)
end
Expand Down
24 changes: 24 additions & 0 deletions test/reducedim.jl
Original file line number Diff line number Diff line change
Expand Up @@ -575,6 +575,30 @@ end
@test B[argmin(B, dims=[2, 3])] == @inferred(minimum(B, dims=[2, 3]))
end

@testset "careful with @inbounds" begin
Base.@propagate_inbounds f(x) = x == 2 ? x[-10000] : x
Base.@propagate_inbounds op(x,y) = x[-10000] + y[-10000]
for (arr, dims) in (([1,1,2], 1), ([1 1 2], 2), ([ones(Int,256);2], 1))
@test_throws BoundsError mapreduce(f, +, arr)
@test_throws BoundsError mapreduce(f, +, arr; dims)
@test_throws BoundsError mapreduce(f, +, arr; dims, init=0)
@test_throws BoundsError mapreduce(identity, op, arr)
try
#=@test_throws BoundsError=# mapreduce(identity, op, arr; dims)
catch ex
@test_broken ex isa BoundsError
end
@test_throws BoundsError mapreduce(identity, op, arr; dims, init=0)

@test_throws BoundsError findmin(f, arr)
@test_throws BoundsError findmin(f, arr; dims)

@test_throws BoundsError mapreduce(f, max, arr)
@test_throws BoundsError mapreduce(f, max, arr; dims)
@test_throws BoundsError mapreduce(f, max, arr; dims, init=0)
end
end

@testset "in-place reductions with mismatched dimensionalities" begin
B = reshape(1:24, 4, 3, 2)
for R in (fill(0, 4), fill(0, 4, 1), fill(0, 4, 1, 1))
Expand Down

0 comments on commit 21ccfc0

Please sign in to comment.