From 5b06e691504452e56bade4781d5473ae1bdf5c6e Mon Sep 17 00:00:00 2001 From: Zachary P Christensen Date: Sat, 3 Jul 2021 03:46:55 -0400 Subject: [PATCH] Implement `preserve` method and add some docs (#2) * Implement `preserve` method and add some docs * Add inline to `preserve` and `ptrcall` option to `LazyPreserve` * version bump * Fixed LazyPreserve{A,Nothing} and `preserve`'s use of it --- Project.toml | 2 +- src/ManualMemory.jl | 128 ++++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 6 ++- 3 files changed, 134 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index e3e3f85..08bc8be 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "ManualMemory" uuid = "d125e4d3-2237-4719-b19c-fa641b8a4667" authors = ["chriselrod and contributors"] -version = "0.1.1" +version = "0.1.2" [compat] julia = "1.5" diff --git a/src/ManualMemory.jl b/src/ManualMemory.jl index 0c8b34e..8c1189f 100644 --- a/src/ManualMemory.jl +++ b/src/ManualMemory.jl @@ -28,6 +28,78 @@ end @generated offsetsize(::Type{T}) where {T} = Base.allocatedinline(T) ? sizeof(T) : sizeof(Int) @inline store!(p::Ptr{T}, v) where {T} = store!(p, convert(T, v)) + +""" + LazyPreserve(x, ptrcall=nothing) + +Used by [`preserve`](@ref) to identify arguments that will be unwrapped with +[`preserve_buffer`](@ref), which is in turn converted to a pointer. If `ptrcall` is +specified conversion to a pointer occurs through a call equivalent to +`ptrcall(preserve_buffer(x), x)`. Otherwise, a call equivalent to +`pointer(preserve_buffer(x))` occurs. +""" +struct LazyPreserve{A,P} + arg::A + ptrcall::P +end +LazyPreserve(arg) = LazyPreserve(arg, nothing) +(p::LazyPreserve)(x) = p.ptrcall(x, p.arg) +(p::LazyPreserve{A,Nothing})(x) where {A} = pointer(x) + + +""" + preserve_buffer(x) + +For structs wrapping arrays, using `GC.@preserve` can trigger heap allocations. +`preserve_buffer` attempts to extract the heap-allocated part. Isolating it by itself +will often allow the heap allocations to be elided. For example: + +```julia +julia> using StaticArrays, BenchmarkTools +julia> # Needed until a release is made featuring https://github.com/JuliaArrays/StaticArrays.jl/commit/a0179213b741c0feebd2fc6a1101a7358a90caed + Base.elsize(::Type{<:MArray{S,T}}) where {S,T} = sizeof(T) +julia> @noinline foo(A) = unsafe_load(A,1) +foo (generic function with 1 method) +julia> function alloc_test_1() + A = view(MMatrix{8,8,Float64}(undef), 2:5, 3:7) + A[begin] = 4 + GC.@preserve A foo(pointer(A)) + end +alloc_test_1 (generic function with 1 method) +julia> function alloc_test_2() + A = view(MMatrix{8,8,Float64}(undef), 2:5, 3:7) + A[begin] = 4 + pb = parent(A) # or `LoopVectorization.preserve_buffer(A)`; `perserve_buffer(::SubArray)` calls `parent` + GC.@preserve pb foo(pointer(A)) + end +alloc_test_2 (generic function with 1 method) +julia> @benchmark alloc_test_1() +BenchmarkTools.Trial: + memory estimate: 544 bytes + allocs estimate: 1 + -------------- + minimum time: 17.227 ns (0.00% GC) + median time: 21.352 ns (0.00% GC) + mean time: 26.151 ns (13.33% GC) + maximum time: 571.130 ns (78.53% GC) + -------------- + samples: 10000 + evals/sample: 998 +julia> @benchmark alloc_test_2() +BenchmarkTools.Trial: + memory estimate: 0 bytes + allocs estimate: 0 + -------------- + minimum time: 3.275 ns (0.00% GC) + median time: 3.493 ns (0.00% GC) + mean time: 3.491 ns (0.00% GC) + maximum time: 4.998 ns (0.00% GC) + -------------- + samples: 10000 + evals/sample: 1000 +``` +""" +@inline preserve_buffer(x::LazyPreserve) = preserve_buffer(x.arg) @inline preserve_buffer(x) = x @inline preserve_buffer(A::AbstractArray) = _preserve_buffer(A, parent(A)) @inline _preserve_buffer(a::A, p::P) where {A,P<:AbstractArray} = _preserve_buffer(p, parent(p)) @@ -101,4 +173,60 @@ end return body end +""" + preserve(op, args...; kwargs...) + +Searches through `args` and `kwargs` for instances of [`LazyPreserve`](@ref), which are +unwrapped using [`preserve_buffer`](@ref) and preserved from garbage collection +(`GC.@preserve`). the resulting buffers are converted to pointers and passed in order to `op`. + +# Examples + +```julia +julia> using ManualMemory: store!, preserve, LazyPreserve + +julia> x = [0 0; 0 0]; + +julia> preserve(store!, LazyPreserve(x), 1) + +julia> x[1] +1 + +``` +""" +preserve(op, args...; kwargs...) = _preserve(op, args, kwargs.data) +@generated function _preserve(op, args::A, kwargs::NamedTuple{syms,K}) where {A,syms,K} + _preserve_expr(A, syms, K) +end +function _preserve_expr(::Type{A}, syms::Tuple{Vararg{Symbol}}, ::Type{K}) where {A,K} + body = Expr(:block, Expr(:meta,:inline)) + call = Expr(:call, :op) + pres = :(GC.@preserve) + @inbounds for i in 1:length(A.parameters) + arg_i = _unwrap_preserve(body, pres, :(getfield(args, $i)), A.parameters[i]) + push!(call.args, arg_i) + end + if length(syms) > 0 + kwargs = Expr(:parameters) + @inbounds for i in 1:length(syms) + arg_i = _unwrap_preserve(body, pres, :(getfield(kwargs, $i)), K.parameters[i]) + push!(call.args, Expr(:kw, syms[i], arg_i)) + end + push!(call.args, kwargs) + end + push!(pres.args, call) + push!(body.args, pres) + return body +end +function _unwrap_preserve(body::Expr, pres::Expr, argexpr::Expr, argtype::Type) + if argtype <: LazyPreserve + bufsym = gensym() + push!(body.args, Expr(:(=), bufsym, Expr(:call, :preserve_buffer, argexpr))) + push!(pres.args, bufsym) + return :($argexpr($bufsym)) + else + return argexpr + end +end + end diff --git a/test/runtests.jl b/test/runtests.jl index 59301e2..a48e388 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,4 @@ -using ManualMemory: MemoryBuffer, load, store! +using ManualMemory: MemoryBuffer, load, store!, LazyPreserve, preserve using Test @testset "ManualMemory.jl" begin @@ -12,6 +12,10 @@ using Test store!(Base.unsafe_convert(Ptr{String}, m), str) @test load(Base.unsafe_convert(Ptr{String}, m)) === str end + + x = [0 0; 0 0]; + preserve(store!, LazyPreserve(x), 1) + @test x[1] === 1 end using ThreadingUtilities