对数组操作 并行化

综合讨论区
1

有一个三维的tensor，我想对每个tensor[i, j, :]序列都进行去重复求其均值，写成函数如下：

function eliminate(tensor::AbstractArray{T, 3}) where T
    dims_1, dims_2, dims_3 = size(tensor)
    mean_array = similar(tensor, dims_1, dims_2)

    @inbounds @views for i in 1: dims_1, j in 1: dims_2
        mean_array[i, j] = unique(tensor[i, j, :]) |> mean
    end

    return mean_array
end

我想对这段代码进行并行化，提高运行速度，第一次接触并行化， 请问有什么好的建议吗？

2

我后来改写了一下：

function eliminate2(tensor::AbstractArray{T, 3}) where T
    dims_1, dims_2, dims_3 = size(tensor)
    mean_array = zeros(T, dims_1, dims_2)

    @inbounds for i in 1: dims_1, j in 1: dims_2
        hash_tab = Set{T}()
        clock = 0
        for k in 1: dims_3
            if tensor[i, j, k] ∉ hash_tab
                push!(hash_tab, tensor[i, j, k])
                mean_array[i, j] += tensor[i, j, k]
            clock += 1
            end
        end
        mean_array[i, j] /= clock
    end

    return mean_array
end

这个版本的函数比上面的版本快了不少，但还是不够。我想说如果对这段代码进行并行化处理的话，是不是可以加速。但是由于涉及到对数组的操作，一头雾水。

3 

-    @inbounds for i in 1: dims_1, j in 1: dims_2
+   Threads.@threads for ind in CartesianIndices((1: dims_1, 1: dims_2))
+        i, j = ind.I

这里你对每个 element 的操作都是独立的，直接上 threads就可以了

1 个赞
4

刚想到了，如果你用 GPU 加速的话，应该也是可以的，先 sort 下，然后写个kernel的函数，遇到重复的就跳过，应该可以轻松实现加速

1 个赞
5

好的，非常感谢，我去尝试一下