如何加速这段FDTD程序？

Question

代码基本是下面一坨了，目前的速度其实还可以了，用了With做宏替换，Compile编译到机器码。
目前感觉还能提速的地方是每个Do循环，因为每个Do循环其实跟 i、j、k 的顺序是无关的，所以在想有没有可能并行起来。但是在Compile内部并行我不知道怎么弄。

Ni = 50;
Nj = 80;
Nk = 50;
di = 5/50;
dj = 2 Pi/Nj;
dk = 5/50;
f = 150 10^6;
μ = 4.0 Pi 1.0 10^-7;
X = 1/f;
ks = 2*Pi^2*f^2;
X = 1/f;
cc = 2.99792458 10^8;
Δt = 
  0.99/cc /Sqrt[1.0/di^2 + 1.0/(0.5*di dj)^2 + 1.0/dk^2];
mat = ConstantArray[0, {Ni + 1, Nj, Nk + 2}];
{EiMat, EjMat, EkMat, HiMat, HjMat, HkMat} = ConstantArray[mat, 6];

compute =
  With[{
    miu = 4.0*Pi*1.0 10^-7,
    ϵ = 1.0/(cc cc μ),
    σ = 10^-6,
    dt = Δt,
    di = di,
    dj = dj,
    dk = dk,
    Ni = Ni,
    Nj = Nj,
    Nk = Nk,
    Δt = Δt,
    ks = ks,
    X = X,
    f = f
    
    },
   Compile[{{t, _Real},
     {matrix, _Real, 4}},
    Module[
     {
      pulse,
      sum,
      CA,
      CB,
      EiMat,
      EjMat,
      EkMat,
      HiMat,
      HjMat,
      HkMat
      },
     {EiMat, EjMat, EkMat, HiMat, HjMat, HkMat} = matrix;
     CA = (2 ϵ - σ Δt)/(
      2 ϵ + σ Δt);
     CB = (2 Δt)/(
      2 ϵ + σ Δt);
     (*Do[*)
     pulse = 2*ks*Sqrt[Exp[1/2/ks]]*Exp[-ks*(t - X)^2]*(t - X);
     (***********************)
     EkMat[[10, 10, 25]] += pulse;
     Do[
      EiMat[[i, j, 
        k]] = (CA EiMat[[i, j, k]] + 
         CB ((HkMat[[i, j, k]] - 
             HkMat[[i, If[j == 1, Nj, j - 1], 
              k]])/((i - 1/2) di dj) - (
            HjMat[[i, j, k]] - HjMat[[i, j, k - 1]])/dk)),
      {i, 1, Ni},
      {j, 1, Nj},
      {k, 2, Nk}
      ];
     
     
     Do[
      EjMat[[i, j, k]] = 
        CA EjMat[[i, j, k]] + 
         CB ((HiMat[[i, j, k]] - HiMat[[i, j, k - 1]])/dk - (
            HkMat[[i, j, k]] - HkMat[[i - 1, j, k]])/di);
      ,
      {i, 2, Ni},
      {j, 1, Nj},
      {k, 2, Nk}];
     
     Do[
      EkMat[[i, j, k]] = 
       CA EkMat[[i, j, k]] + 
        CB (((i - 1/2) HjMat[[i, j, k]] - (i - 3/2) HjMat[[i - 1, j, 
              k]])/((i - 1) di) - (
           HiMat[[i, j, k]] - HiMat[[i, If[j == 1, Nj, j - 1], k]])/(
           di dj (i - 1))),
      {i, 2, Ni},
      {j, 1, Nj},
      {k, 2, Nk}
      ];
     
     (*Ek i=1*)
     Do[
      sum = Sum[HjMat[[1, m, k]], {m, 1, Nj}];
      Do[
       EkMat[[1, j, k]] = CA EkMat[[1, j, k]] + CB (4/(di Nj)) sum,
       {j, 1, Nj}],
      {k, 1, Nk}];
     
     
     Do[
      HiMat[[i, j, k]] = (
        HiMat[[i, j, k]] +
         dt/
          miu ((EjMat[[i, j, k + 1]] - EjMat[[i, j, k]])/dk - (
            EkMat[[i, If[j == Nj, 1, j + 1], k]] - 
             EkMat[[i, j, k]])/((i - 1) di dj))),
      {i, 2, Ni},
      {j, 1, Nj},
      {k, 2, Nk}];
     
     
     Do[
      HkMat[[i, j, k]] =
       HkMat[[i, j, k]] +
        dt/
         miu ((EiMat[[i, If[j == Nj, 1, j + 1], k]] - 
            EiMat[[i, j, k]])/((i - 1/2) di dj) - (
           
           i EjMat[[i + 1, j, k]] - (i - 1) EjMat[[i, j, k]])/((i - 1/
              2) di ));
      HjMat[[i, j, k]] = 
       HjMat[[i, j, k]] + 
        dt/miu ((EkMat[[i + 1, j, k]] - EkMat[[i, j, k]])/di - (
           EiMat[[i, j, k + 1]] - EiMat[[i, j, k]])/dk)
      ,
      {i, 1, Ni},
      {j, 1, Nj},
      {k, 2, Nk}];
     (*,{t,0,50 10^-9,Δt}];*)
     {EiMat, EjMat, EkMat, HiMat, HjMat, HkMat}
     ], CompilationTarget -> "C",
    RuntimeOptions -> "Speed"
    ]
   ];
