[flang] Fold COUNT()

[klausler]

Summary

Complete folding of the intrinsic reduction function COUNT() for all
cases, including partial reductions with DIM= arguments.

[PeteSteinfeld]

Looks great!

flang/lib/Evaluate/CMakeLists.txt

add_flang_library(FortranEvaluate
   expression.cpp
   fold.cpp
   fold-character.cpp
   fold-complex.cpp
   fold-designator.cpp
   fold-integer.cpp
   fold-logical.cpp
   fold-real.cpp
+  fold-reduction.cpp
   formatting.cpp
   host.cpp
   initial-image.cpp
   integer.cpp
   intrinsics.cpp
   intrinsics-library.cpp
   logical.cpp
   real.cpp

flang/lib/Evaluate/fold-implementation.h

if (TotalElementCount(shape) > 0) {
           results.emplace_back(func(std::get<I>(*args)->At(argIndex[I])...));
         }
         (std::get<I>(*args)->IncrementSubscripts(argIndex[I]), ...);
       } while (bounds.IncrementSubscripts(resultIndex));
     }
     // Build and return constant result
     if constexpr (TR::category == TypeCategory::Character) {
       auto len{static_cast<ConstantSubscript>(
-          results.size() ? results[0].length() : 0)};
+          results.empty() ? 0 : results[0].length())};
       return Expr<TR>{Constant<TR>{len, std::move(results), std::move(shape)}};
     } else {
       return Expr<TR>{Constant<TR>{std::move(results), std::move(shape)}};
     }
   }
   return Expr<TR>{std::move(funcRef)};
 }
 
Expr<T> FoldMINorMAX(
   for (auto &arg : funcRef.arguments()) {
     if (auto *cst{Folder<T>{context}.Folding(arg)}) {
       constantArgs.push_back(cst);
     }
   }
   if (constantArgs.size() != funcRef.arguments().size()) {
     return Expr<T>(std::move(funcRef));
   }
-  CHECK(constantArgs.size() > 0);
+  CHECK(!constantArgs.empty());
   Expr<T> result{std::move(*constantArgs[0])};
   for (std::size_t i{1}; i < constantArgs.size(); ++i) {
     Extremum<T> extremum{order, result, Expr<T>{std::move(*constantArgs[i])}};
     result = FoldOperation(context, std::move(extremum));
   }
   return result;
 }
 
