Julia 代码中的内存分配问题

Question

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

\n

import numpy as np\nfrom time import time\n\ndef problem(c):\n    start = time()\n    N = np.array([0, 0])\n    U = np.arange(c)\n    \n    for _ in U:\n        bits = np.bitwise_xor(N[:-1], N[-2::-1])\n        N = np.append(N, np.setdiff1d(U, bits).min())\n\n    return len(*np.where(N==0)), time()-start \n\nproblem(10000)\n

\n

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

\n

function problem(c)\n    N = [0]\n    U = Vector(0:c)\n    \n    for _ in U\n        elems = N[1:length(N)-1]\n        bits = elems .\xe2\x8a\xbb reverse(elems)\n        push!(N, minimum(setdiff(U, bits))) \n    end\n    \n    return sum(N .== 0)\nend\n\n@time problem(10000)\n

\n

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

\n

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

\n

Answer 1

原来的代码可以通过以下方式重写：

\n

function problem2(c)\n    N = zeros(Int, c+2)\n    notseen = falses(c+1)\n\n    for lN in 1:c+1\n        notseen .= true\n        @inbounds for i in 1:lN-1\n            b = N[i] \xe2\x8a\xbb N[lN-i]\n            b <= c && (notseen[b+1] = false)\n        end\n        idx = findfirst(notseen)\n        isnothing(idx) || (N[lN+1] = idx-1)\n    end\n    return count(==(0), N)\nend\n

\n

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

\n

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

\n

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

\n

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

\n

julia> using BenchmarkTools\n\njulia> @btime problem(10000)\n  4.938 s (163884 allocations: 3.25 GiB)\n1475\n\njulia> @btime problem2(10000)\n  76.275 ms (4 allocations: 79.59 KiB)\n1475\n

\n

结果是速度快了 60 倍以上。

\n

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

\n\n

编辑

\n

事实上，更改Vector{Bool}和 删除条件b和c关系（在数学上适用于 的这些值c）可以提供更好的速度：

\n

julia> function problem3(c)\n           N = zeros(Int, c+2)\n           notseen = Vector{Bool}(undef, c+1)\n\n           for lN in 1:c+1\n               notseen .= true\n               @inbounds for i in 1:lN-1\n                   b = N[i] \xe2\x8a\xbb N[lN-i]\n                   notseen[b+1] = false\n               end\n               idx = findfirst(notseen)\n               isnothing(idx) || (N[lN+1] = idx-1)\n           end\n           return count(==(0), N)\n       end\nproblem3 (generic function with 1 method)\n\njulia> @btime problem3(10000)\n  20.714 ms (3 allocations: 88.17 KiB)\n1475\n

\n