Merge vector into master

1 files changed, 36 insertions, 13 deletions

diff --git a/Tools/CodeGen/CalculationFunction.m b/Tools/CodeGen/CalculationFunction.m
index b20587b..d28a8f8 100644
--- a/Tools/CodeGen/CalculationFunction.m
+++ b/Tools/CodeGen/CalculationFunction.m
@@ -198,7 +198,7 @@ localName[x_] :=
 
 definePreDefinitions[pDefs_] :=
   CommentedBlock["Initialize predefined quantities",
-    Map[DeclareAssignVariable["CCTK_REAL", #[[1]], #[[2]]] &, pDefs]];
+    Map[DeclareAssignVariable[DataType[], #[[1]], #[[2]]] &, pDefs]];
 
 (* --------------------------------------------------------------------------
    Equations
@@ -246,11 +246,11 @@ simpCollect[collectList_, eqrhs_, localvar_, debug_] :=
     all];
 
 (* Return a CodeGen block which assigns dest by evaluating expr *)
-assignVariableFromExpression[dest_, expr_, declare_] :=
+assignVariableFromExpression[dest_, expr_, declare_, vectorise_] :=
   Module[{tSym, type, cleanExpr, code},
     tSym = Unique[];
-    type = If[StringMatchQ[ToString[dest], "dir*"], "int", "CCTK_REAL"];
-    cleanExpr = ReplacePowers[expr] /. Kranc`t -> tSym;
+    type = If[StringMatchQ[ToString[dest], "dir*"], "ptrdiff_t", DataType[]];
+    cleanExpr = ReplacePowers[expr, vectorise] /. Kranc`t -> tSym;
 
     If[SOURCELANGUAGE == "C",
       code = If[declare, type <> " ", ""] <> ToString[dest] <> " = " <>
@@ -346,7 +346,7 @@ Options[CreateCalculationFunction] = ThornOptions;
 
 CreateCalculationFunction[calc_, debug_, useCSE_, imp_, opts:OptionsPattern[]] :=
   Module[{gfs, allSymbols, knownSymbols,
-          shorts, eqs, parameters,
+          shorts, eqs, parameters, parameterRules,
           functionName, dsUsed, groups, pddefs, cleancalc, eqLoop, where,
           addToStencilWidth, pDefs, haveCondTextuals, condTextuals},
 
@@ -362,6 +362,8 @@ CreateCalculationFunction[calc_, debug_, useCSE_, imp_, opts:OptionsPattern[]] :
   pDefs = lookup[cleancalc, PreDefinitions];
   haveCondTextuals = mapContains[cleancalc, ConditionalOnTextuals];
 
+  SetDataType[If[OptionValue[UseVectors],"CCTK_REAL_VEC", "CCTK_REAL"]];
+
   VerifyCalculation[cleancalc];
 
   gfs = allGroupVariables[groups];
@@ -386,10 +388,14 @@ CreateCalculationFunction[calc_, debug_, useCSE_, imp_, opts:OptionsPattern[]] :
 
     If[!lookupDefault[cleancalc, NoSimplify, False],
        InfoMessage[InfoFull, "Simplifying equations", eqs];
-       eqs = Simplify[eqs, {r>0}]]];
+       eqs = Simplify[eqs, {r>=0}]]];
 
   InfoMessage[InfoFull, "Equations:"];
 
+  (* Wrap parameters with ToReal *)
+  parameterRules = Map[(#->ToReal[#])&, parameters];
+  eqs = eqs /. parameterRules;
+
   Map[printEq, eqs];
 
   (* Check all the function names *)
@@ -420,7 +426,7 @@ CreateCalculationFunction[calc_, debug_, useCSE_, imp_, opts:OptionsPattern[]] :
        "Calculation is:", cleancalc]];
 
   {
-  DefineFunction[bodyFunctionName, "void", "cGH const * restrict const cctkGH, int const dir, int const face, CCTK_REAL const normal[3], CCTK_REAL const tangentA[3], CCTK_REAL const tangentB[3], int const min[3], int const max[3], int const n_subblock_gfs, CCTK_REAL * restrict const subblock_gfs[]",
+  DefineFunction[bodyFunctionName, "static void", "cGH const * restrict const cctkGH, int const dir, int const face, CCTK_REAL const normal[3], CCTK_REAL const tangentA[3], CCTK_REAL const tangentB[3], int const min[3], int const max[3], int const n_subblock_gfs, CCTK_REAL * restrict const subblock_gfs[]",
   {
     "DECLARE_CCTK_ARGUMENTS;\n",
     "DECLARE_CCTK_PARAMETERS;\n\n",
@@ -440,7 +446,7 @@ CreateCalculationFunction[calc_, debug_, useCSE_, imp_, opts:OptionsPattern[]] :
     CommentedBlock["Include user-supplied include files",
       Map[IncludeFile, lookupDefault[cleancalc, DeclarationIncludes, {}]]],
 
-    InitialiseFDVariables[],
+    InitialiseFDVariables[OptionValue[UseVectors]],
     definePreDefinitions[pDefs],
 
     If[Cases[{pddefs}, SBPDerivative[_], Infinity] != {},
@@ -554,18 +560,18 @@ equationLoop[eqs_, cleancalc_, gfs_, shorts_, incs_, groups_, pddefs_,
     declare = markFirst[First /@ eqsReplaced, Map[localName, gfsInRHS]];
 
     calcCode =
-      MapThread[{assignVariableFromExpression[#1[[1]], #1[[2]], #2], "\n"} &,
+      MapThread[{assignVariableFromExpression[#1[[1]], #1[[2]], #2, OptionValue[UseVectors]], "\n"} &,
        {eqsReplaced, declare}];
 
     GenericGridLoop[functionName,
     {
-      DeclareDerivatives[defsWithoutShorts, eqsOrdered],
+      (* DeclareDerivatives[defsWithoutShorts, eqsOrdered], *)
 
       CommentedBlock["Assign local copies of grid functions",
         Map[DeclareMaybeAssignVariableInLoop[
-              "CCTK_REAL", localName[#], GridName[#],
+              DataType[], localName[#], GridName[#],
               StringMatchQ[ToString[GridName[#]], "eT" ~~ _ ~~ _ ~~ "[" ~~ __ ~~ "]"],
-                "*stress_energy_state"] &,
+                "*stress_energy_state", OptionValue[UseVectors]] &,
             gfsInRHS]],
 
       CommentedBlock["Include user supplied include files",
@@ -580,8 +586,25 @@ equationLoop[eqs_, cleancalc_, gfs_, shorts_, incs_, groups_, pddefs_,
         Map[InfoVariable[#[[1]]] &, (eqs2 /. localMap)],
         ""],
 
+      If[OptionValue[UseVectors], {
+        CommentedBlock["If necessary, store only partial vectors after the first iteration",
+          ConditionalOnParameterTextual["CCTK_REAL_VEC_SIZE > 1 && i<lc_imin",
+            {
+              DeclareAssignVariable["ptrdiff_t", "elt_count", "lc_imin-i"],
+              Map[StoreHighPartialVariableInLoop[GridName[#], localName[#], "elt_count"] &,
+                  gfsInLHS],
+              "continue;\n"
+            }]],
+        CommentedBlock["If necessary, store only partial vectors after the last iteration",
+          ConditionalOnParameterTextual["CCTK_REAL_VEC_SIZE > 1 && i+CCTK_REAL_VEC_SIZE > lc_imax",
+            {
+              DeclareAssignVariable["ptrdiff_t", "elt_count", "lc_imax-i"],
+              Map[StoreLowPartialVariableInLoop[GridName[#], localName[#], "elt_count"] &,
+                  gfsInLHS],
+              "break;\n"
+            }]]}, {}],
       CommentedBlock["Copy local copies back to grid functions",
-        Map[AssignVariableInLoop[GridName[#], localName[#]] &,
+        Map[(If[OptionValue[UseVectors], StoreVariableInLoop, AssignVariableInLoop][GridName[#], localName[#]]) &,
             gfsInLHS]],
 
       If[debugInLoop, Map[InfoVariable[GridName[#]] &, gfsInLHS], ""]}, opts]];