Julia代码中的内存分配问题

Question

我使用Python/Numpy中的一个函数来解决组合博弈论中的一个问题。

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秒。在核心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代码的性能（而不是试图仅仅复制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