Julia 代码中的内存分配问题

Question

我在 Python/Numpy 中使用了一个函数来解决 combinatorial game theory 中的一个问题。

import numpy as np
from time import time

def problem(c):
    start = time()
    N = np.array([0, 0])
    U = np.arange(c)
    
    for _ in U:
        bits = np.bitwise_xor(N[:-1], N[-2::-1])
        N = np.append(N, np.setdiff1d(U, bits).min())

    return len(*np.where(N==0)), time()-start 

problem(10000)

然后我用 Julia 编写它，因为我认为 Julia 使用即时编译会更快。

function problem(c)
    N = [0]
    U = Vector(0:c)
    
    for _ in U
        elems = N[1:length(N)-1]
        bits = elems .⊻ reverse(elems)
        push!(N, minimum(setdiff(U, bits))) 
    end
    
    return sum(N .== 0)
end

@time problem(10000)

但第二个版本要慢得多。对于 c = 10000，Python 版本需要 2.5 秒。在 Core i5 处理器上，Julia 版本需要 4.5 秒。由于 Numpy 操作是用 C 实现的，我想知道 Python 是否确实更快，或者我正在编写一个浪费时间复杂度的函数。

Julia 中的实现分配了大量内存。如何减少分配次数以提高其性能？

Answer 1

原代码可以改写如下：

function problem2(c)
    N = zeros(Int, c+2)
    notseen = falses(c+1)

    for lN in 1:c+1
        notseen .= true
        @inbounds for i in 1:lN-1
            b = N[i] ⊻ N[lN-i]
            b <= c && (notseen[b+1] = false)
        end
        idx = findfirst(notseen)
        isnothing(idx) || (N[lN+1] = idx-1)
    end
    return count(==(0), N)
end

首先检查函数是否产生相同的结果：

julia> problem(10000), problem2(10000)
(1475, 1475)

（我还检查了生成的N 向量是否相同）

现在让我们对这两个函数进行基准测试：

julia> using BenchmarkTools

julia> @btime problem(10000)
  4.938 s (163884 allocations: 3.25 GiB)
1475

julia> @btime problem2(10000)
  76.275 ms (4 allocations: 79.59 KiB)
1475

所以它的速度提高了 60 倍以上。

我为提高性能所做的就是避免分配。在 Julia 中，它既简单又高效。如果代码的任何部分不清楚，请发表评论。请注意，我专注于展示如何提高 Julia 代码的性能（而不是试图仅仅复制 Python 代码，因为 - 正如在原始帖子下所评论的那样 - 进行语言性能比较非常棘手）。我认为最好集中讨论如何使 Julia 代码更快。

---

编辑

确实更改为Vector{Bool} 并删除b 和c 关系上的条件（数学上适用于c 的这些值）提供了更好的速度：

julia> function problem3(c)
           N = zeros(Int, c+2)
           notseen = Vector{Bool}(undef, c+1)

           for lN in 1:c+1
               notseen .= true
               @inbounds for i in 1:lN-1
                   b = N[i] ⊻ N[lN-i]
                   notseen[b+1] = false
               end
               idx = findfirst(notseen)
               isnothing(idx) || (N[lN+1] = idx-1)
           end
           return count(==(0), N)
       end
problem3 (generic function with 1 method)

julia> @btime problem3(10000)
  20.714 ms (3 allocations: 88.17 KiB)
1475