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.

4 files changed, 40 insertions, 2 deletions

diff --git a/Tools/CodeGen/Calculation.m b/Tools/CodeGen/Calculation.m
index c74c44d..6c929a6 100644
--- a/Tools/CodeGen/Calculation.m
+++ b/Tools/CodeGen/Calculation.m
@@ -30,6 +30,7 @@ GetCalculationParameters;
 CalculationStencilSize;
 CalculationOnDevice;
 GetCalculationWhere;
+SplitCalculations;
 
 Begin["`Private`"];
 
@@ -101,6 +102,40 @@ DefFn[
   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
index 5b8109c..5cd0412 100644
--- a/Tools/CodeGen/CalculationFunction.m
+++ b/Tools/CodeGen/CalculationFunction.m
@@ -144,7 +144,7 @@ VerifyCalculation[calc_] :=
          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
index 29bee44..5664372 100644
--- a/Tools/CodeGen/Kranc.m
+++ b/Tools/CodeGen/Kranc.m
@@ -43,7 +43,7 @@ Boundary, Interior, InteriorNoSync, Where, AddToStencilWidth,
 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
index aaea594..09f5ff3 100644
--- a/Tools/CodeGen/KrancThorn.m
+++ b/Tools/CodeGen/KrancThorn.m
@@ -104,6 +104,9 @@ CreateKrancThorn[groupsOrig_, parentDirectory_, thornName_, opts:OptionsPattern[
     cktCheckNamedArgs[{opts}];
 
     calcs = OptionValue[Calculations];
+
+    calcs = SplitCalculations[calcs];
+
     calcs = Map[mapReplaceAdd[#, Shorthands, Join[lookup[#,Shorthands,{}],OptionValue[Shorthands]]] &, calcs];
 
     declaredGroups = OptionValue[DeclaredGroups];