闭包列表的类型稳定性

Question

我正在尝试在 Julia 中设计一些代码，这些代码将采用用户提供的函数列表，并基本上对它们应用一些代数运算。

看来如果这个函数列表是闭包，则不会推断它们的返回值，导致根据@code_warntype 导致类型不稳定的代码。

我尝试为闭包提供返回类型，但似乎无法找到正确的语法。

下面是一个例子：

functions = Function[x -> x]

function f(u)
    ret = zeros(eltype(u), length(u))

    for func in functions
        ret .+= func(u)
    end

    ret
end

运行这个：

u0 = [1.0, 2.0, 3.0]
@code_warntype f(u0)

并获得

Body::Array{Float64,1}
1 ? %1  = (Base.arraylen)(u)::Int64
?   %2  = $(Expr(:foreigncall, :(:jl_alloc_array_1d), Array{Float64,1}, svec(Any, Int64), :(:ccall), 2, Array{Float64,1}, :(%1), :(%1)))::Array{Float64,1}
?   %3  = invoke Base.fill!(%2::Array{Float64,1}, 0.0::Float64)::Array{Float64,1}
?   %4  = Main.functions::Any
?   %5  = (Base.iterate)(%4)::Any
?   %6  = (%5 === nothing)::Bool
?   %7  = (Base.not_int)(%6)::Bool
???       goto #4 if not %7
2 ? %9  = ? (#1 => %5, #3 => %15)::Any
?   %10 = (Core.getfield)(%9, 1)::Any
?   %11 = (Core.getfield)(%9, 2)::Any
?   %12 = (%10)(u)::Any
?   %13 = (Base.broadcasted)(Main.:+, %3, %12)::Any
?         (Base.materialize!)(%3, %13)
?   %15 = (Base.iterate)(%4, %11)::Any
?   %16 = (%15 === nothing)::Bool
?   %17 = (Base.not_int)(%16)::Bool
???       goto #4 if not %17
3 ?       goto #2
4 ?       return %3

那么，如何使这种代码类型稳定？

Answer 1

如果您想要任意函数的类型稳定性，您必须将它们作为元组传递，这允许 julia 提前知道将在哪个阶段应用哪个函数。

function fsequential(u, fs::Fs) where Fs<:Tuple
    ret = similar(u)
    fill!(ret, 0)
    return fsequential!(ret, u, fs...)
end

@inline function fsequential!(ret, u, f::F, fs...) where F
    ret .+= f(u)
    return fsequential!(ret, u, fs...)
end
fsequential!(ret, u) = ret

julia> u0 = [1.0, 2.0, 3.0]
3-element Array{Float64,1}:
 1.0
 2.0
 3.0

julia> fsequential(u0, (identity, x-> x .+ 1))
3-element Array{Float64,1}:
 3.0
 5.0
 7.0

如果你检查这个，@code_warntype你会发现它是可以推断的。

fsequential!是有时被称为“lispy 元组编程”的一个示例，在该示例中，您一次迭代处理一个参数，直到用完所有可变参数。这是一个强大的范例，它允许比for带有数组的-loop更灵活的推理（因为它允许 Julia 为每个“循环迭代”编译单独的代码）。但是，它通常仅在容器中的元素数量相当少时才有用，否则最终会导致编译时间过长。

类型参数F和Fs看起来没有必要，但它们旨在强制 Julia 专门为您传入的特定函数编写代码。

Answer 2

您的代码中存在多个层问题（不幸的是类型稳定性）：

1. functions是一个全局变量，所以从根本上来说你的代码不会是类型稳定的
2. 即使您移动到functions函数定义内部并且它将是一个向量，代码仍然是类型不稳定的，因为容器将具有抽象 eltype（即使您Function之前删除了前缀[，如果您有多个不同的函数，这仍然是正确的）
3. 如果将向量更改为元组（那么集合functions将是类型稳定的），则该函数仍然是类型不稳定的，因为您使用的循环无法在内部推断出返回类型func(u)

解决方案是使用将@generated循环展开为一系列连续应用程序的函数func(u)- 那么您的代码将是类型稳定的。

然而，一般来说，我认为，假设这func(u)是一个昂贵的操作，代码中的类型不稳定应该不会有太大问题，因为最终您无论如何都会将返回值转换func(u)为Float64。

编辑一个@generated版本，以便与蒂姆·霍利提出的版本进行比较。

@generated function fgenerated(u, functions::Tuple{Vararg{Function}})
    expr = :(ret = zeros(eltype(u), size(u)))
    for fun in functions.parameters
        expr = :($expr; ret .+= $(fun.instance)(u))
    end
    return expr
end