Merge pull request #399 from E3SM-Project/bartgol/expression-templates-simplifications

Some simplifications to the expression templates design

Author: bartgol

src/expression/CMakeLists.txt

@@ -9,7 +9,7 @@ target_include_directories(ekat_expression INTERFACE
 
 # Set the PUBLIC_HEADER property
 set (HEADERS
-  ekat_expression_base.hpp
+  ekat_expression_meta.hpp
   ekat_expression_binary_op.hpp
   ekat_expression_compare.hpp
   ekat_expression_conditional.hpp

src/expression/ekat_expression_base.hpp

@@ -1,27 +1,33 @@
 #ifndef EKAT_EXPRESSION_BINARY_OP_HPP
 #define EKAT_EXPRESSION_BINARY_OP_HPP
 
-#include "ekat_expression_base.hpp"
+#include "ekat_expression_meta.hpp"
+
+#include <Kokkos_Core.hpp>
 
 namespace ekat {
 
 enum class BinOp {
   Plus,
   Minus,
   Mult,
-  Div,
-  Max,
-  Min
+  Div
 };
 
 template<typename ELeft, typename ERight, BinOp OP>
-class BinaryExpression : public Expression<BinaryExpression<ELeft,ERight,OP>>{
+class BinaryExpression {
 public:
   static constexpr bool expr_l = is_expr_v<ELeft>;
   static constexpr bool expr_r = is_expr_v<ERight>;
 
+  using eval_left_t  = eval_return_t<ELeft>;
+  using eval_right_t = eval_return_t<ERight>;
+
+  using eval_t = std::common_type_t<eval_left_t,eval_right_t>;
+
+  // Don't create an expression from builtin types, just combine them!
   static_assert (expr_l or expr_r,
-    "[CmpExpression] At least one between ELeft and ERight must be an Expression type.\n");
+    "[BinaryExpression] At least one between ELeft and ERight must be an Expression type.\n");
 
   BinaryExpression (const ELeft& left, const ERight& right)
     : m_left(left)
@@ -30,19 +36,27 @@ class BinaryExpression : public Expression<BinaryExpression<ELeft,ERight,OP>>{
     // Nothing to do here
   }
 
-  int num_indices () const {
-    if constexpr (not expr_l) {
-      return m_right.num_indices();
-    } else if constexpr (not expr_r) {
-      return m_left.num_indices();
+  static constexpr int rank () {
+    if constexpr (expr_l) {
+      if constexpr (expr_r) {
+        static_assert(ELeft::rank()==ERight::rank(),
+          "[BinaryExpression] Error! ELeft and ERight are Expression types of different rank.\n");
+      }
+      return ELeft::rank();
     } else {
-      return std::max(m_left.num_indices(),m_right.num_indices());
+      return ERight::rank();
     }
   }
+  int extent (int i) const {
+    if constexpr (expr_l)
+      return m_left.extent(i);
+    else
+      return m_right.extent(i);
+  }
 
   template<typename... Args>
   KOKKOS_INLINE_FUNCTION
-  auto eval(Args... args) const {
+  eval_t eval(Args... args) const {
     if constexpr (not expr_l) {
       return eval_impl(m_left,m_right.eval(args...));
     } else if constexpr (not expr_r) {
@@ -52,12 +66,10 @@ class BinaryExpression : public Expression<BinaryExpression<ELeft,ERight,OP>>{
     }
   }
 
-  static auto ret_type () { return ELeft::ret_type() + ERight::ret_type(); }
 protected:
 
-  template<typename T1, typename T2>
   KOKKOS_INLINE_FUNCTION
-  auto eval_impl(const T1 l, const T2 r) const {
+  eval_t eval_impl (const eval_left_t& l, const eval_right_t& r) const {
     if constexpr (OP==BinOp::Plus) {
       return l+r;
     } else if constexpr (OP==BinOp::Minus) {
@@ -66,71 +78,21 @@ class BinaryExpression : public Expression<BinaryExpression<ELeft,ERight,OP>>{
       return l*r;
     } else if constexpr (OP==BinOp::Div) {
       return l/r;
-    } else if constexpr (OP==BinOp::Max) {
-      return Kokkos::max(l,r);
-    } else if constexpr (OP==BinOp::Min) {
-      return Kokkos::min(l,r);
+      return Kokkos::min(static_cast<const eval_t&>(l),static_cast<const eval_t&>(r));
     }
   }
 
   ELeft    m_left;
   ERight   m_right;
 };
 
+// Specialize meta utils
 template<typename ELeft, typename ERight, BinOp OP>
 struct is_expr<BinaryExpression<ELeft,ERight,OP>> : std::true_type {};
-
-// Unary minus implemented as -1*expr
-template<typename ERight>
-BinaryExpression<int,ERight,BinOp::Mult>
-operator- (const Expression<ERight>& r)
-{
-  return BinaryExpression<int,ERight,BinOp::Mult>(-1,r.cast());
-}
-
-// Overload arithmetic operators
-template<typename ELeft, typename ERight>
-std::enable_if_t<is_expr_v<ELeft> or is_expr_v<ERight>,BinaryExpression<ELeft,ERight,BinOp::Plus>>
-operator+ (const ELeft& l, const ERight& r)
-{
-  return BinaryExpression<ELeft,ERight,BinOp::Plus>(l,r);
-}
-
-template<typename ELeft, typename ERight>
-std::enable_if_t<is_expr_v<ELeft> or is_expr_v<ERight>,BinaryExpression<ELeft,ERight,BinOp::Minus>>
-operator- (const ELeft& l, const ERight& r)
-{
-  return BinaryExpression<ELeft,ERight,BinOp::Minus>(l,r);
-}
-
-template<typename ELeft, typename ERight>
-std::enable_if_t<is_expr_v<ELeft> or is_expr_v<ERight>,BinaryExpression<ELeft,ERight,BinOp::Mult>>
-operator* (const ELeft& l, const ERight& r)
-{
-  return BinaryExpression<ELeft,ERight,BinOp::Mult>(l,r);
-}
-
-template<typename ELeft, typename ERight>
-std::enable_if_t<is_expr_v<ELeft> or is_expr_v<ERight>,BinaryExpression<ELeft,ERight,BinOp::Div>>
-operator/ (const ELeft& l, const ERight& r)
-{
-  return BinaryExpression<ELeft,ERight,BinOp::Div>(l,r);
-}
-
-// Overload max/min functions
-template<typename ELeft, typename ERight>
-std::enable_if_t<is_expr_v<ELeft> or is_expr_v<ERight>,BinaryExpression<ELeft,ERight,BinOp::Max>>
-max (const ELeft& l, const ERight& r)
-{
-  return BinaryExpression<ELeft,ERight,BinOp::Max>(l,r);
-}
-
-template<typename ELeft, typename ERight>
-std::enable_if_t<is_expr_v<ELeft> or is_expr_v<ERight>,BinaryExpression<ELeft,ERight,BinOp::Min>>
-min (const ELeft& l, const ERight& r)
-{
-  return BinaryExpression<ELeft,ERight,BinOp::Min>(l,r);
-}
+template<typename ELeft, typename ERight, BinOp OP>
+struct eval_return<BinaryExpression<ELeft,ERight,OP>> {
+  using type = typename BinaryExpression<ELeft,ERight,OP>::eval_t;
+};
 
 } // namespace ekat