n = 0; PrintTemporary[Dynamic[n]];
Do[If[Mod[n, 50] == 0, EkMat2 = EkMat; nn = n]; 
 n++; {EiMat, EjMat, EkMat, HiMat, HjMat, HkMat} = 
  compute @@ {t, {EiMat, EjMat, EkMat, HiMat, HjMat, HkMat}},
 {t, 0, 50 10^-9, Δt}]

 

Comments

为了能在外面监视计算过程，50步画一次图什么的。

---

……我觉得你也是知道的不过还是说给其他读者听下：如果是调试阶段这么做倒也无妨，真要提速的话就不能这么做。

---

为什么Compile里第一个Do里，有个除以(i - 1/2) di dj，这里会包含i循环体？难道是我好久没学物理就忘了么。。。

---

这是柱坐标啊，phi方向的步长要用到r方向

---

哦，果然是我的错。忘了柱坐标了。

Answer 1

……这帖一直没人回啊，那我来消灭零回复好了。

就我所知，这帖（就是我在上面评论里给的那帖）已经包含了目前已知的全部（至少是有明确效果的）可以用在FDM中的速度优化技巧，而你也已经在你的代码中使用了其中的大部分，除了Compile`GetElement。这个也是可以方便地用代码重写技术完成的。顺便感觉上你的代码使用代码重写的话应该可以再简洁一些，不过柱坐标下各个方程是没直角坐标这么规整……啊~我不确定，总之你有兴趣可以看看这帖。下面就放代码——啊，先给出你现有代码的测速：

(*时间循环我已放进代码里了，循环50次*)

origin = Hold@With[{miu = 4.0*Pi*1.0 10^-7, ϵ = 1.0/(cc cc μ), σ = 
     10^-6, dt = Δt, di = di, dj = dj, dk = dk, Ni = Ni, Nj = Nj, 
    Nk = Nk, Δt = Δt, ks = ks, X = X, f = f}, 
   Compile[{{matrix, _Real, 4}}, 
    Module[{pulse, sum, CA, CB, EiMat, EjMat, EkMat, HiMat, HjMat, 
      HkMat}, {EiMat, EjMat, EkMat, HiMat, HjMat, HkMat} = matrix;
     CA = (2 ϵ - σ Δt)/(2 ϵ + σ Δt);
     CB = (2 Δt)/(2 ϵ + σ Δt);
     Do[pulse = 2*ks*Sqrt[Exp[1/2/ks]]*Exp[-ks*(t - X)^2]*(t - X);
      (***********************)
      EkMat[[10, 10, 25]] += pulse;
      Do[EiMat[[i, j, k]] = (CA EiMat[[i, j, k]] + 
          CB ((HkMat[[i, j, k]] - 
                HkMat[[i, If[j == 1, Nj, j - 1], k]])/((i - 1/2) di dj) - (HjMat[[i, j, 
                  k]] - HjMat[[i, j, k - 1]])/dk)), {i, 1, Ni}, {j, 1, Nj}, {k, 2, Nk}];
      Do[EjMat[[i, j, k]] = 
         CA EjMat[[i, j, k]] + 
          CB ((HiMat[[i, j, k]] - HiMat[[i, j, k - 1]])/
              dk - (HkMat[[i, j, k]] - HkMat[[i - 1, j, k]])/di);, {i, 2, Ni}, {j, 1, 
        Nj}, {k, 2, Nk}];
      Do[EkMat[[i, j, k]] = 
        CA EkMat[[i, j, k]] + 
         CB (((i - 1/2) HjMat[[i, j, k]] - (i - 3/2) HjMat[[i - 1, j, k]])/((i - 
                 1) di) - (HiMat[[i, j, k]] - 
               HiMat[[i, If[j == 1, Nj, j - 1], k]])/(di dj (i - 1))), {i, 2, Ni}, {j, 1,
         Nj}, {k, 2, Nk}];
      (*Ek i=1*)Do[sum = Sum[HjMat[[1, m, k]], {m, 1, Nj}];
       Do[EkMat[[1, j, k]] = CA EkMat[[1, j, k]] + CB (4/(di Nj)) sum, {j, 1, Nj}], {k, 
        1, Nk}];
      Do[HiMat[[i, j, k]] = (HiMat[[i, j, k]] + 
          dt/miu ((EjMat[[i, j, k + 1]] - EjMat[[i, j, k]])/
              dk - (EkMat[[i, If[j == Nj, 1, j + 1], k]] - 
                EkMat[[i, j, k]])/((i - 1) di dj))), {i, 2, Ni}, {j, 1, Nj}, {k, 2, Nk}];
      Do[HkMat[[i, j, k]] = 
        HkMat[[i, j, k]] + 
         dt/miu ((EiMat[[i, If[j == Nj, 1, j + 1], k]] - 
               
               EiMat[[i, j, k]])/((i - 
                 1/2) di dj) - (i EjMat[[i + 1, j, k]] - (i - 1) EjMat[[i, j, k]])/((i - 
                 1/2) di));
       HjMat[[i, j, k]] = 
        HjMat[[i, j, k]] + 
         dt/miu ((EkMat[[i + 1, j, k]] - EkMat[[i, j, k]])/
             di - (EiMat[[i, j, k + 1]] - EiMat[[i, j, k]])/dk), {i, 1, Ni}, {j, 1, 
        Nj}, {k, 2, Nk}];
      , {t, 0, 50 Δt, Δt}]; {EiMat, EjMat, EkMat, HiMat, 
      HjMat, HkMat}], CompilationTarget -> "C", RuntimeOptions -> "Speed"]];

compute2 = ReleaseHold@origin;
compute2[{EiMat, EjMat, EkMat, HiMat, HjMat, HkMat}]; // AbsoluteTiming

(* {3.026318, Null} *)

然后是把Compile`GetElement用上的版本：

compute3 = ReleaseHold@
   With[{cg = Compile`GetElement}, 
    origin /. Part -> cg /. 
        HoldPattern@(h : Set | AddTo)[cg@a__, b_] :> h[Part@a, b]];

compute3[{EiMat, EjMat, EkMat, HiMat, HjMat, HkMat}]; // AbsoluteTiming

(* {0.886446, Null} *)

执行时间缩到原来的三分之一了嗯。

Comments

很棒的元编程例子。excellent showcase for meta-programming。