Add SplitBy option to calculations to split calculations up by variable

This is not quite perfect yet; scheduling statements which use the original calculation name will not see the new calculation names. We should introduce a group for the original calculation name.
ianhinder · Apr 2, 2012 · e11f056 · e11f056
1 parent 28dad7e
commit e11f056
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 2 deletions.
diff --git a/Tools/CodeGen/Calculation.m b/Tools/CodeGen/Calculation.m
@@ -30,6 +30,7 @@
 CalculationStencilSize;
 CalculationOnDevice;
 GetCalculationWhere;
+SplitCalculations;
 
 Begin["`Private`"];
 
@@ -101,6 +102,40 @@
   CalculationOnDevice[calc_List] :=
   lookupDefault[calc, ExecuteOn, Automatic] === Device];
 
+partialCalculation[calc_, suffix_, updates_, evolVars_] :=
+Module[
+  {name, calc1, replaces, calc2, vars, patterns, eqs, calc3},
+  (* Add suffix to name *)
+  name     = lookup[calc, Name] <> suffix;
+  calc1    = mapReplace[calc, Name, name];
+  (* Replace some entries in the calculation *)
+  replaces = updates //. (lhs_ -> rhs_) -> (mapReplace[#, lhs, rhs]&);
+  calc2 = Apply[Composition, replaces][calc1];
+  (* Remove unnecessary equations *)
+  vars     = Join[evolVars, lookup[calc2, Shorthands]];
+  patterns = Replace[vars, {    Tensor[n_,__]  ->     Tensor[n,__] ,
+                            dot[Tensor[n_,__]] -> dot[Tensor[n,__]]}, 1];
+  eqs      = FilterRules[lookup[calc, Equations], patterns];
+  calc3    = mapReplace[calc2, Equations, eqs];
+  calc3
+];
+
+DefFn[
+  SplitCalculations[calcs_List] :=
+  Flatten[SplitCalculation/@calcs,1]];
+
+DefFn[
+  SplitCalculation[calc_] :=
+  Module[
+    {splitBy = lookup[calc,SplitBy, {}]},
+    If[splitBy === {},
+       {calc},
+       Table[partialCalculation[calc, 
+                                "_"<>StringReplace[ToString[var],{"["->"","]"->"",","->""}],
+                                {},
+                                {var}]~Join~{CachedVariables -> {(* var[[1]] *)}}, (* This is not general *)
+             {var, splitBy}]]]];
+
 End[];
 
 EndPackage[];
diff --git a/Tools/CodeGen/CalculationFunction.m b/Tools/CodeGen/CalculationFunction.m
@@ -144,7 +144,7 @@ indentation after each line break (this will push the line length
          PartialDerivatives, PreDefinitions, Schedule,Equations,
          Shorthands, ConditionalOnKeyword, Before, After,
          ConditionalOnTextuals, Where, ConditionalOnKeywords,
-         CollectList, AllowedSymbols, ApplyBCs, Conditional, CachedVariables};
+         CollectList, AllowedSymbols, ApplyBCs, Conditional, CachedVariables, SplitBy};
 
     usedKeys = Map[First, calc];
     unknownKeys = Complement[usedKeys, allowedKeys];

diff --git a/Tools/CodeGen/Kranc.m b/Tools/CodeGen/Kranc.m
@@ -43,7 +43,7 @@
 Everywhere, normal1, normal2, normal3, INV, SQR, CUB, QAD, dot, pow,
 exp, dt, dx, dy, dz, idx, idy, idz, t, MinMod, VanLeer, BodyFunction,
 CallerFunction, LoopFunction, GFAccessFunction, InitFDVariables,
-MacroPointer, CachedVariables}
+MacroPointer, CachedVariables, SplitBy}
 
 {ConditionalOnKeyword, ConditionalOnKeywords, CollectList, Interior,
 InteriorNoSync, Boundary, BoundaryWithGhosts, Where, PreDefinitions,

diff --git a/Tools/CodeGen/KrancThorn.m b/Tools/CodeGen/KrancThorn.m
@@ -104,6 +104,9 @@ Thorn generation (main entry point for non-tensorial thorns)
     cktCheckNamedArgs[{opts}];
 
     calcs = OptionValue[Calculations];
+
+    calcs = SplitCalculations[calcs];
+
     calcs = Map[mapReplaceAdd[#, Shorthands, Join[lookup[#,Shorthands,{}],OptionValue[Shorthands]]] &, calcs];
 
     declaredGroups = OptionValue[DeclaredGroups];