Allocations with closure and mul!

General Usage Performance
,
1

Hi! I have an issue with allocations using the mul! function and a closure:

using LinearAlgebra, Test

struct operator{T} 
    prod!
end

function operator(A :: AbstractArray{T}) where T 
    prod! = (res, u, α, β) -> mul!(res, A, u, α, β)
    return operator{T}(prod!)
end

import LinearAlgebra.mul!

function mul!(res::Vector{T}, op::operator{T}, u::Vector{T}, α::T, β::T) where T 
    op.prod!(res, u, α, β)
end

function mul!(res::Vector{T}, op::operator{T}, u::Vector{T}) where T 
    op.prod!(res, u, one(T), zero(T))
end

function mul2!(res::Vector{T}, op::operator{T}, u::Vector{T}) where T 
    op.prod!(res, u, 2*one(T), 2*one(T))
end

function test_allocs()
    n = 10
    A = rand(n, n)
    A_op1 = operator(A)
    res1, u1, α1, β1 = rand(n), rand(n), 2.0, 2.0

    ### compile
    A_op1.prod!(res1, u1, β1, α1)
    mul!(res1, A_op1, u1, β1, α1)  
    mul!(res1, A, u1, α1, β1) 
    mul2!(res1, A_op1, u1) 

    ### tests allocs
    allocs1 = @allocated A_op1.prod!(res1, u1, α1, β1) 
    @test allocs1 == 16
    allocs2 = @allocated mul!(res1, A_op1, u1, α1, β1)
    @test allocs2 == 32
    allocs3 = @allocated mul!(res1, A, u1, α1, β1)  # LinearAlgebra mul! function for matrices
    @test allocs3 == 0
    allocs4 = @allocated mul!(res1, A_op1, u1) 
    @test allocs4 == 0
    allocs5 = @allocated mul2!(res1, A_op1, u1) 
    @test allocs5 == 0
end

test_allocs()

All tests passed but I do not know how to remove the allocations for the first and second tests with my type operator{T}. Maybe it comes from α and β because when I removed them with the 3-args mul! and mul2! the allocations disappear, but I could not find a way to solve this problem. Does anybody know how I could fix this?

2

The field prod! of your operator type misses a type annotation and hence seems to be of Any. I assume you wanted to make this prod!::T.

1 Like
3

I do not think that this possible because I want prod! to be a closure (its type could be Function but it does not change the number of allocations). The T in struct operator{T} is only used to indicate that the closure should work with data of type T.

4

If you want to reserve T for data you can just add another type parameter to operator to make prod! type-stable:

struct operator{T,F}
    prod!::F
end
operator{T}(prod!::F) where {T,F} = operator{T,F}(prod!)
5

Thank you!