Skip to content

Commit

Permalink
change IOBuffer to use Memory internally (#53192)
Browse files Browse the repository at this point in the history
An Array is often still allocated on output, but this gives the compiler
a chance to potentially elide that in certain cases.

For measurement, it seems about 10% faster as a string builder:
```
julia> @Btime repr("hello\nworld"^10);
  1.096 μs (10 allocations: 640 bytes) # master
  973.000 ns (9 allocations: 608 bytes) # PR
  994.000 ns (8 allocations: 576 bytes) # also PR, after Revise-ing Base.wrap
```
  • Loading branch information
vtjnash authored Feb 13, 2024
1 parent 81c3474 commit d7b9ac8
Show file tree
Hide file tree
Showing 14 changed files with 230 additions and 137 deletions.
34 changes: 26 additions & 8 deletions base/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3067,7 +3067,8 @@ of [`unsafe_wrap`](@ref) utilizing `Memory` or `MemoryRef` instead of raw pointe
"""
function wrap end

@eval @propagate_inbounds function wrap(::Type{Array}, ref::MemoryRef{T}, dims::NTuple{N, Integer}) where {T, N}
# 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...)
Expand All @@ -3076,18 +3077,35 @@ function wrap end
mem = ccall(:jl_genericmemory_slice, Memory{T}, (Any, Ptr{Cvoid}, Int), mem, ref.ptr_or_offset, len)
ref = MemoryRef(mem)
end
$(Expr(:new, :(Array{T, N}), :ref, :dims))
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."))

function wrap(::Type{Array}, m::Memory{T}, dims::NTuple{N, Integer}) where {T, N}
wrap(Array, MemoryRef(m), dims)
@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
function wrap(::Type{Array}, m::MemoryRef{T}, l::Integer) where {T}
wrap(Array, m, (l,))
@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
function wrap(::Type{Array}, m::Memory{T}, l::Integer) where {T}
wrap(Array, MemoryRef(m), (l,))
@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
Loading

0 comments on commit d7b9ac8

Please sign in to comment.