CUDA.jl：如何优化CuArrays数据拷贝的时间效率？

问题背景 #

我正在用Julia求解受限n体问题（引力场已通过解析法求解），场景为包含太阳、木星及大量小行星/火箭的2D太阳系。借助CUDA.jl的高阶抽象，已实现小行星传播过程的并行化，但每时间步后将所有小行星X、Y位置存入二维数组（每行对应一个时间步的位置）的数据拷贝与存储操作，成为程序性能瓶颈。

当前瓶颈代码 #

目前采用先将CuArray转换为普通Julia数组，再逐元素写入二维数组的方式：

particleXsRAM=Array(particleXs)
particleYsRAM=Array(particleYs)
for count in 1:numberParticles
    particleXHistories[step,count]=particleXsRAM[count]
    particleYHistories[step,count]=particleYsRAM[count]
end

已尝试的无效方案 #

• 创建二维CuArray逐元素拷贝：实际由CPU执行，无性能提升
• 创建CuArray数组：后续绘图耗时过高，性价比低

当前代码虽比非并行版本有性能提升，但未达预期，希望通过CUDA.jl的高阶抽象并行化该存储过程以优化性能。

完整代码示例 #

function MassCoordsCircle(rotation::AbstractFloat,radius::AbstractFloat)
    return radius*cos(rotation),radius*sin(rotation)

end

function FindForcesArr(rocketCoord::AbstractArray,mass1Coord::AbstractFloat,mass2Coord::AbstractFloat,rocketOtherCoord::AbstractArray,mass1OtherCoord::AbstractFloat,mass2OtherCoord::AbstractFloat,mass1::AbstractFloat,mass2::AbstractFloat,G::AbstractFloat)
    return -G*mass1*(rocketCoord.-mass1Coord) ./ sqrt.((rocketCoord.-mass1Coord) .^ 2 .+ (rocketOtherCoord.-mass1OtherCoord) .^2 ) .^ 3 .- G*mass2*(rocketCoord.-mass2Coord) ./ sqrt.((rocketCoord.-mass2Coord) .^ 2 .+ (rocketOtherCoord.-mass2OtherCoord) .^ 2) .^3

end

function TaylorArr(x::AbstractArray,y::AbstractArray,vx::AbstractArray,vy::AbstractArray,mass1::AbstractFloat,mass2::AbstractFloat,dMass1::AbstractFloat,dMass2::AbstractFloat,G::AbstractFloat,mass1Rotation::AbstractFloat,mass2Rotation::AbstractFloat,period::AbstractFloat,dt::AbstractFloat)
    mass1x,mass1y=MassCoords(mass1Rotation,dMass1)
    mass2x,mass2y=MassCoords(mass2Rotation,dMass2)
    ax=FindForcesArr(x,mass1x,mass2x,y,mass1y,mass2y,mass1,mass2,G)
    ay=FindForcesArr(y,mass1y,mass2y,x,mass1x,mass2x,mass1,mass2,G)
    return x .+ dt .* vx .+ dt^2/2 .* ax,y .+ dt.*vy .+ dt^2/2 .* ay,vx .+ dt .* ax,vy .+ dt .* ay
    
end

using CUDA
using Plots
using PyPlot
using PyCall
np=pyimport("numpy")

G=6.6726e-11
#Mass of Sun
mass1=1.989e30
#Mass of Jupiter
mass2=1.898e27
#Sun-Jupiter distance
d_mass1_mass2=7.7834082e11
timeStep=100000.0
numberParticles=2000

d_COM_mass1=d_mass1_mass2*mass2/(mass1+mass2)
d_COM_mass2=d_mass1_mass2*mass1/(mass1+mass2)
MassCoords=MassCoordsCircle
PropagateParticle=TaylorArr

period=sqrt(4*pi^2/(G*(mass1+mass2))*d_mass1_mass2^3)


steps=trunc(Int64,div(period,timeStep)+1)
mass1xHistory=zeros(steps)
mass1yHistory=zeros(steps)
mass2xHistory=zeros(steps)
mass2yHistory=zeros(steps)

