Search code examples
juliamatrix-multiplicationtype-stability

Type stability with container types and matrix-vector multiply in Julia


I am trying to use Julia's A_mul_B! with a container type, something like

# my composite type, contains 2 vectors and 1 matrix of same Float type
type MyContainer{T <: Float}
    z :: Vector
    x :: Matrix
    y :: Vector
    MyContainer(z::Vector{T}, x::Matrix{T}, y::Vector{T}) = new(z,x,y) 
end

I then use an instance of MyContainer with A_mul_B! followed by some arithmetic with the Vector objects:

# only work with single/double precision
typealias Float Union{Float32, Float64}

# function to perform mat-vec multiply
function f{T <: Float}(v::MyContainer{T}) 
    Base.A_mul_B!(v.z, v.x, v.y)
    return sumabs2(v.z) * sumabs2(v.y)
end

As defined, f is curiously not type-stable, even though the constructor itself is type-stable. Is there a place where I can annotate the types of z, x, and y so that A_mul_B! sees them?

Here is a minimal working example:

MyModule.jl

module MyModule

export MyContainer, f

# only work with single/double precision
typealias Float Union{Float32, Float64}

# my composite type, contains 2 vectors and 1 matrix of same Float type
type MyContainer{T <: Float}
    z :: Vector
    x :: Matrix
    y :: Vector
    MyContainer(z::Vector{T}, x::Matrix{T}, y::Vector{T}) = new(z,x,y) 
end

# testing routine initializes all arrays with a single value 
function MyContainer{T <: Float}(n::Int, t::T)
    z = t*ones(T, n)
    x = t*ones(T, (n,n))
    y = t*ones(T, n)
    return MyContainer{eltype(z)}(z, x, y)
end

# function to perform mat-vec multiply
function f{T <: Float}(v::MyContainer{T}) 
    Base.A_mul_B!(v.z, v.x, v.y)
    return sumabs2(v.z) * sumabs2(v.y)
end

end

test.jl

include("MyModule.jl")

function g() 
    # check type stability
    @code_warntype MyModule.MyContainer(10, 1.0) # type-stable
    @code_warntype MyModule.f(v) # red Array{T,1}, Array{T,2}, Any

    # make a container
    v = MyModule.MyContainer(10, 1.0)

    # does type-stability matter for performance?
    @time 1+1 
    MyModule.f(v)
    @time MyModule.f(v) # maybe... note small memory allocation
end

g()

partial output

# omit output of @code_warntype for conciseness
  0.000000 seconds
  0.000001 seconds (3 allocations: 48 bytes)
10000.0

Solution

  • As David Sanders pointed out, the problem is

    type MyContainer{T <: Float}
        z :: Vector
        x :: Matrix
        y :: Vector
        MyContainer(z::Vector{T}, x::Matrix{T}, y::Vector{T}) = new(z,x,y) 
    end
    

    Since Vector and Matrix are abstract types, this type's fields are not concrete-inferrable. The fix is to concretely type them:

    type MyContainer{T <: Float}
        z :: Vector{T}
        x :: Matrix{T}
        y :: Vector{T}
    end