Added @gc_preserve macro.

Project.toml

@@ -1,7 +1,7 @@
 name = "PaddedMatrices"
 uuid = "aaaacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <elrodc@gmail.com>"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -14,7 +14,7 @@ VectorizedRNG = "33b4df10-0173-11e9-2a0c-851a7edac40e"
 
 [compat]
 ArrayInterface = "2.14.4"
-LoopVectorization = "0.9.3"
+LoopVectorization = "0.9.4"
 SLEEFPirates = "0.6"
 VectorizationBase = "0.13.7"
 VectorizedRNG = "0.2.1"

src/PaddedMatrices.jl

@@ -10,7 +10,7 @@ using VectorizationBase, ArrayInterface,
 
 using VectorizationBase: align, gep, AbstractStridedPointer, AbstractSIMDVector, vnoaliasstore!, staticm1,
     static_sizeof, lazymul, vmul, vadd, vsub, StridedPointer, gesp
-using LoopVectorization: maybestaticsize, mᵣ, nᵣ
+using LoopVectorization: maybestaticsize, mᵣ, nᵣ, preserve_buffer
 using ArrayInterface: StaticInt, Zero, One, OptionallyStaticUnitRange, size, strides, offsets, indices,
     static_length, static_first, static_last, axes,
     dense_dims, DenseDims, stride_rank, StrideRank
@@ -21,7 +21,7 @@ using ArrayInterface: StaticInt, Zero, One, OptionallyStaticUnitRange, size, str
 
 # using Parameters: @unpack
 
-export @StrideArray, # @Constant,
+export @StrideArray, @gc_preserve, # @Constant,
     AbstractStrideArray, AbstractStrideVector, AbstractStrideMatrix,
     StrideArray, StrideVector, StrideMatrix,
     PtrArray,# PtrVector, PtrMatrix,

src/blas.jl

@@ -291,7 +291,8 @@ end
     pB = PtrArray(B)
     pC = PtrArray(C)
     M = static_length(Ma); K = static_length(Ka); N = static_length(Na);
-    GC.@preserve C A B begin
+    Cb = preserve_buffer(C); Ab = preserve_buffer(A); Bb = preserve_buffer(B);
+    GC.@preserve Cb Ab Bb begin
         if VectorizationBase.CACHE_SIZE[2] === nothing || ((nᵣ ≥ N) || (contiguousstride1(A) && dontpack(pointer(pA), M, K, stride(A,2), StaticInt{mc}(), StaticInt{kc}(), Tc)))
             loopmul!(pC, pA, pB, α, β, (Ma,Ka,Na))
         elseif VectorizationBase.CACHE_SIZE[3] === nothing || ((contiguousstride1(B) && (kc * nc < K * N)) || firststride(B) < 240)
@@ -845,6 +846,7 @@ end
     TC = promote_type(TA,TB)
     C = StrideArray{TC}(undef, (size(A, StaticInt{1}()),size(B, StaticInt{2}())))
     jmul!(C, A, B)
+    C
 end
 
 @inline extract_λ(a) = a

src/rand.jl

@@ -40,8 +40,9 @@ function rand_expr(expr, args...)
     end
     quote
         array = $array
+        b = array.data
         # we add GC.@preserve to save the array.
-        GC.@preserve array $call
+        GC.@preserve b $call
         array
     end
 end

src/type_declarations.jl

@@ -86,6 +86,42 @@ const StrideMatrix{S,D,T,C,B,R,X,O,A} = StrideArray{S,D,T,2,C,B,R,X,O,A}
 const AbstractStrideVector{S,D,T,C,B,R,X,O} = AbstractStrideArray{S,D,T,1,C,B,R,X,O}
 const AbstractStrideMatrix{S,D,T,C,B,R,X,O} = AbstractStrideArray{S,D,T,2,C,B,R,X,O}
 
+@inline StrideArray(A::AbstractArray) = StrideArray(PtrArray(A), A)
+
+
+@inline LoopVectorization.preserve_buffer(A::MemoryBuffer) = A
+@inline LoopVectorization.preserve_buffer(A::StrideArray) = LoopVectorization.preserve_buffer(A.data)
+@inline LoopVectorization.preserve_buffer(A::PtrArray) = nothing
+
+@inline maybe_ptr_array(A) = A
+@inline maybe_ptr_array(A::AbstractArray) = maybe_ptr_array(ArrayInterface.device(A), A)
+@inline maybe_ptr_array(::ArrayInterface.CPUPointer, A::AbstractArray) = PtrArray(A)
+@inline maybe_ptr_array(_, A::AbstractArray) = A
+
+
+macro gc_preserve(ex)
+    @assert ex.head === :call
+    q = Expr(:block)
+    call = Expr(:call, esc(ex.args[1]))
+    gcp = Expr(:gc_preserve, call)
+    for i ∈ 2:length(ex.args)
+        arg = ex.args[i]
+        A = gensym(:A); buffer = gensym(:buffer);
+        if arg isa Expr && arg.head === :kw
+            push!(call.args, Expr(:kw, arg.args[1], Expr(:call, :maybe_ptr_array, A)))
+            arg = arg.args[2]
+        else
+            push!(call.args, Expr(:call, :maybe_ptr_array, A))
+        end
+        push!(q.args, :($A = $(esc(arg))))
+        push!(q.args, Expr(:(=), buffer, Expr(:call, :preserve_buffer, A)))
+        push!(gcp.args, buffer)
+    end
+    push!(q.args, gcp)
+    q
+end
+
+
 # const FVector{L,T} = StrideVector{Tuple{StaticInt{L}},(true,),T,1,0,(1,)}
 
 # @inline function StrideArray{S,T,N,C,B,R,X,O,D}(ptr, sz, sx, data) where {S,T,N,C,B,R,X,O,D}

test/misc.jl

@@ -1,9 +1,23 @@
 
+@noinline function dostuff(A; B = I)
+    C = A * B
+    s = zero(eltype(C))
+    PaddedMatrices.@avx for i in eachindex(C)
+        s += C[i]
+    end
+    s
+end
+function gc_preserve_test()
+    A = @StrideArray rand(8,8);
+    B = @StrideArray rand(8,8);
+    @gc_preserve dostuff(A, B = B)
+end
 
 @testset "Miscellaneous" begin
     x = @StrideArray rand(127);
     @test maximum(abs, x) == maximum(abs, Array(x))
     y = @StrideArray rand(3);
     @test maximum(abs, y) == maximum(abs, Array(y))
+    @test iszero(@allocated gc_preserve_test())
 end
 

test/runtests.jl

@@ -9,7 +9,7 @@ using Test
 
 @testset "PaddedMatrices.jl" begin
     @test isempty(detect_unbound_args(PaddedMatrices))
+    @time include("misc.jl")
     @time include("matmul_tests.jl")
     @time include("broadcast_tests.jl")
-    @time include("misc.jl")
 end