迭代依赖型MATLAB光束传播模拟代码的加速与集群核心利用问题

问题描述 #

我是一名物理本科生，现有一段用于光学光束传播模拟的MATLAB代码，代码内包含大量for循环，但循环的第n+1次迭代依赖第n次迭代的计算结果，因此无法用parfor替代for循环。我在集群中分别申请28核和56核以batch方式提交任务，发现运行时间无差异，代码仅使用了1个核心。请问若无法使用并行计算，还有哪些加速方法？

补充代码片段（耗时占比最高） #

MATLAB Profiler显示耗时集中在以下片段：

片段1（主时间循环内的双层循环）

for j = 3:Nz
    %u(1,j) = u_o(1,j-2) + (-1i/(2*k_0))*(4*(dz(j-2)+dz(j-1))/( (dr(i-1)+dr(i))/2 )^2)*( u_o(2,j-1) - u_o(1,j-1) ) + ...
    %    (dz(j-2)+dz(j-1))*( -1i*w*u_o(i,j-1)*( n_Kerr*I_o(1,j-1) - n_e*(N(1,j-1)+N_ei) - n_h*(N(1,j-1)+N_hi) )/c - ...
    %        u_o(1,j-1)*( alpha(1,j-1) + beta(1,j-1)*I_o(1,j-1) + sigma_e(1,j-1)*(N(1,j-1)+N_ei) + sigma_h(1,j-1)*(N(1,j-1)+N_hi) )/2 ); 
    for i = 2:(Nr-1)
        u(i,j) = u_o(i,j-2) + (-1i/(2*k_0))*((dz(j-2)+dz(j-1))/( (dr(i-1)+dr(i))/2 )^2)*( (1-1/(2i))*u_o(i-1,j-1) + (1+1/(2i))*u_o(i+1,j-1) - 2*u_o(i,j-1) ) + ...
            (dz(j-2)+dz(j-1))*( -1i*w*u_o(i,j-1)*( n_Kerr*I_o(i,j-1) - n_e*(N(i,j-1)+N_ei) - n_h*(N(i,j-1)+N_hi) )/c - ...
            u_o(i,j-1)*( alpha(i,j-1) + beta(i,j-1)*I_o(i,j-1) + sigma_e(i,j-1)*(N(i,j-1)+N_ei) + sigma_h(i,j-1)*(N(i,j-1)+N_hi) )/2 );      
    end
    u(Nr,j) = u(Nr-1,j); u(1,j) = u(2,j); % Neumann boundary  
end

片段2（主时间循环内的双层循环）

for i= 2:Nr-1
    for j = 2:(Nz-1)
    u(i,j) = u(i,j) + v*(u(i,j-1) + u(i,j-1) - 2*u(i,j))/2;
    end
end

上述循环存在迭代依赖，无法并行化。

加速方法 #

1. 向量化改造 #

MATLAB对矩阵/向量运算的底层优化远优于逐元素循环，即使外层存在迭代依赖，内层循环仍可尝试向量化：

• 针对片段1：固定j时，i的遍历可转化为矩阵切片运算。例如，一次性提取u_o(2:Nr,j-1)、u_o(1:Nr-1,j-1)、u_o(3:Nr+1,j-1)（对应u_o(i,j-1)、u_o(i-1,j-1)、u_o(i+1,j-1)），利用广播机制批量计算所有i=2:Nr-1对应的u(i,j)值，替代内层for循环。
• 针对片段2：对每个i，提取u(i,2:Nz-1)向量，直接通过向量运算完成批量更新，改写为：

  for i=2:Nr-1
      u(i,2:Nz-1) = u(i,2:Nz-1) + v*(u(i,1:Nz-2) - u(i,2:Nz-1));
  end
  

  完全去掉内层循环。

2. 预计算重复表达式 #

循环内大量重复计算的表达式可提前批量计算，减少实时运算量：

• 预计算固定常量：如const1 = -1i/(2*k_0);、const2 = 1/c;、complex_coeff = [1-1/(2i), 1+1/(2i), -2];，循环内直接调用。
• 预计算随j变化的项：如dz_sum = dz(1:Nz-2) + dz(2:Nz-1);，循环到j时直接取dz_sum(j-2)。
• 预计算随i变化的项：如dr_inv_sq = 1 ./ ((dr(1:Nr-1)+dr(2:Nr))/2).^2;，循环到i时取dr_inv_sq(i-1)。
• 预计算物理量中间数组：如提前计算n_term = n_Kerr*I_o - n_e*(N+N_ei) - n_h*(N+N_hi);、loss_term = alpha + beta.*I_o + sigma_e.*(N+N_ei) + sigma_h.*(N+N_hi);，循环内直接引用对应位置的值。

3. 内存与访问优化 #

• 预分配内存：在循环开始前为u等大数组分配内存，避免动态扩容的开销。例如：

  u = zeros(Nr, Nz, 'like', u_o); % 与u_o保持相同的数值类型（复数/实数）
  

• 列优先访问：MATLAB采用列优先存储，尽量让内层循环遍历列维度（即j作为内层），减少缓存失效。比如片段2可尝试交换循环顺序（结合向量化效果更佳）：

  for j=2:Nz-1
      u(2:Nr-1,j) = u(2:Nr-1,j) + v*(u(2:Nr-1,j-1) - u(2:Nr-1,j));
  end
  

4. 启用MATLAB内置优化 #

• 开启JIT编译：MATLAB R2020b及以后版本默认启用，若版本较早可通过feature('JIT', 'on')手动开启，加速循环执行。
• 使用accel加速器：在代码开头添加feature('accel', 'on')，启用MATLAB的底层运算加速。
• 优先使用原生复数运算：避免手动拆分实部虚部计算，MATLAB对原生复数运算的优化更充分。

5. MEX编译核心代码 #

将最耗时的循环片段改写为C/C++，编译为MEX文件调用，可借助CPU的多线程优化内层无依赖循环：

• 针对片段1，外层j循环因依赖需串行，但内层i循环可通过OpenMP实现单节点多线程并行。
• 编译时开启优化：使用mex -O3 -largeArrayDims your_code.cpp命令，-O3启用最高级别的编译器优化，-largeArrayDims支持大数组索引。

6. 数值精度调整 #

• 若模拟允许，将数组类型从double改为single，减少内存占用和计算时间：

  u_o = single(u_o);
  u = single(u);
  

  MATLAB对单精度运算的优化同样成熟，能显著提升速度（需确认精度满足物理模拟要求）。

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