MPI环境下三角矩阵For循环迭代的非均匀负载分割方案问询

问题背景 #

现有一款基于MPI与OpenMP的粒子交互计算程序，需遍历满足对称性（$A_{ij}=A_{ji}$）的粒子交互对称矩阵A。当前采用均匀列分割方式将矩阵分配给不同MPI进程，本地通过OpenMP并行处理单列。但受KD树筛选相关交互及矩阵对称性影响，均匀分割导致严重负载不均——后段进程远早于前段完成计算。

问题：若已知首列的最大相关索引数，能否设计前小后大的非均匀分割策略，实现MPI负载均衡？

核心Fortran伪代码（原实现） #

iter_per_rank = n/n_ranks

if ( mod(n, n_ranks) > 0 ) then
    iter_per_rank = iter_per_rank + 1
end if

first_iter = rank * iter_per_rank + 1
last_iter = first_iter + iter_per_rank - 1

!$omp parallel do schedule(guided) &
!$omp private(i,j,idx)
!$omp reduction(+:...)
do i = last_iter, first_iter, -1
    idx = search%kdtree(array(i))
    do k = 1, size(idx%i%values) -1
        j = idx%i%values(k+1)

        if (i .ge. j) then
            cycle
        end if

        [...] ! 粒子交互计算逻辑

    end do
end do
!$omp end parallel do

call mpi_comms(...)

解决方案：前小后大的非均匀分割策略 #

完全可以通过非均匀分割实现负载均衡，核心思路是根据每列的计算负载（KD树返回的相关交互粒子数）分配列数：由于后段进程负载更低（对应大索引列的交互数更少），让靠前的进程分配更少的列，靠后的进程分配更多的列，使每个进程的总负载尽可能接近。

具体实现步骤

1. 负载分布建模：已知首列最大相关索引数$L_{max}$，结合负载前密后疏的特点，可近似第$i$列的负载为：
   $$load(i) = L_{max} \times \frac{n - i + 1}{n}$$
   （如果能提前跑一次程序统计每列实际负载，替换该模型可获得更精准的分割）

2. 计算目标负载：先估算所有列的总负载总和$Total_L$，每个MPI进程的目标负载为$Target_L = Total_L / n_{ranks}$。

3. 非均匀列划分：从第一列开始累加负载，当累加值接近$Target_L$时，将这段列分配给当前进程；重复此过程直到所有列分配完毕，最后剩余的列直接分配给最后一个进程。

修改后的核心伪代码

! 1. 预计算每列负载（基于首列最大值的近似模型）
allocate(load(n))
L_max = ... ! 已知的首列最大相关索引数
do i = 1, n
    load(i) = L_max * real(n - i + 1)/real(n) ! 负载递减模型，可根据实际调整
end do

! 2. 调用非均匀分割子程序，获取当前进程的列范围
call compute_non_uniform_split(n, n_ranks, rank, first_iter, last_iter, load)

! 3. 原有OpenMP并行计算逻辑保持不变
!$omp parallel do schedule(guided) &
!$omp private(i,j,idx)
!$omp reduction(+:...)
do i = last_iter, first_iter, -1
    idx = search%kdtree(array(i))
    do k = 1, size(idx%i%values) -1
        j = idx%i%values(k+1)

        if (i .ge. j) then
            cycle
        end if

        [...] ! 粒子交互计算逻辑

    end do
end do
!$omp end parallel do

call mpi_comms(...)

非均匀分割子程序实现

subroutine compute_non_uniform_split(n, n_ranks, rank, first, last, load)
    integer, intent(in) :: n, n_ranks, rank
    integer, intent(out) :: first, last
    real, intent(in) :: load(n)
    real :: total_load, target_load, current_load
    integer :: i, current_rank, start_idx

    total_load = sum(load)
    target_load = total_load / n_ranks

    current_load = 0.0
    start_idx = 1
    do current_rank = 0, n_ranks-1
        do i = start_idx, n
            current_load = current_load + load(i)
            ! 最后一个进程直接分配剩余所有列，否则达到目标负载就分割
            if (current_rank == n_ranks-1 .or. current_load >= target_load) then
                if (current_rank == rank) then
                    first = start_idx
                    last = i
                end if
                start_idx = i + 1
                current_load = 0.0
                exit
            end if
        end do
    end do
end subroutine compute_non_uniform_split

注意事项 #

• 如果条件允许，建议先执行一次"负载探测"：让单个进程遍历所有列，统计每列的实际交互数，用真实负载数据替代近似模型，能实现最优的负载均衡。
• 原代码中通过if(i .ge. j) cycle跳过了对称部分的计算，确保只处理矩阵的上三角区域，非均匀分割不会破坏这一逻辑，只需保证每个进程的列范围连续即可。

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