Julia 中容器类型和矩阵向量乘法的类型稳定性

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

我正在尝试将 Julia 的 A_mul_B! 与容器类型一起使用,例如

# 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

然后我将 MyContainer 的实例与 A_mul_B! 一起使用,然后对 Vector 对象进行一些运算:

# 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

根据定义,f 奇怪的是类型不稳定,即使构造函数本身是类型稳定的。有什么地方可以注释 zxy 的类型,以便 A_mul_B! 看到它们吗?

这是一个最小的工作示例:

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()

部分输出

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

正如大卫·桑德斯指出的那样,问题是

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

由于 VectorMatrix 是抽象类型,因此该类型的字段不可具体推断。解决方法是具体键入它们:

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