Skip to content

Commit

Permalink
Add _unsetindex! methods for SubArrays and CartesianIndexes (#5…
Browse files Browse the repository at this point in the history
…3383)

With this, the following (and equivalent calls) work:
```julia
julia> copyto!(view(zeros(BigInt, 2), 1:2), Vector{BigInt}(undef,2))
2-element view(::Vector{BigInt}, 1:2) with eltype BigInt:
 #undef
 #undef

julia> copyto!(view(zeros(BigInt, 2), 1:2), view(Vector{BigInt}(undef,2), 1:2))
2-element view(::Vector{BigInt}, 1:2) with eltype BigInt:
 #undef
 #undef
```

Close #53098. With this, all
the `_unsetindex!` branches in `copyto_unaliased!` work for
`Array`-views, and this makes certain indexing operations vectorize and
speed-up:
```julia
julia> using BenchmarkTools

julia> a = view(rand(100,100), 1:100, 1:100); b = view(similar(a), axes(a)...);

julia> @Btime copyto!($b, $a);
  16.427 μs (0 allocations: 0 bytes) # master
  2.308 μs (0 allocations: 0 bytes) # PR
``` 

Improves (but doesn't resolve)
#40962 and
#53158

```julia
julia> a = rand(40,40); b = rand(40,40);

julia> @Btime $a[1:end,1:end] .= $b;
  5.383 μs (0 allocations: 0 bytes) # v"1.12.0-DEV.16"
  3.194 μs (0 allocations: 0 bytes) # PR
```
ƒ
Co-authored-by: Jameson Nash <vtjnash@gmail.com>
  • Loading branch information
jishnub authored Feb 20, 2024
1 parent ea2b255 commit 1a90409
Show file tree
Hide file tree
Showing 6 changed files with 118 additions and 2 deletions.
15 changes: 14 additions & 1 deletion base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1472,7 +1472,20 @@ function _setindex!(::IndexCartesian, A::AbstractArray, v, I::Vararg{Int,M}) whe
r
end

_unsetindex!(A::AbstractArray, i::Integer) = _unsetindex!(A, to_index(i))
function _unsetindex!(A::AbstractArray, i::Integer...)
@_propagate_inbounds_meta
_unsetindex!(A, map(to_index, i)...)
end

function _unsetindex!(A::AbstractArray{T}, i::Int...) where T
# this provides a fallback method which is a no-op if the element is already unassigned
# such that copying into an uninitialized object generally always will work,
# even if the specific custom array type has not implemented `_unsetindex!`
@inline
@boundscheck checkbounds(A, i...)
allocatedinline(T) || @inbounds(!isassigned(A, i...)) || throw(MethodError(_unsetindex!, (A, i...)))
return A
end

"""
parent(A)
Expand Down
7 changes: 6 additions & 1 deletion base/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,12 @@ function _unsetindex!(A::Array, i::Int)
@inbounds _unsetindex!(GenericMemoryRef(A.ref, i))
return A
end

function _unsetindex!(A::Array, i::Int...)
@inline
@boundscheck checkbounds(A, i...)
@inbounds _unsetindex!(A, _to_linear_index(A, i...))
return A
end

# TODO: deprecate this (aligned_sizeof and/or elsize and/or sizeof(Some{T}) are more correct)
elsize(::Type{A}) where {T,A<:Array{T}} = aligned_sizeof(T)
Expand Down
6 changes: 6 additions & 0 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1610,6 +1610,12 @@ end
end
end

# _unsetindex
@propagate_inbounds function Base._unsetindex!(A::AbstractArray, i::CartesianIndex)
Base._unsetindex!(A, to_indices(A, (i,))...)
return A
end

## permutedims

## Permute array dims ##
Expand Down
19 changes: 19 additions & 0 deletions base/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -410,6 +410,25 @@ function isassigned(V::FastSubArray{<:Any, 1}, i::Int)
r
end

function _unsetindex!(V::FastSubArray, i::Int)
@inline
@boundscheck checkbounds(Bool, V, i)
@inbounds _unsetindex!(V.parent, _reindexlinear(V, i))
return V
end
function _unsetindex!(V::FastSubArray{<:Any,1}, i::Int)
@inline
@boundscheck checkbounds(Bool, V, i)
@inbounds _unsetindex!(V.parent, _reindexlinear(V, i))
return V
end
function _unsetindex!(V::SubArray{T,N}, i::Vararg{Int,N}) where {T,N}
@inline
@boundscheck checkbounds(Bool, V, i...)
@inbounds _unsetindex!(V.parent, reindex(V.indices, i)...)
return V
end

IndexStyle(::Type{<:FastSubArray}) = IndexLinear()
IndexStyle(::Type{<:SubArray}) = IndexCartesian()

Expand Down
31 changes: 31 additions & 0 deletions test/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2019,3 +2019,34 @@ end
@test B == A
end
end

@testset "_unsetindex!" begin
struct MyMatrixUnsetIndexCartInds{T,A<:AbstractMatrix{T}} <: AbstractMatrix{T}
data :: A
end
Base.size(A::MyMatrixUnsetIndexCartInds) = size(A.data)
Base.getindex(M::MyMatrixUnsetIndexCartInds, i::Int, j::Int) = M.data[i,j]
Base.setindex!(M::MyMatrixUnsetIndexCartInds, v, i::Int, j::Int) = setindex!(M.data, v, i, j)
struct MyMatrixUnsetIndexLinInds{T,A<:AbstractMatrix{T}} <: AbstractMatrix{T}
data :: A
end
Base.size(A::MyMatrixUnsetIndexLinInds) = size(A.data)
Base.getindex(M::MyMatrixUnsetIndexLinInds, i::Int) = M.data[i]
Base.setindex!(M::MyMatrixUnsetIndexLinInds, v, i::Int) = setindex!(M.data, v, i)
Base.IndexStyle(::Type{<:MyMatrixUnsetIndexLinInds}) = IndexLinear()

function test_unsetindex(MT)
M = MT(ones(2,2))
M2 = MT(Matrix{BigFloat}(undef, 2,2))
copyto!(M, M2)
@test all(==(1), M)
M3 = MT(Matrix{BigFloat}(undef, 2,2))
for i in eachindex(M3)
@test !isassigned(M3, i)
end
M3 .= 1
@test_throws MethodError copyto!(M3, M2)
end
test_unsetindex(MyMatrixUnsetIndexCartInds)
test_unsetindex(MyMatrixUnsetIndexLinInds)
end
42 changes: 42 additions & 0 deletions test/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1055,4 +1055,46 @@ end
@test !isassigned(v, 1, 2) # inbounds but not assigned
@test !isassigned(v, 3, 3) # out-of-bounds
end

@testset "_unsetindex!" begin
function test_unsetindex(A, B)
copyto!(A, B)
for i in eachindex(A)
@test !isassigned(A, i)
end
end
@testset "dest IndexLinear, src IndexLinear" begin
for p in (fill(BigInt(2)), BigInt[1, 2], BigInt[1 2; 3 4])
A = view(copy(p), ntuple(_->:, ndims(p))...)
B = view(similar(A), ntuple(_->:, ndims(p))...)
test_unsetindex(A, B)
test_unsetindex(p, B)
end
end

@testset "dest IndexLinear, src IndexCartesian" begin
for p in (fill(BigInt(2)), BigInt[1, 2], BigInt[1 2; 3 4])
A = view(copy(p), ntuple(_->:, ndims(p))...)
B = view(similar(A), axes(A)...)
test_unsetindex(A, B)
test_unsetindex(p, B)
end
end

@testset "dest IndexCartesian, src IndexLinear" begin
for p in (fill(BigInt(2)), BigInt[1, 2], BigInt[1 2; 3 4])
A = view(p, axes(p)...)
B = similar(A)
test_unsetindex(A, B)
end
end

@testset "dest IndexCartesian, src IndexCartesian" begin
for p in (fill(BigInt(2)), BigInt[1, 2], BigInt[1 2; 3 4])
A = view(p, axes(p)...)
B = view(similar(A), axes(A)...)
test_unsetindex(A, B)
end
end
end
end

0 comments on commit 1a90409

Please sign in to comment.