From 121226bfd6c93ea7e0c69d544298912ed3de9d70 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 20 Feb 2024 22:35:02 +0530 Subject: [PATCH] Add `_unsetindex!` methods for `SubArray`s and `CartesianIndex`es (#53383) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 https://github.com/JuliaLang/julia/issues/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) https://github.com/JuliaLang/julia/issues/40962 and https://github.com/JuliaLang/julia/issues/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 (cherry picked from commit 1a9040908b9bb89927b2d318da3f9fe1e457abea) --- base/abstractarray.jl | 15 +++++++++++++- base/array.jl | 7 ++++++- base/multidimensional.jl | 6 ++++++ base/subarray.jl | 19 ++++++++++++++++++ test/abstractarray.jl | 31 +++++++++++++++++++++++++++++ test/subarray.jl | 42 ++++++++++++++++++++++++++++++++++++++++ 6 files changed, 118 insertions(+), 2 deletions(-) diff --git a/base/abstractarray.jl b/base/abstractarray.jl index 749ca98a587b0..9042e5f81e462 100644 --- a/base/abstractarray.jl +++ b/base/abstractarray.jl @@ -1473,7 +1473,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) diff --git a/base/array.jl b/base/array.jl index 4b2cd34b2f0db..4fcf29af69b4b 100644 --- a/base/array.jl +++ b/base/array.jl @@ -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) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 146f8c160e8e9..fe29e5c1d8e5f 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -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 ## diff --git a/base/subarray.jl b/base/subarray.jl index 396f7c52bd77a..c15a4d48868fa 100644 --- a/base/subarray.jl +++ b/base/subarray.jl @@ -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() diff --git a/test/abstractarray.jl b/test/abstractarray.jl index 732ffba85ae84..a156953476c43 100644 --- a/test/abstractarray.jl +++ b/test/abstractarray.jl @@ -2059,3 +2059,34 @@ end @test r2[i] == z[j] 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 diff --git a/test/subarray.jl b/test/subarray.jl index 0606ed00e6a7a..ada53d7d91ee3 100644 --- a/test/subarray.jl +++ b/test/subarray.jl @@ -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