particleXs=CUDA.randn(Float64,numberParticles).*1e12
particleYs=CUDA.randn(Float64,numberParticles).*1e12
particleVXs=CUDA.randn(Float64,numberParticles).*15000
particleVYs=CUDA.randn(Float64,numberParticles).*15000

particleXHistories=Array{Float64}(undef,steps,numberParticles)
particleYHistories=Array{Float64}(undef,steps,numberParticles)

time=0
step=1
while step<=steps
    #Sun is 180 degrees out of sync with Jupiter.
    mass1Rotation=time/period*2*pi+pi
    mass2Rotation=time/period*2*pi
    mass1x,mass1y=MassCoords(mass1Rotation,d_COM_mass1)
    mass2x,mass2y=MassCoords(mass2Rotation,d_COM_mass2)

    mass1xHistory[step]=mass1x
    mass1yHistory[step]=mass1y
    mass2xHistory[step]=mass2x
    mass2yHistory[step]=mass2y

    particleXsRAM=Array(particleXs)
    particleYsRAM=Array(particleYs)
    for count in 1:numberParticles
        particleXHistories[step,count]=particleXsRAM[count]
        particleYHistories[step,count]=particleYsRAM[count]

    end

    particleXs,particleYs,particleVXs,particleVYs=PropagateParticle(particleXs,particleYs,particleVXs,particleVYs,mass1,mass2,d_COM_mass1,d_COM_mass2,G,mass1Rotation,mass2Rotation,period,timeStep)

    time=time+timeStep
    step=step+1
        
end

plt.figure(1,figsize=[7,7])
plt.plot(mass1xHistory,mass1yHistory)
plt.plot(mass2xHistory,mass2yHistory)
plt.plot(np.array(particleXHistories),np.array(particleYHistories))
plt.xlim([-9e11,9e11])
plt.ylim([-9e11,9e11])
plt.show()

注：小行星初始位置与速度随机生成，选用Matplotlib而非Plots是因其绘图速度更快。

优化方案 #

1. 利用广播机制替代手动循环 #

直接将CuArray的内容拷贝到CPU数组的对应行，避免逐元素循环的开销：

# 替换原循环代码
particleXHistories[step, :] .= Array(particleXs)
particleYHistories[step, :] .= Array(particleYs)

Julia的广播机制会优化赋值操作，比手动循环效率更高，同时减少代码冗余。

2. 使用copyto!直接拷贝（推荐） #

通过copyto!函数结合view创建的行视图，直接将GPU数据拷贝到CPU数组的指定位置，避免中间数组的创建：

# 替换原循环代码
copyto!(view(particleXHistories, step, :), particleXs)
copyto!(view(particleYHistories, step, :), particleYs)

这种方式减少了一次内存分配（无需创建particleXsRAM中间数组），直接完成设备到主机的内存拷贝，性能提升更明显。

3. 延迟拷贝（显存充足时适用） #

如果GPU显存足够存储所有时间步的历史数据，可以先在GPU上存储历史，模拟完成后一次性拷贝到CPU：

# 初始化GPU上的历史数组（行=时间步，列=粒子）
particleXHistoriesGPU = CuArray{Float64}(undef, steps, numberParticles)
particleYHistoriesGPU = CuArray{Float64}(undef, steps, numberParticles)

# 模拟过程中直接写入GPU数组
particleXHistoriesGPU[step, :] .= particleXs
particleYHistoriesGPU[step, :] .= particleYs

# 模拟结束后一次性拷贝到CPU
particleXHistories = Array(particleXHistoriesGPU)
particleYHistories = Array(particleYHistoriesGPU)

此方法避免了每步的CPU-GPU数据交互，适合粒子数和时间步规模适中的场景；若显存不足，可分批次拷贝（如每100步拷贝一次到CPU）。

内容的提问来源于stack exchange，提问作者TimTheEnchanter