Skip to content

Commit

Permalink
Make reshape and view on Memory produce Arrays and delete wrap (#53896)
Browse files Browse the repository at this point in the history
- Make reshape and view with one based indexing on Memory produce Arrays
- delete wrap

Implements
#53552 (comment)

---------

Co-authored-by: Jameson Nash <vtjnash@gmail.com>
  • Loading branch information
LilithHafner and vtjnash committed Apr 6, 2024
1 parent c707a53 commit 273d91e
Show file tree
Hide file tree
Showing 8 changed files with 63 additions and 80 deletions.
1 change: 0 additions & 1 deletion HISTORY.md
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,6 @@ New library functions
* `copyuntil(out, io, delim)` and `copyline(out, io)` copy data into an `out::IO` stream ([#48273]).
* `eachrsplit(string, pattern)` iterates split substrings right to left.
* `Sys.username()` can be used to return the current user's username ([#51897]).
* `wrap(Array, m::Union{MemoryRef{T}, Memory{T}}, dims)` is the safe counterpart to `unsafe_wrap` ([#52049]).
* `GC.logging_enabled()` can be used to test whether GC logging has been enabled via `GC.enable_logging` ([#51647]).
* `IdSet` is now exported from Base and considered public ([#53262]).

Expand Down
51 changes: 0 additions & 51 deletions base/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3080,54 +3080,3 @@ intersect(r::AbstractRange, v::AbstractVector) = intersect(v, r)
_getindex(v, i)
end
end

"""
wrap(Array, m::Union{Memory{T}, MemoryRef{T}}, dims)
Create an array of size `dims` using `m` as the underlying memory. This can be thought of as a safe version
of [`unsafe_wrap`](@ref) utilizing `Memory` or `MemoryRef` instead of raw pointers.
"""
function wrap end

# validity checking for _wrap calls, separate from allocation of Array so that it can be more likely to inline into the caller
function _wrap(ref::MemoryRef{T}, dims::NTuple{N, Int}) where {T, N}
mem = ref.mem
mem_len = length(mem) + 1 - memoryrefoffset(ref)
len = Core.checked_dims(dims...)
@boundscheck mem_len >= len || invalid_wrap_err(mem_len, dims, len)
if N != 1 && !(ref === GenericMemoryRef(mem) && len === mem_len)
mem = ccall(:jl_genericmemory_slice, Memory{T}, (Any, Ptr{Cvoid}, Int), mem, ref.ptr_or_offset, len)
ref = MemoryRef(mem)
end
return ref
end

@noinline invalid_wrap_err(len, dims, proddims) = throw(DimensionMismatch(
"Attempted to wrap a MemoryRef of length $len with an Array of size dims=$dims, which is invalid because prod(dims) = $proddims > $len, so that the array would have more elements than the underlying memory can store."))

@eval @propagate_inbounds function wrap(::Type{Array}, m::MemoryRef{T}, dims::NTuple{N, Integer}) where {T, N}
dims = convert(Dims, dims)
ref = _wrap(m, dims)
$(Expr(:new, :(Array{T, N}), :ref, :dims))
end

@eval @propagate_inbounds function wrap(::Type{Array}, m::Memory{T}, dims::NTuple{N, Integer}) where {T, N}
dims = convert(Dims, dims)
ref = _wrap(MemoryRef(m), dims)
$(Expr(:new, :(Array{T, N}), :ref, :dims))
end
@eval @propagate_inbounds function wrap(::Type{Array}, m::MemoryRef{T}, l::Integer) where {T}
dims = (Int(l),)
ref = _wrap(m, dims)
$(Expr(:new, :(Array{T, 1}), :ref, :dims))
end
@eval @propagate_inbounds function wrap(::Type{Array}, m::Memory{T}, l::Integer) where {T}
dims = (Int(l),)
ref = _wrap(MemoryRef(m), (l,))
$(Expr(:new, :(Array{T, 1}), :ref, :dims))
end
@eval @propagate_inbounds function wrap(::Type{Array}, m::Memory{T}) where {T}
ref = MemoryRef(m)
dims = (length(m),)
$(Expr(:new, :(Array{T, 1}), :ref, :dims))
end
1 change: 0 additions & 1 deletion base/exports.jl
Original file line number Diff line number Diff line change
Expand Up @@ -460,7 +460,6 @@ export
vcat,
vec,
view,
wrap,
zeros,

# search, find, match and related functions
Expand Down
24 changes: 24 additions & 0 deletions base/genericmemory.jl
Original file line number Diff line number Diff line change
Expand Up @@ -288,3 +288,27 @@ function indcopy(sz::Dims, I::GenericMemory)
src = eltype(I)[I[i][_findin(I[i], i < n ? (1:sz[i]) : (1:s))] for i = 1:n]
dst, src
end

# Wrapping a memory region in an Array
@eval begin # @eval for the Array construction. Block for the docstring.
function reshape(m::GenericMemory{M, T}, dims::Vararg{Int, N}) where {M, T, N}
len = Core.checked_dims(dims...)
length(m) == len || throw(DimensionMismatch("parent has $(length(m)) elements, which is incompatible with size $(dims)"))
ref = MemoryRef(m)
$(Expr(:new, :(Array{T, N}), :ref, :dims))
end

"""
view(m::GenericMemory{M, T}, inds::Union{UnitRange, OneTo})
Create a vector `v::Vector{T}` backed by the specified indices of `m`. It is only safe to
resize `v` if `m` is subseqently not used.
"""
function view(m::GenericMemory{M, T}, inds::Union{UnitRange, OneTo}) where {M, T}
isempty(inds) && return T[] # needed to allow view(Memory{T}(undef, 0), 2:1)
@boundscheck checkbounds(m, inds)
ref = MemoryRef(m, first(inds)) # @inbounds would be safe here but does not help performance.
dims = (length(inds),)
$(Expr(:new, :(Array{T, 1}), :ref, :dims))
end
end
12 changes: 4 additions & 8 deletions base/iobuffer.jl
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ end

# allocate Vector{UInt8}s for IOBuffer storage that can efficiently become Strings
StringMemory(n::Integer) = unsafe_wrap(Memory{UInt8}, _string_n(n))
StringVector(n::Integer) = wrap(Array, StringMemory(n))
StringVector(n::Integer) = view(StringMemory(n), 1:n)::Vector{UInt8}

# IOBuffers behave like Files. They are typically readable and writable. They are seekable. (They can be appendable).

Expand Down Expand Up @@ -456,7 +456,7 @@ function take!(io::IOBuffer)
if nbytes == 0 || io.reinit
data = StringVector(0)
elseif io.writable
data = wrap(Array, MemoryRef(io.data, io.offset + 1), nbytes)
data = view(io.data, io.offset+1:nbytes+io.offset)
else
data = copyto!(StringVector(io.size), 1, io.data, io.offset + 1, nbytes)
end
Expand All @@ -465,7 +465,7 @@ function take!(io::IOBuffer)
if nbytes == 0
data = StringVector(0)
elseif io.writable
data = wrap(Array, MemoryRef(io.data, io.ptr), nbytes)
data = view(io.data, io.ptr:io.ptr+nbytes-1)
else
data = read!(io, data)
end
Expand All @@ -491,11 +491,7 @@ state. This should only be used internally for performance-critical
It might save an allocation compared to `take!` (if the compiler elides the
Array allocation), as well as omits some checks.
"""
_unsafe_take!(io::IOBuffer) =
wrap(Array, io.size == io.offset ?
MemoryRef(Memory{UInt8}()) :
MemoryRef(io.data, io.offset + 1),
io.size - io.offset)
_unsafe_take!(io::IOBuffer) = view(io.data, io.offset+1:io.size)

function write(to::IO, from::GenericIOBuffer)
written::Int = bytesavailable(from)
Expand Down
5 changes: 4 additions & 1 deletion base/strings/string.jl
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,10 @@ String(s::AbstractString) = print_to_string(s)
@assume_effects :total String(s::Symbol) = unsafe_string(unsafe_convert(Ptr{UInt8}, s))

unsafe_wrap(::Type{Memory{UInt8}}, s::String) = ccall(:jl_string_to_genericmemory, Ref{Memory{UInt8}}, (Any,), s)
unsafe_wrap(::Type{Vector{UInt8}}, s::String) = wrap(Array, unsafe_wrap(Memory{UInt8}, s))
function unsafe_wrap(::Type{Vector{UInt8}}, s::String)
mem = unsafe_wrap(Memory{UInt8}, s)
view(mem, eachindex(mem))
end

Vector{UInt8}(s::CodeUnits{UInt8,String}) = copyto!(Vector{UInt8}(undef, length(s)), s)
Vector{UInt8}(s::String) = Vector{UInt8}(codeunits(s))
Expand Down
1 change: 0 additions & 1 deletion doc/src/base/arrays.md
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,6 @@ Base.reshape
Base.dropdims
Base.vec
Base.SubArray
Base.wrap
```

## Concatenation and permutation
Expand Down
48 changes: 31 additions & 17 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3187,23 +3187,37 @@ end
end
end

@testset "Wrapping Memory into Arrays" begin
mem = Memory{Int}(undef, 10) .= 1
memref = MemoryRef(mem)
@test_throws DimensionMismatch wrap(Array, mem, (10, 10))
@test wrap(Array, mem, (5,)) == ones(Int, 5)
@test wrap(Array, mem, 2) == ones(Int, 2)
@test wrap(Array, memref, 10) == ones(Int, 10)
@test wrap(Array, memref, (2,2,2)) == ones(Int,2,2,2)
@test wrap(Array, mem, (5, 2)) == ones(Int, 5, 2)

memref2 = MemoryRef(mem, 3)
@test wrap(Array, memref2, (5,)) == ones(Int, 5)
@test wrap(Array, memref2, 2) == ones(Int, 2)
@test wrap(Array, memref2, (2,2,2)) == ones(Int,2,2,2)
@test wrap(Array, memref2, (3, 2)) == ones(Int, 3, 2)
@test_throws DimensionMismatch wrap(Array, memref2, 9)
@test_throws DimensionMismatch wrap(Array, memref2, 10)
@testset "Wrapping Memory into Arrays with view and reshape" begin
mem::Memory{Int} = Memory{Int}(undef, 10) .= 11:20

@test_throws DimensionMismatch reshape(mem, 10, 10)
@test_throws DimensionMismatch reshape(mem, 5)
@test_throws BoundsError view(mem, 1:10, 1:10)
@test_throws BoundsError view(mem, 1:11)
@test_throws BoundsError view(mem, 3:11)
@test_throws BoundsError view(mem, 0:4)

@test @inferred(view(mem, 1:5))::Vector{Int} == 11:15
@test @inferred(view(mem, 1:2))::Vector{Int} == 11:12
@test @inferred(view(mem, 1:10))::Vector{Int} == 11:20
@test @inferred(view(mem, 3:8))::Vector{Int} == 13:18
@test @inferred(view(mem, 20:19))::Vector{Int} == []
@test @inferred(view(mem, -5:-7))::Vector{Int} == []
@test @inferred(reshape(mem, 5, 2))::Matrix{Int} == reshape(11:20, 5, 2)

empty_mem = Memory{Module}(undef, 0)
@test_throws BoundsError view(empty_mem, 0:1)
@test_throws BoundsError view(empty_mem, 1:2)
@test_throws DimensionMismatch reshape(empty_mem, 1)
@test_throws DimensionMismatch reshape(empty_mem, 1, 2, 3)
@test_throws ArgumentError reshape(empty_mem, 2^16, 2^16, 2^16, 2^16)

@test @inferred(view(empty_mem, 1:0))::Vector{Module} == []
@test @inferred(view(empty_mem, 10:3))::Vector{Module} == []
@test isempty(@inferred(reshape(empty_mem, 0, 7, 1))::Array{Module, 3})

offset_inds = OffsetArrays.IdOffsetRange(values=3:6, indices=53:56)
@test @inferred(view(collect(mem), offset_inds)) == view(mem, offset_inds)
end

@testset "Memory size" begin
Expand Down

0 comments on commit 273d91e

Please sign in to comment.