[flang] Fold DOT_PRODUCT()

[klausler]

Summary

Implement constant folding of the intrinsic function DOT_PRODUCT().

[jeanPerier]

Looks good

flang/include/flang/Evaluate/fold.h

 // an expression, if it has one.  It returns a pointer, which is
 // const-qualified when the expression is so.  The value can be
 // parenthesized.
 template <typename T, typename EXPR>
 auto UnwrapConstantValue(EXPR &expr) -> common::Constify<Constant<T>, EXPR> * {
   if (auto *c{UnwrapExpr<Constant<T>>(expr)}) {
     return c;
   } else {
-    if constexpr (!std::is_same_v<T, SomeDerived>) {
-      if (auto *parens{UnwrapExpr<Parentheses<T>>(expr)}) {
-        return UnwrapConstantValue<T>(parens->left());
-      }
-    }
     return nullptr;
   }
 }
 
 // GetScalarConstantValue() extracts the known scalar constant value of
 // an expression, if it has one.  The value can be parenthesized.
 template <typename T, typename EXPR>
 constexpr auto GetScalarConstantValue(const EXPR &expr)

flang/lib/Evaluate/fold-complex.cpp

if (args.size() > 0 && args[0].has_value()) {
                 ? std::move(*args[1]->UnwrapExpr())
                 : AsGenericExpr(Constant<Part>{Scalar<Part>{}})};
         return Fold(context,
             Expr<T>{
                 ComplexConstructor<KIND>{ToReal<KIND>(context, std::move(re)),
                     ToReal<KIND>(context, std::move(im))}});
       }
     }
+  } else if (name == "dot_product") {
+    return FoldDotProduct<T>(context, std::move(funcRef));
   } else if (name == "merge") {
     return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "product") {
     auto one{Scalar<Part>::FromInteger(value::Integer<8>{1}).value};
     return FoldProduct<T>(context, std::move(funcRef), Scalar<T>{one});
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   }
-  // TODO: dot_product, matmul
+  // TODO: matmul
   return Expr<T>{std::move(funcRef)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Complex, KIND>> FoldOperation(
     FoldingContext &context, ComplexConstructor<KIND> &&x) {
   if (auto array{ApplyElementwise(context, x)}) {
     return *array;

flang/lib/Evaluate/fold-integer.cpp

if (const auto *cx{UnwrapExpr<Expr<SomeInteger>>(args[0])}) {
           [](const auto &kx) {
             return Scalar<typename ResultType<decltype(kx)>::Part>::DIGITS;
           },
           cx->u)};
     }
   } else if (name == "dim") {
     return FoldElementalIntrinsic<T, T, T>(
         context, std::move(funcRef), &Scalar<T>::DIM);
+  } else if (name == "dot_product") {
+    return FoldDotProduct<T>(context, std::move(funcRef));
   } else if (name == "dshiftl" || name == "dshiftr") {
     const auto fptr{
         name == "dshiftl" ? &Scalar<T>::DSHIFTL : &Scalar<T>::DSHIFTR};
     // Third argument can be of any kind. However, it must be smaller or equal
     // than BIT_SIZE. It can be converted to Int4 to simplify.
     return FoldElementalIntrinsic<T, T, T, Int4>(context, std::move(funcRef),
         ScalarFunc<T, T, T, Int4>(
             [&fptr](const Scalar<T> &i, const Scalar<T> &j,

flang/lib/Evaluate/fold-logical.cpp

if (const auto *ix{UnwrapExpr<Expr<SomeInteger>>(args[0])}) {
                             "POS=%jd out of range for BTEST"_err_en_US,
                             static_cast<std::intmax_t>(posVal));
                       }
                       return Scalar<T>{x.BTEST(posVal)};
                     }));
           },
           ix->u);
     }
