Julia 没有使用所有可用的 CPU 线程

Question

我在 Julia 中做了一些计算，发现它运行比 numpy 的对应物慢得多（大约 25 倍！）。

然后我意识到 Julia 只使用了 8 个线程，总共 96 CPU 个线程（48 个物理内核）我的电脑，而 numpy 似乎可以很好地利用 70 个线程。

运行 Julia 与 $julia --thread 96 argument 没有任何区别，即使 julia> Threads.nthreads() returns 96.

此外，结果有点令人失望，我怀疑 Julia 使用所有 96 个线程仍然可能无法与 numpy 的速度相匹配。

这是 Julia 代码。我只是用 julia> @time calc_fid(mat_a, mat_b) 测量时间，这给了我 90 秒的平均值。

using Statistics
using LinearAlgebra

function calc(A::Array{Float32,2}, B::Array{Float32,2})
    μ_A = mean(A, dims=2)
    μ_B = mean(B, dims=2)
    σ_A = cov(A, dims=2)
    σ_B = cov(B, dims=2)

    ssdiff = sum((μ_A - μ_B).^2)
    covmean = sqrt(σ_A * σ_B)
    res = ssdiff + tr(σ_A .+ σ_B .- 2.0 * covmean)
    
    return res
end

这是平均需要 3.5 秒的 numpy 代码。用 time.perf_counter()

测量

import numpy as np

from numpy import cov
from numpy import trace
from scipy.linalg import sqrtm

def calc(A, B):
    mu_A = A.mean(axis=0)
    mu_B = B.mean(axis=0)
    sigma_A = cov(A, rowvar=False)
    sigma_B = cov(B, rowvar=False)

    ssdiff = np.sum((mu_A - mu_B) ** 2.0)
    covmean = sqrtm(sigma_A.dot(sigma_B))
    res = ssdiff + trace(sigma_A + sigma_B - 2.0 * covmean)

    return res

任何 suggestion/explanation 将不胜感激！

Answer 1

您的 Julia 示例代码实际上并未使用 Julia 的多线程功能。

在 Julia 中，如果你想要多线程，你通常需要显式地启用它。 mean 或 cov 之类的函数不会仅仅凭借以多线程启动 Julia 就自动成为多线程；如果您希望这些函数是多线程的，您将需要编写它们的多线程版本，或者使用已经这样做的包。

在您的 Julia 示例代码中，唯一会像当前编写的那样完全多线程的是线性代数，因为它会回退到 BLAS（numpy 确实很可能这样做），它有自己的多线程系统完全独立于 Julia 的多线程（BLAS 是用 Fortran 编写的）。不过，这也可能没有使用所有 96 个线程，并且在您的情况下几乎可以肯定默认为 8。可以使用

检查 Julia 中 BLAS 使用的线程数

using LinearAlgebra
BLAS.get_num_threads()

或类似地设置为

BLAS.set_num_threads(n)

Julia 没有使用所有可用的 CPU 线程

Julia is not using all of the available CPU threads

performance

benchmarking

multithreading

numpy

julia