Factor out gridfunction access (phi[index]) function

4 files changed, 18 insertions, 13 deletions

diff --git a/Tools/CodeGen/CaKernel.m b/Tools/CodeGen/CaKernel.m
index 39a5f1f..3b2e159 100644
--- a/Tools/CodeGen/CaKernel.m
+++ b/Tools/CodeGen/CaKernel.m
@@ -77,7 +77,8 @@ DefFn[CaKernelCode[calc_List,opts___] :=
 
     calc2 = Join[calc, {BodyFunction -> (codeBlock[kernel, #] &), 
                        CallerFunction -> False,
-                       LoopFunction -> (codeBlock[kernel<>"_Computations", #] &)}];
+                       LoopFunction -> (codeBlock[kernel<>"_Computations", #] &),
+                       GFAccessFunction -> ({"I3D(",Riffle[{#,0,0,0},","],")"} &)}];
 
     CreateCalculationFunction[calc2,opts]]];
 
diff --git a/Tools/CodeGen/CalculationFunction.m b/Tools/CodeGen/CalculationFunction.m
index 364c4ab..dea7331 100644
--- a/Tools/CodeGen/CalculationFunction.m
+++ b/Tools/CodeGen/CalculationFunction.m
@@ -382,10 +382,11 @@ DefFn[
           shorts, eqs, parameters, parameterRules, odeGroups,
           functionName, dsUsed, groups, pddefs, cleancalc, eqLoop, where,
           addToStencilWidth, pDefs, haveCondTextuals, condTextuals, calc,
-          kernelCall,debug,imp},
+          kernelCall,debug,imp,gridName},
 
   debug = OptionValue[Debug];
   imp = lookup[calcp, Implementation];
+  gridName = lookup[calcp, GFAccessFunction];
 
   functionName = ToString@lookup[calcp, Name];
   bodyFunctionName = functionName <> "_Body";
@@ -600,10 +601,12 @@ DefFn[
           loopFunction, gfsInBoth, gfsDifferentiated,
           gfsDifferentiatedAndOnLHS, declare, eqsReplaced,
           arraysInRHS, arraysInLHS, arraysOnlyInRHS, odeVars,
-          generateEquationCode, groupedIfs, IfThenGroup, noSimplify},
+          generateEquationCode, groupedIfs, IfThenGroup, noSimplify,gridName},
 
     InfoMessage[InfoFull, "Equation loop"];
 
+    gridName = Function[x,FlattenBlock[lookup[cleancalc, GFAccessFunction][x]]];
+
     rhss = Map[#[[2]] &, eqs];
     lhss = Map[#[[1]] &, eqs];
 
@@ -738,7 +741,7 @@ DefFn[
          code
       ];
 
-    assignLocalGridFunctions[gs_, useVectors_, useJacobian_] := assignLocalFunctions[gs, useVectors, useJacobian, GridName];
+    assignLocalGridFunctions[gs_, useVectors_, useJacobian_] := assignLocalFunctions[gs, useVectors, useJacobian, gridName];
     assignLocalArrayFunctions[gs_] := assignLocalFunctions[gs, False, False, ArrayName];
 
     (* separate grid and array variables *)
@@ -782,7 +785,7 @@ DefFn[
 
       Which[OptionValue[UseOpenCL],
               CommentedBlock["Copy local copies back to grid functions",
-                Map[StorePartialVariableInLoop[GridName[#], localName[#]] &, gfsInLHS]],
+                Map[StorePartialVariableInLoop[gridName[#], localName[#]] &, gfsInLHS]],
             OptionValue[UseVectors],
               {
                 CommentedBlock["If necessary, store only partial vectors after the first iteration",
@@ -790,30 +793,30 @@ DefFn[
                   {
                     DeclareAssignVariable["ptrdiff_t", "elt_count_lo", "lc_imin-i"],
                     DeclareAssignVariable["ptrdiff_t", "elt_count_hi", "lc_imax-i"],
-                    Map[StoreMiddlePartialVariableInLoop[GridName[#], localName[#], "elt_count_lo", "elt_count_hi"] &, gfsInLHS],
+                    Map[StoreMiddlePartialVariableInLoop[gridName[#], localName[#], "elt_count_lo", "elt_count_hi"] &, gfsInLHS],
                     "break;\n"
                   }]],
                 CommentedBlock["If necessary, store only partial vectors after the first iteration",
                   ConditionalOnParameterTextual["CCTK_REAL_VEC_SIZE > 1 && CCTK_BUILTIN_EXPECT(i < lc_imin, 0)",
                   {
                     DeclareAssignVariable["ptrdiff_t", "elt_count", "lc_imin-i"],
-                    Map[StoreHighPartialVariableInLoop[GridName[#], localName[#], "elt_count"] &, gfsInLHS],
+                    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 && CCTK_BUILTIN_EXPECT(i+CCTK_REAL_VEC_SIZE > lc_imax, 0)",
                   {
                     DeclareAssignVariable["ptrdiff_t", "elt_count", "lc_imax-i"],
-                    Map[StoreLowPartialVariableInLoop[GridName[#], localName[#], "elt_count"] &, gfsInLHS],
+                    Map[StoreLowPartialVariableInLoop[gridName[#], localName[#], "elt_count"] &, gfsInLHS],
                     "break;\n"
                   }]],
-                Map[StoreVariableInLoop[GridName[#], localName[#]] &, gfsInLHS]
+                Map[StoreVariableInLoop[gridName[#], localName[#]] &, gfsInLHS]
               },
             True,
               CommentedBlock["Copy local copies back to grid functions",
-                Map[AssignVariableInLoop[GridName[#], localName[#]] &, gfsInLHS]]],
+                Map[AssignVariableInLoop[gridName[#], localName[#]] &, gfsInLHS]]],
 
-      If[debugInLoop, Map[InfoVariable[GridName[#]] &, gfsInLHS], ""]}, opts]}]];
+      If[debugInLoop, Map[InfoVariable[gridName[#]] &, gfsInLHS], ""]}, opts]}]];
 
 (* Unsorted *)
 
diff --git a/Tools/CodeGen/Kranc.m b/Tools/CodeGen/Kranc.m
index 345bc3d..04bb427 100644
--- a/Tools/CodeGen/Kranc.m
+++ b/Tools/CodeGen/Kranc.m
@@ -40,7 +40,7 @@ LoopPreIncludes, GroupImplementations, PartialDerivatives, NoSimplify,
 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}
+CallerFunction, LoopFunction, GFAccessFunction}
 
 {ConditionalOnKeyword, ConditionalOnKeywords, CollectList, Interior,
 InteriorNoSync, Boundary, BoundaryWithGhosts, Where, PreDefinitions,
diff --git a/Tools/CodeGen/Thorn.m b/Tools/CodeGen/Thorn.m
index f41c6ad..920dc35 100644
--- a/Tools/CodeGen/Thorn.m
+++ b/Tools/CodeGen/Thorn.m
@@ -540,7 +540,8 @@ CreateSetterSource[calcs_, debug_, include_,
 
    calc = Join[calc, {BodyFunction -> bodyFunction, 
                       CallerFunction -> True,
-                      LoopFunction -> (GenericGridLoop[lookup[calc,Name],#,opts] &)}];
+                      LoopFunction -> (GenericGridLoop[lookup[calc,Name],#,opts] &),
+                     GFAccessFunction -> ({#,"[","index","]"} &)}];
 
    CreateCalculationFunction[calc, opts]}];