在 Julia 中通过递归调用减少 JIT 时间

Question

我有一个递归函数，它操作整数的二叉树，实现为一对嵌套的对或整数。我的函数创建了一棵具有不同结构的新树，并递归调用自身直到满足某些条件。我发现的问题是代码第一次是运行时，JIT 编译所有可能的函数签名需要很长时间；之后运行就好了。

这是最小的工作示例：

my_tree = ((((6 => 7) => (6 => 7)) => ((7 => 7) => (0 => 7))) => (((8 => 7) => (7 => 7)) => ((8 => 8) => (8 => 0)))) => ((((2 => 4) => 7) => (6 => (0 => 5))) => (((6 => 8) => (2 => 8)) => ((2 => 1) => (4 => 5))))

function tree_reduce(tree::Pair)
    left, right = tree
    left isa Pair && (left = tree_reduce(left))
    right isa Pair && (right = tree_reduce(right))
    return left + right
end

@show my_tree
@show tree_reduce(my_tree)

using MethodAnalysis
methods = methodinstances(tree_reduce)
@show length(methods)

虽然这个例子在感知上并不慢，但它仍然生成了 9 个方法实例：

tree_reduce(::Pair{Pair{Pair{Int64, Int64}, Pair{Int64, Int64}}, Pair{Pair{Int64, Int64}, Pair{Int64, Int64}}})
tree_reduce(::Pair{Pair{Int64, Int64}, Pair{Int64, Int64}})
tree_reduce(::Pair{Int64, Int64})
tree_reduce(::Pair{Pair{Pair{Pair{Int64, Int64}, Int64}, Pair{Int64, Pair{Int64, Int64}}}, Pair{Pair{Pair{Int64, Int64}, Pair{Int64, Int64}}, Pair{Pair{Int64, Int64}, Pair{Int64, Int64}}}})
etc ...

有没有办法避免这种情况/预编译/加速/编写通用函数/运行在解释模式下使用特定（部分）函数？我准备好让代码的整体性能稍微好一点，而不是在第一次顶级调用 tree_reduce.

时让它运行更快

Answer 1

是的，Base.@nospecialize:


function tree_reduce2(Base.@nospecialize(tree::Pair))
    left, right = tree
    left isa Pair && (left = tree_reduce(left))
    right isa Pair && (right = tree_reduce(right))
    return left + right
end

@show my_tree
@show tree_reduce2(my_tree)

using MethodAnalysis
methods = methodinstances(tree_reduce2)
@show length(methods) #1

查看 docs of the macro 了解更多使用方法。

Answer 2

@nospecialize 是一个选项，但我认为在这种情况下，一个不捕获其类型的整个数据结构的单独数据结构是有序的。 Pair 确实适用于强类型的事物对，而不适用于大型嵌套结构。

julia> abstract type BiTree{T} end
julia> struct Branch{T} <: BiTree{T} 
           left::BiTree{T}
           right::BiTree{T}
       end

julia> struct Leaf{T} <: BiTree{T}
           value::T
       end

julia> Base.foldl(f, b::Branch) = f(foldl(f, b.left), foldl(f, b.right))

julia> Base.foldl(f, l::Leaf) = f(l.value)

julia> (→)(l::Branch, r::Branch) = Branch(l, r) # just for illustration
→ (generic function with 1 method)

julia> (→)(l, r::Branch) = Branch(Leaf(l), r)
→ (generic function with 2 methods)

julia> (→)(l::Branch, r) = Branch(l, Leaf(r))
→ (generic function with 3 methods)

julia> (→)(l, r) = Branch(Leaf(l), Leaf(r))
→ (generic function with 4 methods)

julia> my_tree = ((((6 → 7) → (6 → 7)) → ((7 → 7) → (0 → 7))) → (((8 → 7) → (7 → 7)) → ((8 → 8) → (8 → 0)))) → ((((2 → 4) → 7) → (6 → (0 → 5))) → (((6 → 8) → (2 → 8)) → ((2 → 1) → (4 → 5))));

julia> typeof(my_tree)
Branch{Int64}

julia> foldl(+, my_tree)
160

这还有一个好处，就是您可以在不破坏任何东西的情况下重载其他方法，例如打印或索引。

在 Julia 中通过递归调用减少 JIT 时间

Reducing JIT time with recursive calls in Julia

optimization

jit

julia