WIP: Lazy stencil expressions for kernel fusion

examples/hasegawa-wakatani/hw.cxx

@@ -1,8 +1,14 @@
 
+#include <bout/bout_types.hxx>
 #include <bout/derivs.hxx>
+#include <bout/field3d.hxx>
 #include <bout/invert_laplace.hxx>
 #include <bout/physicsmodel.hxx>
 #include <bout/smoothing.hxx>
+#include <bout/stencil_expr.hxx>
+#include <bout/unused.hxx>
+
+#include <memory>
 
 class HW : public PhysicsModel {
 private:
@@ -109,10 +115,10 @@ class HW : public PhysicsModel {
       nonzonal_phi -= averageY(DC(phi));
     }
 
-    ddt(n) =
-        -bracket(phi, n, bm) + alpha * (nonzonal_phi - nonzonal_n) - kappa * DDZ(phi);
+    ddt(n) = -bracket_arakawa(phi, n) + alpha * (nonzonal_phi - nonzonal_n)
+             - kappa * DDZ_stencil(phi);
 
-    ddt(vort) = -bracket(phi, vort, bm) + alpha * (nonzonal_phi - nonzonal_n);
+    ddt(vort) = -bracket_arakawa(phi, vort) + alpha * (nonzonal_phi - nonzonal_n);
 
     return 0;
   }

examples/performance/ddz/ddz.cxx

