WIP: Lazy stencil expressions for kernel fusion

examples/hasegawa-wakatani/hw.cxx

@@ -3,6 +3,7 @@
 #include <bout/invert_laplace.hxx>
 #include <bout/physicsmodel.hxx>
 #include <bout/smoothing.hxx>
+#include <bout/stencil_expr.hxx>
 
 class HW : public PhysicsModel {
 private:
@@ -110,9 +111,9 @@ class HW : public PhysicsModel {
     }
 
     ddt(n) =
-        -bracket(phi, n, bm) + alpha * (nonzonal_phi - nonzonal_n) - kappa * DDZ(phi);
+        -bracket_arakawa(phi, n) + alpha * (nonzonal_phi - nonzonal_n) - kappa * DDZ(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/derivs.hxx

@@ -31,6 +31,7 @@
 
 #include "bout/field2d.hxx"
 #include "bout/field3d.hxx"
+#include "bout/stencil_expr.hxx"
 #include "bout/vector2d.hxx"
 #include "bout/vector3d.hxx"
 
@@ -102,22 +103,6 @@ Coordinates::FieldMetric DDY(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
-///
-///   \f$\partial / \partial z\f$
-///
-/// @param[in] f       The field to be differentiated
-/// @param[in] outloc  The cell location where the result is desired. If
-///                    staggered grids is not enabled then this has no effect
-///                    If not given, defaults to CELL_DEFAULT
-/// @param[in] method  Differencing method to use. This overrides the 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");
-
 /// Calculate first partial derivative in Z
 ///
 ///   \f$\partial / \partial z\f$
@@ -147,7 +132,11 @@ Coordinates::FieldMetric DDZ(const Field2D& f, CELL_LOC outloc = CELL_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,
-             const std::string& method = "DEFAULT",
+             DIFF_METHOD method = DIFF_DEFAULT,
+             const std::string& region = "RGN_NOBNDRY");
+
+/// Compatibility overload for string-based callers.
+Vector3D DDZ(const Vector3D& f, CELL_LOC outloc, const std::string& method,
              const std::string& region = "RGN_NOBNDRY");
 
 /// Calculate first partial derivative in Z
@@ -163,7 +152,11 @@ Vector3D DDZ(const Vector3D& f, CELL_LOC outloc = CELL_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,
-             const std::string& method = "DEFAULT",
+             DIFF_METHOD method = DIFF_DEFAULT,
+             const std::string& region = "RGN_NOBNDRY");
+
+/// Compatibility overload for string-based callers.
+Vector2D DDZ(const Vector2D& f, CELL_LOC outloc, const std::string& method,
              const std::string& region = "RGN_NOBNDRY");
 
 ////////// SECOND DERIVATIVES //////////

include/bout/field.hxx

@@ -539,53 +539,127 @@ inline BoutReal mean(const BinaryExpr<ResT, L, R, Func>& f, bool allpe = false,
   return bout::reduce::Mean::finalize(state);
 }
 
-/// Exponent: pow(lhs, lhs) is \p lhs raised to the power of \p rhs
-///
-/// This loops over the entire domain, including guard/boundary cells by
-/// default (can be changed using the \p rgn argument)
-/// If CHECK >= 3 then the result will be checked for non-finite numbers
-template <typename T, typename = bout::utils::EnableIfField<T>>
-T pow(const T& lhs, const T& rhs, const std::string& rgn = "RGN_ALL") {
-
-  ASSERT1(areFieldsCompatible(lhs, rhs));
+class Field3DParallel;
+class FieldPerp;
 
-  T result{emptyFrom(lhs)};
+namespace bout::detail {
+template <typename T>
+struct expression_result {
+  using type = std::decay_t<T>;
+};
 
-  BOUT_FOR(i, result.getRegion(rgn)) { result[i] = ::pow(lhs[i], rhs[i]); }
+template <typename ResT, typename L, typename R, typename Func>
+struct expression_result<BinaryExpr<ResT, L, R, Func>> {
+  using type = ResT;
+};
 
-  checkData(result);
-  return result;
-}
+template <typename T>
+using expression_result_t = typename expression_result<std::decay_t<T>>::type;
 
-template <typename T, typename = bout::utils::EnableIfField<T>>
-T pow(const T& lhs, BoutReal rhs, const std::string& rgn = "RGN_ALL") {
+template <typename T>
+inline constexpr bool is_expression_field_v =
+    is_expr_field2d_v<T> || is_expr_field3d_v<T> || is_expr_fieldperp_v<T>;
 
-  // Check if the inputs are allocated
-  checkData(lhs);
-  checkData(rhs);
+template <typename L, typename R>
+inline constexpr bool is_same_expression_rank_v =
+    (is_expr_field2d_v<L> && is_expr_field2d_v<R>)
+    || (is_expr_field3d_v<L> && is_expr_field3d_v<R>)
+    || (is_expr_fieldperp_v<L> && is_expr_fieldperp_v<R>);
 
-  T result{emptyFrom(lhs)};
+template <typename T>
+std::optional<int> getPerpYIndex(const T& value) {
+  if constexpr (std::is_same_v<std::decay_t<T>, ::FieldPerp>) {
+    return value.getIndex();
+  } else {
+    return std::nullopt;
+  }
+}
 
-  BOUT_FOR(i, result.getRegion(rgn)) { result[i] = ::pow(lhs[i], rhs); }
+template <typename ResT, typename L, typename R, typename Func>
+std::optional<int> getPerpYIndex(const BinaryExpr<ResT, L, R, Func>& expr) {
+  if constexpr (std::is_same_v<ResT, ::FieldPerp>) {
+    return expr.getIndex();
+  } else {
+    return std::nullopt;
+  }
+}
 
-  checkData(result);
-  return result;
+template <typename Expr>
+std::optional<size_t> getExpressionRegionID(const Expr& expr,
+                                            const std::string& region_name) {
+  std::optional<size_t> region_id{};
+  if (expr.getMesh()->hasRegion3D(region_name)) {
+    region_id = expr.getMesh()->getRegionID(region_name);
+  }
+  return expr.getMesh()->getCommonRegion(region_id, expr.getRegionID());
 }
 
-template <typename T, typename = bout::utils::EnableIfField<T>>
-T pow(BoutReal lhs, const T& rhs, const std::string& rgn = "RGN_ALL") {
+template <typename L, typename R>
+std::optional<size_t> getExpressionRegionID(const L& lhs, const R& rhs,
+                                            const std::string& region_name) {
+  return lhs.getMesh()->getCommonRegion(getExpressionRegionID(lhs, region_name),
+                                        rhs.getRegionID());
+}
+} // namespace bout::detail
 
-  // Check if the inputs are allocated
-  checkData(lhs);
-  checkData(rhs);
+/// Exponent: pow(lhs, lhs) is \p lhs raised to the power of \p rhs
+///
+/// This loops over the entire domain, including guard/boundary cells by
+/// default (can be changed using the \p rgn argument)
+/// If CHECK >= 3 then the result will be checked for non-finite numbers
+template <typename L, typename R, typename ResT = bout::detail::expression_result_t<L>,
+          typename = std::enable_if_t<bout::detail::is_same_expression_rank_v<L, R>>>
+auto pow(const L& lhs, const R& rhs, const std::string& rgn = "RGN_ALL") {
 
-  // Define and allocate the output result
-  T result{emptyFrom(rhs)};
+  if constexpr (bout::utils::is_Field_v<std::decay_t<L>>
+                && bout::utils::is_Field_v<std::decay_t<R>>) {
+    ASSERT1(areFieldsCompatible(lhs, rhs));
+  } else {
+    ASSERT1_EXPR_COMPATIBLE(lhs, rhs);
+  }
+
+  return BinaryExpr<ResT, L, R, bout::op::Pow>{
+      static_cast<typename L::View>(lhs),
+      static_cast<typename R::View>(rhs),
+      bout::op::Pow{},
+      lhs.getMesh(),
+      lhs.getLocation(),
+      lhs.getDirections(),
+      bout::detail::getExpressionRegionID(lhs, rhs, rgn),
+      lhs.getMesh()->template getRegion<ResT>(rgn),
+      bout::detail::getPerpYIndex(lhs)};
+}
 
-  BOUT_FOR(i, result.getRegion(rgn)) { result[i] = ::pow(lhs, rhs[i]); }
+template <typename T, typename ResT = bout::detail::expression_result_t<T>,
+          typename = std::enable_if_t<bout::detail::is_expression_field_v<T>>>
+auto pow(const T& lhs, BoutReal rhs, const std::string& rgn = "RGN_ALL") {
+
+  return BinaryExpr<ResT, T, Constant<BoutReal>, bout::op::Pow>{
+      static_cast<typename T::View>(lhs),
+      static_cast<typename Constant<BoutReal>::View>(rhs),
+      bout::op::Pow{},
+      lhs.getMesh(),
+      lhs.getLocation(),
+      lhs.getDirections(),
+      bout::detail::getExpressionRegionID(lhs, rgn),
+      lhs.getMesh()->template getRegion<ResT>(rgn),
+      bout::detail::getPerpYIndex(lhs)};
+}
 
-  checkData(result);
-  return result;
+template <typename T, typename ResT = bout::detail::expression_result_t<T>,
+          typename = std::enable_if_t<bout::detail::is_expression_field_v<T>>>
+auto pow(BoutReal lhs, const T& rhs, const std::string& rgn = "RGN_ALL") {
+
+  return BinaryExpr<ResT, Constant<BoutReal>, T, bout::op::Pow>{
+      static_cast<typename Constant<BoutReal>::View>(lhs),
+      static_cast<typename T::View>(rhs),
+      bout::op::Pow{},
+      rhs.getMesh(),
+      rhs.getLocation(),
+      rhs.getDirections(),
+      bout::detail::getExpressionRegionID(rhs, rgn),
+      rhs.getMesh()->template getRegion<ResT>(rgn),
+      bout::detail::getPerpYIndex(rhs)};
 }
 
 /*!
@@ -604,29 +678,6 @@ T pow(BoutReal lhs, const T& rhs, const std::string& rgn = "RGN_ALL") {
  * result for non-finite numbers
  *
  */
-class Field3DParallel;
-class FieldPerp;
-
-namespace bout::detail {
-template <typename T>
-std::optional<int> getPerpYIndex(const T& value) {
-  if constexpr (std::is_same_v<std::decay_t<T>, ::FieldPerp>) {
-    return value.getIndex();
-  } else {
-    return std::nullopt;
-  }
-}
-
-template <typename ResT, typename L, typename R, typename Func>
-std::optional<int> getPerpYIndex(const BinaryExpr<ResT, L, R, Func>& expr) {
-  if constexpr (std::is_same_v<ResT, ::FieldPerp>) {
-    return expr.getIndex();
-  } else {
-    return std::nullopt;
-  }
-}
-} // namespace bout::detail
-
 #ifdef FIELD_FUNC
 #error This macro has already been defined
 #else

include/bout/field3d.hxx

@@ -900,7 +900,24 @@ inline auto operator-(const Field3D& f) {
 /// This loops over the entire domain, including guard/boundary cells by
 /// default (can be changed using the \p rgn argument).
 /// If CHECK >= 3 then the result will be checked for non-finite numbers
-Field3D pow(const Field3D& lhs, const Field2D& rhs, const std::string& rgn = "RGN_ALL");
+template <typename L, typename R>
+std::enable_if_t<is_expr_field3d_v<L> && is_expr_field2d_v<R>,
+                 BinaryExpr<Field3D, L, R, bout::op::Pow>>
+pow(const L& lhs, const R& rhs, const std::string& rgn = "RGN_ALL") {
+  ASSERT1_EXPR_COMPATIBLE(lhs, rhs);
+  auto regionID = bout::detail::getExpressionRegionID(lhs, rhs, rgn);
+  int mesh_nz = lhs.getMesh()->LocalNz;
+  return BinaryExpr<Field3D, L, R, bout::op::Pow>{
+      static_cast<typename L::View>(lhs),
+      static_cast<typename R::View>(rhs).setScale(1, mesh_nz),
+      bout::op::Pow{},
+      lhs.getMesh(),
+      lhs.getLocation(),
+      lhs.getDirections(),
+      regionID,
+      (regionID.has_value() ? lhs.getMesh()->getRegion(regionID.value())
+                            : lhs.getMesh()->getRegion(rgn))};
+}
 FieldPerp pow(const Field3D& lhs, const FieldPerp& rhs,
               const std::string& rgn = "RGN_ALL");
 

include/bout/fieldops.hxx

@@ -118,6 +118,16 @@ struct Div {
     return a / b;
   }
 };
+struct Pow {
+  template <typename LView, typename RView>
+  BOUT_HOST_DEVICE BOUT_FORCEINLINE BoutReal operator()(int idx, const LView& L,
+                                                        const RView& R) const {
+    return ::pow(L(idx), R(idx));
+  }
+  BOUT_HOST_DEVICE BOUT_FORCEINLINE BoutReal operator()(BoutReal a, BoutReal b) const {
+    return ::pow(a, b);
+  }
+};
 struct IfElse {
   bool condition;