+  } else if (name == "dot_product") {
+    return FoldDotProduct<T>(context, std::move(funcRef));
   } else if (name == "extends_type_of") {
     // Type extension testing with EXTENDS_TYPE_OF() ignores any type
     // parameters. Returns a constant truth value when the result is known now.
     if (args[0] && args[1]) {
       auto t0{args[0]->GetType()};
       auto t1{args[1]->GetType()};
       if (t0 && t1) {
         if (auto result{t0->ExtendsTypeOf(*t1)}) {
if (args[0] && args[1]) {
       name == "__builtin_ieee_support_io" ||
       name == "__builtin_ieee_support_nan" ||
       name == "__builtin_ieee_support_sqrt" ||
       name == "__builtin_ieee_support_standard" ||
       name == "__builtin_ieee_support_subnormal" ||
       name == "__builtin_ieee_support_underflow_control") {
     return Expr<T>{true};
   }
-  // TODO: dot_product, is_iostat_end,
+  // TODO: is_iostat_end,
   // is_iostat_eor, logical, matmul, out_of_range,
   // parity
   return Expr<T>{std::move(funcRef)};
 }
 
 template <typename T>
 Expr<LogicalResult> FoldOperation(
     FoldingContext &context, Relational<T> &&relation) {

flang/lib/Evaluate/fold-real.cpp

return FoldElementalIntrinsic<T, T>(context, std::move(funcRef),
           return y.value;
         }));
   } else if (name == "dim") {
     return FoldElementalIntrinsic<T, T, T>(context, std::move(funcRef),
         ScalarFunc<T, T, T>(
             [](const Scalar<T> &x, const Scalar<T> &y) -> Scalar<T> {
               return x.DIM(y).value;
             }));
+  } else if (name == "dot_product") {
+    return FoldDotProduct<T>(context, std::move(funcRef));
   } else if (name == "dprod") {
     if (auto scalars{GetScalarConstantArguments<T, T>(context, args)}) {
       return Fold(context,
           Expr<T>{Multiply<T>{
               Expr<T>{std::get<0>(*scalars)}, Expr<T>{std::get<1>(*scalars)}}});
     }
   } else if (name == "epsilon") {
     return Expr<T>{Scalar<T>::EPSILON()};

flang/lib/Evaluate/fold-reduction.h

 //===-- lib/Evaluate/fold-reduction.h -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-// TODO: DOT_PRODUCT, NORM2, PARITY
+// TODO: NORM2, PARITY
 
 #ifndef FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 #define FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 
 #include "fold-implementation.h"
 
 namespace Fortran::evaluate {
 
-// Fold and validate a DIM= argument.  Returns true (with &dim empty)
-// when DIM= is not present or (with &dim set) when DIM= is present, constant,
-// and valid.  Returns false, possibly with an error message, when
-// DIM= is present but either not constant or not valid.
+// DOT_PRODUCT
+template <typename T>
+static Expr<T> FoldDotProduct(
+    FoldingContext &context, FunctionRef<T> &&funcRef) {
+  using Element = typename Constant<T>::Element;
+  auto args{funcRef.arguments()};
+  CHECK(args.size() == 2);
+  Folder<T> folder{context};
+  Constant<T> *va{folder.Folding(args[0])};
+  Constant<T> *vb{folder.Folding(args[1])};
+  if (va && vb) {
+    CHECK(va->Rank() == 1 && vb->Rank() == 1);
+    if (va->size() != vb->size()) {
+      context.messages().Say(
+          "Vector arguments to DOT_PRODUCT have distinct extents %zd and %zd"_err_en_US,
+          va->size(), vb->size());
+      return MakeInvalidIntrinsic(std::move(funcRef));
+    }
+    Element sum{};
+    bool overflow{false};
+    if constexpr (T::category == TypeCategory::Complex) {
+      std::vector<Element> conjugates;
+      for (const Element &x : va->values()) {
+        conjugates.emplace_back(x.CONJG());
+      }
+      Constant<T> conjgA{
+          std::move(conjugates), ConstantSubscripts{va->shape()}};
+      Expr<T> products{Fold(
+          context, Expr<T>{std::move(conjgA)} * Expr<T>{Constant<T>{*vb}})};
+      Constant<T> &cProducts{DEREF(UnwrapConstantValue<T>(products))};
+      Element correction; // Use Kahan summation for greater precision.
+      const auto &rounding{context.targetCharacteristics().roundingMode()};
+      for (const Element &x : cProducts.values()) {
+        auto next{correction.Add(x, rounding)};
+        overflow |= next.flags.test(RealFlag::Overflow);
+        auto added{sum.Add(next.value, rounding)};
+        overflow |= added.flags.test(RealFlag::Overflow);
+        correction = added.value.Subtract(sum, rounding)
+                         .value.Subtract(next.value, rounding)
+                         .value;
+        sum = std::move(added.value);
+      }
+    } else if constexpr (T::category == TypeCategory::Logical) {
+      Expr<T> conjunctions{Fold(context,
+          Expr<T>{LogicalOperation<T::kind>{LogicalOperator::And,
+              Expr<T>{Constant<T>{*va}}, Expr<T>{Constant<T>{*vb}}}})};
+      Constant<T> &cConjunctions{DEREF(UnwrapConstantValue<T>(conjunctions))};
+      for (const Element &x : cConjunctions.values()) {
+        if (x.IsTrue()) {
+          sum = Element{true};
+          break;
+        }
+      }
+    } else if constexpr (T::category == TypeCategory::Integer) {
+      Expr<T> products{
+          Fold(context, Expr<T>{Constant<T>{*va}} * Expr<T>{Constant<T>{*vb}})};
+      Constant<T> &cProducts{DEREF(UnwrapConstantValue<T>(products))};
+      for (const Element &x : cProducts.values()) {
+        auto next{sum.AddSigned(x)};
+        overflow |= next.overflow;
+        sum = std::move(next.value);
+      }
+    } else { // T::category == TypeCategory::Real
+      Expr<T> products{
+          Fold(context, Expr<T>{Constant<T>{*va}} * Expr<T>{Constant<T>{*vb}})};
+      Constant<T> &cProducts{DEREF(UnwrapConstantValue<T>(products))};
+      Element correction; // Use Kahan summation for greater precision.
+      const auto &rounding{context.targetCharacteristics().roundingMode()};
+      for (const Element &x : cProducts.values()) {
+        auto next{correction.Add(x, rounding)};
+        overflow |= next.flags.test(RealFlag::Overflow);
+        auto added{sum.Add(next.value, rounding)};
+        overflow |= added.flags.test(RealFlag::Overflow);
+        correction = added.value.Subtract(sum, rounding)
+                         .value.Subtract(next.value, rounding)
+                         .value;
+        sum = std::move(added.value);
+      }
+    }
+    if (overflow) {
+      context.messages().Say(
+          "DOT_PRODUCT of %s data overflowed during computation"_warn_en_US,
+          T::AsFortran());
+    }
+    return Expr<T>{Constant<T>{std::move(sum)}};
+  }
+  return Expr<T>{std::move(funcRef)};
+}
+
+// Fold and validate a DIM= argument.  Returns false on error.
 bool CheckReductionDIM(std::optional<int> &dim, FoldingContext &,
     ActualArguments &, std::optional<int> dimIndex, int rank);
 
 // Fold and validate a MASK= argument.  Return null on error, absent MASK=, or
 // non-constant MASK=.
 Constant<LogicalResult> *GetReductionMASK(
     std::optional<ActualArgument> &maskArg, const ConstantSubscripts &shape,
     FoldingContext &);
if (std::optional<Constant<T>> array{
               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
     bool overflow{false};
     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
       if constexpr (T::category == TypeCategory::Integer) {
         auto sum{element.AddSigned(array->At(at))};
         overflow |= sum.overflow;
         element = sum.value;
       } else { // Real & Complex: use Kahan summation
-        auto next{array->At(at).Add(correction)};
+        const auto &rounding{context.targetCharacteristics().roundingMode()};
+        auto next{array->At(at).Add(correction, rounding)};
         overflow |= next.flags.test(RealFlag::Overflow);
-        auto sum{element.Add(next.value)};
+        auto sum{element.Add(next.value, rounding)};
         overflow |= sum.flags.test(RealFlag::Overflow);
         // correction = (sum - element) - next; algebraically zero
-        correction =
-            sum.value.Subtract(element).value.Subtract(next.value).value;
+        correction = sum.value.Subtract(element, rounding)
+                         .value.Subtract(next.value, rounding)
+                         .value;
         element = sum.value;
       }
     }};
     if (overflow) {
       context.messages().Say(
           "SUM() of %s data overflowed"_warn_en_US, T::AsFortran());
     } else {
       return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
     }
   }
   return Expr<T>{std::move(ref)};
 }
 
 } // namespace Fortran::evaluate
 #endif // FORTRAN_EVALUATE_FOLD_REDUCTION_H_

flang/test/Evaluate/fold-dot.f90

+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Tests folding of DOT_PRODUCT()
+module m
+  logical, parameter :: test_i4a = dot_product([(j,j=1,10)],[(j,j=1,10)]) == sum([(j*j,j=1,10)])
+  logical, parameter :: test_r4a = dot_product([(1.*j,j=1,10)],[(j,j=1,10)]) == sum([(j*j,j=1,10)])
+  logical, parameter :: test_z4a = dot_product([((j,j),j=1,10)],[((j,j),j=1,10)]) == sum([(((j,-j)*(j,j)),j=1,10)])
+  logical, parameter :: test_l4a = .not. dot_product([logical::],[logical::])
+  logical, parameter :: test_l4b = .not. dot_product([(j==2,j=1,10)], [(j==3,j=1,10)])
+  logical, parameter :: test_l4c = dot_product([(j==4,j=1,10)], [(j==4,j=1,10)])
+end