@@ -65,17 +65,9 @@ int main(int argc, char** argv) {
   // Nested loops over block data
   ITERATOR_TEST_BLOCK("DDZ Default", result = DDZ(a););
 
-  ITERATOR_TEST_BLOCK("DDZ C2", result = DDZ(a, CELL_DEFAULT, "DIFF_C2"););
+  ITERATOR_TEST_BLOCK("DDZ C2", result = DDZ(a, CELL_DEFAULT, DIFF_C2););
 
-  ITERATOR_TEST_BLOCK("DDZ C4", result = DDZ(a, CELL_DEFAULT, "DIFF_C4"););
-
-  ITERATOR_TEST_BLOCK("DDZ S2", result = DDZ(a, CELL_DEFAULT, "DIFF_S2"););
-
-  ITERATOR_TEST_BLOCK("DDZ W2", result = DDZ(a, CELL_DEFAULT, "DIFF_W2"););
-
-  ITERATOR_TEST_BLOCK("DDZ W3", result = DDZ(a, CELL_DEFAULT, "DIFF_W3"););
-
-  ITERATOR_TEST_BLOCK("DDZ FFT", result = DDZ(a, CELL_DEFAULT, "DIFF_FFT"););
+  ITERATOR_TEST_BLOCK("DDZ C4", result = DDZ(a, CELL_DEFAULT, DIFF_C4););
 
   if (profileMode) {
     int nthreads = 0;

include/bout/coordinates.hxx

@@ -80,6 +80,12 @@ public:
               FieldMetric g_22, FieldMetric g_33, FieldMetric g_12, FieldMetric g_13,
               FieldMetric g_23, FieldMetric ShiftTorsion, FieldMetric IntShiftTorsion);
 
+  ~Coordinates();
+  Coordinates(const Coordinates&) = delete;
+  Coordinates& operator=(const Coordinates&) = delete;
+  Coordinates(Coordinates&&) noexcept = default;
+  Coordinates& operator=(Coordinates&&) noexcept = default;
+
   /// Add variables to \p output_options, for post-processing
   void outputVars(Options& output_options);
 

include/bout/coordinates_accessor.hxx

@@ -54,7 +54,8 @@ struct CoordinatesAccessor {
     d1_dx,
     d1_dy,
     d1_dz, // Grid spacing non-uniformity
-    J,     // Jacobian
+    IntShiftTorsion,
+    J, // Jacobian
     B,
     Byup,
     Bydown, // Magnetic field magnitude
@@ -119,6 +120,7 @@ struct CoordinatesAccessor {
 
   COORD_FN(dx, dy, dz);
   COORD_FN(d1_dx, d1_dy, d1_dz);
+  COORD_FN(IntShiftTorsion);
   COORD_FN(J);
   COORD_FN(B, Byup, Bydown);
   COORD_FN(G1, G3);

include/bout/deriv_store.hxx

@@ -4,10 +4,9 @@
  * Definition of derivative methods storage class
  *
  **************************************************************************
- * Copyright 2018
- *    D.Dickinson, P.Hill, B.Dudson
+ * Copyright 2018 - 2026 BOUT++ contributors
  *
- * Contact: Ben Dudson, bd512@york.ac.uk
+ * Contact: Ben Dudson, dudson2@llnl.gov
  *
  * This file is part of BOUT++.
  *
@@ -26,22 +25,24 @@
  *
  **************************************************************************/
 
-#ifndef __DERIV_STORE_HXX__
-#define __DERIV_STORE_HXX__
+#ifndef DERIV_STORE_HXX
+#define DERIV_STORE_HXX
 
 #include <cstddef>
 #include <functional>
 #include <map>
 #include <set>
+#include <string>
 #include <type_traits>
 #include <unordered_map>
 
-#include "bout/field3d.hxx"
-#include <bout/scorepwrapper.hxx>
-
 #include <bout/bout_types.hxx>
 #include <bout/boutexception.hxx>
+#include <bout/field3d.hxx>
 #include <bout/options.hxx>
+#include <bout/output.hxx>
+#include <bout/scorepwrapper.hxx>
+#include <bout/utils.hxx>
 
 /// Here we have a templated singleton that is used to store DerivativeFunctions
 /// for all types of derivatives. It is templated on the FieldType (2D or 3D) as
@@ -100,9 +101,8 @@ struct DerivativeStore {
     auto key = getKey(direction, stagger, toString(derivType));
     if (isEmpty(key)) {
       return std::set<std::string>{};
-    } else {
-      return registeredMethods.at(key);
     }
+    return registeredMethods.at(key);
   };
 
   /// Outputs a list of all registered method names for the
@@ -308,21 +308,23 @@ struct DerivativeStore {
   };
 
   void initialise(Options* options) {
+    defaultMethods.clear();
+    setDefaults();
 
     // To replicate the existing behaviour we first search for a section called
     //"dd?" and if the option isn't in there we search a section called "diff"
-    auto backupSection = options->getSection("diff");
+    auto* backupSection = options->getSection("diff");
 
-    std::map<DIRECTION, std::string> directions = {{DIRECTION::X, "ddx"},
-                                                   {DIRECTION::Y, "ddy"},
-                                                   {DIRECTION::YOrthogonal, "ddy"},
-                                                   {DIRECTION::Z, "ddz"}};
+    const std::map<DIRECTION, std::string> directions = {{DIRECTION::X, "ddx"},
+                                                         {DIRECTION::Y, "ddy"},
+                                                         {DIRECTION::YOrthogonal, "ddy"},
+                                                         {DIRECTION::Z, "ddz"}};
 
-    std::map<DERIV, std::string> derivTypes = {{DERIV::Standard, "first"},
-                                               {DERIV::StandardSecond, "second"},
-                                               {DERIV::StandardFourth, "fourth"},
-                                               {DERIV::Upwind, "upwind"},
-                                               {DERIV::Flux, "flux"}};
+    const std::map<DERIV, std::string> derivTypes = {{DERIV::Standard, "first"},
+                                                     {DERIV::StandardSecond, "second"},
+                                                     {DERIV::StandardFourth, "fourth"},
+                                                     {DERIV::Upwind, "upwind"},
+                                                     {DERIV::Flux, "flux"}};
 
     for (const auto& direction : directions) {
       for (const auto& deriv : derivTypes) {
@@ -342,7 +344,7 @@ struct DerivativeStore {
         //-------------------------------------------------------------
 
         // The direction specific section to consider
-        auto specificSection = options->getSection(direction.second);
+        auto* specificSection = options->getSection(direction.second);
 
         // Find the appropriate value for theDefault either from
         // the input file or if not found then use the value in
@@ -355,6 +357,11 @@ struct DerivativeStore {
           backupSection->get(derivName, theDefault, "");
         }
 
+        if (uppercase(theDefault) == toString(DIFF_DEFAULT)) {
+          throw BoutException("Default derivative options must resolve to a concrete "
+                              "method");
+        }
+
         // Now we have the default method we should store it in defaultMethods
         theDefault = uppercase(theDefault);
         defaultMethods[getKey(theDirection, STAGGER::None, theDerivTypeString)] =
@@ -377,6 +384,11 @@ struct DerivativeStore {
           specificSection->get(derivName, theDefault, "");
         }
 
+        if (uppercase(theDefault) == toString(DIFF_DEFAULT)) {
+          throw BoutException("Default derivative options must resolve to a concrete "
+                              "method");
+        }
+
         // Now we have the default method we should store it in defaultMethods
         theDefault = uppercase(theDefault);
         defaultMethods[getKey(theDirection, STAGGER::L2C, theDerivTypeString)] =

include/bout/derivs.hxx

@@ -33,6 +33,7 @@
 #include "bout/field2d.hxx"
 #include "bout/field3d.hxx"
 #include "bout/metric_tensor.hxx"
+#include "bout/stencil_expr.hxx"
 #include "bout/vector2d.hxx"
 #include "bout/vector3d.hxx"
 
@@ -53,10 +54,6 @@
 ///                    If not given, defaults to DIFF_DEFAULT
 /// @param[in] region  What region is expected to be calculated
 ///                    If not given, defaults to RGN_NOBNDRY
-Field3D DDX(const Field3D& f, CELL_LOC outloc = CELL_DEFAULT,
-            const std::string& method = "DEFAULT",
-            const std::string& region = "RGN_NOBNDRY");
-
 /// Calculate first partial derivative in X
 ///
 ///   \f$\partial / \partial x\f$
@@ -124,9 +121,9 @@ bout::FieldMetric DDY(const Field2D& f, CELL_LOC outloc = CELL_DEFAULT,
 ///                    If not given, defaults to DIFF_DEFAULT
 /// @param[in] region  What region is expected to be calculated
 ///                    If not given, defaults to RGN_NOBNDRY
-Field3D DDZ(const Field3D& f, CELL_LOC outloc = CELL_DEFAULT,
-            const std::string& method = "DEFAULT",
-            const std::string& region = "RGN_NOBNDRY");
+bout::FieldMetric DDZ(const Field2D& f, CELL_LOC outloc = CELL_DEFAULT,
+                      const std::string& method = "DEFAULT",
+                      const std::string& region = "RGN_NOBNDRY");
 
 /// Calculate first partial derivative in Z
 ///
@@ -140,9 +137,20 @@ Field3D DDZ(const Field3D& f, CELL_LOC outloc = CELL_DEFAULT,
 ///                    If not given, defaults to DIFF_DEFAULT
 /// @param[in] region  What region is expected to be calculated
 ///                    If not given, defaults to RGN_NOBNDRY
-bout::FieldMetric DDZ(const Field2D& f, CELL_LOC outloc = CELL_DEFAULT,
-                      const std::string& method = "DEFAULT",
-                      const std::string& region = "RGN_NOBNDRY");
+Field3D DDZ(const Field3D& f, CELL_LOC outloc = CELL_DEFAULT,
+            DIFF_METHOD method = DIFF_DEFAULT, const std::string& region = "RGN_NOBNDRY");
+
+/// Compatibility overload for string-based callers.
+Field3D DDZ(const Field3D& f, CELL_LOC outloc, const std::string& method,
+            const std::string& region = "RGN_NOBNDRY");
+
+/// Calculate the Z derivative using the FFT-based implementation.
+///
+/// This is kept separate from the stencil-expression `DDZ_stencil` overload
+/// because FFT differentiation cannot be represented by the stencil-expression
+/// machinery.
+Field3D DDZ_FFT(const Field3D& f, CELL_LOC outloc = CELL_DEFAULT,
+                const std::string& region = "RGN_NOBNDRY");
 
 /// Calculate first partial derivative in Z
 ///
@@ -156,6 +164,11 @@ bout::FieldMetric DDZ(const Field2D& f, CELL_LOC outloc = CELL_DEFAULT,
 ///                    If not given, defaults to DIFF_DEFAULT
 /// @param[in] region  What region is expected to be calculated
 ///                    If not given, defaults to RGN_NOBNDRY
+Vector3D DDZ(const Vector3D& f, CELL_LOC outloc = CELL_DEFAULT,
+             DIFF_METHOD method = DIFF_DEFAULT,
+             const std::string& region = "RGN_NOBNDRY");
+
+/// Compatibility overload for string-based callers.
 Vector3D DDZ(const Vector3D& v, CELL_LOC outloc = CELL_DEFAULT,
              const std::string& method = "DEFAULT",
              const std::string& region = "RGN_NOBNDRY");
@@ -172,6 +185,11 @@ Vector3D DDZ(const Vector3D& v, CELL_LOC outloc = CELL_DEFAULT,
 ///                    If not given, defaults to DIFF_DEFAULT
 /// @param[in] region  What region is expected to be calculated
 ///                    If not given, defaults to RGN_NOBNDRY
+Vector2D DDZ(const Vector2D& f, CELL_LOC outloc = CELL_DEFAULT,
+             DIFF_METHOD method = DIFF_DEFAULT,
+             const std::string& region = "RGN_NOBNDRY");
+
+/// Compatibility overload for string-based callers.
 Vector2D DDZ(const Vector2D& v, CELL_LOC outloc = CELL_DEFAULT,
              const std::string& method = "DEFAULT",
              const std::string& region = "RGN_NOBNDRY");