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

Posted 2023-03-06

【中文标题】在 Julia 中使用递归调用减少 JIT 时间

【英文标题】：Reducing JIT time with recursive calls in Julia

【发布时间】：2022-01-21 07:30:16

【问题描述】：

我有一个递归函数，它操作整数的二叉树，实现为一对嵌套的对或整数。我的函数创建一个具有不同结构的新树，并递归调用自身，直到满足某些条件。我发现的问题是代码第一次运行时，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(::PairPairPairInt64, Int64, PairInt64, Int64, PairPairInt64, Int64, PairInt64, Int64)
tree_reduce(::PairPairInt64, Int64, PairInt64, Int64)
tree_reduce(::PairInt64, Int64)
tree_reduce(::PairPairPairPairInt64, Int64, Int64, PairInt64, PairInt64, Int64, PairPairPairInt64, Int64, PairInt64, Int64, PairPairInt64, Int64, PairInt64, Int64)
etc ...

有没有办法避免这种情况/预编译/加速/编写通用函数/在解释模式下运行特定（部分）函数？我准备让代码的整体性能稍微起作用，同时让它在第一次***调用 tree_reduce 时运行得更快。

【参考方案1】：

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

julia> abstract type BiTreeT end
julia> struct BranchT <: BiTreeT 
           left::BiTreeT
           right::BiTreeT
       end

julia> struct LeafT <: BiTreeT
           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)
BranchInt64

julia> foldl(+, my_tree)
160

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

【讨论】：

很好的答案，但我认为你一开始就错过了abstract type BiTreeT end。

这个答案已经足够好了，我将它（非常轻微地）改编为***.com/questions/70479843/…中的答案

【参考方案2】：

是的，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 了解更多使用方法。

【讨论】：

我想知道是否值得为基本案例PairT,U where T<:Integer,U<:Integer 提供一个单独的专门化方法。它没有isa Pair 行，并且可以针对 T,U 类型进行优化。

在我的书中，这种带有isa 的成语只是要求几个方法定义。您只是手动实现类型分派。