有好心人教教我怎么使用@threads加速我的代码吗？

综合讨论区 Julia入门
1

我在建模的过程中遇到类似下面这些代码的问题：

using SparseArrays, Base.Threads
function t_s(i::Int)
	length_I = round(Int, i + rand())
	o1 = fill(i, length_I)
	o2 = fill(i^2, length_I)
	o3 = fill(i^3, length_I)
	sleep(2.5e-7)
	return o1, o2, o3
end
function t_a(x::Int)
	o1 = Vector{Int}(undef, x^2)
	o2 = Vector{Int}(undef, x^2)
	o3 = Vector{Int}(undef, x^2)
	m::Int = 0
	for i in 1:x
		o1s, o2s, o3s = t_s(i)
		length_I = length(o1s)
		o1[m+1:m+length_I] = o1s
		o2[m+1:m+length_I] = o2s
		o3[m+1:m+length_I] = o3s
		m += length_I
	end
	o1 = o1[1:m]
	o2 = o2[1:m]
	o3 = o3[1:m]
	sparse(o1, o2, o3)
end
t_a(10000)

我想用@threads加速，但不知道怎么写，有人教教我吗

2

Julia中文文档-多线程

并行计算的最大线程数=CPU的核心数，因为只有核心可以并行计算，同一核上的多线程是不能并行的。（https://blog.csdn.net/qq_41607336/article/details/118180220）

查一下你的CPU有几个核，然后以最大核数启动julia，并将for i in 1:x替换成Threads.@threads for i in 1:x运行即可。参考顶部的教程。

在for循环内添加println("Working on thread $(Threads.threadid())")来展示这次循环在哪个线程上工作。

如果你使用vscode，在设置内搜索julia: Num Threads，点击在设置中编辑，修改"julia.NumThreads": 0,这里的数字即可。

另参考 Julia中文文档-并行计算

此外你有责任确保程序没有数据竞争，如果你不遵守该要求，就得不到正确的结果。在下述代码中，o1,o2,o3被多个线程抢用，所以这段代码在多线程时可能需要修改。

o1[m+1:m+length_I] = o1s
o2[m+1:m+length_I] = o2s
o3[m+1:m+length_I] = o3s

可以将o1分配成一个矩阵，每个线程用一列：o1 = Matrix{Int}(undef, x^2, x)

o1[m+1:m+length_I, i] = o1s
3

感谢您的回复。根据你的提示，我修改的t_a函数，（我同时修改了t_s和t_a的逻辑，以便更清晰地检查输出）

function t_s(i::Int)
	#length_I = round(Int, 3 + rand())
	length_I = 3
	o1 = fill(i, length_I)
	o2 = fill(i^2, length_I)
	o3 = fill(i^3, length_I)
	sleep(2.5e-3)
	return o1, o2, o3
end
function t_a(x::Int)
	o1 = zeros(Int, 4x)
	o2 = zeros(Int, 4x)
	o3 = zeros(Int, 4x)
	m::Int = 0
	for i in 1:x
		o1s, o2s, o3s = t_s(i)
		length_I = length(o1s)
		o1[m+1:m+length_I] = o1s
		o2[m+1:m+length_I] = o2s
		o3[m+1:m+length_I] = o3s
		m += length_I
	end
	o1 = o1[1:m]
	o2 = o2[1:m]
	o3 = o3[1:m]
	return o1, o2, o3
end

上面是原问题，下面的t_a1函数是根据您的提示写的：

function t_a1(x::Int)
    o1 = zeros(Int, 4, x)
    o2 = zeros(Int, 4, x)
    o3 = zeros(Int, 4, x)
    @threads for i in 1:x
        o1s, o2s, o3s = t_s(i)
        length_I = length(o1s)
        o1[1:length_I, i] = o1s
        o2[1:length_I, i] = o2s
        o3[1:length_I, i] = o3s
    end
	nz_idx = findall(i -> i != 0, o1)
    return o1[nz_idx], o2[nz_idx], o3[nz_idx]
end

下面是测试：

julia> using Base.Threads, BenchmarkTools

julia> nthreads()
8

julia> @btime t_a(3)
  12.221 ms (42 allocations: 1.88 KiB)
([1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 1, 1, 4, 4, 4, 9, 9, 9], [1, 1, 1, 8, 8, 8, 27, 27, 27])

julia> @btime t_a1(3)
  3.426 ms (89 allocations: 6.61 KiB)
([1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 1, 1, 4, 4, 4, 9, 9, 9], [1, 1, 1, 8, 8, 8, 27, 27, 27])

可以看到确实完成任务，并且有加速。
但我觉得，不把o1、o2、o3初始化成矩阵是不是也可以啊？即：

function t_a2(x::Int)
	o1 = zeros(Int, 4x)
	o2 = zeros(Int, 4x)
	o3 = zeros(Int, 4x)
    m = Atomic{Int}(0)
    @threads for i in 1:x
        o1s, o2s, o3s = t_s(i)
        length_I = length(o1s)
        old_m = atomic_add!(m, length_I)
        o1[old_m+1:old_m+length_I] = o1s
        o2[old_m+1:old_m+length_I] = o2s
        o3[old_m+1:old_m+length_I] = o3s
    end
    return o1[1:m[]], o2[1:m[]], o3[1:m[]]
end

julia> @btime t_a2(3)
  3.297 ms (85 allocations: 6.27 KiB)
([1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 1, 1, 4, 4, 4, 9, 9, 9], [1, 1, 1, 8, 8, 8, 27, 27, 27])

好像也完成了任务。

4

使用 原子操作自然是更好的，而且可以发现原子操作的3.297 ms (85 allocations: 6.27 KiB) < 矩阵法 3.426 ms (89 allocations: 6.61 KiB)。（给你发帖前我还没看原子操作那节，刚看了下发现确实不错）
感觉这个话题挺有意思的，楼主的循环是最大3次，t_a1(3)，但是却分配了8线程，有没有可能分配3线程更快？
我用三线程运行了楼主的代码，@btime的结果为

6.555 ms (44 allocations: 3.84 KiB)

用8线程运行：

9.434 ms (69 allocations: 7.00 KiB)

我猜用可以被for循环整除的线程效率会更高一些。

5

顺带一提，threads 适用于 cpu 运算密集型任务。如果这里 sleep 省略的是 IO 密集型任务，比如读写下载，那可以用 async 协程并发来加速。

1 个赞