Expr<T> FoldArray(ArrayConstructor<T> &&array) {
     return Expr<T>{std::move(array)};
   }
 
 private:
   bool FoldArray(const common::CopyableIndirection<Expr<T>> &expr) {
     Expr<T> folded{Fold(context_, common::Clone(expr.value()))};
     if (const auto *c{UnwrapConstantValue<T>(folded)}) {
       // Copy elements in Fortran array element order
-      if (c->size() > 0) {
+      if (!c->empty()) {
         ConstantSubscripts index{c->lbounds()};
         do {
           elements_.emplace_back(c->At(index));
         } while (c->IncrementSubscripts(index));
       }
       return true;
     } else {
       return false;
bool ArrayConstructorIsFlat(const ArrayConstructorValues<T> &values) {
   }
   return true;
 }
 
 template <typename T>
 std::optional<Expr<T>> AsFlatArrayConstructor(const Expr<T> &expr) {
   if (const auto *c{UnwrapConstantValue<T>(expr)}) {
     ArrayConstructor<T> result{expr};
-    if (c->size() > 0) {
+    if (!c->empty()) {
       ConstantSubscripts at{c->lbounds()};
       do {
         result.Push(Expr<T>{Constant<T>{c->At(at)}});
       } while (c->IncrementSubscripts(at));
     }
     return std::make_optional<Expr<T>>(std::move(result));
   } else if (const auto *a{UnwrapExpr<ArrayConstructor<T>>(expr)}) {
     if (ArrayConstructorIsFlat(*a)) {

flang/lib/Evaluate/fold-integer.cpp

if (int rank{array->Rank()}; rank > 0) {
           }
         }
       }
     }
   }
   return Expr<T>{std::move(funcRef)};
 }
 
+// COUNT()
+template <typename T>
+static Expr<T> FoldCount(FoldingContext &context, FunctionRef<T> &&ref) {
+  static_assert(T::category == TypeCategory::Integer);
+  ActualArguments &arg{ref.arguments()};
+  if (const Constant<LogicalResult> *mask{arg.empty()
+              ? nullptr
+              : Folder<LogicalResult>{context}.Folding(arg[0])}) {
+    std::optional<ConstantSubscript> dim;
+    if (arg.size() > 1 && arg[1]) {
+      dim = CheckDIM(context, arg[1], mask->Rank());
+      if (!dim) {
+        mask = nullptr;
+      }
+    }
+    if (mask) {
+      auto accumulator{[&](Scalar<T> &element, const ConstantSubscripts &at) {
+        if (mask->At(at).IsTrue()) {
+          element = element.AddSigned(Scalar<T>{1}).value;
+        }
+      }};
+      return Expr<T>{DoReduction<T>(*mask, dim, Scalar<T>{}, accumulator)};
+    }
+  }
+  return Expr<T>{std::move(ref)};
+}
+
 // for IALL, IANY, & IPARITY
 template <typename T>
 static Expr<T> FoldBitReduction(FoldingContext &context, FunctionRef<T> &&ref,
     Scalar<T> (Scalar<T>::*operation)(const Scalar<T> &) const,
     Scalar<T> identity) {
   static_assert(T::category == TypeCategory::Integer);
-  using Element = Scalar<T>;
   std::optional<ConstantSubscript> dim;
   if (std::optional<Constant<T>> array{
           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
-    auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
+    auto accumulator{[&](Scalar<T> &element, const ConstantSubscripts &at) {
       element = (element.*operation)(array->At(at));
     }};
-    return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+    return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
   }
   return Expr<T>{std::move(ref)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
     FoldingContext &context,
     FunctionRef<Type<TypeCategory::Integer, KIND>> &&funcRef) {
if (const auto *cx{UnwrapExpr<Expr<SomeReal>>(args[0])}) {
                         "%s intrinsic folding overflow"_en_US, name);
                   }
                   return y.value;
                 }));
           },
           cx->u);
     }
   } else if (name == "count") {
-    if (!args[1]) { // TODO: COUNT(x,DIM=d)
-      if (const auto *constant{UnwrapConstantValue<LogicalResult>(args[0])}) {
-        std::int64_t result{0};
-        for (const auto &element : constant->values()) {
-          if (element.IsTrue()) {
-            ++result;
-          }
-        }
-        return Expr<T>{result};
-      }
-    }
+    return FoldCount<T>(context, std::move(funcRef));
   } else if (name == "digits") {
     if (const auto *cx{UnwrapExpr<Expr<SomeInteger>>(args[0])}) {
       return Expr<T>{std::visit(
           [](const auto &kx) {
             return Scalar<ResultType<decltype(kx)>>::DIGITS;
           },
           cx->u)};
     } else if (const auto *cx{UnwrapExpr<Expr<SomeReal>>(args[0])}) {

flang/lib/Evaluate/fold-logical.cpp

static Expr<T> FoldAllAny(FoldingContext &context, FunctionRef<T> &&ref,
   using Element = Scalar<T>;
   std::optional<ConstantSubscript> dim;
   if (std::optional<Constant<T>> array{
           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
               /*ARRAY(MASK)=*/0, /*DIM=*/1)}) {
     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
       element = (element.*operation)(array->At(at));
     }};
-    return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+    return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
   }
   return Expr<T>{std::move(ref)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Logical, KIND>> FoldIntrinsicFunction(
     FoldingContext &context,
     FunctionRef<Type<TypeCategory::Logical, KIND>> &&funcRef) {

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: ALL, ANY, COUNT, DOT_PRODUCT, FINDLOC, IALL, IANY, IPARITY,
-// NORM2, MAXLOC, MINLOC, PARITY, PRODUCT, SUM
+// TODO: DOT_PRODUCT, FINDLOC, NORM2, MAXLOC, MINLOC, PARITY
 
 #ifndef FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 #define FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 
 #include "fold-implementation.h"
 
 namespace Fortran::evaluate {
 
+// Folds & validates a DIM= actual argument.
+std::optional<ConstantSubscript> CheckDIM(
+    FoldingContext &, std::optional<ActualArgument> &, int rank);
+
 // Common preprocessing for reduction transformational intrinsic function
 // folding.  If the intrinsic can have DIM= &/or MASK= arguments, extract
 // and check them.  If a MASK= is present, apply it to the array data and
 // substitute identity values for elements corresponding to .FALSE. in
 // the mask.  If the result is present, the intrinsic call can be folded.
 template <typename T>
 static std::optional<Constant<T>> ProcessReductionArgs(FoldingContext &context,
     ActualArguments &arg, std::optional<ConstantSubscript> &dim,
     const Scalar<T> &identity, int arrayIndex,
     std::optional<std::size_t> dimIndex = std::nullopt,
     std::optional<std::size_t> maskIndex = std::nullopt) {
   if (arg.empty()) {
     return std::nullopt;
   }
   Constant<T> *folded{Folder<T>{context}.Folding(arg[arrayIndex])};
   if (!folded || folded->Rank() < 1) {
     return std::nullopt;
   }
   if (dimIndex && arg.size() >= *dimIndex + 1 && arg[*dimIndex]) {
-    if (auto *dimConst{
-            Folder<SubscriptInteger>{context}.Folding(arg[*dimIndex])}) {
-      if (auto dimScalar{dimConst->GetScalarValue()}) {
-        dim.emplace(dimScalar->ToInt64());
-        if (*dim < 1 || *dim > folded->Rank()) {
-          context.messages().Say(
-              "DIM=%jd is not valid for an array of rank %d"_err_en_US,
-              static_cast<std::intmax_t>(*dim), folded->Rank());
-          dim.reset();
-        }
-      }
-    }
+    dim = CheckDIM(context, arg[*dimIndex], folded->Rank());
     if (!dim) {
       return std::nullopt;
     }
   }
   if (maskIndex && arg.size() >= *maskIndex + 1 && arg[*maskIndex]) {
     if (Constant<LogicalResult> *
         mask{Folder<LogicalResult>{context}.Folding(arg[*maskIndex])}) {
       if (CheckConformance(context.messages(), AsShape(folded->shape()),
if (maskIndex && arg.size() >= *maskIndex + 1 && arg[*maskIndex]) {
     }
   } else {
     return Constant<T>{*folded};
   }
 }
 
 // Generalized reduction to an array of one dimension fewer (w/ DIM=)
 // or to a scalar (w/o DIM=).
-template <typename T, typename ACCUMULATOR>
-static Constant<T> DoReduction(const Constant<T> &array,
+template <typename T, typename ACCUMULATOR, typename ARRAY>
+static Constant<T> DoReduction(const Constant<ARRAY> &array,
     std::optional<ConstantSubscript> &dim, const Scalar<T> &identity,
     ACCUMULATOR &accumulator) {
   ConstantSubscripts at{array.lbounds()};
   std::vector<typename Constant<T>::Element> elements;
   ConstantSubscripts resultShape; // empty -> scalar
   if (dim) { // DIM= is present, so result is an array
     resultShape = array.shape();
     resultShape.erase(resultShape.begin() + (*dim - 1));
auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
           Expr<T>{Constant<T>{array->At(at)}}, Expr<T>{Constant<T>{element}})};
       auto folded{GetScalarConstantValue<LogicalResult>(
           test.Rewrite(context, std::move(test)))};
       CHECK(folded.has_value());
       if (folded->IsTrue()) {
         element = array->At(at);
       }
     }};
-    return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+    return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
   }
   return Expr<T>{std::move(ref)};
 }
 
 // PRODUCT
 template <typename T>
 static Expr<T> FoldProduct(
     FoldingContext &context, FunctionRef<T> &&ref, Scalar<T> identity) {
auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
         overflow |= prod.flags.test(RealFlag::Overflow);
         element = prod.value;
       }
     }};
     if (overflow) {
       context.messages().Say(
           "PRODUCT() of %s data overflowed"_en_US, T::AsFortran());
     } else {
-      return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+      return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
     }
   }
   return Expr<T>{std::move(ref)};
 }
 
 // SUM
 template <typename T>
 static Expr<T> FoldSum(FoldingContext &context, FunctionRef<T> &&ref) {
auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
             sum.value.Subtract(element).value.Subtract(next.value).value;
         element = sum.value;
       }
     }};
     if (overflow) {
       context.messages().Say(
           "SUM() of %s data overflowed"_en_US, T::AsFortran());
     } else {
-      return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+      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/lib/Evaluate/fold-reduction.cpp

+//===-- lib/Evaluate/fold-reduction.cpp -----------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "fold-reduction.h"
+
+namespace Fortran::evaluate {
+
+std::optional<ConstantSubscript> CheckDIM(
+    FoldingContext &context, std::optional<ActualArgument> &arg, int rank) {
+  if (arg) {
+    if (auto *dimConst{Folder<SubscriptInteger>{context}.Folding(arg)}) {
+      if (auto dimScalar{dimConst->GetScalarValue()}) {
+        auto dim{dimScalar->ToInt64()};
+        if (dim >= 1 && dim <= rank) {
+          return {dim};
+        } else {
+          context.messages().Say(
+              "DIM=%jd is not valid for an array of rank %d"_err_en_US,
+              static_cast<std::intmax_t>(dim), rank);
+        }
+      }
+    }
+  }
+  return std::nullopt;
+}
+
+} // namespace Fortran::evaluate

flang/test/Evaluate/folding29.f90

+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Tests folding of COUNT()
+module m
+  logical, parameter :: arr(3,4) = reshape([(modulo(j, 2) == 1, j = 1, size(arr))], shape(arr))
+  logical, parameter :: test_1 = count([1, 2, 3, 2, 1] < [(j, j=1, 5)]) == 2
+  logical, parameter :: test_2 = count(arr) == 6
+  logical, parameter :: test_3 = all(count(arr, dim=1) == [2, 1, 2, 1])
+  logical, parameter :: test_4 = all(count(arr, dim=2) == [2, 2, 2])
+  logical, parameter :: test_5 = count(logical(arr, kind=1)) == 6
+  logical, parameter :: test_6 = count(logical(arr, kind=2)) == 6
+end module